cmd/compile: add basic wasmexport support

This CL adds a compiler directive go:wasmexport, which applies to
a Go function and makes it an exported function of the Wasm module
being built, so it can be called directly from the host. As
proposed in #65199, parameter and result types are limited to
32-bit and 64-bit integers and floats, and there can be at most
one result.

As the Go and Wasm calling conventions are different, for a
wasmexport function we generate a wrapper function does the ABI
conversion at compile time.

Currently this CL only adds basic support. In particular,
- it only supports executable mode, i.e. the Go wasm module calls
  into the host via wasmimport, which then calls back to Go via
  wasmexport. Library (c-shared) mode is not implemented yet.
- only supports wasip1, not js.
- if the exported function unwinds stacks (goroutine switch, stack
growth, etc.), it probably doesn't work.

TODO: support stack unwinding, c-shared mode, js.

For #65199.

Change-Id: Id1777c2d44f7d51942c1caed3173c0a82f120cc4
Reviewed-on: https://go-review.googlesource.com/c/go/+/603055
Reviewed-by: Than McIntosh <thanm@golang.org>
Reviewed-by: Randy Reddig <randy.reddig@fastly.com>
Reviewed-by: David Chase <drchase@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

4 files changed, 136 insertions, 2 deletions

diff --git a/src/cmd/internal/obj/link.go b/src/cmd/internal/obj/link.go
index 27626d9deb..a3e4a0d309 100644
--- a/src/cmd/internal/obj/link.go
+++ b/src/cmd/internal/obj/link.go
@@ -500,6 +500,7 @@ type FuncInfo struct {
 	FuncInfoSym *LSym
 
 	WasmImport *WasmImport
+	WasmExport *WasmExport
 
 	sehUnwindInfoSym *LSym
 }
@@ -665,9 +666,9 @@ func (wi *WasmImport) Read(b []byte) {
 // parameters and results translated into WASM types based on the Go function
 // declaration.
 type WasmFuncType struct {
-	// Params holds the imported function parameter fields.
+	// Params holds the function parameter fields.
 	Params []WasmField
-	// Results holds the imported function result fields.
+	// Results holds the function result fields.
 	Results []WasmField
 }
 
@@ -724,6 +725,27 @@ func (ft *WasmFuncType) Read(b []byte) {
 	}
 }
 
