Implement closures and bound methods

compiler.go

@@ -214,6 +214,7 @@ func (c *Compiler) Parse(mainPath string, buildTags []string) error {
 	c.ir.SimpleDCE()                   // remove most dead code
 	c.ir.AnalyseCallgraph()            // set up callgraph
 	c.ir.AnalyseInterfaceConversions() // determine which types are converted to an interface
+	c.ir.AnalyseFunctionPointers()     // determine which function pointer signatures need context
 	c.ir.AnalyseBlockingRecursive()    // make all parents of blocking calls blocking (transitively)
 	c.ir.AnalyseGoCalls()              // check whether we need a scheduler
 
@@ -542,8 +543,18 @@ func (c *Compiler) getLLVMType(goType types.Type) (llvm.Type, error) {
 			}
 			paramTypes = append(paramTypes, subType)
 		}
-		// make a function pointer of it
-		return llvm.PointerType(llvm.FunctionType(returnType, paramTypes, false), 0), nil
+		var ptr llvm.Type
+		if c.ir.SignatureNeedsContext(typ) {
+			// make a closure type (with a function pointer type inside):
+			// {context, funcptr}
+			paramTypes = append(paramTypes, c.i8ptrType)
+			ptr = llvm.PointerType(llvm.FunctionType(returnType, paramTypes, false), 0)
+			ptr = c.ctx.StructType([]llvm.Type{c.i8ptrType, ptr}, false)
+		} else {
+			// make a simple function pointer
+			ptr = llvm.PointerType(llvm.FunctionType(returnType, paramTypes, false), 0)
+		}
+		return ptr, nil
 	case *types.Slice:
 		elemType, err := c.getLLVMType(typ.Elem())
 		if err != nil {
@@ -703,6 +714,12 @@ func (c *Compiler) parseFuncDecl(f *Function) (*Frame, error) {
 		frame.params[param] = i
 	}
 
+	if c.ir.FunctionNeedsContext(f) {
+		// This function gets an extra parameter: the context pointer (for
+		// closures and bound methods). Add it as an extra paramter here.
+		paramTypes = append(paramTypes, c.i8ptrType)
+	}
+
 	fnType := llvm.FunctionType(retType, paramTypes, false)
 
 	name := f.LinkName()
@@ -781,7 +798,13 @@ func (c *Compiler) getInterpretedValue(value Value) (llvm.Value, error) {
 			}
 			return getZeroValue(llvmType)
 		}
-		return c.ir.GetFunction(value.Elem).llvmFn, nil
+		fn := c.ir.GetFunction(value.Elem)
+		ptr := fn.llvmFn
+		if c.ir.SignatureNeedsContext(fn.fn.Signature) {
+			// Create closure value: {context, function pointer}
+			ptr = llvm.ConstStruct([]llvm.Value{llvm.ConstPointerNull(c.i8ptrType), ptr}, false)
+		}
+		return ptr, nil
 
 	case *GlobalValue:
 		zero := llvm.ConstInt(llvm.Int32Type(), 0, false)
@@ -1001,6 +1024,54 @@ func (c *Compiler) parseFunc(frame *Frame) error {
 		frame.locals[param] = llvmParam
 	}
 
+	// Load free variables from the context. This is a closure (or bound
+	// method).
+	if len(frame.fn.fn.FreeVars) != 0 {
+		if !c.ir.FunctionNeedsContext(frame.fn) {
+			panic("free variables on function without context")
+		}
+		c.builder.SetInsertPointAtEnd(frame.blocks[frame.fn.fn.Blocks[0]])
+		context := frame.fn.llvmFn.Param(len(frame.fn.fn.Params))
+
+		// Determine the context type. It's a struct containing all variables.
+		freeVarTypes := make([]llvm.Type, 0, len(frame.fn.fn.FreeVars))
+		for _, freeVar := range frame.fn.fn.FreeVars {
+			typ, err := c.getLLVMType(freeVar.Type())
+			if err != nil {
+				return err
+			}
+			freeVarTypes = append(freeVarTypes, typ)
+		}
+		contextType := llvm.StructType(freeVarTypes, false)
+
+		// Get a correctly-typed pointer to the context.
+		contextAlloc := llvm.Value{}
+		if c.targetData.TypeAllocSize(contextType) <= c.targetData.TypeAllocSize(c.i8ptrType) {
+			// Context stored directly in pointer. Load it using an alloca.
+			contextRawAlloc := c.builder.CreateAlloca(llvm.PointerType(c.i8ptrType, 0), "")
+			contextRawValue := c.builder.CreateBitCast(context, llvm.PointerType(c.i8ptrType, 0), "")
+			c.builder.CreateStore(contextRawValue, contextRawAlloc)
+			contextAlloc = c.builder.CreateBitCast(contextRawAlloc, llvm.PointerType(contextType, 0), "")
+		} else {
+			// Context stored in the heap. Bitcast the passed-in pointer to the
+			// correct pointer type.
+			contextAlloc = c.builder.CreateBitCast(context, llvm.PointerType(contextType, 0), "")
+		}
+
+		// Load each free variable from the context.
+		// A free variable is always a pointer when this is a closure, but it
+		// can be another type when it is a wrapper for a bound method (these
+		// wrappers are generated by the ssa package).
+		for i, freeVar := range frame.fn.fn.FreeVars {
+			indices := []llvm.Value{
+				llvm.ConstInt(llvm.Int32Type(), 0, false),
+				llvm.ConstInt(llvm.Int32Type(), uint64(i), false),
+			}
+			gep := c.builder.CreateInBoundsGEP(contextAlloc, indices, "")
+			frame.locals[freeVar] = c.builder.CreateLoad(gep, "")
+		}
+	}
+
 	if frame.blocking {
 		// Coroutine initialization.
 		c.builder.SetInsertPointAtEnd(frame.blocks[frame.fn.fn.Blocks[0]])
@@ -1372,7 +1443,7 @@ func (c *Compiler) parseBuiltin(frame *Frame, args []ssa.Value, callName string)
 	}
 }
 
