compiler: ResolveImportedFunction should infer the typeID from the importing module

internal/engine/interpreter/interpreter.go

@@ -487,7 +487,7 @@ func (e *engine) setLabelAddress(op *uint64, label label, labelAddressResolution
 }
 
 // ResolveImportedFunction implements wasm.ModuleEngine.
-func (e *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
+func (e *moduleEngine) ResolveImportedFunction(index, descFunc, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
 	imported := importedModuleEngine.(*moduleEngine)
 	e.functions[index] = imported.functions[indexInImportedModule]
 }

internal/engine/wazevo/module_engine.go

@@ -237,20 +237,20 @@ func (m *moduleEngine) putLocalMemory() {
 }
 
 // ResolveImportedFunction implements wasm.ModuleEngine.
-func (m *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
+func (m *moduleEngine) ResolveImportedFunction(index, descFunc, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
 	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)
+		m.ResolveImportedFunction(index, descFunc, imported.indexInModule, imported.me)
 		return // Recursively resolve the imported function.
 	}
 
 	offset := importedME.parent.functionOffsets[indexInImportedModule]
-	typeID := getTypeIDOf(indexInImportedModule, importedME.module)
+	typeID := m.module.TypeIDs[descFunc]
 	executable := &importedME.parent.executable[offset]
 	// Write functionInstance.
 	binary.LittleEndian.PutUint64(m.opaque[executableOffset:], uint64(uintptr(unsafe.Pointer(executable))))
@@ -261,28 +261,6 @@ func (m *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm
 	m.importedFunctions[index] = importedFunction{me: importedME, indexInModule: indexInImportedModule}
 }
 
-func getTypeIDOf(funcIndex wasm.Index, m *wasm.ModuleInstance) wasm.FunctionTypeID {
-	source := m.Source
-
-	var typeIndex wasm.Index
-	if funcIndex >= source.ImportFunctionCount {
-		funcIndex -= source.ImportFunctionCount
-		typeIndex = source.FunctionSection[funcIndex]
-	} else {
-		var cnt wasm.Index
-		for i := range source.ImportSection {
-			if source.ImportSection[i].Type == wasm.ExternTypeFunc {
-				if cnt == funcIndex {
-					typeIndex = source.ImportSection[i].DescFunc
-					break
-				}
-				cnt++
-			}
-		}
-	}
-	return m.TypeIDs[typeIndex]
-}
-
 // ResolveImportedMemory implements wasm.ModuleEngine.
 func (m *moduleEngine) ResolveImportedMemory(importedModuleEngine wasm.ModuleEngine) {
 	importedME := importedModuleEngine.(*moduleEngine)

internal/engine/wazevo/module_engine_test.go

@@ -152,28 +152,19 @@ func TestModuleEngine_setupOpaque(t *testing.T) {
 }
 
 func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
-	const begin = 5000
-	m := &moduleEngine{
-		opaque:            make([]byte, 10000),
-		importedFunctions: make([]importedFunction, 4),
-		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
-			ImportedFunctionsBegin: begin,
-		}},
-	}
-
 	var op1, op2 byte = 0xaa, 0xbb
-	im1 := &moduleEngine{
+	importing := &moduleEngine{
 		opaquePtr: &op1,
 		parent: &compiledModule{
 			executables:     &executables{executable: make([]byte, 1000)},
 			functionOffsets: []int{1, 5, 10},
 		},
 		module: &wasm.ModuleInstance{
-			TypeIDs: []wasm.FunctionTypeID{0, 0, 0, 0, 111, 222, 333},
+			TypeIDs: []wasm.FunctionTypeID{0, 0, 0, 888, 111, 222, 333},
 			Source:  &wasm.Module{FunctionSection: []wasm.Index{4, 5, 6}},
 		},
 	}
-	im2 := &moduleEngine{
+	imported := &moduleEngine{
 		opaquePtr: &op2,
 		parent: &compiledModule{
 			executables:     &executables{executable: make([]byte, 1000)},
@@ -185,21 +176,31 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 		},
 	}
 
-	m.ResolveImportedFunction(0, 0, im1)
-	m.ResolveImportedFunction(1, 0, im2)
-	m.ResolveImportedFunction(2, 2, im1)
-	m.ResolveImportedFunction(3, 1, im1)
+	const begin = 5000
+	m := &moduleEngine{
+		opaque:            make([]byte, 10000),
+		importedFunctions: make([]importedFunction, 4),
+		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
+			ImportedFunctionsBegin: begin,
+		}},
+		module: importing.module,
+	}
+
+	m.ResolveImportedFunction(0, 4, 0, importing)
+	m.ResolveImportedFunction(1, 3, 0, imported)
+	m.ResolveImportedFunction(2, 6, 2, importing)
+	m.ResolveImportedFunction(3, 5, 1, importing)
 
 	for i, tc := range []struct {
 		index      int
 		op         *byte
 		executable *byte
 		expTypeID  wasm.FunctionTypeID
 	}{
-		{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},
+		{index: 0, op: &op1, executable: &importing.parent.executable[1], expTypeID: 111},
+		{index: 1, op: &op2, executable: &imported.parent.executable[50], expTypeID: 888},
+		{index: 2, op: &op1, executable: &importing.parent.executable[10], expTypeID: 333},
+		{index: 3, op: &op1, executable: &importing.parent.executable[5], expTypeID: 222},
 	} {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
 			buf := m.opaque[begin+wazevoapi.FunctionInstanceSize*tc.index:]
@@ -216,15 +217,6 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 }
 
 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,