+// WasmExport represents a WebAssembly (WASM) exported function with
+// parameters and results translated into WASM types based on the Go function
+// declaration.
+type WasmExport struct {
+	WasmFuncType
+
+	WrappedSym *LSym // the wrapped Go function
+	AuxSym     *LSym // aux symbol to pass metadata to the linker
+}
+
+func (we *WasmExport) CreateAuxSym() {
+	var b bytes.Buffer
+	we.WasmFuncType.Write(&b)
+	p := b.Bytes()
+	we.AuxSym = &LSym{
+		Type: objabi.SDATA, // doesn't really matter
+		P:    append([]byte(nil), p...),
+		Size: int64(len(p)),
+	}
+}
+
 type WasmField struct {
 	Type WasmFieldType
 	// Offset holds the frame-pointer-relative locations for Go's stack-based
diff --git a/src/cmd/internal/obj/objfile.go b/src/cmd/internal/obj/objfile.go
index cbdc5a3486..de38349930 100644
--- a/src/cmd/internal/obj/objfile.go
+++ b/src/cmd/internal/obj/objfile.go
@@ -361,6 +361,9 @@ func (w *writer) Sym(s *LSym) {
 	if s.ABIWrapper() {
 		flag2 |= goobj.SymFlagABIWrapper
 	}
+	if s.Func() != nil && s.Func().WasmExport != nil {
+		flag2 |= goobj.SymFlagWasmExport
+	}
 	if strings.HasPrefix(name, "gofile..") {
 		name = filepath.ToSlash(name)
 	}
@@ -627,6 +630,9 @@ func (w *writer) Aux(s *LSym) {
 			}
 			w.aux1(goobj.AuxWasmImport, fn.WasmImport.AuxSym)
 		}
+		if fn.WasmExport != nil {
+			w.aux1(goobj.AuxWasmType, fn.WasmExport.AuxSym)
+		}
 	} else if v := s.VarInfo(); v != nil {
 		if v.dwarfInfoSym != nil && v.dwarfInfoSym.Size != 0 {
 			w.aux1(goobj.AuxDwarfInfo, v.dwarfInfoSym)
@@ -737,6 +743,9 @@ func nAuxSym(s *LSym) int {
 			}
 			n++
 		}
+		if fn.WasmExport != nil {
+			n++
+		}
 	} else if v := s.VarInfo(); v != nil {
 		if v.dwarfInfoSym != nil && v.dwarfInfoSym.Size != 0 {
 			n++
@@ -801,6 +810,9 @@ func genFuncInfoSyms(ctxt *Link) {
 		if wi := fn.WasmImport; wi != nil {
 			auxsyms = append(auxsyms, wi.AuxSym)
 		}
+		if we := fn.WasmExport; we != nil {
+			auxsyms = append(auxsyms, we.AuxSym)
+		}
 		for _, s := range auxsyms {
 			if s == nil || s.Size == 0 {
 				continue
diff --git a/src/cmd/internal/obj/sym.go b/src/cmd/internal/obj/sym.go
index d2e61832ba..943be3c38c 100644
--- a/src/cmd/internal/obj/sym.go
+++ b/src/cmd/internal/obj/sym.go
@@ -462,6 +462,9 @@ func (ctxt *Link) traverseFuncAux(flag traverseFlag, fsym *LSym, fn func(parent
 	if wi := fninfo.WasmImport; wi != nil {
 		auxsyms = append(auxsyms, wi.AuxSym)
 	}
+	if we := fninfo.WasmExport; we != nil {
+		auxsyms = append(auxsyms, we.AuxSym)
+	}
 	for _, s := range auxsyms {
 		if s == nil || s.Size == 0 {
 			continue
diff --git a/src/cmd/internal/obj/wasm/wasmobj.go b/src/cmd/internal/obj/wasm/wasmobj.go
index dcbf35e886..4b5324cc56 100644
--- a/src/cmd/internal/obj/wasm/wasmobj.go
+++ b/src/cmd/internal/obj/wasm/wasmobj.go
@@ -196,6 +196,8 @@ func preprocess(ctxt *obj.Link, s *obj.LSym, newprog obj.ProgAlloc) {
 		// isn't run. We don't want the frame expansion code because our function
 		// body is just the code to translate and call the imported function.
 		framesize = 0
+	} else if s.Func().WasmExport != nil {
+		genWasmExportWrapper(s, appendp)
 	} else if s.Func().Text.From.Sym.Wrapper() {
 		// if g._panic != nil && g._panic.argp == FP {
 		//   g._panic.argp = bottom-of-frame
@@ -900,6 +902,95 @@ func genWasmImportWrapper(s *obj.LSym, appendp func(p *obj.Prog, as obj.As, args
 	p = appendp(p, obj.ARET)
 }
 
+// Generate function body for wasmexport wrapper function.
+func genWasmExportWrapper(s *obj.LSym, appendp func(p *obj.Prog, as obj.As, args ...obj.Addr) *obj.Prog) {
+	we := s.Func().WasmExport
+	we.CreateAuxSym()
+	p := s.Func().Text
+	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 {
+		p = appendp(p, AGet, regAddr(REG_SP))
+		p = appendp(p, AGet, regAddr(REG_R0+int16(i)))
+		switch f.Type {
+		case obj.WasmI32:
+			p = appendp(p, AI32Store, constAddr(f.Offset))
+		case obj.WasmI64:
+			p = appendp(p, AI64Store, constAddr(f.Offset))
+		case obj.WasmF32:
+			p = appendp(p, AF32Store, constAddr(f.Offset))
+		case obj.WasmF64:
+			p = appendp(p, AF64Store, constAddr(f.Offset))
+		case obj.WasmPtr:
+			p = appendp(p, AI64ExtendI32U)
+			p = appendp(p, AI64Store, constAddr(f.Offset))
+		default:
+			panic("bad param type")
+		}
+	}
+
+	// Call the Go function.
+	// XXX maybe use ACALL and let later phase expand? But we don't use PC_B. Maybe we should?
+	// Go calling convention expects we push a return PC before call.
+	// SP -= 8
+	p = appendp(p, AGet, regAddr(REG_SP))
+	p = appendp(p, AI32Const, constAddr(8))
+	p = appendp(p, AI32Sub)
+	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{
+		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, 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
+	p = appendp(p, AIf)
+	p = appendp(p, obj.AUNDEF)
+	p = appendp(p, AEnd)
+
+	// Load result
+	if len(we.Results) > 1 {
+		panic("invalid results type")
+	} else if len(we.Results) == 1 {
+		p = appendp(p, AGet, regAddr(REG_SP))
+		f := we.Results[0]
+		switch f.Type {
+		case obj.WasmI32:
+			p = appendp(p, AI32Load, constAddr(f.Offset))
+		case obj.WasmI64:
+			p = appendp(p, AI64Load, constAddr(f.Offset))
+		case obj.WasmF32:
+			p = appendp(p, AF32Load, constAddr(f.Offset))
+		case obj.WasmF64:
+			p = appendp(p, AF64Load, constAddr(f.Offset))
+		case obj.WasmPtr:
+			p = appendp(p, AI64Load, constAddr(f.Offset))
+			p = appendp(p, AI32WrapI64)
+		default:
+			panic("bad result type")
+		}
+	}
+
+	// Epilogue. Cannot use ARET as we don't follow Go calling convention.
+	// SP += framesize
+	p = appendp(p, AGet, regAddr(REG_SP))
+	p = appendp(p, AI32Const, constAddr(framesize))
+	p = appendp(p, AI32Add)
+	p = appendp(p, ASet, regAddr(REG_SP))
+	p = appendp(p, AReturn)
+}
+
 func constAddr(value int64) obj.Addr {
 	return obj.Addr{Type: obj.TYPE_CONST, Offset: value}
 }
@@ -991,6 +1082,12 @@ func assemble(ctxt *obj.Link, s *obj.LSym, newprog obj.ProgAlloc) {
 		// no locals
 		useAssemblyRegMap()
 	default:
+		if s.Func().WasmExport != nil {
+			// no local SP, not following Go calling convention
+			useAssemblyRegMap()
+			break
+		}
+
 		// Normal calling convention: PC_B as WebAssembly parameter. First local variable is local SP cache.
 		regVars[REG_PC_B-MINREG] = &regVar{false, 0}
 		hasLocalSP = true