Repository: tetratelabs/wazero
Pull request #1630 (2023-08-15 02:01:57)

wazevo: adds support for host functions
#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi.go b/internal/engine/wazevo/backend/isa/arm64/abi.go
index 556c934a..e115e0fe 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi.go
@@ -76,19 +76,23 @@ func (a *abiImpl) init(sig *ssa.Signature) {
 	if len(a.rets) < len(sig.Results) {
 		a.rets = make([]backend.ABIArg, len(sig.Results))
 	}
+	a.rets = a.rets[:len(sig.Results)]
 	a.retStackSize = a.setABIArgs(a.rets, sig.Results)
 	if argsNum := len(sig.Params); len(a.args) < argsNum {
 		a.args = make([]backend.ABIArg, argsNum)
 	}
+	a.args = a.args[:len(sig.Params)]
 	a.argStackSize = a.setABIArgs(a.args, sig.Params)
 
 	// Gather the real registers usages in arg/return.
+	a.retRealRegs = a.retRealRegs[:0]
 	for i := range a.rets {
 		r := &a.rets[i]
 		if r.Kind == backend.ABIArgKindReg {
 			a.retRealRegs = append(a.retRealRegs, r.Reg)
 		}
 	}
+	a.argRealRegs = a.argRealRegs[:0]
 	for i := range a.args {
 		arg := &a.args[i]
 		if arg.Kind == backend.ABIArgKindReg {
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi_go_call.go b/internal/engine/wazevo/backend/isa/arm64/abi_go_call.go
index e4df1872..d48b3853 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi_go_call.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi_go_call.go
@@ -13,7 +13,7 @@ var calleeSavedRegistersPlusLinkRegSorted = []regalloc.VReg{
 }
 
 // CompileGoFunctionTrampoline implements backend.Machine.
-func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature) {
+func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature, needModuleContextPtr bool) {
 	cur := m.allocateInstr()
 	cur.asNop0()
 	m.rootInstr = cur
@@ -26,6 +26,25 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 
 	// Next, we need to store all the arguments to the execution context.
 
+	argBegin := 1 // Skips exec context by default.
+	if needModuleContextPtr {
+		offset := wazevoapi.ExecutionContextOffsets.GoFunctionCallCalleeModuleContextOpaque.I64()
+		if !offsetFitsInAddressModeKindRegUnsignedImm12(64, offset) {
+			panic("BUG: too large or un-aligned offset for goFunctionCallCalleeModuleContextOpaque in execution context")
+		}
+
+		// Module context is always the second argument.
+		moduleCtrPtr := x1VReg
+		store := m.allocateInstr()
+		amode := addressMode{kind: addressModeKindRegUnsignedImm12, rn: execCtrPtr, imm: offset}
+		store.asStore(operandNR(moduleCtrPtr), amode, 64)
+		store.prev = cur
+		cur.next = store
+		cur = store
+
+		argBegin++
+	}
+
 	stackPtrReg := x15VReg // Caller save, so we can use it for whatever we want.
 	imm12Op, ok := asImm12Operand(wazevoapi.ExecutionContextOffsets.GoFunctionCallStackBegin.U64())
 	if !ok {
@@ -39,7 +58,7 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 	cur = goCallStackPtrCalc
 
 	abi := m.getOrCreateABIImpl(sig)
-	for _, arg := range abi.args[1:] { // Skips exec context.
+	for _, arg := range abi.args[argBegin:] {
 		if arg.Kind == backend.ABIArgKindReg {
 			store := m.allocateInstr()
 			v := arg.Reg
@@ -90,6 +109,7 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 			movTmpToReg := m.allocateInstr()
 
 			mode := addressMode{kind: addressModeKindPostIndex, rn: stackPtrReg}
+			tmpRegVRegVec := v17VReg // Caller save, so we can use it for whatever we want.
 			switch r.Type {
 			case ssa.TypeI32:
 				mode.imm = 8 // We use uint64 for all basic types, except SIMD v128.
@@ -101,12 +121,12 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 				movTmpToReg.asMove64(r.Reg, tmpRegVReg)
 			case ssa.TypeF32:
 				mode.imm = 8 // We use uint64 for all basic types, except SIMD v128.
-				loadIntoTmp.asFpuLoad(operandNR(tmpRegVReg), mode, 32)
-				movTmpToReg.asFpuMov64(r.Reg, tmpRegVReg)
+				loadIntoTmp.asFpuLoad(operandNR(tmpRegVRegVec), mode, 32)
+				movTmpToReg.asFpuMov64(r.Reg, tmpRegVRegVec)
 			case ssa.TypeF64:
 				mode.imm = 8 // We use uint64 for all basic types, except SIMD v128.
-				loadIntoTmp.asFpuLoad(operandNR(tmpRegVReg), mode, 64)
-				movTmpToReg.asFpuMov64(r.Reg, tmpRegVReg)
+				loadIntoTmp.asFpuLoad(operandNR(tmpRegVRegVec), mode, 64)
+				movTmpToReg.asFpuMov64(r.Reg, tmpRegVRegVec)
 			}
 			loadIntoTmp.prev = cur
 			cur.next = loadIntoTmp
@@ -181,16 +201,19 @@ func (m *machine) restoreRegistersInExecutionContext(cur *instruction, regs []re
 }
 
 func (m *machine) setExitCode(cur *instruction, execCtr regalloc.VReg, exitCode wazevoapi.ExitCode) *instruction {
+	constReg := x17VReg // caller-saved, so we can use it.
+
+	m.pendingInstructions = m.pendingInstructions[:0]
+	m.lowerConstantI32(x17VReg, int32(exitCode))
+	for _, instr := range m.pendingInstructions {
+		instr.prev = cur
+		cur.next = instr
+		cur = instr
+	}
+
 	// Set the exit status on the execution context.
-	// 	movz tmp, #wazevoapi.ExitCodeGrowStack
-	// 	str tmp, [exec_context]
-	loadStatusConst := m.allocateInstrAfterLowering()
-	loadStatusConst.asMOVZ(tmpRegVReg, uint64(exitCode), 0, true)
-	loadStatusConst.prev = cur
-	cur.next = loadStatusConst
-	cur = loadStatusConst
 	setExistStatus := m.allocateInstrAfterLowering()
-	setExistStatus.asStore(operandNR(tmpRegVReg),
+	setExistStatus.asStore(operandNR(constReg),
 		addressMode{
 			kind: addressModeKindRegUnsignedImm12,
 			rn:   execCtr, imm: wazevoapi.ExecutionContextOffsets.ExitCodeOffset.I64(),
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine.go b/internal/engine/wazevo/backend/isa/arm64/machine.go
index cc17b5a5..0cd712af 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine.go
@@ -134,6 +134,7 @@ func (m *machine) Reset() {
 	m.regAllocFn.reset()
 	m.spillSlotSize = 0
 	m.unresolvedAddressModes = m.unresolvedAddressModes[:0]
+	m.rootInstr = nil
 }
 
 // InitializeABI implements backend.Machine InitializeABI.
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
index dda78451..383d40f0 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
@@ -331,6 +331,7 @@ func (m *machine) insertStackBoundsCheck(requiredStackSize int64, cur *instructi
 	// Set the required stack size and set it to the exec context.
 	{
 		// First load the requiredStackSize into the temporary register,
+		m.pendingInstructions = m.pendingInstructions[:0]
 		m.lowerConstantI64(tmpRegVReg, requiredStackSize)
 		// lowerConstantI64 adds instructions into m.pendingInstructions,
 		// so we manually link them together.
diff --git a/internal/engine/wazevo/backend/machine.go b/internal/engine/wazevo/backend/machine.go
index f9bbaeee..2ebd4e2f 100644
--- a/internal/engine/wazevo/backend/machine.go
+++ b/internal/engine/wazevo/backend/machine.go
@@ -97,6 +97,6 @@ type (
 		Encode()
 
 		// CompileGoFunctionTrampoline compiles the trampoline function  to call a Go function of the given exit code and signature.
-		CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature)
+		CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature, needModuleContextPtr bool)
 	}
 )
diff --git a/internal/engine/wazevo/backend/regalloc/assign.go b/internal/engine/wazevo/backend/regalloc/assign.go
index 995f6893..df120eaf 100644
--- a/internal/engine/wazevo/backend/regalloc/assign.go
+++ b/internal/engine/wazevo/backend/regalloc/assign.go
@@ -79,7 +79,8 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
 		return
 	case 1:
 	default:
-		panic("multiple defs (== call instruction) should be special cased")
+		// multiple defs (== call instruction) can be special cased, and no need to assign (already real regs following the calling convention.
+		return
 	}
 
 	d := defs[0]
diff --git a/internal/engine/wazevo/call_engine.go b/internal/engine/wazevo/call_engine.go
index 700715e4..f5da650c 100644
--- a/internal/engine/wazevo/call_engine.go
+++ b/internal/engine/wazevo/call_engine.go
@@ -61,11 +61,15 @@ type (
 		// savedRegisters is the opaque spaces for save/restore registers.
 		// We want to align 16 bytes for each register, so we use [64][2]uint64.
 		savedRegisters [64][2]uint64
+		// goFunctionCallCalleeModuleContextOpaque is the pointer to the target Go function's moduleContextOpaque.
+		goFunctionCallCalleeModuleContextOpaque uintptr
 		// goFunctionCallStack is used to pass/receive parameters/results for Go function calls.
-		goFunctionCallStack [128]uint64
+		goFunctionCallStack [goFunctionCallStackSize]uint64
 	}
 )
 
+const goFunctionCallStackSize = 128
+
 var initialStackSize uint64 = 512
 
 func (c *callEngine) init() {
@@ -111,7 +115,7 @@ func (c *callEngine) CallWithStack(ctx context.Context, paramResultStack []uint6
 
 	entrypoint(c.executable, c.execCtxPtr, c.parent.opaquePtr, paramResultPtr, c.stackTop)
 	for {
-		switch c.execCtx.exitCode {
+		switch ec := c.execCtx.exitCode; ec & wazevoapi.ExitCodeMask {
 		case wazevoapi.ExitCodeOK:
 			return nil
 		case wazevoapi.ExitCodeGrowStack:
@@ -126,7 +130,7 @@ func (c *callEngine) CallWithStack(ctx context.Context, paramResultStack []uint6
 		case wazevoapi.ExitCodeMemoryOutOfBounds:
 			return wasmruntime.ErrRuntimeOutOfBoundsMemoryAccess
 		case wazevoapi.ExitCodeGrowMemory:
-			mod := moduleInstanceFromPtr(c.execCtx.callerModuleContextPtr)
+			mod := c.callerModuleInstance()
 			mem := mod.MemoryInstance
 			argRes := &c.execCtx.goFunctionCallStack[0]
 			if res, ok := mem.Grow(uint32(*argRes)); !ok {
@@ -145,14 +149,27 @@ func (c *callEngine) CallWithStack(ctx context.Context, paramResultStack []uint6
 			}
 			c.execCtx.exitCode = wazevoapi.ExitCodeOK
 			afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr, c.execCtx.stackPointerBeforeGoCall)
+		case wazevoapi.ExitCodeCallGoFunction:
+			index := wazevoapi.GoFunctionIndexFromExitCode(ec)
+			f := hostModuleGoFuncFromOpaque[api.GoFunction](index, c.execCtx.goFunctionCallCalleeModuleContextOpaque)
+			f.Call(ctx, c.execCtx.goFunctionCallStack[:])
+			c.execCtx.exitCode = wazevoapi.ExitCodeOK
+			afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr, c.execCtx.stackPointerBeforeGoCall)
+		case wazevoapi.ExitCodeCallGoModuleFunction:
+			index := wazevoapi.GoFunctionIndexFromExitCode(ec)
+			f := hostModuleGoFuncFromOpaque[api.GoModuleFunction](index, c.execCtx.goFunctionCallCalleeModuleContextOpaque)
+			mod := c.callerModuleInstance()
+			f.Call(ctx, mod, c.execCtx.goFunctionCallStack[:])
+			c.execCtx.exitCode = wazevoapi.ExitCodeOK
+			afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr, c.execCtx.stackPointerBeforeGoCall)
 		default:
 			panic("BUG")
 		}
 	}
 }
 
-func moduleInstanceFromPtr(ptr *byte) *wasm.ModuleInstance {
-	return *(**wasm.ModuleInstance)(unsafe.Pointer(ptr))
+func (c *callEngine) callerModuleInstance() *wasm.ModuleInstance {
+	return *(**wasm.ModuleInstance)(unsafe.Pointer(c.execCtx.callerModuleContextPtr))
 }
 
 func opaqueViewFromPtr(ptr uintptr) []byte {
diff --git a/internal/engine/wazevo/engine.go b/internal/engine/wazevo/engine.go
index 6fd38a1f..19c6bbbd 100644
--- a/internal/engine/wazevo/engine.go
+++ b/internal/engine/wazevo/engine.go
@@ -36,17 +36,21 @@ type (
 
 	// compiledModule is a compiled variant of a wasm.Module and ready to be used for instantiation.
 	compiledModule struct {
-		executable       []byte
-		functionOffsets  []compiledFunctionOffset
+		executable      []byte
+		functionOffsets []compiledFunctionOffset
+
+		// The followings are only available for non host modules.
+
 		offsets          wazevoapi.ModuleContextOffsetData
 		builtinFunctions *builtinFunctions
 	}
 
 	// compiledFunctionOffset tells us that where in the executable a function begins.
 	compiledFunctionOffset struct {
-		// offset is the beggining of the function.
+		// offset is the beginning of the function.
 		offset int
 		// goPreambleSize is the size of Go preamble of the function.
+		// This is only needed for non host modules.
 		goPreambleSize int
 	}
 )
@@ -62,10 +66,14 @@ func NewEngine(_ context.Context, _ api.CoreFeatures, _ filecache.Cache) wasm.En
 }
 
 // CompileModule implements wasm.Engine.
-func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []experimental.FunctionListener, ensureTermination bool) error {
+func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []experimental.FunctionListener, _ bool) error {
 	e.rels = e.rels[:0]
 	cm := &compiledModule{offsets: wazevoapi.NewModuleContextOffsetData(module)}
 
+	if module.IsHostModule {
+		return e.compileHostModule(module)
+	}
+
 	importedFns, localFns := int(module.ImportFunctionCount), len(module.FunctionSection)
 	if importedFns+localFns == 0 {
 		e.addCompiledModule(module, cm)
@@ -181,6 +189,97 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, _ []exper
 	return nil
 }
 
+func (e *engine) compileHostModule(module *wasm.Module) error {
+	machine := newMachine()
+	be := backend.NewCompiler(machine, ssa.NewBuilder())
+
+	num := len(module.CodeSection)
+	cm := &compiledModule{}
+	cm.functionOffsets = make([]compiledFunctionOffset, num)
+	totalSize := 0 // Total binary size of the executable.
+	bodies := make([][]byte, num)
+	var sig ssa.Signature
+	for i := range module.CodeSection {
+		totalSize = (totalSize + 15) &^ 15
+		cm.functionOffsets[i].offset = totalSize
+
+		typIndex := module.FunctionSection[i]
+		typ := &module.TypeSection[typIndex]
+		if typ.ParamNumInUint64 >= goFunctionCallStackSize || typ.ResultNumInUint64 >= goFunctionCallStackSize {
+			return fmt.Errorf("too many params or results for a host function (maximum %d): %v",
+				goFunctionCallStackSize, typ)
+		}
+
+		// We can relax until the index fits together in ExitCode as we do in wazevoapi.ExitCodeCallGoModuleFunctionWithIndex.
+		// However, 1 << 16 should be large enough for a real use case.
+		const hostFunctionNumMaximum = 1 << 16
+		if i >= hostFunctionNumMaximum {
+			return fmt.Errorf("too many host functions (maximum %d)", hostFunctionNumMaximum)
+		}
+
+		sig.ID = ssa.SignatureID(typIndex) // This is important since we reuse the `machine` which caches the ABI based on the SignatureID.
+		sig.Params = append(sig.Params[:0],
+			ssa.TypeI64, // First argument must be exec context.
+			ssa.TypeI64, // The second argument is the moduleContextOpaque of this host module.
+		)
+		for _, t := range typ.Params {
+			sig.Params = append(sig.Params, frontend.WasmTypeToSSAType(t))
+		}
+
+		sig.Results = sig.Results[:0]
+		for _, t := range typ.Results {
+			sig.Results = append(sig.Results, frontend.WasmTypeToSSAType(t))
+		}
+
+		c := &module.CodeSection[i]
+		if c.GoFunc == nil {
+			panic("BUG: GoFunc must be set for host module")
+		}
+
+		var exitCode wazevoapi.ExitCode
+		fn := c.GoFunc
+		switch fn.(type) {
+		case api.GoModuleFunction:
+			exitCode = wazevoapi.ExitCodeCallGoModuleFunctionWithIndex(i)
+		case api.GoFunction:
+			exitCode = wazevoapi.ExitCodeCallGoFunctionWithIndex(i)
+		}
+
+		be.Init(false)
+		machine.CompileGoFunctionTrampoline(exitCode, &sig, true)
+		be.Encode()
+		body := be.Buf()
+
+		// TODO: optimize as zero copy.
+		copied := make([]byte, len(body))
+		copy(copied, body)
+		bodies[i] = copied
+		totalSize += len(body)
+	}
+
+	// Allocate executable memory and then copy the generated machine code.
+	executable, err := platform.MmapCodeSegment(totalSize)
+	if err != nil {
+		panic(err)
+	}
+	cm.executable = executable
+
+	for i, b := range bodies {
+		offset := cm.functionOffsets[i]
+		copy(executable[offset.offset:], b)
+	}
+
+	if runtime.GOARCH == "arm64" {
+		// On arm64, we cannot give all of rwx at the same time, so we change it to exec.
+		if err = platform.MprotectRX(executable); err != nil {
+			return err
+		}
+	}
+
+	e.compiledModules[module.ID] = cm
+	return nil
+}
+
 // Close implements wasm.Engine.
 func (e *engine) Close() (err error) {
 	e.mux.Lock()
@@ -228,10 +327,15 @@ func (e *engine) NewModuleEngine(m *wasm.Module, mi *wasm.ModuleInstance) (wasm.
 	me.parent = compiled
 	me.module = mi
 
-	if size := compiled.offsets.TotalSize; size != 0 {
-		opaque := make([]byte, size)
-		me.opaque = opaque
-		me.opaquePtr = &opaque[0]
+	if m.IsHostModule {
+		me.opaque = buildHostModuleOpaque(m)
+		me.opaquePtr = &me.opaque[0]
+	} else {
+		if size := compiled.offsets.TotalSize; size != 0 {
+			opaque := make([]byte, size)
+			me.opaque = opaque
+			me.opaquePtr = &opaque[0]
+		}
 	}
 	return me, nil
 }
@@ -244,7 +348,7 @@ func (e *engine) compileBuiltinFunctions() {
 		machine.CompileGoFunctionTrampoline(wazevoapi.ExitCodeGrowMemory, &ssa.Signature{
 			Params:  []ssa.Type{ssa.TypeI32 /* exec context */, ssa.TypeI32},
 			Results: []ssa.Type{ssa.TypeI32},
-		})
+		}, false)
 		be.Encode()
 		src := be.Buf()
 
diff --git a/internal/engine/wazevo/frontend/frontend.go b/internal/engine/wazevo/frontend/frontend.go
index 1310e51b..752073fd 100644
--- a/internal/engine/wazevo/frontend/frontend.go
+++ b/internal/engine/wazevo/frontend/frontend.go
@@ -64,10 +64,10 @@ func NewFrontendCompiler(m *wasm.Module, ssaBuilder ssa.Builder, offset *wazevoa
 		sig.Params[0] = executionContextPtrTyp
 		sig.Params[1] = moduleContextPtrTyp
 		for j, typ := range wasmSig.Params {
-			sig.Params[j+2] = wasmToSSA(typ)
+			sig.Params[j+2] = WasmTypeToSSAType(typ)
 		}
 		for j, typ := range wasmSig.Results {
-			sig.Results[j] = wasmToSSA(typ)
+			sig.Results[j] = WasmTypeToSSAType(typ)
 		}
 		c.signatures[wasmSig] = sig
 		c.ssaBuilder.DeclareSignature(sig)
@@ -134,7 +134,7 @@ func (c *Compiler) LowerToSSA() error {
 	builder.AnnotateValue(c.moduleCtxPtrValue, "module_ctx")
 
 	for i, typ := range c.wasmFunctionTyp.Params {
-		st := wasmToSSA(typ)
+		st := WasmTypeToSSAType(typ)
 		variable := builder.DeclareVariable(st)
 		value := entryBlock.AddParam(builder, st)
 		builder.DefineVariable(variable, value, entryBlock)
@@ -156,7 +156,7 @@ func (c *Compiler) localVariable(index wasm.Index) ssa.Variable {
 func (c *Compiler) declareWasmLocals(entry ssa.BasicBlock) {
 	localCount := wasm.Index(len(c.wasmFunctionTyp.Params))
 	for i, typ := range c.wasmFunctionLocalTypes {
-		st := wasmToSSA(typ)
+		st := WasmTypeToSSAType(typ)
 		variable := c.ssaBuilder.DeclareVariable(st)
 		c.wasmLocalToVariable[wasm.Index(i)+localCount] = variable
 
@@ -227,8 +227,8 @@ func (c *Compiler) declareWasmGlobal(typ wasm.ValueType, mutable bool) {
 	}
 }
 
-// wasmToSSA converts wasm.ValueType to ssa.Type.
-func wasmToSSA(vt wasm.ValueType) ssa.Type {
+// WasmTypeToSSAType converts wasm.ValueType to ssa.Type.
+func WasmTypeToSSAType(vt wasm.ValueType) ssa.Type {
 	switch vt {
 	case wasm.ValueTypeI32:
 		return ssa.TypeI32
@@ -246,7 +246,7 @@ func wasmToSSA(vt wasm.ValueType) ssa.Type {
 // addBlockParamsFromWasmTypes adds the block parameters to the given block.
 func (c *Compiler) addBlockParamsFromWasmTypes(tps []wasm.ValueType, blk ssa.BasicBlock) {
 	for _, typ := range tps {
-		st := wasmToSSA(typ)
+		st := WasmTypeToSSAType(typ)
 		blk.AddParam(c.ssaBuilder, st)
 	}
 }
diff --git a/internal/engine/wazevo/hostmodule.go b/internal/engine/wazevo/hostmodule.go
new file mode 100644
index 00000000..b75d9811
--- /dev/null
+++ b/internal/engine/wazevo/hostmodule.go
@@ -0,0 +1,46 @@
+package wazevo
+
+import (
+	"encoding/binary"
+	"reflect"
+	"unsafe"
+
+	"github.com/tetratelabs/wazero/internal/wasm"
+)
+
+func buildHostModuleOpaque(m *wasm.Module) moduleContextOpaque {
+	size := len(m.CodeSection) * 16
+	ret := make(moduleContextOpaque, size)
+
+	var offset int
+	for i := range m.CodeSection {
+		goFn := m.CodeSection[i].GoFunc
+		writeIface(goFn, ret[offset:])
+		offset += 16
+	}
+	return ret
+}
+
+func hostModuleGoFuncFromOpaque[T any](index int, opaqueBegin uintptr) T {
+	offset := uintptr(index * 16)
+	ptr := opaqueBegin + offset
+
+	var opaqueViewOverFunction []byte
+	sh := (*reflect.SliceHeader)(unsafe.Pointer(&opaqueViewOverFunction))
+	sh.Data = ptr
+	sh.Len = 16
+	sh.Cap = 16
+	return readIface(opaqueViewOverFunction).(T)
+}
+
+func writeIface(goFn interface{}, buf []byte) {
+	goFnIface := *(*[2]uint64)(unsafe.Pointer(&goFn))
+	binary.LittleEndian.PutUint64(buf, goFnIface[0])
+	binary.LittleEndian.PutUint64(buf[8:], goFnIface[1])
+}
+
+func readIface(buf []byte) interface{} {
+	b := binary.LittleEndian.Uint64(buf)
+	s := binary.LittleEndian.Uint64(buf[8:])
+	return *(*interface{})(unsafe.Pointer(&[2]uint64{b, s}))
+}
diff --git a/internal/engine/wazevo/module_engine.go b/internal/engine/wazevo/module_engine.go
index d5a1035c..d0f97dca 100644
--- a/internal/engine/wazevo/module_engine.go
+++ b/internal/engine/wazevo/module_engine.go
@@ -26,9 +26,11 @@ type (
 	// Internally, the buffer is structured as follows:
 	//
 	// 	type moduleContextOpaque struct {
+	// 	    moduleInstance                            *wasm.ModuleInstance
 	// 	    localMemoryBufferPtr                      *byte                (optional)
 	// 	    localMemoryLength                         uint64               (optional)
 	// 	    importedMemoryInstance                    *wasm.MemoryInstance (optional)
+	// 	    importedMemoryOwnerOpaqueCtx              *byte                (optional)
 	// 	    importedFunctions [importedFunctions] struct { the total size depends on # of imported functions.
 	// 	        executable      *byte
 	// 	        opaqueCtx       *moduleContextOpaque
@@ -38,6 +40,8 @@ type (
 	// 	}
 	//
 	// See wazevoapi.NewModuleContextOffsetData for the details of the offsets.
+	//
+	// Note that for host modules, the structure is entirely different. See buildHostModuleOpaque.
 	moduleContextOpaque []byte
 )
 
@@ -121,6 +125,11 @@ func (m *moduleEngine) ResolveImportedMemory(importedModuleEngine wasm.ModuleEng
 	importedME := importedModuleEngine.(*moduleEngine)
 	inst := importedME.module
 
+	if importedME.parent.offsets.ImportedMemoryBegin >= 0 {
+		// This case can be resolved by recursively resolving the owner.
+		panic("TODO: support re-exported memory import")
+	}
+
 	offset := m.parent.offsets.ImportedMemoryBegin
 	b := uint64(uintptr(unsafe.Pointer(inst.MemoryInstance)))
 	binary.LittleEndian.PutUint64(m.opaque[offset:], b)
@@ -129,7 +138,9 @@ func (m *moduleEngine) ResolveImportedMemory(importedModuleEngine wasm.ModuleEng
 
 // DoneInstantiation implements wasm.ModuleEngine.
 func (m *moduleEngine) DoneInstantiation() {
-	m.setupOpaque()
+	if !m.module.Source.IsHostModule {
+		m.setupOpaque()
+	}
 }
 
 // LookupFunction implements wasm.ModuleEngine.
diff --git a/internal/engine/wazevo/wazevoapi/exitcode.go b/internal/engine/wazevo/wazevoapi/exitcode.go
index bf6eca72..608650d4 100644
--- a/internal/engine/wazevo/wazevoapi/exitcode.go
+++ b/internal/engine/wazevo/wazevoapi/exitcode.go
@@ -9,8 +9,16 @@ const (
 	ExitCodeGrowMemory
 	ExitCodeUnreachable
 	ExitCodeMemoryOutOfBounds
+	// ExitCodeCallGoModuleFunction is an exit code for a call to an api.GoModuleFunction.
+	ExitCodeCallGoModuleFunction
+	// ExitCodeCallGoFunction is an exit code for a call to an api.GoFunction.
+	ExitCodeCallGoFunction
+
+	exitCodeMax
 )
 
+const ExitCodeMask = 0xff
+
 // String implements fmt.Stringer.
 func (e ExitCode) String() string {
 	switch e {
@@ -18,6 +26,10 @@ func (e ExitCode) String() string {
 		return "ok"
 	case ExitCodeGrowStack:
 		return "grow_stack"
+	case ExitCodeCallGoModuleFunction:
+		return "call_go_module_function"
+	case ExitCodeCallGoFunction:
+		return "call_go_function"
 	case ExitCodeUnreachable:
 		return "unreachable"
 	case ExitCodeMemoryOutOfBounds:
@@ -25,3 +37,15 @@ func (e ExitCode) String() string {
 	}
 	panic("TODO")
 }
+
+func ExitCodeCallGoModuleFunctionWithIndex(index int) ExitCode {
+	return ExitCodeCallGoModuleFunction | ExitCode(index<<8)
+}
+
+func ExitCodeCallGoFunctionWithIndex(index int) ExitCode {
+	return ExitCodeCallGoFunction | ExitCode(index<<8)
+}
+
+func GoFunctionIndexFromExitCode(exitCode ExitCode) int {
+	return int(exitCode >> 8)
+}
diff --git a/internal/engine/wazevo/wazevoapi/offsetdata.go b/internal/engine/wazevo/wazevoapi/offsetdata.go
index c1d44b0a..d9538970 100644
--- a/internal/engine/wazevo/wazevoapi/offsetdata.go
+++ b/internal/engine/wazevo/wazevoapi/offsetdata.go
@@ -3,18 +3,19 @@ package wazevoapi
 import "github.com/tetratelabs/wazero/internal/wasm"
 
 var ExecutionContextOffsets = ExecutionContextOffsetData{
-	ExitCodeOffset:              0,
-	CallerModuleContextPtr:      8,
-	OriginalFramePointer:        16,
-	OriginalStackPointer:        24,
-	GoReturnAddress:             32,
-	StackBottomPtr:              40,
-	GoCallReturnAddress:         48,
-	StackPointerBeforeGrow:      56,
-	StackGrowRequiredSize:       64,
-	MemoryGrowTrampolineAddress: 72,
-	SavedRegistersBegin:         80,
-	GoFunctionCallStackBegin:    1104,
+	ExitCodeOffset:                          0,
+	CallerModuleContextPtr:                  8,
+	OriginalFramePointer:                    16,
+	OriginalStackPointer:                    24,
+	GoReturnAddress:                         32,
+	StackBottomPtr:                          40,
+	GoCallReturnAddress:                     48,
+	StackPointerBeforeGrow:                  56,
+	StackGrowRequiredSize:                   64,
+	MemoryGrowTrampolineAddress:             72,
+	SavedRegistersBegin:                     80,
+	GoFunctionCallCalleeModuleContextOpaque: 1104,
+	GoFunctionCallStackBegin:                1112,
 }
 
 // ExecutionContextOffsetData allows the compilers to get the information about offsets to the fields of wazevo.executionContext,
@@ -42,6 +43,8 @@ type ExecutionContextOffsetData struct {
 	MemoryGrowTrampolineAddress Offset
 	// GoCallReturnAddress is an offset of the first element of `savedRegisters` field in wazevo.executionContext
 	SavedRegistersBegin Offset
+	// GoFunctionCallCalleeModuleContextOpaque is an offset of `goFunctionCallCalleeModuleContextOpaque` field in wazevo.executionContext
+	GoFunctionCallCalleeModuleContextOpaque Offset
 	// GoFunctionCallStackBegin is an offset of the first element of `goFunctionCallStack` field in wazevo.executionContext
 	GoFunctionCallStackBegin Offset
 }


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go b/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go
index a4997a93..ec5bfa44 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go
@@ -27,11 +27,169 @@ func Test_calleeSavedRegistersPlusLinkRegSorted(t *testing.T) {
 
 func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
 	for _, tc := range []struct {
-		name     string
-		exitCode wazevoapi.ExitCode
-		sig      *ssa.Signature
-		exp      string
+		name                 string
+		exitCode             wazevoapi.ExitCode
+		sig                  *ssa.Signature
+		needModuleContextPtr bool
+		exp                  string
 	}{
+		{
+			name:     "go call",
+			exitCode: wazevoapi.ExitCodeCallGoFunctionWithIndex(100),
+			sig: &ssa.Signature{
+				Params:  []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeF64},
+				Results: []ssa.Type{ssa.TypeI32, ssa.TypeI64, ssa.TypeF32, ssa.TypeF64},
+			},
+			needModuleContextPtr: true,
+			exp: `
+	str x18, [x0, #0x50]
+	str x19, [x0, #0x60]
+	str x20, [x0, #0x70]
+	str x21, [x0, #0x80]
+	str x22, [x0, #0x90]
+	str x23, [x0, #0xa0]
+	str x24, [x0, #0xb0]
+	str x25, [x0, #0xc0]
+	str x26, [x0, #0xd0]
+	str x28, [x0, #0xe0]
+	str x30, [x0, #0xf0]
+	str q18, [x0, #0x100]
+	str q19, [x0, #0x110]
+	str q20, [x0, #0x120]
+	str q21, [x0, #0x130]
+	str q22, [x0, #0x140]
+	str q23, [x0, #0x150]
+	str q24, [x0, #0x160]
+	str q25, [x0, #0x170]
+	str q26, [x0, #0x180]
+	str q27, [x0, #0x190]
+	str q28, [x0, #0x1a0]
+	str q29, [x0, #0x1b0]
+	str q30, [x0, #0x1c0]
+	str q31, [x0, #0x1d0]
+	str x1, [x0, #0x450]
+	add x15, x0, #0x458
+	str d0, [x15], #0x8
+	movz w17, #0x6406, LSL 0
+	str w17, [x0]
+	mov x27, sp
+	str x27, [x0, #0x38]
+	adr x27, #0x1c
+	str x27, [x0, #0x30]
+	exit_sequence w0
+	ldr x18, [x0, #0x50]
+	ldr x19, [x0, #0x60]
+	ldr x20, [x0, #0x70]
+	ldr x21, [x0, #0x80]
+	ldr x22, [x0, #0x90]
+	ldr x23, [x0, #0xa0]
+	ldr x24, [x0, #0xb0]
+	ldr x25, [x0, #0xc0]
+	ldr x26, [x0, #0xd0]
+	ldr x28, [x0, #0xe0]
+	ldr x30, [x0, #0xf0]
+	ldr q18, [x0, #0x100]
+	ldr q19, [x0, #0x110]
+	ldr q20, [x0, #0x120]
+	ldr q21, [x0, #0x130]
+	ldr q22, [x0, #0x140]
+	ldr q23, [x0, #0x150]
+	ldr q24, [x0, #0x160]
+	ldr q25, [x0, #0x170]
+	ldr q26, [x0, #0x180]
+	ldr q27, [x0, #0x190]
+	ldr q28, [x0, #0x1a0]
+	ldr q29, [x0, #0x1b0]
+	ldr q30, [x0, #0x1c0]
+	ldr q31, [x0, #0x1d0]
+	add x15, x0, #0x458
+	ldr w27, [x15], #0x8
+	mov w0, w27
+	ldr x27, [x15], #0x8
+	mov x1, x27
+	ldr s17, [x15], #0x8
+	mov q0.8b, q17.8b
+	ldr d17, [x15], #0x8
+	mov q1.8b, q17.8b
+	ret
+`,
+		},
+		{
+			name:     "go call",
+			exitCode: wazevoapi.ExitCodeCallGoFunctionWithIndex(100),
+			sig: &ssa.Signature{
+				Params:  []ssa.Type{ssa.TypeI64, ssa.TypeI64, ssa.TypeF64, ssa.TypeF64, ssa.TypeI32, ssa.TypeI32},
+				Results: []ssa.Type{},
+			},
+			needModuleContextPtr: true,
+			exp: `
+	str x18, [x0, #0x50]
+	str x19, [x0, #0x60]
+	str x20, [x0, #0x70]
+	str x21, [x0, #0x80]
+	str x22, [x0, #0x90]
+	str x23, [x0, #0xa0]
+	str x24, [x0, #0xb0]
+	str x25, [x0, #0xc0]
+	str x26, [x0, #0xd0]
+	str x28, [x0, #0xe0]
+	str x30, [x0, #0xf0]
+	str q18, [x0, #0x100]
+	str q19, [x0, #0x110]
+	str q20, [x0, #0x120]
+	str q21, [x0, #0x130]
+	str q22, [x0, #0x140]
+	str q23, [x0, #0x150]
+	str q24, [x0, #0x160]
+	str q25, [x0, #0x170]
+	str q26, [x0, #0x180]
+	str q27, [x0, #0x190]
+	str q28, [x0, #0x1a0]
+	str q29, [x0, #0x1b0]
+	str q30, [x0, #0x1c0]
+	str q31, [x0, #0x1d0]
+	str x1, [x0, #0x450]
+	add x15, x0, #0x458
+	str d0, [x15], #0x8
+	str d1, [x15], #0x8
+	str x2, [x15], #0x8
+	str x3, [x15], #0x8
+	movz w17, #0x6406, LSL 0
+	str w17, [x0]
+	mov x27, sp
+	str x27, [x0, #0x38]
+	adr x27, #0x1c
+	str x27, [x0, #0x30]
+	exit_sequence w0
+	ldr x18, [x0, #0x50]
+	ldr x19, [x0, #0x60]
+	ldr x20, [x0, #0x70]
+	ldr x21, [x0, #0x80]
+	ldr x22, [x0, #0x90]
+	ldr x23, [x0, #0xa0]
+	ldr x24, [x0, #0xb0]
+	ldr x25, [x0, #0xc0]
+	ldr x26, [x0, #0xd0]
+	ldr x28, [x0, #0xe0]
+	ldr x30, [x0, #0xf0]
+	ldr q18, [x0, #0x100]
+	ldr q19, [x0, #0x110]
+	ldr q20, [x0, #0x120]
+	ldr q21, [x0, #0x130]
+	ldr q22, [x0, #0x140]
+	ldr q23, [x0, #0x150]
+	ldr q24, [x0, #0x160]
+	ldr q25, [x0, #0x170]
+	ldr q26, [x0, #0x180]
+	ldr q27, [x0, #0x190]
+	ldr q28, [x0, #0x1a0]
+	ldr q29, [x0, #0x1b0]
+	ldr q30, [x0, #0x1c0]
+	ldr q31, [x0, #0x1d0]
+	add x15, x0, #0x458
+	ret
+`,
+		},
 		{
 			name:     "grow memory",
 			exitCode: wazevoapi.ExitCodeGrowMemory,
@@ -65,10 +223,10 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
 	str q29, [x0, #0x1b0]
 	str q30, [x0, #0x1c0]
 	str q31, [x0, #0x1d0]
-	add x15, x0, #0x450
+	add x15, x0, #0x458
 	str x1, [x15], #0x8
-	movz x27, #0x2, LSL 0
-	str w27, [x0]
+	orr w17, wzr, #0x2
+	str w17, [x0]
 	mov x27, sp
 	str x27, [x0, #0x38]
 	adr x27, #0x1c
@@ -99,7 +257,7 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
 	ldr q29, [x0, #0x1b0]
 	ldr q30, [x0, #0x1c0]
 	ldr q31, [x0, #0x1d0]
-	add x15, x0, #0x450
+	add x15, x0, #0x458
 	ldr w27, [x15], #0x8
 	mov w0, w27
 	ret
@@ -108,7 +266,7 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			_, _, m := newSetupWithMockContext()
-			m.CompileGoFunctionTrampoline(tc.exitCode, tc.sig)
+			m.CompileGoFunctionTrampoline(tc.exitCode, tc.sig, tc.needModuleContextPtr)
 
 			fmt.Println(m.Format())
 			require.Equal(t, tc.exp, m.Format())
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go
index 40b92ad2..a5fc2a38 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go
@@ -217,8 +217,8 @@ func TestMachine_insertStackBoundsCheck(t *testing.T) {
 	str q31, [x0, #0x2c0]
 	mov x27, sp
 	str x27, [x0, #0x38]
-	movz x27, #0x1, LSL 0
-	str w27, [x0]
+	orr w17, wzr, #0x1
+	str w17, [x0]
 	movz x27, #0xfff0, LSL 0
 	str x27, [x0, #0x40]
 	adr x27, #0x1c
@@ -315,8 +315,8 @@ func TestMachine_insertStackBoundsCheck(t *testing.T) {
 	str q31, [x0, #0x2c0]
 	mov x27, sp
 	str x27, [x0, #0x38]
-	movz x27, #0x1, LSL 0
-	str w27, [x0]
+	orr w17, wzr, #0x1
+	str w17, [x0]
 	orr x27, xzr, #0x10
 	str x27, [x0, #0x40]
 	adr x27, #0x1c
diff --git a/internal/engine/wazevo/backend/machine_test.go b/internal/engine/wazevo/backend/machine_test.go
index db585d17..f328dad8 100644
--- a/internal/engine/wazevo/backend/machine_test.go
+++ b/internal/engine/wazevo/backend/machine_test.go
@@ -24,10 +24,8 @@ type mockMachine struct {
 	rinfo                  *regalloc.RegisterInfo
 }
 
-func (m mockMachine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *ssa.Signature) {
-	// TODO implement me
-	panic("implement me")
-}
+// CompileGoFunctionTrampoline implements Machine.CompileGoFunctionTrampoline.
+func (m mockMachine) CompileGoFunctionTrampoline(wazevoapi.ExitCode, *ssa.Signature, bool) {}
 
 // Encode implements Machine.Encode.
 func (m mockMachine) Encode() {
@@ -36,8 +34,7 @@ func (m mockMachine) Encode() {
 }
 
 // ResolveRelocations implements Machine.ResolveRelocations.
-func (m mockMachine) ResolveRelocations(refToBinaryOffset map[ssa.FuncRef]int, binary []byte, relocations []RelocationInfo) {
-}
+func (m mockMachine) ResolveRelocations(map[ssa.FuncRef]int, []byte, []RelocationInfo) {}
 
 // SetupPrologue implements Machine.SetupPrologue.
 func (m mockMachine) SetupPrologue() {}
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 63b8676d..11b01783 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -17,15 +17,23 @@ import (
 	"github.com/tetratelabs/wazero/internal/wasm"
 )
 
+const (
+	i32 = wasm.ValueTypeI32
+	i64 = wasm.ValueTypeI64
+	f32 = wasm.ValueTypeF32
+	f64 = wasm.ValueTypeF64
+)
+
 func TestE2E(t *testing.T) {
 	type callCase struct {
 		params, expResults []uint64
 		expErr             string
 	}
 	for _, tc := range []struct {
-		name        string
-		imported, m *wasm.Module
-		calls       []callCase
+		name           string
+		imported, m    *wasm.Module
+		needHostModule bool
+		calls          []callCase
 	}{
 		{
 			name: "swap", m: testcases.SwapParamAndReturn.Module,
@@ -224,3 +232,79 @@ func configureWazevo(config wazero.RuntimeConfig) {
 	// Insert the wazevo implementation.
 	cm.newEngine = wazevo.NewEngine
 }
+
+func TestE2E_host_functions(t *testing.T) {
+	config := wazero.NewRuntimeConfigCompiler()
+
+	// Configure the new optimizing backend!
+	configureWazevo(config)
+
+	ctx := context.Background()
+	r := wazero.NewRuntimeWithConfig(ctx, config)
+	defer func() {
+		require.NoError(t, r.Close(ctx))
+	}()
+
+	var expectedMod api.Module
+
+	b := r.NewHostModuleBuilder("env")
+	b.NewFunctionBuilder().WithFunc(func(ctx2 context.Context, d float64) float64 {
+		require.Equal(t, ctx, ctx2)
+		fmt.Printf("%#x\n", math.Float64bits(d))
+		require.Equal(t, 35.0, d)
+		return math.Sqrt(d)
+	}).Export("root")
+	b.NewFunctionBuilder().WithFunc(func(ctx2 context.Context, mod api.Module, a uint32, b uint64, c float32, d float64) (uint32, uint64, float32, float64) {
+		require.Equal(t, expectedMod, mod)
+		require.Equal(t, ctx, ctx2)
+		require.Equal(t, uint32(2), a)
+		require.Equal(t, uint64(100), b)
+		require.Equal(t, float32(15.0), c)
+		require.Equal(t, 35.0, d)
+		return a * a, b * b, c * c, d * d
+	}).Export("square")
+
+	_, err := b.Instantiate(ctx)
+	require.NoError(t, err)
+
+	m := &wasm.Module{
+		ImportFunctionCount: 2,
+		ImportSection: []wasm.Import{
+			{Module: "env", Name: "root", Type: wasm.ExternTypeFunc, DescFunc: 0},
+			{Module: "env", Name: "square", Type: wasm.ExternTypeFunc, DescFunc: 1},
+		},
+		TypeSection: []wasm.FunctionType{
+			{Results: []wasm.ValueType{f64}, Params: []wasm.ValueType{f64}},
+			{Results: []wasm.ValueType{i32, i64, f32, f64}, Params: []wasm.ValueType{i32, i64, f32, f64}},
+			{Results: []wasm.ValueType{i32, i64, f32, f64, f64}, Params: []wasm.ValueType{i32, i64, f32, f64}},
+		},
+		FunctionSection: []wasm.Index{2},
+		CodeSection: []wasm.Code{{
+			Body: []byte{
+				wasm.OpcodeLocalGet, 0, wasm.OpcodeLocalGet, 1, wasm.OpcodeLocalGet, 2, wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeCall, 1,
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeCall, 0,
+				wasm.OpcodeEnd,
+			},
+		}},
+		ExportSection: []wasm.Export{{Name: "f", Type: wasm.ExternTypeFunc, Index: 2}},
+	}
+
+	compiled, err := r.CompileModule(ctx, binaryencoding.EncodeModule(m))
+	require.NoError(t, err)
+
+	inst, err := r.InstantiateModule(ctx, compiled, wazero.NewModuleConfig())
+	require.NoError(t, err)
+
+	expectedMod = inst
+
+	f := inst.ExportedFunction("f")
+
+	res, err := f.Call(ctx, []uint64{2, 100, uint64(math.Float32bits(15.0)), math.Float64bits(35.0)}...)
+	require.NoError(t, err)
+	require.Equal(t, []uint64{
+		2 * 2, 100 * 100, uint64(math.Float32bits(15.0 * 15.0)), math.Float64bits(35.0 * 35.0),
+		math.Float64bits(math.Sqrt(35.0)),
+	}, res)
+}
diff --git a/internal/engine/wazevo/hostmodule_test.go b/internal/engine/wazevo/hostmodule_test.go
new file mode 100644
index 00000000..0e38290d
--- /dev/null
+++ b/internal/engine/wazevo/hostmodule_test.go
@@ -0,0 +1,22 @@
+package wazevo
+
+import (
+	"context"
+	"testing"
+
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+func Test_writeIface_readIface(t *testing.T) {
+	buf := make([]byte, 100)
+
+	var called bool
+	var goFn api.GoFunction = api.GoFunc(func(context.Context, []uint64) {
+		called = true
+	})
+	writeIface(goFn, buf)
+	got := readIface(buf).(api.GoFunction)
+	got.Call(context.Background(), nil)
+	require.True(t, called)
+}
diff --git a/internal/engine/wazevo/module_engine_test.go b/internal/engine/wazevo/module_engine_test.go
index 88c4cf24..3e2b529a 100644
--- a/internal/engine/wazevo/module_engine_test.go
+++ b/internal/engine/wazevo/module_engine_test.go
@@ -69,7 +69,10 @@ func TestModuleEngine_setupOpaque(t *testing.T) {
 				require.Equal(t, expLen, actualLen)
 			}
 			if tc.offset.ImportedMemoryBegin >= 0 {
-				imported := &moduleEngine{opaque: []byte{1, 2, 3}, module: &wasm.ModuleInstance{MemoryInstance: tc.m.MemoryInstance}}
+				imported := &moduleEngine{
+					opaque: []byte{1, 2, 3}, module: &wasm.ModuleInstance{MemoryInstance: tc.m.MemoryInstance},
+					parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{ImportedMemoryBegin: -1}},
+				}
 				imported.opaquePtr = &imported.opaque[0]
 				m.ResolveImportedMemory(imported)
 
diff --git a/internal/engine/wazevo/wazevo_test.go b/internal/engine/wazevo/wazevo_test.go
index 812d40b3..790395b4 100644
--- a/internal/engine/wazevo/wazevo_test.go
+++ b/internal/engine/wazevo/wazevo_test.go
@@ -61,5 +61,6 @@ func Test_ExecutionContextOffsets(t *testing.T) {
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.savedRegisters))%16, wazevoapi.Offset(0),
 		"SavedRegistersBegin must be aligned to 16 bytes")
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.savedRegisters)), offsets.SavedRegistersBegin)
+	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.goFunctionCallCalleeModuleContextOpaque)), offsets.GoFunctionCallCalleeModuleContextOpaque)
 	require.Equal(t, wazevoapi.Offset(unsafe.Offsetof(execCtx.goFunctionCallStack)), offsets.GoFunctionCallStackBegin)
 }
diff --git a/internal/engine/wazevo/wazevoapi/exitcode_test.go b/internal/engine/wazevo/wazevoapi/exitcode_test.go
new file mode 100644
index 00000000..ef7e317c
--- /dev/null
+++ b/internal/engine/wazevo/wazevoapi/exitcode_test.go
@@ -0,0 +1,11 @@
+package wazevoapi
+
+import (
+	"testing"
+
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+func TestExitCode_withinByte(t *testing.T) {
+	require.True(t, exitCodeMax < ExitCodeMask) //nolint
+}