@@ -245,21 +237,31 @@ func TestModuleEngine_ResolveImportedFunction_recursive(t *testing.T) {
 		},
 		importedFunctions: []importedFunction{{me: imported, indexInModule: 0}},
 		module: &wasm.ModuleInstance{
-			TypeIDs: []wasm.FunctionTypeID{0, 222, 0},
+			TypeIDs: []wasm.FunctionTypeID{999, 222, 0},
 			Source:  &wasm.Module{FunctionSection: []wasm.Index{1}},
 		},
 	}
 
-	m.ResolveImportedFunction(0, 0, importing)
-	m.ResolveImportedFunction(1, 1, importing)
+	const begin = 5000
+	m := &moduleEngine{
+		opaque:            make([]byte, 10000),
+		importedFunctions: make([]importedFunction, 4),
+		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
+			ImportedFunctionsBegin: begin,
+		}},
+		module: importing.module,
+	}
+
+	m.ResolveImportedFunction(0, 0, 0, importing)
+	m.ResolveImportedFunction(1, 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: 0, op: &importedOp, executable: &imported.parent.executable[10], expTypeID: 999},
 		{index: 1, op: &importingOp, executable: &importing.parent.executable[500], expTypeID: 222},
 	} {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
@@ -312,26 +314,6 @@ func Test_functionInstance_offsets(t *testing.T) {
 	require.Equal(t, typeID, ptr+16)
 }
 
-func Test_getTypeIDOf(t *testing.T) {
-	m := &wasm.ModuleInstance{
-		TypeIDs: []wasm.FunctionTypeID{111, 222, 333, 444},
-		Source: &wasm.Module{
-			ImportFunctionCount: 1,
-			ImportSection: []wasm.Import{
-				{Type: wasm.ExternTypeMemory},
-				{Type: wasm.ExternTypeTable},
-				{Type: wasm.ExternTypeFunc, DescFunc: 3},
-			},
-			FunctionSection: []wasm.Index{2, 1, 0},
-		},
-	}
-
-	require.Equal(t, wasm.FunctionTypeID(444), getTypeIDOf(0, m))
-	require.Equal(t, wasm.FunctionTypeID(333), getTypeIDOf(1, m))
-	require.Equal(t, wasm.FunctionTypeID(222), getTypeIDOf(2, m))
-	require.Equal(t, wasm.FunctionTypeID(111), getTypeIDOf(3, m))
-}
-
 func Test_newAlignedOpaque(t *testing.T) {
 	for i := 0; i < 100; i++ {
 		s := 16 * (i + 10)

internal/wasm/engine.go

@@ -44,9 +44,10 @@ type ModuleEngine interface {
 
 	// ResolveImportedFunction is used to add imported functions needed to make this ModuleEngine fully functional.
 	// 	- `index` is the function Index of this imported function.
+	// 	- `descFunc` is the type Index in Module.TypeSection of this imported function. It corresponds to Import.DescFunc.
 	// 	- `indexInImportedModule` is the function Index of the imported function in the imported module.
 	//	- `importedModuleEngine` is the ModuleEngine for the imported ModuleInstance.
-	ResolveImportedFunction(index, indexInImportedModule Index, importedModuleEngine ModuleEngine)
+	ResolveImportedFunction(index, descFunc, indexInImportedModule Index, importedModuleEngine ModuleEngine)
 
 	// ResolveImportedMemory is called when this module imports a memory from another module.
 	ResolveImportedMemory(importedModuleEngine ModuleEngine)

internal/wasm/store.go

@@ -446,7 +446,7 @@ func (m *ModuleInstance) resolveImports(ctx context.Context, module *Module) (er
 					return
 				}
 
-				m.Engine.ResolveImportedFunction(i.IndexPerType, imported.Index, importedModule.Engine)
+				m.Engine.ResolveImportedFunction(i.IndexPerType, i.DescFunc, imported.Index, importedModule.Engine)
 			case ExternTypeTable:
 				expected := i.DescTable
 				importedTable := importedModule.Tables[imported.Index]

internal/wasm/store_test.go

@@ -341,7 +341,7 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 				importedModuleName: {{Type: ExternTypeFunc, Module: importedModuleName, Name: "fn", DescFunc: 0}},
 				"non-exist":        {{Name: "fn", DescFunc: 0}},
 			},
-		}, importingModuleName, nil, nil)
+		}, importingModuleName, nil, []FunctionTypeID{0})
 		require.EqualError(t, err, "module[non-exist] not instantiated")
 	})
 
@@ -485,7 +485,7 @@ func (e *mockModuleEngine) FunctionInstanceReference(i Index) Reference {
 }
 
 // ResolveImportedFunction implements the same method as documented on wasm.ModuleEngine.
-func (e *mockModuleEngine) ResolveImportedFunction(index, importedIndex Index, _ ModuleEngine) {
+func (e *mockModuleEngine) ResolveImportedFunction(index, _, importedIndex Index, _ ModuleEngine) {
 	e.resolveImportsCalled[index] = importedIndex
 }
 
@@ -742,7 +742,7 @@ func Test_resolveImports(t *testing.T) {
 				},
 			}
 
-			m := &ModuleInstance{Engine: &mockModuleEngine{resolveImportsCalled: map[Index]Index{}}, s: s, Source: module}
+			m := &ModuleInstance{Engine: &mockModuleEngine{resolveImportsCalled: map[Index]Index{}}, s: s, Source: module, TypeIDs: []FunctionTypeID{0, 1}}
 			err := m.resolveImports(context.Background(), module)
 			require.NoError(t, err)