Repository: tetratelabs/wazero
Pull request #1628 (2023-08-14 06:19:00)

wazevo: adds support for imported memory
#1496 


Patch:
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index 95b26885..77df6994 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -656,6 +656,9 @@ func (e *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm
 	e.functions[index] = imported.functions[indexInImportedModule]
 }
 
+// ResolveImportedMemory implements wasm.ModuleEngine.
+func (e *moduleEngine) ResolveImportedMemory(wasm.ModuleEngine) {}
+
 // FunctionInstanceReference implements the same method as documented on wasm.ModuleEngine.
 func (e *moduleEngine) FunctionInstanceReference(funcIndex wasm.Index) wasm.Reference {
 	return uintptr(unsafe.Pointer(&e.functions[funcIndex]))
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index 81cb6b95..f21dd192 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -442,6 +442,9 @@ func (e *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm
 	e.functions[index] = imported.functions[indexInImportedModule]
 }
 
+// ResolveImportedMemory implements wasm.ModuleEngine.
+func (e *moduleEngine) ResolveImportedMemory(wasm.ModuleEngine) {}
+
 // DoneInstantiation implements wasm.ModuleEngine.
 func (e *moduleEngine) DoneInstantiation() {}
 
diff --git a/internal/engine/wazevo/call_engine.go b/internal/engine/wazevo/call_engine.go
index 283d1cbc..700715e4 100644
--- a/internal/engine/wazevo/call_engine.go
+++ b/internal/engine/wazevo/call_engine.go
@@ -2,6 +2,7 @@ package wazevo
 
 import (
 	"context"
+	"encoding/binary"
 	"reflect"
 	"unsafe"
 
@@ -60,7 +61,7 @@ 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
-
+		// goFunctionCallStack is used to pass/receive parameters/results for Go function calls.
 		goFunctionCallStack [128]uint64
 	}
 )
