wazevo: fixes re-exported function imports (#1708)

internal/engine/wazevo/call_engine.go

@@ -158,6 +158,11 @@ func (c *callEngine) callWithStack(ctx context.Context, paramResultStack []uint6
 				err = c.parent.module.FailIfClosed()
 			}
 		}
+
+		if err != nil {
+			// Ensures that we can reuse this callEngine even after an error.
+			c.execCtx.exitCode = wazevoapi.ExitCodeOK
+		}
 	}()
 
 	entrypoint(c.preambleExecutable, c.executable, c.execCtxPtr, c.parent.opaquePtr, paramResultPtr, c.stackTop)

internal/engine/wazevo/engine.go

@@ -53,8 +53,8 @@ type (
 	// compiledModule is a compiled variant of a wasm.Module and ready to be used for instantiation.
 	compiledModule struct {
 		executable []byte
-		// functionOffsets maps a local function index to compiledFunctionOffset.
-		functionOffsets []compiledFunctionOffset
+		// functionOffsets maps a local function index to the offset in the executable.
+		functionOffsets []int
 		parent          *engine
 		module          *wasm.Module
 		entryPreambles  []*byte // indexed-correlated with the type index.
@@ -64,22 +64,8 @@ type (
 		offsets         wazevoapi.ModuleContextOffsetData
 		sharedFunctions *sharedFunctions
 	}
-
-	// compiledFunctionOffset tells us that where in the executable a function begins.
-	compiledFunctionOffset struct {
-		// 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
-	}
 )
 
-// nativeBegin returns the offset of the beginning of the function in the executable after the Go preamble if any.
-func (c compiledFunctionOffset) nativeBegin() int {
-	return c.offset + c.goPreambleSize
-}
-
 var _ wasm.Engine = (*engine)(nil)
 
 // NewEngine returns the implementation of wasm.Engine.
@@ -150,7 +136,7 @@ func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listene
 	be := backend.NewCompiler(ctx, machine, ssaBuilder)
 
 	totalSize := 0 // Total binary size of the executable.
-	cm.functionOffsets = make([]compiledFunctionOffset, localFns)
+	cm.functionOffsets = make([]int, localFns)
 	bodies := make([][]byte, localFns)
 	for i := range module.CodeSection {
 		if wazevoapi.DeterministicCompilationVerifierEnabled {
@@ -175,8 +161,7 @@ func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listene
 
 		// Align 16-bytes boundary.
 		totalSize = (totalSize + 15) &^ 15
-		compiledFuncOffset := &cm.functionOffsets[i]
-		compiledFuncOffset.offset = totalSize
+		cm.functionOffsets[i] = totalSize
 
 		fref := frontend.FunctionIndexToFuncRef(fidx)
 		e.refToBinaryOffset[fref] = totalSize
@@ -204,7 +189,7 @@ func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listene
 
 	for i, b := range bodies {
 		offset := cm.functionOffsets[i]
-		copy(executable[offset.offset:], b)
+		copy(executable[offset:], b)
 	}
 
 	// Resolve relocations for local function calls.
@@ -289,13 +274,13 @@ func (e *engine) compileHostModule(ctx context.Context, module *wasm.Module) (*c
 
 	num := len(module.CodeSection)
 	cm := &compiledModule{module: module}
-	cm.functionOffsets = make([]compiledFunctionOffset, num)
+	cm.functionOffsets = make([]int, 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
+		cm.functionOffsets[i] = totalSize
 
 		typIndex := module.FunctionSection[i]
 		typ := &module.TypeSection[typIndex]
@@ -360,7 +345,7 @@ func (e *engine) compileHostModule(ctx context.Context, module *wasm.Module) (*c
 
 	for i, b := range bodies {
 		offset := cm.functionOffsets[i]
-		copy(executable[offset.offset:], b)
+		copy(executable[offset:], b)
 	}
 
 	if runtime.GOARCH == "arm64" {
@@ -553,7 +538,7 @@ func (cm *compiledModule) functionIndexOf(addr uintptr) wasm.Index {
 	addr -= uintptr(unsafe.Pointer(&cm.executable[0]))
 	offset := cm.functionOffsets
 	index := sort.Search(len(offset), func(i int) bool {
-		return offset[i].offset > int(addr)
+		return offset[i] > int(addr)
 	})
 	index -= 1
 	if index < 0 {

internal/engine/wazevo/engine_test.go

@@ -165,13 +165,8 @@ func TestCompiledModule_functionIndexOf(t *testing.T) {
 	}
 
 	cm := &compiledModule{
-		executable: executable,
-		functionOffsets: []compiledFunctionOffset{
-			{offset: 0, goPreambleSize: 100},
-			{offset: 500, goPreambleSize: 200},
-			{offset: 1000, goPreambleSize: 0},
-			{offset: 2000, goPreambleSize: 0},
-		},
+		executable:      executable,
+		functionOffsets: []int{0, 500, 1000, 2000},
 	}
 	require.Equal(t, wasm.Index(0), cm.functionIndexOf(executableAddr))
 	require.Equal(t, wasm.Index(0), cm.functionIndexOf(executableAddr+499))

internal/engine/wazevo/module_engine.go

@@ -133,7 +133,7 @@ func (m *moduleEngine) NewFunction(index wasm.Index) api.Function {
 
 	ce := &callEngine{
 		indexInModule:          index,
-		executable:             &p.executable[offset.offset],
+		executable:             &p.executable[offset],
 		parent:                 m,
 		preambleExecutable:     m.parent.entryPreambles[typIndex],
 		sizeOfParamResultSlice: sizeOfParamResultSlice,
@@ -152,10 +152,17 @@ func (m *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm
 	executableOffset, moduleCtxOffset, typeIDOffset := m.parent.offsets.ImportedFunctionOffset(index)
 	importedME := importedModuleEngine.(*moduleEngine)
 
+	if int(indexInImportedModule) >= len(importedME.importedFunctions) {
+		indexInImportedModule -= wasm.Index(len(importedME.importedFunctions))
+	} else {
+		imported := &importedME.importedFunctions[indexInImportedModule]
+		m.ResolveImportedFunction(index, imported.indexInModule, imported.me)
+		return // Recursively resolve the imported function.
+	}
+
 	offset := importedME.parent.functionOffsets[indexInImportedModule]
 	typeID := getTypeIDOf(indexInImportedModule, importedME.module)
-	// When calling imported function from the machine code, we need to skip the Go preamble.
-	executable := &importedME.parent.executable[offset.nativeBegin()]
+	executable := &importedME.parent.executable[offset]
 	// Write functionInstance.
 	binary.LittleEndian.PutUint64(m.opaque[executableOffset:], uint64(uintptr(unsafe.Pointer(executable))))
 	binary.LittleEndian.PutUint64(m.opaque[moduleCtxOffset:], uint64(uintptr(unsafe.Pointer(importedME.opaquePtr))))
@@ -222,7 +229,7 @@ func (m *moduleEngine) FunctionInstanceReference(funcIndex wasm.Index) wasm.Refe
 	}
 	localIndex := funcIndex - m.module.Source.ImportFunctionCount
 	p := m.parent
-	executable := &p.executable[p.functionOffsets[localIndex].nativeBegin()]
+	executable := &p.executable[p.functionOffsets[localIndex]]
 	typeID := m.module.TypeIDs[m.module.Source.FunctionSection[localIndex]]
 
 	lf := &functionInstance{

internal/engine/wazevo/module_engine_test.go

@@ -126,7 +126,7 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 		opaquePtr: &op1,
 		parent: &compiledModule{
 			executable:      make([]byte, 1000),
-			functionOffsets: []compiledFunctionOffset{{offset: 1, goPreambleSize: 4}, {offset: 5, goPreambleSize: 4}, {offset: 10, goPreambleSize: 4}},
+			functionOffsets: []int{1, 5, 10},
 		},
 		module: &wasm.ModuleInstance{
 			TypeIDs: []wasm.FunctionTypeID{0, 0, 0, 0, 111, 222, 333},
@@ -137,7 +137,7 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 		opaquePtr: &op2,
 		parent: &compiledModule{
 			executable:      make([]byte, 1000),
-			functionOffsets: []compiledFunctionOffset{{offset: 50, goPreambleSize: 4}},
+			functionOffsets: []int{50, 4},
 		},
 		module: &wasm.ModuleInstance{
 			TypeIDs: []wasm.FunctionTypeID{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 999},
@@ -156,10 +156,65 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 		executable *byte
 		expTypeID  wasm.FunctionTypeID
 	}{
-		{index: 0, op: &op1, executable: &im1.parent.executable[1+4], expTypeID: 111},
-		{index: 1, op: &op2, executable: &im2.parent.executable[50+4], expTypeID: 999},
-		{index: 2, op: &op1, executable: &im1.parent.executable[10+4], expTypeID: 333},
-		{index: 3, op: &op1, executable: &im1.parent.executable[5+4], expTypeID: 222},
+		{index: 0, op: &op1, executable: &im1.parent.executable[1], expTypeID: 111},
+		{index: 1, op: &op2, executable: &im2.parent.executable[50], expTypeID: 999},
+		{index: 2, op: &op1, executable: &im1.parent.executable[10], expTypeID: 333},
+		{index: 3, op: &op1, executable: &im1.parent.executable[5], expTypeID: 222},
+	} {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			buf := m.opaque[begin+wazevoapi.FunctionInstanceSize*tc.index:]
+			actualExecutable := binary.LittleEndian.Uint64(buf)
+			actualOpaquePtr := binary.LittleEndian.Uint64(buf[8:])
+			actualTypeID := binary.LittleEndian.Uint64(buf[16:])
+			expExecutable := uint64(uintptr(unsafe.Pointer(tc.executable)))
+			expOpaquePtr := uint64(uintptr(unsafe.Pointer(tc.op)))
+			require.Equal(t, expExecutable, actualExecutable)
+			require.Equal(t, expOpaquePtr, actualOpaquePtr)
+			require.Equal(t, uint64(tc.expTypeID), actualTypeID)
+		})
+	}
+}
+
+func TestModuleEngine_ResolveImportedFunction_recursive(t *testing.T) {
+	const begin = 5000
+	m := &moduleEngine{
+		opaque:            make([]byte, 10000),
+		importedFunctions: make([]importedFunction, 4),
+		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
+			ImportedFunctionsBegin: begin,
+		}},
+	}
+
+	var importingOp, importedOp byte = 0xaa, 0xbb
+	imported := &moduleEngine{
+		opaquePtr: &importedOp,
+		parent:    &compiledModule{executable: make([]byte, 50), functionOffsets: []int{10}},
+		module: &wasm.ModuleInstance{
+			TypeIDs: []wasm.FunctionTypeID{111},
+			Source:  &wasm.Module{FunctionSection: []wasm.Index{0}},
+		},
+	}
+	importing := &moduleEngine{
+		opaquePtr:         &importingOp,
+		parent:            &compiledModule{executable: make([]byte, 1000), functionOffsets: []int{500}},
+		importedFunctions: []importedFunction{{me: imported, indexInModule: 0}},
+		module: &wasm.ModuleInstance{
+			TypeIDs: []wasm.FunctionTypeID{0, 222, 0},
+			Source:  &wasm.Module{FunctionSection: []wasm.Index{1}},
+		},
+	}
+
+	m.ResolveImportedFunction(0, 0, importing)
+	m.ResolveImportedFunction(1, 1, importing)
+
+	for i, tc := range []struct {
+		index      int
+		op         *byte
+		executable *byte
+		expTypeID  wasm.FunctionTypeID
+	}{
+		{index: 0, op: &importedOp, executable: &imported.parent.executable[10], expTypeID: 111},
+		{index: 1, op: &importingOp, executable: &importing.parent.executable[500], expTypeID: 222},
 	} {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
 			buf := m.opaque[begin+wazevoapi.FunctionInstanceSize*tc.index:]

internal/integration_test/engine/adhoc_test.go

@@ -46,7 +46,7 @@ var tests = map[string]testCase{
 	"host function with numeric parameter":              {f: testHostFunctionNumericParameter},
 	"close module with in-flight calls":                 {f: testCloseInFlight},
 	"multiple instantiation from same source":           {f: testMultipleInstantiation},
-	"exported function that grows memory":               {f: testMemOps, wazevoSkip: true},
+	"exported function that grows memory":               {f: testMemOps},
 	"import functions with reference type in signature": {f: testReftypeImports, wazevoSkip: true},
 	"overflow integer addition":                         {f: testOverflow},
 	"un-signed extend global":                           {f: testGlobalExtend},