cmd/internal/obj/wasm: handle stack unwinding in wasmexport

CL 603055 added basic support of wasmexport. This CL follows it
and adds stack unwinding handling. If the wasmexport Go function
returns normally, we directly return to the host. If the Go
function unwinds the stack (e.g. goroutine switch, stack growth),
we need to run a PC loop to call functions on the new stack,
similar to wasm_pc_f_loop. One difference is that when the
wasmexport function returns normally, we need to exit the loop and
return to the host.

Now a wasmimport function can call back into the Go via wasmexport.
During the callback the stack could have moved. The wasmimport
code needs to read a new SP after the host function returns,
instead of assuming the SP doesn't change.

For #65199.

Change-Id: I62c1cde1c46f7eb72625892dea41e8137b361891
Reviewed-on: https://go-review.googlesource.com/c/go/+/603836
Reviewed-by: Michael Knyszek <mknyszek@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Achille Roussel <achille.roussel@gmail.com>

src/cmd/compile/internal/ssagen/abi.go

@@ -397,9 +397,16 @@ func GenWasmExportWrapper(wrapped *ir.Func) {
 
 	pp := objw.NewProgs(fn, 0)
 	defer pp.Free()
+	// TEXT. Has a frame to pass args on stack to the Go function.
 	pp.Text.To.Type = obj.TYPE_TEXTSIZE
 	pp.Text.To.Val = int32(0)
 	pp.Text.To.Offset = types.RoundUp(ft.ArgWidth(), int64(types.RegSize))
+	// No locals. (Callee's args are covered in the callee's stackmap.)
+	p := pp.Prog(obj.AFUNCDATA)
+	p.From.SetConst(rtabi.FUNCDATA_LocalsPointerMaps)
+	p.To.Type = obj.TYPE_MEM
+	p.To.Name = obj.NAME_EXTERN
+	p.To.Sym = base.Ctxt.Lookup("no_pointers_stackmap")
 	pp.Flush()
 	// Actual code geneneration is in cmd/internal/obj/wasm.
 }

src/cmd/internal/obj/wasm/wasmobj.go