-func (c *Compiler) parseFunctionCall(frame *Frame, args []ssa.Value, llvmFn llvm.Value, blocking bool, parentHandle llvm.Value) (llvm.Value, error) {
+func (c *Compiler) parseFunctionCall(frame *Frame, args []ssa.Value, llvmFn, context llvm.Value, blocking bool, parentHandle llvm.Value) (llvm.Value, error) {
 	var params []llvm.Value
 	if blocking {
 		if parentHandle.IsNil() {
@@ -1391,6 +1462,12 @@ func (c *Compiler) parseFunctionCall(frame *Frame, args []ssa.Value, llvmFn llvm
 		params = append(params, val)
 	}
 
+	if !context.IsNil() {
+		// This function takes a context parameter.
+		// Add it to the end of the parameter list.
+		params = append(params, context)
+	}
+
 	if frame.blocking && llvmFn.Name() == "runtime.Sleep" {
 		// Set task state to TASK_STATE_SLEEP and set the duration.
 		c.builder.CreateCall(c.mod.NamedFunction("runtime.sleepTask"), []llvm.Value{frame.taskHandle, params[0]}, "")
@@ -1443,10 +1520,22 @@ func (c *Compiler) parseCall(frame *Frame, instr *ssa.CallCommon, parentHandle l
 		if err != nil {
 			return llvm.Value{}, err
 		}
+
 		llvmFnType, err := c.getLLVMType(instr.Method.Type())
 		if err != nil {
 			return llvm.Value{}, err
 		}
+		if c.ir.SignatureNeedsContext(instr.Method.Type().(*types.Signature)) {
+			// This is somewhat of a hack.
+			// getLLVMType() has created a closure type for us, but we don't
+			// actually want a closure type as an interface call can never be a
+			// closure call. So extract the function pointer type from the
+			// closure.
+			// This happens because somewhere the same function signature is
+			// used in a closure or bound method.
+			llvmFnType = llvmFnType.Subtypes()[1]
+		}
+
 		values := []llvm.Value{
 			itf,
 			llvm.ConstInt(llvm.Int16Type(), uint64(c.ir.MethodNum(instr.Method)), false),
@@ -1454,6 +1543,7 @@ func (c *Compiler) parseCall(frame *Frame, instr *ssa.CallCommon, parentHandle l
 		fn := c.builder.CreateCall(c.mod.NamedFunction("runtime.interfaceMethod"), values, "invoke.func")
 		fnCast := c.builder.CreateBitCast(fn, llvmFnType, "invoke.func.cast")
 		receiverValue := c.builder.CreateExtractValue(itf, 1, "invoke.func.receiver")
+
 		args := []llvm.Value{receiverValue}
 		for _, arg := range instr.Args {
 			val, err := c.parseExpr(frame, arg)
@@ -1462,16 +1552,21 @@ func (c *Compiler) parseCall(frame *Frame, instr *ssa.CallCommon, parentHandle l
 			}
 			args = append(args, val)
 		}
+		if c.ir.SignatureNeedsContext(instr.Method.Type().(*types.Signature)) {
+			// This function takes an extra context parameter. An interface call
+			// cannot also be a closure but we have to supply the nil pointer
+			// anyway.
+			args = append(args, llvm.ConstPointerNull(c.i8ptrType))
+		}
+
 		// TODO: blocking methods (needs analysis)
 		return c.builder.CreateCall(fnCast, args, ""), nil
 	}
 
-	// Regular function, builtin, or function pointer.
-	switch call := instr.Value.(type) {
-	case *ssa.Builtin:
-		return c.parseBuiltin(frame, instr.Args, call.Name())
-	case *ssa.Function:
-		if call.Name() == "Asm" && len(instr.Args) == 1 {
+	// Try to call the function directly for trivially static calls.
+	fn := instr.StaticCallee()
+	if fn != nil {
+		if fn.Name() == "Asm" && len(instr.Args) == 1 {
 			// Magic function: insert inline assembly instead of calling it.
 			if named, ok := instr.Args[0].Type().(*types.Named); ok && named.Obj().Name() == "__asm" {
 				fnType := llvm.FunctionType(llvm.VoidType(), []llvm.Type{}, false)
@@ -1480,20 +1575,52 @@ func (c *Compiler) parseCall(frame *Frame, instr *ssa.CallCommon, parentHandle l
 				return c.builder.CreateCall(target, nil, ""), nil
 			}
 		}
-		targetFunc := c.ir.GetFunction(call)
-		name := targetFunc.LinkName()
-		llvmFn := c.mod.NamedFunction(name)
-		if llvmFn.IsNil() {
-			return llvm.Value{}, errors.New("undefined function: " + name)
+		targetFunc := c.ir.GetFunction(fn)
+		if targetFunc.llvmFn.IsNil() {
+			return llvm.Value{}, errors.New("undefined function: " + targetFunc.LinkName())
 		}
-		return c.parseFunctionCall(frame, instr.Args, llvmFn, targetFunc.blocking, parentHandle)
+		var context llvm.Value
+		if c.ir.FunctionNeedsContext(targetFunc) {
+			// This function call is to a (potential) closure, not a regular
+			// function. See whether it is a closure and if so, call it as such.
+			// Else, supply a dummy nil pointer as the last parameter.
+			var err error
+			if mkClosure, ok := instr.Value.(*ssa.MakeClosure); ok {
+				// closure is {context, function pointer}
+				closure, err := c.parseExpr(frame, mkClosure)
+				if err != nil {
+					return llvm.Value{}, err
+				}
+				context = c.builder.CreateExtractValue(closure, 0, "")
+			} else {
+				context, err = getZeroValue(c.i8ptrType)
+				if err != nil {
+					return llvm.Value{}, err
+				}
+			}
+		}
+		return c.parseFunctionCall(frame, instr.Args, targetFunc.llvmFn, context, targetFunc.blocking, parentHandle)
+	}
+
+	// Builtin or function pointer.
+	switch call := instr.Value.(type) {
+	case *ssa.Builtin:
+		return c.parseBuiltin(frame, instr.Args, call.Name())
 	default: // function pointer
 		value, err := c.parseExpr(frame, instr.Value)
 		if err != nil {
 			return llvm.Value{}, err
 		}
 		// TODO: blocking function pointers (needs analysis)
-		return c.parseFunctionCall(frame, instr.Args, value, false, parentHandle)
+		var context llvm.Value
+		if c.ir.SignatureNeedsContext(instr.Signature()) {
+			// 'value' is a closure, not a raw function pointer.
+			// Extract the function pointer and the context pointer.
+			// closure: {context, function pointer}
+			context = c.builder.CreateExtractValue(value, 0, "")
+			value = c.builder.CreateExtractValue(value, 1, "")
+		}
+		return c.parseFunctionCall(frame, instr.Args, value, context, false, parentHandle)
 	}
 }
 
@@ -1583,7 +1710,17 @@ func (c *Compiler) parseExpr(frame *Frame, expr ssa.Value) (llvm.Value, error) {
 		}
 		return c.builder.CreateGEP(val, indices, ""), nil
 	case *ssa.Function:
-		return c.mod.NamedFunction(c.ir.GetFunction(expr).LinkName()), nil
+		fn := c.ir.GetFunction(expr)
+		ptr := fn.llvmFn
+		if c.ir.FunctionNeedsContext(fn) {
+			// Create closure for function pointer.
+			// Closure is: {context, function pointer}
+			ptr = llvm.ConstStruct([]llvm.Value{
+				llvm.ConstPointerNull(c.i8ptrType),
+				ptr,
+			}, false)
+		}
+		return ptr, nil
 	case *ssa.Global:
 		if strings.HasPrefix(expr.Name(), "__cgofn__cgo_") || strings.HasPrefix(expr.Name(), "_cgo_") {
 			// Ignore CGo global variables which we don't use.
@@ -1731,6 +1868,10 @@ func (c *Compiler) parseExpr(frame *Frame, expr ssa.Value) (llvm.Value, error) {
 		default:
 			panic("unknown lookup type: " + expr.String())
 		}
+
+	case *ssa.MakeClosure:
+		return c.parseMakeClosure(frame, expr)
+
 	case *ssa.MakeInterface:
 		val, err := c.parseExpr(frame, expr.X)
 		if err != nil {
@@ -2249,6 +2390,85 @@ func (c *Compiler) parseConvert(typeFrom, typeTo types.Type, value llvm.Value) (
 	}
 }
 
+func (c *Compiler) parseMakeClosure(frame *Frame, expr *ssa.MakeClosure) (llvm.Value, error) {
+	if len(expr.Bindings) == 0 {
+		panic("unexpected: MakeClosure without bound variables")
+	}
+	f := c.ir.GetFunction(expr.Fn.(*ssa.Function))
+	if !c.ir.FunctionNeedsContext(f) {
+		// Maybe AnalyseFunctionPointers didn't run?
+		panic("MakeClosure on function signature without context")
+	}
+
+	// Collect all bound variables.
+	boundVars := make([]llvm.Value, 0, len(expr.Bindings))
+	boundVarTypes := make([]llvm.Type, 0, len(expr.Bindings))
+	for _, binding := range expr.Bindings {
+		// The context stores the bound variables.
+		llvmBoundVar, err := c.parseExpr(frame, binding)
+		if err != nil {
+			return llvm.Value{}, err
+		}
+		boundVars = append(boundVars, llvmBoundVar)
+		boundVarTypes = append(boundVarTypes, llvmBoundVar.Type())
+	}
+	contextType := llvm.StructType(boundVarTypes, false)
+
+	// Allocate memory for the context.
+	contextAlloc := llvm.Value{}
+	contextHeapAlloc := llvm.Value{}
+	if c.targetData.TypeAllocSize(contextType) <= c.targetData.TypeAllocSize(c.i8ptrType) {
+		// Context fits in a pointer - e.g. when it is a pointer. Store it
+		// directly in the stack after a convert.
+		// Because contextType is a struct and we have to cast it to a *i8,
+		// store it in an alloca first for bitcasting (store+bitcast+load).
+		contextAlloc = c.builder.CreateAlloca(contextType, "")
+	} else {
+		// Context is bigger than a pointer, so allocate it on the heap.
+		size := c.targetData.TypeAllocSize(contextType)
+		sizeValue := llvm.ConstInt(c.uintptrType, size, false)
+		contextHeapAlloc = c.builder.CreateCall(c.allocFunc, []llvm.Value{sizeValue}, "")
+		contextAlloc = c.builder.CreateBitCast(contextHeapAlloc, llvm.PointerType(contextType, 0), "")
+	}
+
+	// Store all bound variables in the alloca or heap pointer.
+	for i, boundVar := range boundVars {
+		indices := []llvm.Value{
+			llvm.ConstInt(llvm.Int32Type(), 0, false),
+			llvm.ConstInt(llvm.Int32Type(), uint64(i), false),
+		}
+		gep := c.builder.CreateInBoundsGEP(contextAlloc, indices, "")
+		c.builder.CreateStore(boundVar, gep)
+	}
+
+	context := llvm.Value{}
+	if c.targetData.TypeAllocSize(contextType) <= c.targetData.TypeAllocSize(c.i8ptrType) {
+		// Load value (as *i8) from the alloca.
+		contextAlloc = c.builder.CreateBitCast(contextAlloc, llvm.PointerType(c.i8ptrType, 0), "")
+		context = c.builder.CreateLoad(contextAlloc, "")
+	} else {
+		// Get the original heap allocation pointer, which already is an
+		// *i8.
+		context = contextHeapAlloc
+	}
+
+	// Get the function signature type, which is a closure type.
+	// A closure is a tuple of {context, function pointer}.
+	typ, err := c.getLLVMType(f.fn.Signature)
+	if err != nil {
+		return llvm.Value{}, err
+	}
+
+	// Create the closure, which is a struct: {context, function pointer}.
+	closure, err := getZeroValue(typ)
+	if err != nil {
+		return llvm.Value{}, err
+	}
+	closure = c.builder.CreateInsertValue(closure, f.llvmFn, 1, "")
+	closure = c.builder.CreateInsertValue(closure, context, 0, "")
+	return closure, nil
+}
+
 func (c *Compiler) parseMakeInterface(val llvm.Value, typ types.Type, isConst bool) (llvm.Value, error) {
 	var itfValue llvm.Value
 	size := c.targetData.TypeAllocSize(val.Type())