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/compiler.go

@@ -1059,32 +1059,47 @@ func (c *Compiler) parseInstr(frame *Frame, instr ssa.Instruction) {
 	case *ssa.Defer:
 		c.emitDefer(frame, instr)
 	case *ssa.Go:
-		if instr.Call.IsInvoke() {
-			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)
-
 		// 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() && c.selectScheduler() != "tasks" {
-			// For coroutine scheduling, this is only required when calling an
-			// external function.
-			// For tasks, because all params are stored in a single object, no
-			// unnecessary parameters should be stored anyway.
-			params = append(params, llvm.Undef(c.i8ptrType)) // context parameter
-			params = append(params, llvm.Undef(c.i8ptrType)) // parent coroutine handle
-		}
 
-		c.emitStartGoroutine(calleeFn.LLVMFn, params)
+		// Start a new goroutine.
+		if callee := instr.Call.StaticCallee(); callee != nil {
+			// Static callee is known. This makes it easier to start a new
+			// goroutine.
+			calleeFn := c.ir.GetFunction(callee)
+			if !calleeFn.IsExported() && c.selectScheduler() != "tasks" {
+				// For coroutine scheduling, this is only required when calling
+				// an external function.
+				// For tasks, because all params are stored in a single object,
+				// no unnecessary parameters should be stored anyway.
+				params = append(params, llvm.Undef(c.i8ptrType))            // context parameter
+				params = append(params, llvm.ConstPointerNull(c.i8ptrType)) // parent coroutine handle
+			}
+			c.emitStartGoroutine(calleeFn.LLVMFn, params)
+		} else if !instr.Call.IsInvoke() {
+			// This is a function pointer.
+			// At the moment, two extra params are passed to the newly started
+			// goroutine:
+			//   * The function context, for closures.
+			//   * The parent handle (for coroutines) or the function pointer
+			//     itself (for tasks).
+			funcPtr, context := c.decodeFuncValue(c.getValue(frame, instr.Call.Value), instr.Call.Value.Type().(*types.Signature))
+			params = append(params, context) // context parameter
+			switch c.selectScheduler() {
+			case "coroutines":
+				params = append(params, llvm.ConstPointerNull(c.i8ptrType)) // parent coroutine handle
+			case "tasks":
+				params = append(params, funcPtr)
+			default:
+				panic("unknown scheduler type")
+			}
+			c.emitStartGoroutine(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.emitStartGoroutine, 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,10 +32,15 @@ const (
 // funcImplementation picks an appropriate func value implementation for the
 // target.
 func (c *Compiler) funcImplementation() funcValueImplementation {
-	if c.GOARCH == "wasm" {
+	// Always pick the switch implementation, as it allows the use of blocking
+	// inside a function that is used as a func value.
+	switch c.selectScheduler() {
+	case "coroutines":
 		return funcValueSwitch
-	} else {
+	case "tasks":
 		return funcValueDoubleword
+	default:
+		panic("unknown scheduler type")
 	}
 }