Skip to content

Commit

Permalink
compiler: ResolveImportedFunction should infer the typeID from the im…
Browse files Browse the repository at this point in the history
…porting module (#2314)

Signed-off-by: Edoardo Vacchi <[email protected]>
  • Loading branch information
evacchi authored Sep 16, 2024
1 parent 5848888 commit b468ada
Show file tree
Hide file tree
Showing 6 changed files with 45 additions and 84 deletions.
2 changes: 1 addition & 1 deletion internal/engine/interpreter/interpreter.go
Original file line number Diff line number Diff line change
Expand Up @@ -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]
}
Expand Down
28 changes: 3 additions & 25 deletions internal/engine/wazevo/module_engine.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))))
Expand All @@ -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)
Expand Down
88 changes: 35 additions & 53 deletions internal/engine/wazevo/module_engine_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)},
Expand All @@ -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:]
Expand All @@ -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,
Expand All @@ -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) {
Expand Down Expand Up @@ -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)
Expand Down
3 changes: 2 additions & 1 deletion internal/wasm/engine.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion internal/wasm/store.go
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
6 changes: 3 additions & 3 deletions internal/wasm/store_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
})

Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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)

Expand Down

0 comments on commit b468ada

Please sign in to comment.