@@ -125,9 +125,10 @@ var Linkwasm = obj.LinkArch{
 }
 
 var (
-	morestack       *obj.LSym
-	morestackNoCtxt *obj.LSym
-	sigpanic        *obj.LSym
+	morestack             *obj.LSym
+	morestackNoCtxt       *obj.LSym
+	sigpanic              *obj.LSym
+	wasm_pc_f_loop_export *obj.LSym
 )
 
 const (
@@ -147,6 +148,7 @@ func instinit(ctxt *obj.Link) {
 	morestack = ctxt.Lookup("runtime.morestack")
 	morestackNoCtxt = ctxt.Lookup("runtime.morestack_noctxt")
 	sigpanic = ctxt.LookupABI("runtime.sigpanic", obj.ABIInternal)
+	wasm_pc_f_loop_export = ctxt.Lookup("wasm_pc_f_loop_export")
 }
 
 func preprocess(ctxt *obj.Link, s *obj.LSym, newprog obj.ProgAlloc) {
@@ -825,13 +827,6 @@ func genWasmImportWrapper(s *obj.LSym, appendp func(p *obj.Prog, as obj.As, args
 			// https://github.com/WebAssembly/multi-value/blob/master/proposals/multi-value/Overview.md
 			panic("invalid results type") // impossible until multi-value proposal has landed
 		}
-		if len(wi.Results) == 1 {
-			// If we have a result (rather than returning nothing at all), then
-			// we'll write the result to the Go stack relative to the current stack pointer.
-			// We cache the current stack pointer value on the wasm stack here and then use
-			// it after the Call instruction to store the result.
-			p = appendp(p, AGet, regAddr(REG_SP))
-		}
 		for _, f := range wi.Params {
 			// Each load instructions will consume the value of sp on the stack, so
 			// we need to read sp for each param. WASM appears to not have a stack dup instruction
@@ -878,20 +873,38 @@ func genWasmImportWrapper(s *obj.LSym, appendp func(p *obj.Prog, as obj.As, args
 			// to by 8 bytes to account for the return address on the Go stack.
 			storeOffset := f.Offset + 8
 
-			// This code is paired the code above that reads the stack pointer onto the wasm
-			// stack. We've done this so we have a consistent view of the sp value as it might
-			// be manipulated by the call and we want to ignore that manipulation here.
+			// We need to push SP on the Wasm stack for the Store instruction, which needs to
+			// be pushed before the value (call result). So we pop the value into a register,
+			// push SP, and push the value back.
+			// We cannot get the SP onto the stack before the call, as if the host function
+			// calls back into Go, the Go stack may have moved.
 			switch f.Type {
 			case obj.WasmI32:
-				p = appendp(p, AI32Store, constAddr(storeOffset))
+				p = appendp(p, AI64ExtendI32U) // the register is 64-bit, so we have to extend
+				p = appendp(p, ASet, regAddr(REG_R0))
+				p = appendp(p, AGet, regAddr(REG_SP))
+				p = appendp(p, AGet, regAddr(REG_R0))
+				p = appendp(p, AI64Store32, constAddr(storeOffset))
 			case obj.WasmI64:
+				p = appendp(p, ASet, regAddr(REG_R0))
+				p = appendp(p, AGet, regAddr(REG_SP))
+				p = appendp(p, AGet, regAddr(REG_R0))
 				p = appendp(p, AI64Store, constAddr(storeOffset))
 			case obj.WasmF32:
+				p = appendp(p, ASet, regAddr(REG_F0))
+				p = appendp(p, AGet, regAddr(REG_SP))
+				p = appendp(p, AGet, regAddr(REG_F0))
 				p = appendp(p, AF32Store, constAddr(storeOffset))
 			case obj.WasmF64:
+				p = appendp(p, ASet, regAddr(REG_F16))
+				p = appendp(p, AGet, regAddr(REG_SP))
+				p = appendp(p, AGet, regAddr(REG_F16))
 				p = appendp(p, AF64Store, constAddr(storeOffset))
 			case obj.WasmPtr:
 				p = appendp(p, AI64ExtendI32U)
+				p = appendp(p, ASet, regAddr(REG_R0))
+				p = appendp(p, AGet, regAddr(REG_SP))
+				p = appendp(p, AGet, regAddr(REG_R0))
 				p = appendp(p, AI64Store, constAddr(storeOffset))
 			default:
 				panic("bad result type")
@@ -907,10 +920,13 @@ func genWasmExportWrapper(s *obj.LSym, appendp func(p *obj.Prog, as obj.As, args
 	we := s.Func().WasmExport
 	we.CreateAuxSym()
 	p := s.Func().Text
+	framesize := p.To.Offset
+	for p.Link != nil && p.Link.As == obj.AFUNCDATA {
+		p = p.Link
+	}
 	if p.Link != nil {
 		panic("wrapper functions for WASM export should not have a body")
 	}
-	framesize := p.To.Offset
 
 	// Store args
 	for i, f := range we.Params {
@@ -943,20 +959,25 @@ func genWasmExportWrapper(s *obj.LSym, appendp func(p *obj.Prog, as obj.As, args
 	p = appendp(p, ASet, regAddr(REG_SP))
 	// write return address to Go stack
 	p = appendp(p, AGet, regAddr(REG_SP))
-	p = appendp(p, AI64Const, obj.Addr{
+	retAddr := obj.Addr{
 		Type:   obj.TYPE_ADDR,
 		Name:   obj.NAME_EXTERN,
 		Sym:    s, // PC_F
 		Offset: 1, // PC_B=1, past the prologue, so we have the right SP delta
-	})
+	}
+	p = appendp(p, AI64Const, retAddr)
 	p = appendp(p, AI64Store, constAddr(0))
 	// Set PC_B parameter to function entry
 	p = appendp(p, AI32Const, constAddr(0))
 	p = appendp(p, ACall, obj.Addr{Type: obj.TYPE_MEM, Name: obj.NAME_EXTERN, Sym: we.WrappedSym})
-	// return value is on the top of the stack, indicating whether to unwind the Wasm stack
-	// TODO: handle stack unwinding
+	// Return value is on the top of the stack, indicating whether to unwind the Wasm stack.
+	// In the unwinding case, we call wasm_pc_f_loop_export to handle stack switch and rewinding,
+	// until a normal return (non-unwinding) back to this function.
 	p = appendp(p, AIf)
-	p = appendp(p, obj.AUNDEF)
+	p = appendp(p, AI32Const, retAddr)
+	p = appendp(p, AI32Const, constAddr(16))
+	p = appendp(p, AI32ShrU)
+	p = appendp(p, ACall, obj.Addr{Type: obj.TYPE_MEM, Name: obj.NAME_EXTERN, Sym: wasm_pc_f_loop_export})
 	p = appendp(p, AEnd)
 
 	// Load result
@@ -1008,6 +1029,7 @@ var notUsePC_B = map[string]bool{
 	"wasm_export_resume":      true,
 	"wasm_export_getsp":       true,
 	"wasm_pc_f_loop":          true,
+	"wasm_pc_f_loop_export":   true,
 	"gcWriteBarrier":          true,
 	"runtime.gcWriteBarrier1": true,
 	"runtime.gcWriteBarrier2": true,
@@ -1062,6 +1084,9 @@ func assemble(ctxt *obj.Link, s *obj.LSym, newprog obj.ProgAlloc) {
 		"wasm_pc_f_loop", "runtime.wasmDiv", "runtime.wasmTruncS", "runtime.wasmTruncU", "memeqbody":
 		varDecls = []*varDecl{}
 		useAssemblyRegMap()
+	case "wasm_pc_f_loop_export":
+		varDecls = []*varDecl{{count: 2, typ: i32}}
+		useAssemblyRegMap()
 	case "memchr", "memcmp":
 		varDecls = []*varDecl{{count: 2, typ: i32}}
 		useAssemblyRegMap()

src/cmd/link/internal/wasm/asm.go

@@ -73,6 +73,7 @@ var wasmFuncTypes = map[string]*wasmFuncType{
 	"wasm_export_resume":      {Params: []byte{}},                                         //
 	"wasm_export_getsp":       {Results: []byte{I32}},                                     // sp
 	"wasm_pc_f_loop":          {Params: []byte{}},                                         //
+	"wasm_pc_f_loop_export":   {Params: []byte{I32}},                                      // pc_f
 	"runtime.wasmDiv":         {Params: []byte{I64, I64}, Results: []byte{I64}},           // x, y -> x/y
 	"runtime.wasmTruncS":      {Params: []byte{F64}, Results: []byte{I64}},                // x -> int(x)
 	"runtime.wasmTruncU":      {Params: []byte{F64}, Results: []byte{I64}},                // x -> uint(x)

src/runtime/asm_wasm.s

@@ -554,5 +554,57 @@ TEXT wasm_pc_f_loop(SB),NOSPLIT,$0
 
 	Return
 
+// wasm_pc_f_loop_export is like wasm_pc_f_loop, except that this takes an
+// argument (on Wasm stack) that is a PC_F, and the loop stops when we get
+// to that PC in a normal return (not unwinding).
+// This is for handling an wasmexport function when it needs to switch the
+// stack.
+TEXT wasm_pc_f_loop_export(SB),NOSPLIT,$0
+	Get PAUSE
+	I32Eqz
+outer:
+	If
+		// R1 is whether a function return normally (0) or unwinding (1).
+		// Start with unwinding.
+		I32Const $1
+		Set R1
+	loop:
+		Loop
+			// Get PC_F & PC_B from -8(SP)
+			Get SP
+			I32Const $8
+			I32Sub
+			I32Load16U $2 // PC_F
+			Tee R2
+
+			Get R0
+			I32Eq
+			If // PC_F == R0, we're at the stop PC
+				Get R1
+				I32Eqz
+				// Break if it is a normal return
+				BrIf outer // actually jump to after the corresponding End
+			End
+
+			Get SP
+			I32Const $8
+			I32Sub
+			I32Load16U $0 // PC_B
+
+			Get R2 // PC_F
+			CallIndirect $0
+			Set R1 // save return/unwinding state for next iteration
+
+			Get PAUSE
+			I32Eqz
+			BrIf loop
+		End
+	End
+
+	I32Const $0
+	Set PAUSE
+
+	Return
+
 TEXT wasm_export_lib(SB),NOSPLIT,$0
 	UNDEF