compiler: add support for 'go' on func values

This commit allows starting a new goroutine directly from a func value,
not just when the static callee is known.

This is necessary to support the whole time package, not just the
commonly used subset that was compiled with the SimpleDCE pass enabled.

Author: aykevl

compiler/calls.go

@@ -12,6 +12,19 @@ import (
 // a struct contains more fields, it is passed as a struct without expanding.
 const MaxFieldsPerParam = 3
 
+// createGoCall starts a new goroutine with the provided function pointer and parameters.
+// Because a go statement doesn't return anything, return undef.
+func (c *Compiler) createGoCall(funcPtr llvm.Value, params []llvm.Value) llvm.Value {
+	// We roundtrip through runtime.makeGoroutine as a signal (to find these
+	// calls) and to break any optimizations LLVM will try to do: they are
+	// invalid if we called this as a regular function to be updated later.
+	calleeValue := c.builder.CreatePtrToInt(funcPtr, c.uintptrType, "")
+	calleeValue = c.createRuntimeCall("makeGoroutine", []llvm.Value{calleeValue}, "")
+	calleeValue = c.builder.CreateIntToPtr(calleeValue, funcPtr.Type(), "")
+	c.createCall(calleeValue, params, "")
+	return llvm.Undef(funcPtr.Type().ElementType().ReturnType())
+}
+
 // Shortcut: create a call to runtime.<fnName> with the given arguments.
 func (c *Compiler) createRuntimeCall(fnName string, args []llvm.Value, name string) llvm.Value {
 	runtimePkg := c.ir.Program.ImportedPackage("runtime")

compiler/compiler.go

@@ -1034,32 +1034,33 @@ func (c *Compiler) parseInstr(frame *Frame, instr ssa.Instruction) {
 			c.addError(instr.Pos(), "todo: go on method receiver")
 			return
 		}
-		callee := instr.Call.StaticCallee()
-		if callee == nil {
-			c.addError(instr.Pos(), "todo: go on non-direct function (function pointer, etc.)")
-			return
-		}
-		calleeFn := c.ir.GetFunction(callee)
-
-		// Mark this function as a 'go' invocation and break invalid
-		// interprocedural optimizations. For example, heap-to-stack
-		// transformations are not sound as goroutines can outlive their parent.
-		calleeType := calleeFn.LLVMFn.Type()
-		calleeValue := c.builder.CreatePtrToInt(calleeFn.LLVMFn, c.uintptrType, "")
-		calleeValue = c.createRuntimeCall("makeGoroutine", []llvm.Value{calleeValue}, "")
-		calleeValue = c.builder.CreateIntToPtr(calleeValue, calleeType, "")
 
 		// Get all function parameters to pass to the goroutine.
 		var params []llvm.Value
 		for _, param := range instr.Call.Args {
 			params = append(params, c.getValue(frame, param))
 		}
-		if !calleeFn.IsExported() {
-			params = append(params, llvm.Undef(c.i8ptrType)) // context parameter
-			params = append(params, llvm.Undef(c.i8ptrType)) // parent coroutine handle
-		}
 
-		c.createCall(calleeValue, params, "")
+		// Mark this function as a 'go' invocation and break invalid
+		// interprocedural optimizations. For example, heap-to-stack
+		// transformations are not sound as goroutines can outlive their parent.
+		if callee := instr.Call.StaticCallee(); callee != nil {
+			// Static callee is known: this is a regular function call.
+			calleeFn := c.ir.GetFunction(callee)
+			if !calleeFn.IsExported() {
+				params = append(params, llvm.Undef(c.i8ptrType))     // context parameter
+				params = append(params, c.getZeroValue(c.i8ptrType)) // parent coroutine handle
+			}
+			c.createGoCall(calleeFn.LLVMFn, params)
+		} else if !instr.Call.IsInvoke() {
+			// Start a new goroutine by calling a function pointer.
+			funcPtr, context := c.decodeFuncValue(c.getValue(frame, instr.Call.Value), instr.Call.Value.Type().(*types.Signature))
+			params = append(params, context)                     // context parameter
+			params = append(params, c.getZeroValue(c.i8ptrType)) // parent coroutine handle
+			c.createGoCall(funcPtr, params)
+		} else {
+			c.addError(instr.Pos(), "todo: go on interface call")
+		}
 	case *ssa.If:
 		cond := c.getValue(frame, instr.Cond)
 		block := instr.Block()

compiler/func-lowering.go

@@ -152,11 +152,11 @@ func (c *Compiler) LowerFuncValues() {
 				// There are multiple functions used in a func value that
 				// implement this signature.
 				// What we'll do is transform the following:
-				//     rawPtr := runtime.getFuncPtr(fn)
-				//     if func.rawPtr == nil {
+				//     rawPtr := runtime.getFuncPtr(func.ptr)
+				//     if rawPtr == nil {
 				//         runtime.nilPanic()
 				//     }
-				//     result := func.rawPtr(...args, func.context)
+				//     result := rawPtr(...args, func.context)
 				// into this:
 				//     if false {
 				//         runtime.nilPanic()
@@ -175,95 +175,111 @@ func (c *Compiler) LowerFuncValues() {
 
 				// Remove some casts, checks, and the old call which we're going
 				// to replace.
-				var funcCall llvm.Value
-				for _, inttoptr := range getUses(getFuncPtrCall) {
-					if inttoptr.IsAIntToPtrInst().IsNil() {
+				for _, callIntPtr := range getUses(getFuncPtrCall) {
+					if !callIntPtr.IsACallInst().IsNil() && callIntPtr.CalledValue().Name() == "runtime.makeGoroutine" {
+						for _, inttoptr := range getUses(callIntPtr) {
+							if inttoptr.IsAIntToPtrInst().IsNil() {
+								panic("expected a inttoptr")
+							}
+							for _, use := range getUses(inttoptr) {
+								c.addFuncLoweringSwitch(funcID, use, c.createGoCall, functions)
+								use.EraseFromParentAsInstruction()
+							}
+							inttoptr.EraseFromParentAsInstruction()
+						}
+						callIntPtr.EraseFromParentAsInstruction()
+						continue
+					}
+					if callIntPtr.IsAIntToPtrInst().IsNil() {
 						panic("expected inttoptr")
 					}
-					for _, ptrUse := range getUses(inttoptr) {
+					for _, ptrUse := range getUses(callIntPtr) {
 						if !ptrUse.IsABitCastInst().IsNil() {
 							for _, bitcastUse := range getUses(ptrUse) {
-								if bitcastUse.IsACallInst().IsNil() || bitcastUse.CalledValue().Name() != "runtime.isnil" {
+								if bitcastUse.IsACallInst().IsNil() || bitcastUse.CalledValue().IsAFunction().IsNil() {
+									panic("expected a call instruction")
+								}
+								switch bitcastUse.CalledValue().Name() {
+								case "runtime.isnil":
+									bitcastUse.ReplaceAllUsesWith(llvm.ConstInt(c.ctx.Int1Type(), 0, false))
+									bitcastUse.EraseFromParentAsInstruction()
+								default:
 									panic("expected a call to runtime.isnil")
 								}
-								bitcastUse.ReplaceAllUsesWith(llvm.ConstInt(c.ctx.Int1Type(), 0, false))
-								bitcastUse.EraseFromParentAsInstruction()
 							}
-							ptrUse.EraseFromParentAsInstruction()
-						} else if !ptrUse.IsACallInst().IsNil() && ptrUse.CalledValue() == inttoptr {
-							if !funcCall.IsNil() {
-								panic("multiple calls on a single runtime.getFuncPtr")
-							}
-							funcCall = ptrUse
+						} else if !ptrUse.IsACallInst().IsNil() && ptrUse.CalledValue() == callIntPtr {
+							c.addFuncLoweringSwitch(funcID, ptrUse, func(funcPtr llvm.Value, params []llvm.Value) llvm.Value {
+								return c.builder.CreateCall(funcPtr, params, "")
+							}, functions)
 						} else {
 							panic("unexpected getFuncPtrCall")
 						}
+						ptrUse.EraseFromParentAsInstruction()
 					}
+					callIntPtr.EraseFromParentAsInstruction()
 				}
-				if funcCall.IsNil() {
-					panic("expected exactly one call use of a runtime.getFuncPtr")
-				}
-
-				// The block that cannot be reached with correct funcValues (to
-				// help the optimizer).
-				c.builder.SetInsertPointBefore(funcCall)
-				defaultBlock := llvm.AddBasicBlock(funcCall.InstructionParent().Parent(), "func.default")
-				c.builder.SetInsertPointAtEnd(defaultBlock)
-				c.builder.CreateUnreachable()
+				getFuncPtrCall.EraseFromParentAsInstruction()
+			}
+		}
+	}
+}
 
-				// Create the switch.
-				c.builder.SetInsertPointBefore(funcCall)
-				sw := c.builder.CreateSwitch(funcID, defaultBlock, len(functions)+1)
+// addFuncLoweringSwitch creates a new switch on a function ID and inserts calls
+// to the newly created direct calls. The funcID is the number to switch on,
+// call is the call instruction to replace, and createCall is the callback that
+// actually creates the new call. By changing createCall to something other than
+// c.builder.CreateCall, instead of calling a function it can start a new
+// goroutine for example.
+func (c *Compiler) addFuncLoweringSwitch(funcID, call llvm.Value, createCall func(funcPtr llvm.Value, params []llvm.Value) llvm.Value, functions funcWithUsesList) {
+	// The block that cannot be reached with correct funcValues (to help the
+	// optimizer).
+	c.builder.SetInsertPointBefore(call)
+	defaultBlock := llvm.AddBasicBlock(call.InstructionParent().Parent(), "func.default")
+	c.builder.SetInsertPointAtEnd(defaultBlock)
+	c.builder.CreateUnreachable()
 
-				// Split right after the switch. We will need to insert a few
-				// basic blocks in this gap.
-				nextBlock := c.splitBasicBlock(sw, llvm.NextBasicBlock(sw.InstructionParent()), "func.next")
+	// Create the switch.
+	c.builder.SetInsertPointBefore(call)
+	sw := c.builder.CreateSwitch(funcID, defaultBlock, len(functions)+1)
 
-				// The 0 case, which is actually a nil check.
-				nilBlock := llvm.InsertBasicBlock(nextBlock, "func.nil")
-				c.builder.SetInsertPointAtEnd(nilBlock)
-				c.createRuntimeCall("nilPanic", nil, "")
-				c.builder.CreateUnreachable()
-				sw.AddCase(llvm.ConstInt(c.uintptrType, 0, false), nilBlock)
+	// Split right after the switch. We will need to insert a few basic blocks
+	// in this gap.
+	nextBlock := c.splitBasicBlock(sw, llvm.NextBasicBlock(sw.InstructionParent()), "func.next")
 
-				// Gather the list of parameters for every call we're going to
-				// make.
-				callParams := make([]llvm.Value, funcCall.OperandsCount()-1)
-				for i := range callParams {
-					callParams[i] = funcCall.Operand(i)
-				}
+	// The 0 case, which is actually a nil check.
+	nilBlock := llvm.InsertBasicBlock(nextBlock, "func.nil")
+	c.builder.SetInsertPointAtEnd(nilBlock)
+	c.createRuntimeCall("nilPanic", nil, "")
+	c.builder.CreateUnreachable()
+	sw.AddCase(llvm.ConstInt(c.uintptrType, 0, false), nilBlock)
 
-				// If the call produces a value, we need to get it using a PHI
-				// node.
-				phiBlocks := make([]llvm.BasicBlock, len(functions))
-				phiValues := make([]llvm.Value, len(functions))
-				for i, fn := range functions {
-					// Insert a switch case.
-					bb := llvm.InsertBasicBlock(nextBlock, "func.call"+strconv.Itoa(fn.id))
-					c.builder.SetInsertPointAtEnd(bb)
-					result := c.builder.CreateCall(fn.funcPtr, callParams, "")
-					c.builder.CreateBr(nextBlock)
-					sw.AddCase(llvm.ConstInt(c.uintptrType, uint64(fn.id), false), bb)
-					phiBlocks[i] = bb
-					phiValues[i] = result
-				}
-				// Create the PHI node so that the call result flows into the
-				// next block (after the split). This is only necessary when the
-				// call produced a value.
-				if funcCall.Type().TypeKind() != llvm.VoidTypeKind {
-					c.builder.SetInsertPointBefore(nextBlock.FirstInstruction())
-					phi := c.builder.CreatePHI(funcCall.Type(), "")
-					phi.AddIncoming(phiValues, phiBlocks)
-					funcCall.ReplaceAllUsesWith(phi)
-				}
+	// Gather the list of parameters for every call we're going to make.
+	callParams := make([]llvm.Value, call.OperandsCount()-1)
+	for i := range callParams {
+		callParams[i] = call.Operand(i)
+	}
 
-				// Finally, remove the old instructions.
-				funcCall.EraseFromParentAsInstruction()
-				for _, inttoptr := range getUses(getFuncPtrCall) {
-					inttoptr.EraseFromParentAsInstruction()
-				}
-				getFuncPtrCall.EraseFromParentAsInstruction()
-			}
-		}
+	// If the call produces a value, we need to get it using a PHI
+	// node.
+	phiBlocks := make([]llvm.BasicBlock, len(functions))
+	phiValues := make([]llvm.Value, len(functions))
+	for i, fn := range functions {
+		// Insert a switch case.
+		bb := llvm.InsertBasicBlock(nextBlock, "func.call"+strconv.Itoa(fn.id))
+		c.builder.SetInsertPointAtEnd(bb)
+		result := createCall(fn.funcPtr, callParams)
+		c.builder.CreateBr(nextBlock)
+		sw.AddCase(llvm.ConstInt(c.uintptrType, uint64(fn.id), false), bb)
+		phiBlocks[i] = bb
+		phiValues[i] = result
+	}
+	// Create the PHI node so that the call result flows into the
+	// next block (after the split). This is only necessary when the
+	// call produced a value.
+	if call.Type().TypeKind() != llvm.VoidTypeKind {
+		c.builder.SetInsertPointBefore(nextBlock.FirstInstruction())
+		phi := c.builder.CreatePHI(call.Type(), "")
+		phi.AddIncoming(phiValues, phiBlocks)
+		call.ReplaceAllUsesWith(phi)
 	}
 }

compiler/func.go

@@ -32,7 +32,7 @@ const (
 // funcImplementation picks an appropriate func value implementation for the
 // target.
 func (c *Compiler) funcImplementation() funcValueImplementation {
-	if c.GOARCH == "wasm" {
+	if c.GOARCH == "wasm" || true {
 		return funcValueSwitch
 	} else {
 		return funcValueDoubleword

testdata/coroutines.go

@@ -28,6 +28,19 @@ func main() {
 	var printer Printer
 	printer = &myPrinter{}
 	printer.Print()
+
+	sleepFuncValue(func(x int) {
+		time.Sleep(1 * time.Millisecond)
+		println("slept inside func pointer", x)
+	})
+	time.Sleep(1 * time.Millisecond)
+	n := 20
+	sleepFuncValue(func(x int) {
+		time.Sleep(1 * time.Millisecond)
+		println("slept inside closure, with value:", n, x)
+	})
+
+	time.Sleep(2 * time.Millisecond)
 }
 
 func sub() {
@@ -47,6 +60,10 @@ func delayedValue() int {
 	return 42
 }
 
+func sleepFuncValue(fn func(int)) {
+	go fn(8)
+}
+
 func nowait() {
 	println("non-blocking goroutine")
 }
@@ -55,7 +72,7 @@ type Printer interface {
 	Print()
 }
 
-type myPrinter struct{
+type myPrinter struct {
 }
 
 func (i *myPrinter) Print() {

testdata/coroutines.txt

@@ -11,3 +11,5 @@ value produced after some time: 42
 non-blocking goroutine
 done with non-blocking goroutine
 async interface method call
+slept inside func pointer 8
+slept inside closure, with value: 20 8