@@ -132,13 +133,14 @@ func (c *callEngine) CallWithStack(ctx context.Context, paramResultStack []uint6
 				*argRes = uint64(0xffffffff) // = -1 in signed 32-bit integer.
 			} else {
 				*argRes = uint64(res)
-				if mod.Source.MemorySection != nil {
-					var opaque []byte
-					sh := (*reflect.SliceHeader)(unsafe.Pointer(&opaque))
-					sh.Data = uintptr(unsafe.Pointer(c.execCtx.callerModuleContextPtr))
-					sh.Len = 24
-					sh.Cap = 24
-					putLocalMemory(opaque, 8, mem)
+				calleeOpaque := opaqueViewFromPtr(uintptr(unsafe.Pointer(c.execCtx.callerModuleContextPtr)))
+				if mod.Source.MemorySection != nil { // Local memory.
+					putLocalMemory(calleeOpaque, 8 /* local memory begins at 8 */, mem)
+				} else {
+					// Imported memory's owner at offset 16 of the callerModuleContextPtr.
+					opaquePtr := uintptr(binary.LittleEndian.Uint64(calleeOpaque[16:]))
+					importedMemOwner := opaqueViewFromPtr(opaquePtr)
+					putLocalMemory(importedMemOwner, 8 /* local memory begins at 8 */, mem)
 				}
 			}
 			c.execCtx.exitCode = wazevoapi.ExitCodeOK
@@ -153,6 +155,15 @@ func moduleInstanceFromPtr(ptr *byte) *wasm.ModuleInstance {
 	return *(**wasm.ModuleInstance)(unsafe.Pointer(ptr))
 }
 
+func opaqueViewFromPtr(ptr uintptr) []byte {
+	var opaque []byte
+	sh := (*reflect.SliceHeader)(unsafe.Pointer(&opaque))
+	sh.Data = ptr
+	sh.Len = 24
+	sh.Cap = 24
+	return opaque
+}
+
 const callStackCeiling = uintptr(5000000) // in uint64 (8 bytes) == 40000000 bytes in total == 40mb.
 
 // growStack grows the stack, and returns the new stack pointer.
diff --git a/internal/engine/wazevo/frontend/frontend.go b/internal/engine/wazevo/frontend/frontend.go
index f59a7d95..1310e51b 100644
--- a/internal/engine/wazevo/frontend/frontend.go
+++ b/internal/engine/wazevo/frontend/frontend.go
@@ -76,7 +76,7 @@ func NewFrontendCompiler(m *wasm.Module, ssaBuilder ssa.Builder, offset *wazevoa
 	c.memoryGrowSig = ssa.Signature{
 		ID: ssa.SignatureID(len(m.TypeSection)),
 		// Takes execution context and the page size to grow.
-		Params: []ssa.Type{ssa.TypeI32, ssa.TypeI32},
+		Params: []ssa.Type{ssa.TypeI64, ssa.TypeI32},
 		// Returns the new size.
 		Results: []ssa.Type{ssa.TypeI32},
 	}
@@ -181,7 +181,7 @@ func (c *Compiler) declareWasmLocals(entry ssa.BasicBlock) {
 }
 
 func (c *Compiler) declareNecessaryVariables() {
-	c.needMemory = len(c.m.ImportedMemories()) > 0 || c.m.MemorySection != nil
+	c.needMemory = c.m.ImportMemoryCount > 0 || c.m.MemorySection != nil
 	if c.needMemory {
 		c.memoryBaseVariable = c.ssaBuilder.DeclareVariable(ssa.TypeI64)
 		c.memoryLenVariable = c.ssaBuilder.DeclareVariable(ssa.TypeI64)
diff --git a/internal/engine/wazevo/frontend/lower.go b/internal/engine/wazevo/frontend/lower.go
index 0e568949..277ae1d1 100644
--- a/internal/engine/wazevo/frontend/lower.go
+++ b/internal/engine/wazevo/frontend/lower.go
@@ -456,7 +456,16 @@ func (c *Compiler) lowerOpcode(op wasm.Opcode) {
 
 		var memSizeInBytes ssa.Value
 		if c.offset.LocalMemoryBegin < 0 {
-			panic("TODO: imported memory")
+			loadMemInstPtr := builder.AllocateInstruction()
+			loadMemInstPtr.AsLoad(c.moduleCtxPtrValue, c.offset.ImportedMemoryBegin.U32(), ssa.TypeI64)
+			builder.InsertInstruction(loadMemInstPtr)
+			memInstPtr := loadMemInstPtr.Return()
+
+			loadBufSizePtr := builder.AllocateInstruction()
+			loadBufSizePtr.AsLoad(memInstPtr, memoryInstanceBufSizeOffset, ssa.TypeI64)
+			builder.InsertInstruction(loadBufSizePtr)
+			memSizeInBytes = loadBufSizePtr.Return()
+
 		} else {
 			load := builder.AllocateInstruction()
 			load.AsLoad(c.moduleCtxPtrValue, c.offset.LocalMemoryLen().U32(), ssa.TypeI32)
@@ -845,7 +854,23 @@ func (c *Compiler) lowerOpcode(op wasm.Opcode) {
 		// TODO: maybe this can be optimized out if this is in-module function calls. Investigate later.
 		c.storeCallerModuleContext()
 
-		typIndex := c.m.FunctionSection[fnIndex]
+		var typIndex wasm.Index
+		if fnIndex < c.m.ImportFunctionCount {
+			var fi int
+			for i := range c.m.ImportSection {
+				imp := &c.m.ImportSection[i]
+				if imp.Type == wasm.ExternTypeFunc {
+					if fi == int(fnIndex) {
+						typIndex = imp.DescFunc
+						break
+					}
+					fi++
+				}
+			}
+		} else {
+			fnIndex -= c.m.ImportFunctionCount
+			typIndex = c.m.FunctionSection[fnIndex]
+		}
 		typ := &c.m.TypeSection[typIndex]
 
 		// TODO: reuse slice?
@@ -957,32 +982,43 @@ func (c *Compiler) getWasmGlobalValue(index wasm.Index, forceLoad bool) ssa.Valu
 	return ret
 }
 
-func (c *Compiler) getMemoryBaseValue(forceReload bool) ssa.Value {
-	if c.offset.LocalMemoryBegin < 0 {
-		panic("TODO: imported memory")
-	}
+const (
+	memoryInstanceBufOffset     = 0
+	memoryInstanceBufSizeOffset = memoryInstanceBufOffset + 8
+)
 
-	variable := c.memoryBaseVariable
+func (c *Compiler) getMemoryBaseValue(forceReload bool) ssa.Value {
 	builder := c.ssaBuilder
+	variable := c.memoryBaseVariable
 	if !forceReload {
 		if v := builder.FindValue(variable); v.Valid() {
 			return v
 		}
 	}
 
-	load := builder.AllocateInstruction()
-	load.AsLoad(c.moduleCtxPtrValue, c.offset.LocalMemoryBase().U32(), ssa.TypeI64)
-	builder.InsertInstruction(load)
-	ret := load.Return()
+	var ret ssa.Value
+	if c.offset.LocalMemoryBegin < 0 {
+		loadMemInstPtr := builder.AllocateInstruction()
+		loadMemInstPtr.AsLoad(c.moduleCtxPtrValue, c.offset.ImportedMemoryBegin.U32(), ssa.TypeI64)
+		builder.InsertInstruction(loadMemInstPtr)
+		memInstPtr := loadMemInstPtr.Return()
+
+		loadBufPtr := builder.AllocateInstruction()
+		loadBufPtr.AsLoad(memInstPtr, memoryInstanceBufOffset, ssa.TypeI64)
+		builder.InsertInstruction(loadBufPtr)
+		ret = loadBufPtr.Return()
+	} else {
+		load := builder.AllocateInstruction()
+		load.AsLoad(c.moduleCtxPtrValue, c.offset.LocalMemoryBase().U32(), ssa.TypeI64)
+		builder.InsertInstruction(load)
+		ret = load.Return()
+	}
+
 	builder.DefineVariableInCurrentBB(variable, ret)
 	return ret
 }
 
 func (c *Compiler) getMemoryLenValue(forceReload bool) ssa.Value {
-	if c.offset.LocalMemoryBegin < 0 {
-		panic("TODO: imported memory")
-	}
-
 	variable := c.memoryLenVariable
 	builder := c.ssaBuilder
 	if !forceReload {
@@ -990,10 +1026,26 @@ func (c *Compiler) getMemoryLenValue(forceReload bool) ssa.Value {
 			return v
 		}
 	}
-	load := builder.AllocateInstruction()
-	load.AsExtLoad(ssa.OpcodeUload32, c.moduleCtxPtrValue, c.offset.LocalMemoryLen().U32(), true)
-	builder.InsertInstruction(load)
-	ret := load.Return()
+
+	var ret ssa.Value
+	if c.offset.LocalMemoryBegin < 0 {
+		loadMemInstPtr := builder.AllocateInstruction()
+		loadMemInstPtr.AsLoad(c.moduleCtxPtrValue, c.offset.ImportedMemoryBegin.U32(), ssa.TypeI64)
+		builder.InsertInstruction(loadMemInstPtr)
+		memInstPtr := loadMemInstPtr.Return()
+
+		loadBufSizePtr := builder.AllocateInstruction()
+		loadBufSizePtr.AsLoad(memInstPtr, memoryInstanceBufSizeOffset, ssa.TypeI64)
+		builder.InsertInstruction(loadBufSizePtr)
+
+		ret = loadBufSizePtr.Return()
+	} else {
+		load := builder.AllocateInstruction()
+		load.AsExtLoad(ssa.OpcodeUload32, c.moduleCtxPtrValue, c.offset.LocalMemoryLen().U32(), true)
+		builder.InsertInstruction(load)
+		ret = load.Return()
+	}
+
 	builder.DefineVariableInCurrentBB(variable, ret)
 	return ret
 }
diff --git a/internal/engine/wazevo/module_engine.go b/internal/engine/wazevo/module_engine.go
index 7e73997a..d5a1035c 100644
--- a/internal/engine/wazevo/module_engine.go
+++ b/internal/engine/wazevo/module_engine.go
@@ -61,10 +61,7 @@ func (m *moduleEngine) setupOpaque() {
 		putLocalMemory(opaque, lm, inst.MemoryInstance)
 	}
 
-	if im := offsets.ImportedMemoryBegin; im >= 0 {
-		b := uint64(uintptr(unsafe.Pointer(&inst.MemoryInstance)))
-		binary.LittleEndian.PutUint64(opaque[im:], b)
-	}
+	// Note: imported memory is resolved in ResolveImportedFunction.
 
 	// Note: imported functions are resolved in ResolveImportedFunction.
 
@@ -119,6 +116,17 @@ func (m *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm
 	binary.LittleEndian.PutUint64(m.opaque[moduleCtx:], uint64(uintptr(unsafe.Pointer(importedME.opaquePtr))))
 }
 
+// ResolveImportedMemory implements wasm.ModuleEngine.
+func (m *moduleEngine) ResolveImportedMemory(importedModuleEngine wasm.ModuleEngine) {
+	importedME := importedModuleEngine.(*moduleEngine)
+	inst := importedME.module
+
+	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))))
+}
+
 // DoneInstantiation implements wasm.ModuleEngine.
 func (m *moduleEngine) DoneInstantiation() {
 	m.setupOpaque()
diff --git a/internal/engine/wazevo/wazevoapi/offsetdata.go b/internal/engine/wazevo/wazevoapi/offsetdata.go
index 9941198c..c1d44b0a 100644
--- a/internal/engine/wazevo/wazevoapi/offsetdata.go
+++ b/internal/engine/wazevo/wazevoapi/offsetdata.go
@@ -51,7 +51,9 @@ type ExecutionContextOffsetData struct {
 type ModuleContextOffsetData struct {
 	TotalSize int
 	ModuleInstanceOffset,
-	LocalMemoryBegin, ImportedMemoryBegin, ImportedFunctionsBegin,
+	LocalMemoryBegin,
+	ImportedMemoryBegin,
+	ImportedFunctionsBegin,
 	GlobalsBegin Offset
 }
 
@@ -109,18 +111,18 @@ func NewModuleContextOffsetData(m *wasm.Module) ModuleContextOffsetData {
 	if m.MemorySection != nil {
 		ret.LocalMemoryBegin = offset
 		// buffer base + memory size.
-		const localMemorySizeInOpaqueVMContext = 16
-		offset += localMemorySizeInOpaqueVMContext
+		const localMemorySizeInOpaqueModuleContext = 16
+		offset += localMemorySizeInOpaqueModuleContext
 	} else {
 		// Indicates that there's no local memory
 		ret.LocalMemoryBegin = -1
 	}
 
 	if m.ImportMemoryCount > 0 {
-		// *wasm.MemoryInstance
-		const importedMemorySizeInOpaqueVMCContext = 8
+		// *wasm.MemoryInstance + imported memory's owner (moduleContextOpaque)
+		const importedMemorySizeInOpaqueModuleContext = 16
 		ret.ImportedMemoryBegin = offset
-		offset += importedMemorySizeInOpaqueVMCContext
+		offset += importedMemorySizeInOpaqueModuleContext
 	} else {
 		// Indicates that there's no imported memory
 		ret.ImportedMemoryBegin = -1
diff --git a/internal/wasm/engine.go b/internal/wasm/engine.go
index c0516308..26d196b0 100644
--- a/internal/wasm/engine.go
+++ b/internal/wasm/engine.go
@@ -48,6 +48,9 @@ type ModuleEngine interface {
 	//	- `importedModuleEngine` is the ModuleEngine for the imported ModuleInstance.
 	ResolveImportedFunction(index, indexInImportedModule Index, importedModuleEngine ModuleEngine)
 
+	// ResolveImportedMemory is called when this module imports a memory from another module.
+	ResolveImportedMemory(importedModuleEngine ModuleEngine)
+
 	// LookupFunction returns the api.Function created from the function in the function table at the given offset.
 	LookupFunction(t *TableInstance, typeId FunctionTypeID, tableOffset Index) (api.Function, error)
 
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index b27466cd..d056133b 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -451,6 +451,7 @@ func (m *ModuleInstance) resolveImports(module *Module) (err error) {
 					return
 				}
 				m.MemoryInstance = importedMemory
+				m.Engine.ResolveImportedMemory(importedModule.Engine)
 			case ExternTypeGlobal:
 				expected := i.DescGlobal
 				importedGlobal := importedModule.Globals[imported.Index]


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 8cc28aa3..63b8676d 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -139,6 +139,12 @@ func TestE2E(t *testing.T) {
 			m:     testcases.MemorySizeGrow.Module,
 			calls: []callCase{{expResults: []uint64{1, 2, 0xffffffff}}},
 		},
+		{
+			name:     "imported_memory_grow",
+			imported: testcases.ImportedMemoryGrow.Imported,
+			m:        testcases.ImportedMemoryGrow.Module,
+			calls:    []callCase{{expResults: []uint64{1, 1, 11, 11}}},
+		},
 	} {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index d83cefa8..ba5f17c8 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -1249,6 +1249,78 @@ blk0: (exec_ctx:i64, module_ctx:i64)
 	v16:i64 = Load module_ctx, 0x20
 	v17:f64 = Load v16, 0x8
 	Jump blk_ret, v3, v5, v7, v9, v11, v13, v15, v17
+`,
+		},
+		{
+			name: "imported_memory_grow",
+			m:    testcases.ImportedMemoryGrow.Module,
+			exp: `
+signatures:
+	sig0: i64i64_i32
+	sig2: i64i32_i32
+
+blk0: (exec_ctx:i64, module_ctx:i64)
+	Store module_ctx, exec_ctx, 0x8
+	v2:i64 = Load module_ctx, 0x18
+	v3:i64 = Load module_ctx, 0x20
+	v4:i32 = CallIndirect v2:sig0, exec_ctx, v3
+	v5:i64 = Load module_ctx, 0x8
+	v6:i64 = Load v5, 0x0
+	v7:i64 = Load module_ctx, 0x8
+	v8:i64 = Load v7, 0x8
+	v9:i64 = Load module_ctx, 0x8
+	v10:i64 = Load v9, 0x8
+	v11:i32 = Iconst_32 0x10
+	v12:i64 = Ushr v10, v11
+	v13:i32 = Iconst_32 0xa
+	Store module_ctx, exec_ctx, 0x8
+	v14:i64 = Load exec_ctx, 0x48
+	v15:i32 = CallIndirect v14:sig2, exec_ctx, v13
+	v16:i64 = Load module_ctx, 0x8
+	v17:i64 = Load v16, 0x0
+	v18:i64 = Load module_ctx, 0x8
+	v19:i64 = Load v18, 0x8
+	Store module_ctx, exec_ctx, 0x8
+	v20:i64 = Load module_ctx, 0x18
+	v21:i64 = Load module_ctx, 0x20
+	v22:i32 = CallIndirect v20:sig0, exec_ctx, v21
+	v23:i64 = Load module_ctx, 0x8
+	v24:i64 = Load v23, 0x0
+	v25:i64 = Load module_ctx, 0x8
+	v26:i64 = Load v25, 0x8
+	v27:i64 = Load module_ctx, 0x8
+	v28:i64 = Load v27, 0x8
+	v29:i32 = Iconst_32 0x10
+	v30:i64 = Ushr v28, v29
+	Jump blk_ret, v4, v12, v22, v30
+`,
+			expAfterOpt: `
+signatures:
+	sig0: i64i64_i32
+	sig2: i64i32_i32
+
+blk0: (exec_ctx:i64, module_ctx:i64)
+	Store module_ctx, exec_ctx, 0x8
+	v2:i64 = Load module_ctx, 0x18
+	v3:i64 = Load module_ctx, 0x20
+	v4:i32 = CallIndirect v2:sig0, exec_ctx, v3
+	v9:i64 = Load module_ctx, 0x8
+	v10:i64 = Load v9, 0x8
+	v11:i32 = Iconst_32 0x10
+	v12:i64 = Ushr v10, v11
+	v13:i32 = Iconst_32 0xa
+	Store module_ctx, exec_ctx, 0x8
+	v14:i64 = Load exec_ctx, 0x48
+	v15:i32 = CallIndirect v14:sig2, exec_ctx, v13
+	Store module_ctx, exec_ctx, 0x8
+	v20:i64 = Load module_ctx, 0x18
+	v21:i64 = Load module_ctx, 0x20
+	v22:i32 = CallIndirect v20:sig0, exec_ctx, v21
+	v27:i64 = Load module_ctx, 0x8
+	v28:i64 = Load v27, 0x8
+	v29:i32 = Iconst_32 0x10
+	v30:i64 = Ushr v28, v29
+	Jump blk_ret, v4, v12, v22, v30
 `,
 		},
 		{
@@ -1256,7 +1328,7 @@ blk0: (exec_ctx:i64, module_ctx:i64)
 			m:    testcases.MemorySizeGrow.Module,
 			exp: `
 signatures:
-	sig1: i32i32_i32
+	sig1: i64i32_i32
 
 blk0: (exec_ctx:i64, module_ctx:i64)
 	v2:i32 = Iconst_32 0x1
@@ -1278,7 +1350,7 @@ blk0: (exec_ctx:i64, module_ctx:i64)
 `,
 			expAfterOpt: `
 signatures:
-	sig1: i32i32_i32
+	sig1: i64i32_i32
 
 blk0: (exec_ctx:i64, module_ctx:i64)
 	v2:i32 = Iconst_32 0x1
diff --git a/internal/engine/wazevo/frontend/lower_test.go b/internal/engine/wazevo/frontend/lower_test.go
index 9d4983aa..8fa928b1 100644
--- a/internal/engine/wazevo/frontend/lower_test.go
+++ b/internal/engine/wazevo/frontend/lower_test.go
@@ -11,4 +11,6 @@ import (
 func TestGlobalInstanceValueOffset(t *testing.T) {
 	var globalInstance wasm.GlobalInstance
 	require.Equal(t, int(unsafe.Offsetof(globalInstance.Val)), globalInstanceValueOffset)
+	var memInstance wasm.MemoryInstance
+	require.Equal(t, int(unsafe.Offsetof(memInstance.Buffer)), memoryInstanceBufOffset)
 }
diff --git a/internal/engine/wazevo/module_engine_test.go b/internal/engine/wazevo/module_engine_test.go
index 46ac3f43..88c4cf24 100644
--- a/internal/engine/wazevo/module_engine_test.go
+++ b/internal/engine/wazevo/module_engine_test.go
@@ -2,6 +2,7 @@ package wazevo
 
 import (
 	"encoding/binary"
+	"runtime"
 	"strconv"
 	"testing"
 	"unsafe"
@@ -68,9 +69,17 @@ 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.opaquePtr = &imported.opaque[0]
+				m.ResolveImportedMemory(imported)
+
 				actualPtr := uintptr(binary.LittleEndian.Uint64(m.opaque[tc.offset.ImportedMemoryBegin:]))
-				expPtr := uintptr(unsafe.Pointer(&tc.m.MemoryInstance))
+				expPtr := uintptr(unsafe.Pointer(tc.m.MemoryInstance))
 				require.Equal(t, expPtr, actualPtr)
+
+				actualOpaquePtr := uintptr(binary.LittleEndian.Uint64(m.opaque[tc.offset.ImportedMemoryBegin+8:]))
+				require.Equal(t, uintptr(unsafe.Pointer(imported.opaquePtr)), actualOpaquePtr)
+				runtime.KeepAlive(imported)
 			}
 			if tc.offset.GlobalsBegin >= 0 {
 				for i, g := range tc.m.Globals {
diff --git a/internal/engine/wazevo/testcases/testcases.go b/internal/engine/wazevo/testcases/testcases.go
index 30a929ed..cba795e4 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -845,6 +845,45 @@ var (
 		},
 	}
 
+	ImportedMemoryGrow = TestCase{
+		Name: "imported_memory_grow",
+		Imported: &wasm.Module{
+			ExportSection: []wasm.Export{
+				{Name: "mem", Type: wasm.ExternTypeMemory, Index: 0},
+				{Name: "size", Type: wasm.ExternTypeFunc, Index: 0},
+			},
+			MemorySection:   &wasm.Memory{Min: 1},
+			TypeSection:     []wasm.FunctionType{v_i32},
+			FunctionSection: []wasm.Index{0},
+			CodeSection:     []wasm.Code{{Body: []byte{wasm.OpcodeMemorySize, 0, wasm.OpcodeEnd}}},
+			DataSection:     []wasm.DataSegment{{OffsetExpression: constExprI32(0), Init: maskedBuf(int(wasm.MemoryPageSize))}},
+			NameSection:     &wasm.NameSection{ModuleName: "env"},
+		},
+		Module: &wasm.Module{
+			ImportMemoryCount:   1,
+			ImportFunctionCount: 1,
+			ImportSection: []wasm.Import{
+				{Module: "env", Name: "mem", Type: wasm.ExternTypeMemory, DescMem: &wasm.Memory{Min: 1}},
+				{Module: "env", Name: "size", Type: wasm.ExternTypeFunc, DescFunc: 0},
+			},
+			TypeSection:     []wasm.FunctionType{v_i32, {Results: []wasm.ValueType{i32, i32, i32, i32}}},
+			ExportSection:   []wasm.Export{{Name: ExportName, Type: wasm.ExternTypeFunc, Index: 1}},
+			FunctionSection: []wasm.Index{1},
+			CodeSection: []wasm.Code{
+				{Body: []byte{
+					wasm.OpcodeCall, 0, // Call imported size function. --> 1
+					wasm.OpcodeMemorySize, 0, // --> 1.
+					wasm.OpcodeI32Const, 10,
+					wasm.OpcodeMemoryGrow, 0,
+					wasm.OpcodeDrop,
+					wasm.OpcodeCall, 0, // Call imported size function. --> 11.
+					wasm.OpcodeMemorySize, 0, // --> 11.
+					wasm.OpcodeEnd,
+				}},
+			},
+		},
+	}
+
 	GlobalsGet = TestCase{
 		Name: "globals_get",
 		Module: &wasm.Module{
diff --git a/internal/engine/wazevo/wazevoapi/offsetdata_test.go b/internal/engine/wazevo/wazevoapi/offsetdata_test.go
index a0245f99..ef2e2f74 100644
--- a/internal/engine/wazevo/wazevoapi/offsetdata_test.go
+++ b/internal/engine/wazevo/wazevoapi/offsetdata_test.go
@@ -43,7 +43,7 @@ func TestNewModuleContextOffsetData(t *testing.T) {
 				ImportedMemoryBegin:    8,
 				ImportedFunctionsBegin: -1,
 				GlobalsBegin:           -1,
-				TotalSize:              16,
+				TotalSize:              24,
 			},
 		},
 		{
@@ -63,9 +63,9 @@ func TestNewModuleContextOffsetData(t *testing.T) {
 			exp: ModuleContextOffsetData{
 				LocalMemoryBegin:       -1,
 				ImportedMemoryBegin:    8,
-				ImportedFunctionsBegin: 16,
+				ImportedFunctionsBegin: 24,
 				GlobalsBegin:           -1,
-				TotalSize:              176,
+				TotalSize:              184,
 			},
 		},
 		{
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 60ac45a5..195095fb 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -412,6 +412,7 @@ type mockModuleEngine struct {
 	callFailIndex        int
 	functionRefs         map[Index]Reference
 	resolveImportsCalled map[Index]Index
+	importedMemModEngine ModuleEngine
 }
 
 type mockCallEngine struct {
@@ -466,6 +467,11 @@ func (e *mockModuleEngine) ResolveImportedFunction(index, importedIndex Index, _
 	e.resolveImportsCalled[index] = importedIndex
 }
 
+// ResolveImportedMemory implements the same method as documented on wasm.ModuleEngine.
+func (e *mockModuleEngine) ResolveImportedMemory(imp ModuleEngine) {
+	e.importedMemModEngine = imp
+}
+
 // NewFunction implements the same method as documented on wasm.ModuleEngine.
 func (e *mockModuleEngine) NewFunction(index Index) api.Function {
 	return &mockCallEngine{index: index, callFailIndex: e.callFailIndex}
@@ -808,14 +814,16 @@ func Test_resolveImports(t *testing.T) {
 			max := uint32(10)
 			memoryInst := &MemoryInstance{Max: max}
 			s := newStore()
+			importedME := &mockModuleEngine{}
 			s.nameToModule[moduleName] = &ModuleInstance{
 				MemoryInstance: memoryInst,
 				Exports: map[string]*Export{name: {
 					Type: ExternTypeMemory,
 				}},
 				ModuleName: moduleName,
+				Engine:     importedME,
 			}
-			m := &ModuleInstance{s: s}
+			m := &ModuleInstance{s: s, Engine: &mockModuleEngine{resolveImportsCalled: map[Index]Index{}}}
 			err := m.resolveImports(&Module{
 				ImportPerModule: map[string][]*Import{
 					moduleName: {{Module: moduleName, Name: name, Type: ExternTypeMemory, DescMem: &Memory{Max: max}}},
@@ -823,6 +831,7 @@ func Test_resolveImports(t *testing.T) {
 			})
 			require.NoError(t, err)
 			require.Equal(t, m.MemoryInstance, memoryInst)
+			require.Equal(t, importedME, m.Engine.(*mockModuleEngine).importedMemModEngine)
 		})
 		t.Run("minimum size mismatch", func(t *testing.T) {
 			importMemoryType := &Memory{Min: 2, Cap: 2}