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>

1 files changed, 45 insertions, 20 deletions

diff --git a/src/cmd/internal/obj/wasm/wasmobj.go b/src/cmd/internal/obj/wasm/wasmobj.go
index 4b5324cc56..20ed142812 100644
--- a/src/cmd/internal/obj/wasm/wasmobj.go
+++ b/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()