nativespawn.go raw

   1  package compiler
   2  
   3  // Native spawn channel routing. Phase A added pipeBoundFd/pipeChanID
   4  // metadata to the channel struct plus pipe-routing helpers. Phase B
   5  // wires the compiler so that:
   6  //
   7  //   1. createSpawn (native path) calls runtime.MarkParentPipeChannel on
   8  //      every channel argument right after spawnDomain returns. The
   9  //      runtime helper looks up the parent fd from LastSpawnedParentFd.
  10  //   2. setupNativeSpawnTargetChannels (mirror of the wasm helper) emits
  11  //      runtime.MarkChildPipeChannel at the entry of every spawn-target
  12  //      function. The helper internally checks ChildPipeFd >= 0 so that
  13  //      calling the function NOT through spawn leaves its parameters as
  14  //      ordinary local channels.
  15  //   3. createChanSend / createChanRecv check b.pipeChannels and emit a
  16  //      pipe-routed call (PipeChanSendCodec / PipeChanRecvCodec) that
  17  //      invokes the user-defined Codec methods to serialize the chan
  18  //      element, instead of the in-runtime queue path.
  19  //
  20  // Channels are multiplexed over a single socketpair fd by chanID. A
  21  // stable layout is fixed by parameter order: the i-th channel parameter
  22  // of the spawn target gets chanID = i+1. ChanID 0 is reserved for
  23  // pipe control messages (currently just close-of-channel notifications).
  24  
  25  import (
  26  	"go/types"
  27  
  28  	"golang.org/x/tools/go/ssa"
  29  	"tinygo.org/x/go-llvm"
  30  )
  31  
  32  type nativeSpawnTarget struct {
  33  	fn            *ssa.Function
  34  	chanParamIdxs []int    // indices of channel-typed params, declaration order
  35  	chanIDs       []uint16 // aligned with chanParamIdxs, 1-based
  36  }
  37  
  38  func (c *compilerContext) scanNativeSpawnTargets(pkg *ssa.Package) {
  39  	if c.isWasmContext() {
  40  		return
  41  	}
  42  	if c.nativeSpawnIndex == nil {
  43  		c.nativeSpawnIndex = make(map[*ssa.Function]*nativeSpawnTarget)
  44  	}
  45  	// Build the entire program's SSA so we can scan EVERY package for
  46  	// spawn() call sites. A spawn target function defined in package A
  47  	// but invoked from package B needs the child-side prologue
  48  	// (MarkChildPipeChannel) emitted when A is compiled — but the
  49  	// spawn site that marks the function as a target lives in B's SSA
  50  	// (which is lazy-built by default and may not yet exist when A is
  51  	// being compiled).
  52  	//
  53  	// program.Build() is idempotent and fast on already-built packages,
  54  	// so the cost is paid once per compilation context.
  55  	c.program.Build()
  56  	seen := map[*ssa.Function]bool{}
  57  	scanned := map[*ssa.Function]bool{}
  58  	// Walk every package, then for each pkg member walk the SSA
  59  	// reachable from it. Methods aren't in pkg.Members directly — they
  60  	// belong to type definitions — so for *ssa.Type members we have to
  61  	// pull the method set via prog.MethodValue. Function values
  62  	// referenced from anonymous functions / closures get picked up via
  63  	// fn.AnonFuncs.
  64  	for _, p := range c.program.AllPackages() {
  65  		for _, mem := range p.Members {
  66  			switch m := mem.(type) {
  67  			case *ssa.Function:
  68  				c.scanFnDeep(m, seen, scanned)
  69  			case *ssa.Type:
  70  				c.scanTypeMethods(m.Type(), seen, scanned)
  71  			}
  72  		}
  73  	}
  74  	_ = pkg
  75  }
  76  
  77  // scanFnDeep scans fn, its anonymous funcs, and its direct callees.
  78  // Following direct callees is required to find spawn() sites inside
  79  // package init() bodies: source-level `func init()` declarations compile
  80  // to SSA functions named init$1, init$2, etc., which are NOT package
  81  // Members — they are only reachable as callees of the synthetic init
  82  // wrapper. Without callee traversal, any spawn() call inside an init()
  83  // is invisible to the scanner, so the target's channels are never
  84  // registered and MarkChildPipeChannel is never emitted in the child
  85  // process, leaving the child's channel.pipeBoundFd == 0 and breaking
  86  // all channel I/O across that spawn boundary.
  87  func (c *compilerContext) scanFnDeep(fn *ssa.Function, seen, scanned map[*ssa.Function]bool) {
  88  	if scanned[fn] {
  89  		return
  90  	}
  91  	scanned[fn] = true
  92  	c.scanFnForNativeSpawn(fn, seen)
  93  	for _, anon := range fn.AnonFuncs {
  94  		c.scanFnDeep(anon, seen, scanned)
  95  	}
  96  	// Follow direct static callees so init$N bodies are reachable.
  97  	for _, blk := range fn.Blocks {
  98  		for _, instr := range blk.Instrs {
  99  			if call, ok := instr.(*ssa.Call); ok {
 100  				if callee, ok := call.Call.Value.(*ssa.Function); ok {
 101  					c.scanFnDeep(callee, seen, scanned)
 102  				}
 103  			}
 104  		}
 105  	}
 106  }
 107  
 108  // scanTypeMethods walks every method on a named type (and on *T) and
 109  // scans its body. The Go SSA program knows methods via MethodValue
 110  // once the program has been Built.
 111  func (c *compilerContext) scanTypeMethods(t types.Type, seen, scanned map[*ssa.Function]bool) {
 112  	mset := c.program.MethodSets.MethodSet(t)
 113  	for i := 0; i < mset.Len(); i++ {
 114  		fn := c.program.MethodValue(mset.At(i))
 115  		if fn != nil {
 116  			c.scanFnDeep(fn, seen, scanned)
 117  		}
 118  	}
 119  	mset = c.program.MethodSets.MethodSet(types.NewPointer(t))
 120  	for i := 0; i < mset.Len(); i++ {
 121  		fn := c.program.MethodValue(mset.At(i))
 122  		if fn != nil {
 123  			c.scanFnDeep(fn, seen, scanned)
 124  		}
 125  	}
 126  }
 127  
 128  func (c *compilerContext) scanFnForNativeSpawn(fn *ssa.Function, seen map[*ssa.Function]bool) {
 129  	if len(fn.Blocks) == 0 {
 130  		return
 131  	}
 132  	for _, blk := range fn.Blocks {
 133  		for _, instr := range blk.Instrs {
 134  			call, ok := instr.(*ssa.Call)
 135  			if !ok {
 136  				continue
 137  			}
 138  			bi, ok := call.Call.Value.(*ssa.Builtin)
 139  			if !ok || bi.Name() != "spawn" {
 140  				continue
 141  			}
 142  			target := extractWasmSpawnTargetFn(call.Call.Args)
 143  			if target == nil || seen[target] {
 144  				continue
 145  			}
 146  			seen[target] = true
 147  			sig := target.Signature
 148  			var chanIdxs []int
 149  			var chanIDs []uint16
 150  			var nextID uint16 = 1
 151  			for i := 0; i < sig.Params().Len(); i++ {
 152  				if _, isChan := sig.Params().At(i).Type().Underlying().(*types.Chan); isChan {
 153  					chanIdxs = append(chanIdxs, i)
 154  					chanIDs = append(chanIDs, nextID)
 155  					nextID++
 156  				}
 157  			}
 158  			t := &nativeSpawnTarget{
 159  				fn:            target,
 160  				chanParamIdxs: chanIdxs,
 161  				chanIDs:       chanIDs,
 162  			}
 163  			c.nativeSpawnTargets = append(c.nativeSpawnTargets, t)
 164  			c.nativeSpawnIndex[target] = t
 165  		}
 166  	}
 167  }
 168  
 169  // isWasmContext returns true on the wasm target. Native spawn handling
 170  // is suppressed for wasm because the wasm path uses SAB channels.
 171  func (c *compilerContext) isWasmContext() bool {
 172  	return c.GOOS == "js" && len(c.Triple) >= 4 && c.Triple[:4] == "wasm"
 173  }
 174  
 175  func (c *compilerContext) nativeSpawnTargetFor(fn *ssa.Function) *nativeSpawnTarget {
 176  	if c.nativeSpawnIndex == nil {
 177  		return nil
 178  	}
 179  	return c.nativeSpawnIndex[fn]
 180  }
 181  
 182  // setupNativeSpawnTargetChannels emits the child-side prologue for a
 183  // spawn target function. Linear, no branching: a single call to
 184  // runtime.MarkChildPipeChannel(param, chanID) per channel parameter.
 185  // The helper internally no-ops when ChildPipeFd < 0.
 186  func (b *builder) setupNativeSpawnTargetChannels() {
 187  	if b.isWasmTarget() || b.fn == nil {
 188  		return
 189  	}
 190  	target := b.nativeSpawnTargetFor(b.fn)
 191  	if target == nil {
 192  		return
 193  	}
 194  	if b.pipeChannels == nil {
 195  		b.pipeChannels = make(map[ssa.Value]bool)
 196  	}
 197  
 198  	runtimePkg := b.program.ImportedPackage("runtime")
 199  	mark := runtimePkg.Members["MarkChildPipeChannel"].(*ssa.Function)
 200  	markType, markFn := b.getFunction(mark)
 201  	i16 := b.ctx.Int16Type()
 202  
 203  	for i, paramIdx := range target.chanParamIdxs {
 204  		if paramIdx >= len(b.fn.Params) {
 205  			continue
 206  		}
 207  		param := b.fn.Params[paramIdx]
 208  		llvmParam, ok := b.locals[param]
 209  		if !ok {
 210  			continue
 211  		}
 212  		chanID := target.chanIDs[i]
 213  		b.CreateCall(markType, markFn, []llvm.Value{
 214  			llvmParam,
 215  			llvm.ConstInt(i16, uint64(chanID), false),
 216  			llvm.Undef(b.dataPtrType),
 217  		}, "")
 218  		b.pipeChannels[param] = true
 219  	}
 220  
 221  	// Signal the parent that the child is ready.
 222  	readyFn := runtimePkg.Members["PipeChildReady"].(*ssa.Function)
 223  	readyType, readyVal := b.getFunction(readyFn)
 224  	b.CreateCall(readyType, readyVal, []llvm.Value{llvm.Undef(b.dataPtrType)}, "")
 225  }
 226  
 227  // emitNativeSpawnChannelMark emits the parent-side MarkParentPipeChannel
 228  // calls right after spawnDomain returns, plus tracks each channel SSA
 229  // value in b.pipeChannels so subsequent chan ops route via pipe.
 230  //
 231  // chanArgs lists the (ssa.Value, llvm.Value) pairs for every channel
 232  // argument in the spawned function's parameter order. Aligns with the
 233  // target's chanIDs.
 234  func (b *builder) emitNativeSpawnChannelMark(target *nativeSpawnTarget, chanArgs []spawnChanArg) {
 235  	if len(chanArgs) == 0 {
 236  		return
 237  	}
 238  	runtimePkg := b.program.ImportedPackage("runtime")
 239  	mark := runtimePkg.Members["MarkParentPipeChannel"].(*ssa.Function)
 240  	markType, markFn := b.getFunction(mark)
 241  
 242  	if b.pipeChannels == nil {
 243  		b.pipeChannels = make(map[ssa.Value]bool)
 244  	}
 245  	i16 := b.ctx.Int16Type()
 246  	for i, ca := range chanArgs {
 247  		chanID := target.chanIDs[i]
 248  		b.CreateCall(markType, markFn, []llvm.Value{
 249  			ca.llvm,
 250  			llvm.ConstInt(i16, uint64(chanID), false),
 251  			llvm.Undef(b.dataPtrType),
 252  		}, "")
 253  		b.pipeChannels[ca.ssa] = true
 254  	}
 255  }
 256  
 257  type spawnChanArg struct {
 258  	ssa  ssa.Value
 259  	llvm llvm.Value
 260  }
 261