wazevo: adds support for imported-exported memory (#1643)

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

internal/engine/wazevo/backend/isa/arm64/lower_instr.go

@@ -7,8 +7,6 @@ package arm64
 // and merge the multiple instructions if possible. It can be considered as "N:1" instruction selection.
 
 import (
-	"fmt"
-
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/wazevoapi"
@@ -379,7 +377,6 @@ func (m *machine) lowerFcmpToFlag(x, y ssa.Value) {
 	rm := m.getOperand_NR(m.compiler.ValueDefinition(y), extModeNone)
 	cmp := m.allocateInstr()
 	cmp.asFpuCmp(rn, rm, x.Type().Bits() == 64)
-	fmt.Println(x.Type())
 	m.insert(cmp)
 }
 

internal/engine/wazevo/e2e_test.go

@@ -427,3 +427,54 @@ func TestE2E_stores(t *testing.T) {
 		})
 	}
 }
+
+func TestE2E_reexported_memory(t *testing.T) {
+	m1 := &wasm.Module{
+		ExportSection: []wasm.Export{{Name: "mem", Type: wasm.ExternTypeMemory, Index: 0}},
+		MemorySection: &wasm.Memory{Min: 1},
+		NameSection:   &wasm.NameSection{ModuleName: "m1"},
+	}
+	m2 := &wasm.Module{
+		ImportMemoryCount: 1,
+		ExportSection:     []wasm.Export{{Name: "mem2", Type: wasm.ExternTypeMemory, Index: 0}},
+		ImportSection:     []wasm.Import{{Module: "m1", Name: "mem", Type: wasm.ExternTypeMemory, DescMem: &wasm.Memory{Min: 1}}},
+		NameSection:       &wasm.NameSection{ModuleName: "m2"},
+	}
+	m3 := &wasm.Module{
+		ImportMemoryCount: 1,
+		ImportSection:     []wasm.Import{{Module: "m2", Name: "mem2", Type: wasm.ExternTypeMemory, DescMem: &wasm.Memory{Min: 1}}},
+		TypeSection:       []wasm.FunctionType{{Results: []wasm.ValueType{i32}}},
+		ExportSection:     []wasm.Export{{Name: testcases.ExportedFunctionName, Type: wasm.ExternTypeFunc, Index: 0}},
+		FunctionSection:   []wasm.Index{0},
+		CodeSection:       []wasm.Code{{Body: []byte{wasm.OpcodeI32Const, 10, wasm.OpcodeMemoryGrow, 0, wasm.OpcodeEnd}}},
+	}
+
+	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))
+	}()
+
+	m1Inst, err := r.Instantiate(ctx, binaryencoding.EncodeModule(m1))
+	require.NoError(t, err)
+
+	m2Inst, err := r.Instantiate(ctx, binaryencoding.EncodeModule(m2))
+	require.NoError(t, err)
+
+	m3Inst, err := r.Instantiate(ctx, binaryencoding.EncodeModule(m3))
+	require.NoError(t, err)
+
+	f := m3Inst.ExportedFunction(testcases.ExportedFunctionName)
+	result, err := f.Call(ctx)
+	require.NoError(t, err)
+	require.Equal(t, uint64(1), result[0])
+	mem := m1Inst.Memory()
+	require.Equal(t, mem, m3Inst.Memory())
+	require.Equal(t, mem, m2Inst.Memory())
+	require.Equal(t, uint32(11), mem.Size()/65536)
+}

internal/engine/wazevo/module_engine.go

@@ -177,15 +177,19 @@ 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")
+	var memInstPtr uint64
+	var memOwnerOpaquePtr uint64
+	if offs := importedME.parent.offsets; offs.ImportedMemoryBegin >= 0 {
+		offset := offs.ImportedMemoryBegin
+		memInstPtr = binary.LittleEndian.Uint64(importedME.opaque[offset:])
+		memOwnerOpaquePtr = binary.LittleEndian.Uint64(importedME.opaque[offset+8:])
+	} else {
+		memInstPtr = uint64(uintptr(unsafe.Pointer(inst.MemoryInstance)))
+		memOwnerOpaquePtr = uint64(uintptr(unsafe.Pointer(importedME.opaquePtr)))
 	}
-
 	offset := m.parent.offsets.ImportedMemoryBegin
-	b := uint64(uintptr(unsafe.Pointer(inst.MemoryInstance)))
-	binary.LittleEndian.PutUint64(m.opaque[offset:], b)
-	binary.LittleEndian.PutUint64(m.opaque[offset+8:], uint64(uintptr(unsafe.Pointer(importedME.opaquePtr))))
+	binary.LittleEndian.PutUint64(m.opaque[offset:], memInstPtr)
+	binary.LittleEndian.PutUint64(m.opaque[offset+8:], memOwnerOpaquePtr)
 }
 
 // DoneInstantiation implements wasm.ModuleEngine.

internal/engine/wazevo/module_engine_test.go

@@ -175,6 +175,28 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 	}
 }
 
+func TestModuleEngine_ResolveImportedMemory_reexported(t *testing.T) {
+	m := &moduleEngine{
+		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
+			ImportedMemoryBegin: 50,
+		}},
+		opaque: make([]byte, 100),
+	}
+
+	importedME := &moduleEngine{
+		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
+			ImportedMemoryBegin: 1000,
+		}},
+		opaque: make([]byte, 2000),
+	}
+	binary.LittleEndian.PutUint64(importedME.opaque[1000:], 0x1234567890abcdef)
+	binary.LittleEndian.PutUint64(importedME.opaque[1000+8:], 0xabcdef1234567890)
+
+	m.ResolveImportedMemory(importedME)
+	require.Equal(t, uint64(0x1234567890abcdef), binary.LittleEndian.Uint64(m.opaque[50:]))
+	require.Equal(t, uint64(0xabcdef1234567890), binary.LittleEndian.Uint64(m.opaque[50+8:]))
+}
+
 func Test_functionInstance_offsets(t *testing.T) {
 	var fi functionInstance
 	require.Equal(t, wazevoapi.FunctionInstanceSize, int(unsafe.Sizeof(fi)))