Repository: tetratelabs/wazero
Pull request #275 (2022-02-22 05:51:50)

Adds ModuleExports.Memory and renames Memory.Len to Memory.Size
This adds a function to get an exported memory by name, allowing end
users to check ahead of time if memory writes might fail. This also
renames memory.len to size to as that's what the spec calls it.


codefromthecrypt: @mathetake  I could use a hand debugging these jit failures.. I triple checked jit.TestVerifyOffsetValue is checking every explicit offset field.

```
--- FAIL: TestJIT (1.37s)
    --- FAIL: TestJIT/linking.wast (0.01s)
        --- FAIL: TestJIT/linking.wast/assert_return/line:387 (0.00s)
            spec_test.go:436: 
                	Error Trace:	spec_test.go:436
                	            				spec_test.go:314
                	Error:      	Not equal: 
                	            	expected: 0x68
                	            	actual  : 0x0
                	Test:       	TestJIT/linking.wast/assert_return/line:387
                	Messages:   	linking.wast:387 assert_return invoke get memory[0] ([]) in module $Ms
        --- FAIL: TestJIT/linking.wast/assert_return/line:388 (0.00s)
            spec_test.go:310: 
                	Error Trace:	spec_test.go:310
                	Error:      	Received unexpected error:
                	            	wasm runtime error: invalid table access
                	            	wasm backtrace:
                	            		0: unknown
                	Test:       	TestJIT/linking.wast/assert_return/line:388
                	Messages:   	linking.wast:388 assert_return invoke get table[0] ([]) in module $Ms
--- FAIL: TestInterpreter (0.91s)
    --- FAIL: TestInterpreter/linking.wast (0.00s)
        --- FAIL: TestInterpreter/linking.wast/assert_return/line:387 (0.00s)
            spec_test.go:436: 
                	Error Trace:	spec_test.go:436
                	            				spec_test.go:314
                	Error:      	Not equal: 
                	            	expected: 0x68
                	            	actual  : 0x0
                	Test:       	TestInterpreter/linking.wast/assert_return/line:387
                	Messages:   	linking.wast:387 assert_return invoke get memory[0] ([]) in module $Ms
        --- FAIL: TestInterpreter/linking.wast/assert_return/line:388 (0.00s)
            spec_test.go:310: 
                	Error Trace:	spec_test.go:310
                	Error:      	Received unexpected error:
                	            	wasm runtime error: invalid table access
                	            	wasm backtrace:
                	            		0: unknown
                	Test:       	TestInterpreter/linking.wast/assert_return/line:388
                	Messages:   	linking.wast:388 assert_return invoke get table[0] ([]) in module $Ms
FAIL
FAIL	github.com/tetratelabs/wazero/tests/spectest	3.349s

```
codefromthecrypt: So the linking.wast module looks like this:
```
(module $Ms
  (type $t (func (result i32)))
  (memory (export "memory") 1)
  (table (export "table") 1 funcref)
  (func (export "get memory[0]") (type $t)
    (i32.load8_u (i32.const 0))
  )
  (func (export "get table[0]") (type $t)
    (call_indirect (type $t) (i32.const 0))
  )
)
```

This failed because the value returned was zero
```
(assert_return (invoke $Ms "get memory[0]") (i32.const 104))  ;; 'h'
```

This failed with "invalid table access" which is the error that suggests an offset problem
```
(assert_return (invoke $Ms "get table[0]") (i32.const 0xdead))
```
mathetake: i see the failure is not jit specific so I hope you can get your hand dirty :D
codefromthecrypt: >> i see the failure is not jit specific so I hope you can get your hand dirty :D

you are right.. I missed that this also fails in the interpreter. will look closer..
codefromthecrypt: still not seeing what is broken.. if you can't I guess will have to start trying to recreate what the linking.wast tests it trying to do in a normal unit test (because the lifecycle is weird in the spec tests, it is harder to see what's going on)
mathetake: you can do it :D
codefromthecrypt: thanks for the look. I'll try to isolate the two tests to find the bug, then backfill a test for it in the main unit test tree
codefromthecrypt: finally :D


Patch:
diff --git a/internal/wasi/wasi.go b/internal/wasi/wasi.go
index e47359ce..2abf6236 100644
--- a/internal/wasi/wasi.go
+++ b/internal/wasi/wasi.go
@@ -187,7 +187,7 @@ type SnapshotPreview1 interface {
 	//
 	// * argv - is the offset to begin writing argument offsets in uint32 little-endian encoding.
 	//   * ArgsSizesGet result argc * 4 bytes are written to this offset
-	// * argvBuf - is the offset to write the null terminated arguments to wasm.MemoryInstance
+	// * argvBuf - is the offset to write the null terminated arguments to ctx.Memory
 	//   * ArgsSizesGet result argv_buf_size bytes are written to this offset
 	//
 	// For example, if ArgsSizesGet wrote argc=2 and argvBufSize=5 for arguments: "a" and "bc"
@@ -215,8 +215,8 @@ type SnapshotPreview1 interface {
 	// corresponding sizes in uint32 little-endian encoding. If either are invalid due to memory constraints, this
 	// returns ErrnoFault.
 	//
-	// * resultArgc - is the offset to write the argument count to wasm.Memory
-	// * resultArgvBufSize - is the offset to write the null-terminated argument length to wasm.Memory
+	// * resultArgc - is the offset to write the argument count to ctx.Memory
+	// * resultArgvBufSize - is the offset to write the null-terminated argument length to ctx.Memory
 	//
 	// For example, if Args are []string{"a","bc"} and
 	//   ArgsSizesGet parameters resultArgc=1 and resultArgvBufSize=6, we expect `ctx.Memory` to contain:
@@ -243,7 +243,7 @@ type SnapshotPreview1 interface {
 	//
 	// * environ - is the offset to begin writing environment variables offsets in uint32 little-endian encoding.
 	//   * EnvironSizesGet result environc * 4 bytes are written to this offset
-	// * environBuf - is the offset to write the environment variables to wasm.Memory
+	// * environBuf - is the offset to write the environment variables to ctx.Memory
 	//   * the format is the same as os.Environ, null terminated "key=val" entries
 	//   * EnvironSizesGet result environBufSize bytes are written to this offset
 	//
@@ -271,8 +271,8 @@ type SnapshotPreview1 interface {
 	// corresponding sizes in uint32 little-endian encoding. If either are invalid due to memory constraints, this
 	// returns ErrnoFault.
 	//
-	// * resultEnvironc - is the offset to write the environment variable count to wasm.Memory
-	// * resultEnvironBufSize - is the offset to write the null-terminated environment variable length to wasm.Memory
+	// * resultEnvironc - is the offset to write the environment variable count to ctx.Memory
+	// * resultEnvironBufSize - is the offset to write the null-terminated environment variable length to ctx.Memory
 	//
 	// For example, if Environ is []string{"a=b","b=cd"} and
 	//   EnvironSizesGet parameters are resultEnvironc=1 and resultEnvironBufSize=6, we expect `ctx.Memory` to contain:
@@ -299,7 +299,7 @@ type SnapshotPreview1 interface {
 	//
 	// * id - The clock id for which to return the time.
 	// * precision - The maximum lag (exclusive) that the returned time value may have, compared to its actual value.
-	// * resultTimestamp - the offset to write the timestamp to wasm.Memory
+	// * resultTimestamp - the offset to write the timestamp to ctx.Memory
 	//   * the timestamp is epoch nanoseconds encoded as a uint64 little-endian encoding.
 	//
 	// For example, if time.Now returned exactly midnight UTC 2022-01-01 (1640995200000000000), and
@@ -481,7 +481,7 @@ type SnapshotPreview1 interface {
 
 	// RandomGet is the WASI function named FunctionRandomGet that write random data in buffer (rand.Read()).
 	//
-	// * buf - is the wasm.Memory offset to write random values
+	// * buf - is the ctx.Memory offset to write random values
 	// * bufLen - size of random data in bytes
 	//
 	// For example, if underlying random source was seeded like `rand.NewSource(42)`, we expect `ctx.Memory` to contain:
diff --git a/internal/wasm/host.go b/internal/wasm/host.go
index 6fdb4b6b..9da927dd 100644
--- a/internal/wasm/host.go
+++ b/internal/wasm/host.go
@@ -16,19 +16,20 @@ func NewModuleContext(ctx context.Context, engine Engine, instance *ModuleInstan
 	return &ModuleContext{
 		ctx:    ctx,
 		engine: engine,
-		memory: instance.Memory,
+		memory: instance.MemoryInstance,
 		Module: instance,
 	}
 }
 
 // ModuleContext implements wasm.ModuleContext and wasm.Module
 type ModuleContext struct {
-	// ctx is the default context, exposed as wasm.ModuleContext Context
-	ctx    context.Context
+	// ctx is returned by Context and overridden WithContext
+	ctx context.Context
+	// engine is used to implement function.Call
 	engine Engine
-	// Module is exported for spectests.callFunction
+	// Module is exported for spectests
 	Module *ModuleInstance
-	// memory is exposed as wasm.ModuleContext Memory
+	// memory is returned by Memory and overridden WithMemory
 	memory publicwasm.Memory
 }
 
@@ -50,16 +51,17 @@ func (c *ModuleContext) WithMemory(memory *MemoryInstance) *ModuleContext {
 	return c
 }
 
+// Context implements wasm.ModuleContext Context
 func (c *ModuleContext) Context() context.Context {
 	return c.ctx
 }
 
-// Memory implements wasm.ModuleExports Memory
+// Memory implements wasm.ModuleContext Memory
 func (c *ModuleContext) Memory() publicwasm.Memory {
 	return c.memory
 }
 
-// Function implements wasm.ModuleExports Function
+// Function implements wasm.ModuleContext Function
 func (c *ModuleContext) Function(name string) publicwasm.Function {
 	exp, err := c.Module.GetExport(name, ExportKindFunc)
 	if err != nil {
@@ -89,11 +91,15 @@ func (s *Store) ExportHostFunctions(moduleName string, nameToGoFunc map[string]i
 		return nil, err
 	}
 
-	ret := HostExports{NameToFunctionInstance: make(map[string]*FunctionInstance, len(nameToGoFunc))}
+	exportCount := len(nameToGoFunc)
+
+	hostModule := &ModuleInstance{Name: moduleName, Exports: make(map[string]*ExportInstance, exportCount)}
+	s.ModuleInstances[moduleName] = hostModule
+	ret := HostExports{NameToFunctionInstance: make(map[string]*FunctionInstance, exportCount)}
 	for name, goFunc := range nameToGoFunc {
 		if hf, err := NewGoFunc(name, goFunc); err != nil {
 			return nil, err
-		} else if function, err := s.AddHostFunction(moduleName, hf); err != nil {
+		} else if function, err := s.AddHostFunction(hostModule, hf); err != nil {
 			return nil, err
 		} else {
 			ret.NameToFunctionInstance[name] = function
@@ -135,19 +141,19 @@ func (g *HostExports) Function(name string) publicwasm.HostFunction {
 	return f.call
 }
 
-// Len implements wasm.ModuleContext Len
-func (m *MemoryInstance) Len() uint32 {
+// Size implements wasm.ModuleContext Size
+func (m *MemoryInstance) Size() uint32 {
 	return uint32(len(m.Buffer))
 }
 
-// hasLen returns true if Len is sufficient for sizeInBytes at the given offset.
-func (m *MemoryInstance) hasLen(offset uint32, sizeInBytes uint32) bool {
-	return uint64(offset+sizeInBytes) <= uint64(m.Len()) // uint64 prevents overflow on add
+// hasSize returns true if Len is sufficient for sizeInBytes at the given offset.
+func (m *MemoryInstance) hasSize(offset uint32, sizeInBytes uint32) bool {
+	return uint64(offset+sizeInBytes) <= uint64(m.Size()) // uint64 prevents overflow on add
 }
 
 // ReadUint32Le implements wasm.ModuleContext ReadUint32Le
 func (m *MemoryInstance) ReadUint32Le(offset uint32) (uint32, bool) {
-	if !m.hasLen(offset, 4) {
+	if !m.hasSize(offset, 4) {
 		return 0, false
 	}
 	return binary.LittleEndian.Uint32(m.Buffer[offset : offset+4]), true
@@ -164,7 +170,7 @@ func (m *MemoryInstance) ReadFloat32Le(offset uint32) (float32, bool) {
 
 // ReadUint64Le implements wasm.ModuleContext ReadUint64Le
 func (m *MemoryInstance) ReadUint64Le(offset uint32) (uint64, bool) {
-	if !m.hasLen(offset, 8) {
+	if !m.hasSize(offset, 8) {
 		return 0, false
 	}
 	return binary.LittleEndian.Uint64(m.Buffer[offset : offset+8]), true
@@ -181,7 +187,7 @@ func (m *MemoryInstance) ReadFloat64Le(offset uint32) (float64, bool) {
 
 // Read implements wasm.ModuleContext Read
 func (m *MemoryInstance) Read(offset, byteCount uint32) ([]byte, bool) {
-	if !m.hasLen(offset, byteCount) {
+	if !m.hasSize(offset, byteCount) {
 		return nil, false
 	}
 	return m.Buffer[offset : offset+byteCount], true
@@ -189,7 +195,7 @@ func (m *MemoryInstance) Read(offset, byteCount uint32) ([]byte, bool) {
 
 // WriteUint32Le implements wasm.ModuleContext WriteUint32Le
 func (m *MemoryInstance) WriteUint32Le(offset, v uint32) bool {
-	if !m.hasLen(offset, 4) {
+	if !m.hasSize(offset, 4) {
 		return false
 	}
 	binary.LittleEndian.PutUint32(m.Buffer[offset:], v)
@@ -203,7 +209,7 @@ func (m *MemoryInstance) WriteFloat32Le(offset uint32, v float32) bool {
 
 // WriteUint64Le implements wasm.ModuleContext WriteUint64Le
 func (m *MemoryInstance) WriteUint64Le(offset uint32, v uint64) bool {
-	if !m.hasLen(offset, 8) {
+	if !m.hasSize(offset, 8) {
 		return false
 	}
 	binary.LittleEndian.PutUint64(m.Buffer[offset:], v)
@@ -217,7 +223,7 @@ func (m *MemoryInstance) WriteFloat64Le(offset uint32, v float64) bool {
 
 // Write implements wasm.ModuleContext Write
 func (m *MemoryInstance) Write(offset uint32, val []byte) bool {
-	if !m.hasLen(offset, uint32(len(val))) {
+	if !m.hasSize(offset, uint32(len(val))) {
 		return false
 	}
 	copy(m.Buffer[offset:], val)
@@ -230,8 +236,8 @@ var NoopMemory = &noopMemory{}
 type noopMemory struct {
 }
 
-// Len implements wasm.ModuleContext Len
-func (m *noopMemory) Len() uint32 {
+// Size implements wasm.ModuleContext Size
+func (m *noopMemory) Size() uint32 {
 	return 0
 }
 
diff --git a/internal/wasm/interpreter/interpreter.go b/internal/wasm/interpreter/interpreter.go
index 80c16eb4..15450b70 100644
--- a/internal/wasm/interpreter/interpreter.go
+++ b/internal/wasm/interpreter/interpreter.go
@@ -517,8 +517,9 @@ func (it *interpreter) callHostFunc(ctx *wasm.ModuleContext, f *interpreterFunct
 		in[i] = val
 	}
 
+	// A host function is invoked with the calling frame's memory, which may be different if in another module.
 	if len(it.frames) > 0 {
-		ctx = ctx.WithMemory(it.frames[len(it.frames)-1].f.funcInstance.ModuleInstance.Memory)
+		ctx = ctx.WithMemory(it.frames[len(it.frames)-1].f.funcInstance.ModuleInstance.MemoryInstance)
 	}
 
 	// Handle any special parameter zero
@@ -548,7 +549,7 @@ func (it *interpreter) callHostFunc(ctx *wasm.ModuleContext, f *interpreterFunct
 func (it *interpreter) callNativeFunc(ctx *wasm.ModuleContext, f *interpreterFunction) {
 	frame := &interpreterFrame{f: f}
 	moduleInst := f.funcInstance.ModuleInstance
-	memoryInst := moduleInst.Memory
+	memoryInst := moduleInst.MemoryInstance
 	globals := moduleInst.Globals
 	var table *wasm.TableInstance
 	if len(moduleInst.Tables) > 0 {
diff --git a/internal/wasm/jit/engine.go b/internal/wasm/jit/engine.go
index 3576585e..7007de42 100644
--- a/internal/wasm/jit/engine.go
+++ b/internal/wasm/jit/engine.go
@@ -83,15 +83,15 @@ type (
 		// This is only used by JIT code so mark this as nolint.
 		moduleInstanceAddress uintptr //nolint
 
-		// globalElement0Address is the address of the first element in the globa slice,
+		// globalElement0Address is the address of the first element in the global slice,
 		// i.e. &ModuleInstance.Globals[0] as uintptr.
 		globalElement0Address uintptr
-		// memoryElement0Address is the address of the first element in the globa slice,
-		// i.e. &ModuleInstance.Memory.Buffer[0] as uintptr.
+		// memoryElement0Address is the address of the first element in the global slice,
+		// i.e. &ModuleInstance.MemoryInstance.Buffer[0] as uintptr.
 		memoryElement0Address uintptr
-		// memorySliceLen is the length of the memory buffer, i.e. len(ModuleInstance.Memory.Buffer).
+		// memorySliceLen is the length of the memory buffer, i.e. len(ModuleInstance.MemoryInstance.Buffer).
 		memorySliceLen uint64
-		// tableElement0Address is the address of the first item in the globa slice,
+		// tableElement0Address is the address of the first item in the global slice,
 		// i.e. &ModuleInstance.Tables[0].Table[0] as uintptr.
 		tableElement0Address uintptr
 		// tableSliceLen is the length of the memory buffer, i.e. len(ModuleInstance.Tables[0].Table).
@@ -184,7 +184,7 @@ const (
 	engineGlobalContextCallFrameStackElement0AddressOffset    = 16
 	engineGlobalContextCallFrameStackLenOffset                = 24
 	engineGlobalContextCallFrameStackPointerOffset            = 32
-	engineGlobalContextPreviouscallFrameStackPointer          = 40
+	engineGlobalContextPreviousCallFrameStackPointer          = 40
 	engineGlobalContextCompiledFunctionsElement0AddressOffset = 48
 
 	// Offsets for engine.moduleContext.
@@ -538,8 +538,9 @@ jitentry:
 			fn := e.compiledFunctions[e.exitContext.functionCallAddress]
 			callerCompiledFunction := e.callFrameAt(1).compiledFunction
 			saved := e.globalContext.previousCallFrameStackPointer
+			// A host function is invoked with the calling frame's memory, which may be different if in another module.
 			e.execHostFunction(fn.source.FunctionKind, fn.source.HostFunction,
-				ctx.WithMemory(callerCompiledFunction.source.ModuleInstance.Memory),
+				ctx.WithMemory(callerCompiledFunction.source.ModuleInstance.MemoryInstance),
 			)
 			e.globalContext.previousCallFrameStackPointer = saved
 			goto jitentry
@@ -547,7 +548,7 @@ jitentry:
 			switch e.exitContext.functionCallAddress {
 			case builtinFunctionAddressMemoryGrow:
 				callerCompiledFunction := e.callFrameTop().compiledFunction
-				e.builtinFunctionMemoryGrow(callerCompiledFunction.source.ModuleInstance.Memory)
+				e.builtinFunctionMemoryGrow(callerCompiledFunction.source.ModuleInstance.MemoryInstance)
 			case builtinFunctionAddressGrowValueStack:
 				callerCompiledFunction := e.callFrameTop().compiledFunction
 				e.builtinFunctionGrowValueStack(callerCompiledFunction.stackPointerCeil)
diff --git a/internal/wasm/jit/jit_amd64.go b/internal/wasm/jit/jit_amd64.go
index 9d42a468..fc65ed79 100644
--- a/internal/wasm/jit/jit_amd64.go
+++ b/internal/wasm/jit/jit_amd64.go
@@ -4874,7 +4874,7 @@ func (c *amd64Compiler) returnFunction() error {
 	cmpWithPreviousCallStackPointer.To.Reg = decrementedCallFrameStackPointerRegister
 	cmpWithPreviousCallStackPointer.From.Type = obj.TYPE_MEM
 	cmpWithPreviousCallStackPointer.From.Reg = reservedRegisterForEngine
-	cmpWithPreviousCallStackPointer.From.Offset = engineGlobalContextPreviouscallFrameStackPointer
+	cmpWithPreviousCallStackPointer.From.Offset = engineGlobalContextPreviousCallFrameStackPointer
 	c.addInstruction(cmpWithPreviousCallStackPointer)
 
 	jmpIfNotPreviousCallStackPointer := c.newProg()
@@ -5259,7 +5259,7 @@ func (c *amd64Compiler) initializeReservedStackBasePointer() {
 }
 
 func (c *amd64Compiler) initializeReservedMemoryPointer() {
-	if c.f.ModuleInstance.Memory != nil {
+	if c.f.ModuleInstance.MemoryInstance != nil {
 		setupMemoryRegister := c.newProg()
 		setupMemoryRegister.As = x86.AMOVQ
 		setupMemoryRegister.To.Type = obj.TYPE_REG
@@ -5481,7 +5481,7 @@ func (c *amd64Compiler) initializeModuleContext() error {
 	// Note: if there's memory instruction in the function, memory instance must be non-nil.
 	// That is ensured by function validation at module instantiation phase, and that's
 	// why it is ok to skip the initialization if the module's memory instance is nil.
-	if c.f.ModuleInstance.Memory != nil {
+	if c.f.ModuleInstance.MemoryInstance != nil {
 		getMemoryInstanceAddress := c.newProg()
 		getMemoryInstanceAddress.As = x86.AMOVQ
 		getMemoryInstanceAddress.To.Type = obj.TYPE_REG
diff --git a/internal/wasm/jit/jit_arm64.go b/internal/wasm/jit/jit_arm64.go
index 3c84b9a0..ce076557 100644
--- a/internal/wasm/jit/jit_arm64.go
+++ b/internal/wasm/jit/jit_arm64.go
@@ -517,7 +517,7 @@ func (c *arm64Compiler) compileReturnFunction() error {
 
 	// Next we compare the decremented call frame stack pointer with the engine.precviousCallFrameStackPointer.
 	c.compileMemoryToRegisterInstruction(arm64.AMOVD,
-		reservedRegisterForEngine, engineGlobalContextPreviouscallFrameStackPointer,
+		reservedRegisterForEngine, engineGlobalContextPreviousCallFrameStackPointer,
 		tmpReg,
 	)
 	c.compileTwoRegistersToNoneInstruction(arm64.ACMP, callFramePointerReg, tmpReg)
@@ -3338,7 +3338,7 @@ func (c *arm64Compiler) compileReservedStackBasePointerRegisterInitialization()
 }
 
 func (c *arm64Compiler) compileReservedMemoryRegisterInitialization() {
-	if c.f.ModuleInstance.Memory != nil {
+	if c.f.ModuleInstance.MemoryInstance != nil {
 		// "reservedRegisterForMemory = engine.MemoryElement0Address"
 		c.compileMemoryToRegisterInstruction(
 			arm64.AMOVD,
@@ -3404,7 +3404,7 @@ func (c *arm64Compiler) compileModuleContextInitialization() error {
 	// Note: if there's memory instruction in the function, memory instance must be non-nil.
 	// That is ensured by function validation at module instantiation phase, and that's
 	// why it is ok to skip the initialization if the module's memory instance is nil.
-	if c.f.ModuleInstance.Memory != nil {
+	if c.f.ModuleInstance.MemoryInstance != nil {
 		// "tmpX = moduleInstance.Memory"
 		c.compileMemoryToRegisterInstruction(
 			arm64.AMOVD,
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index 7e784470..a3639b84 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -87,9 +87,11 @@ type (
 		Exports   map[string]*ExportInstance
 		Functions []*FunctionInstance
 		Globals   []*GlobalInstance
-		Memory    *MemoryInstance
-		Tables    []*TableInstance
-		Types     []*TypeInstance
+		// MemoryInstance is set when Module.MemorySection had a memory, regardless of whether it was exported.
+		// Note: This avoids the name "Memory" which is an interface method name.
+		MemoryInstance *MemoryInstance
+		Tables         []*TableInstance
+		Types          []*TypeInstance
 	}
 
 	// ExportInstance represents an exported instance in a Store.
@@ -248,7 +250,7 @@ func NewStore(ctx context.Context, engine Engine) *Store {
 	}
 }
 
-func (s *Store) Instantiate(module *Module, name string) (*ModuleContext, error) {
+func (s *Store) Instantiate(module *Module, name string) (*ModuleExports, error) {
 	if err := s.requireModuleUnused(name); err != nil {
 		return nil, err
 	}
@@ -336,19 +338,46 @@ func (s *Store) Instantiate(module *Module, name string) (*ModuleContext, error)
 			return nil, fmt.Errorf("module[%s] start function failed: %w", name, err)
 		}
 	}
-	return modCtx, nil
+	return &ModuleExports{s, modCtx}, nil
 }
 
 // ModuleExports implements wasm.Store ModuleExports
-func (s *Store) ModuleExports(moduleName string) (publicwasm.ModuleExports, bool) {
-	e, ok := s.ModuleContexts[moduleName]
-	return e, ok
+func (s *Store) ModuleExports(moduleName string) publicwasm.ModuleExports {
+	if m, ok := s.ModuleContexts[moduleName]; !ok {
+		return nil
+	} else {
+		return &ModuleExports{s, m}
+	}
+}
+
+// ModuleExports implements wasm.ModuleExports
+type ModuleExports struct {
+	s *Store
+	// Context is exported for /wasi.go
+	Context *ModuleContext
+}
+
+// Function implements wasm.ModuleExports Function
+func (m *ModuleExports) Function(name string) publicwasm.Function {
+	exp, err := m.Context.Module.GetExport(name, ExportKindFunc)
+	if err != nil {
+		return nil
+	}
+	return (&function{c: m.Context, f: exp.Function}).Call
+}
+
+// Memory implements wasm.ModuleExports Memory
+func (m *ModuleExports) Memory(name string) publicwasm.Memory {
+	exp, err := m.Context.Module.GetExport(name, ExportKindMemory)
+	if err != nil {
+		return nil
+	}
+	return exp.Memory
 }
 
 // HostExports implements wasm.Store HostExports
-func (s *Store) HostExports(moduleName string) (publicwasm.HostExports, bool) {
-	e, ok := s.hostExports[moduleName]
-	return e, ok
+func (s *Store) HostExports(moduleName string) publicwasm.HostExports {
+	return s.hostExports[moduleName]
 }
 
 func (s *Store) getExport(moduleName string, name string, kind ExportKind) (exp *ExportInstance, err error) {
@@ -450,7 +479,7 @@ func (s *Store) applyTableImport(target *ModuleInstance, tableTypePtr *TableType
 }
 
 func (s *Store) applyMemoryImport(target *ModuleInstance, memoryTypePtr *MemoryType, externModuleExportIsntance *ExportInstance) error {
-	if target.Memory != nil {
+	if target.MemoryInstance != nil {
 		// The current Wasm spec doesn't allow multiple memories.
 		return fmt.Errorf("multiple memories are not supported")
 	} else if memoryTypePtr == nil {
@@ -467,7 +496,7 @@ func (s *Store) applyMemoryImport(target *ModuleInstance, memoryTypePtr *MemoryT
 			return fmt.Errorf("incompatible memory imports: maximum size mismatch")
 		}
 	}
-	target.Memory = memory
+	target.MemoryInstance = memory
 	return nil
 }
 
@@ -652,22 +681,22 @@ func (s *Store) buildFunctionInstances(module *Module, target *ModuleInstance) (
 func (s *Store) buildMemoryInstances(module *Module, target *ModuleInstance) (rollbackFuncs []func(), err error) {
 	// Allocate memory instances.
 	for _, memSec := range module.MemorySection {
-		if target.Memory != nil {
+		if target.MemoryInstance != nil {
 			// This case the memory instance is already imported,
 			// and the current Wasm spec doesn't allow multiple memories.
 			return rollbackFuncs, fmt.Errorf("multiple memories not supported")
 		}
-		target.Memory = &MemoryInstance{
+		target.MemoryInstance = &MemoryInstance{
 			Buffer: make([]byte, memoryPagesToBytesNum(memSec.Min)),
 			Min:    memSec.Min,
 			Max:    memSec.Max,
 		}
-		s.Memories = append(s.Memories, target.Memory)
+		s.Memories = append(s.Memories, target.MemoryInstance)
 	}
 
 	// Initialize the memory instance according to the Data section.
 	for _, d := range module.DataSection {
-		if target.Memory == nil {
+		if target.MemoryInstance == nil {
 			return rollbackFuncs, fmt.Errorf("unknown memory")
 		} else if d.MemoryIndex != 0 {
 			return rollbackFuncs, fmt.Errorf("memory index must be zero")
@@ -696,7 +725,7 @@ func (s *Store) buildMemoryInstances(module *Module, target *ModuleInstance) (ro
 			return rollbackFuncs, fmt.Errorf("memory size out of limit %d * 64Ki", int(*(module.MemorySection[d.MemoryIndex].Max)))
 		}
 
-		memoryInst := target.Memory
+		memoryInst := target.MemoryInstance
 		if size > uint64(len(memoryInst.Buffer)) {
 			return rollbackFuncs, fmt.Errorf("out of bounds memory access")
 		}
@@ -802,10 +831,10 @@ func (s *Store) buildExportInstances(module *Module, target *ModuleInstance) (ro
 			}
 			ei = &ExportInstance{Kind: exp.Kind, Global: target.Globals[index]}
 		case ExportKindMemory:
-			if index != 0 || target.Memory == nil {
+			if index != 0 || target.MemoryInstance == nil {
 				return nil, fmt.Errorf("unknown memory for export[%s]", name)
 			}
-			ei = &ExportInstance{Kind: exp.Kind, Memory: target.Memory}
+			ei = &ExportInstance{Kind: exp.Kind, Memory: target.MemoryInstance}
 		case ExportKindTable:
 			if index >= uint32(len(target.Tables)) {
 				return nil, fmt.Errorf("unknown table for export[%s]", name)
@@ -851,16 +880,14 @@ func DecodeBlockType(types []*TypeInstance, r io.Reader) (*FunctionType, uint64,
 // exists for this module and name it is ignored rather than overwritten.
 //
 // Note: The wasm.Memory of the fn will be from the importing module.
-// Note: The ModuleInstance of this host function is lazy created and only includes exported functions and their types.
-func (s *Store) AddHostFunction(moduleName string, hf *GoFunc) (*FunctionInstance, error) {
-	m := s.getModuleInstance(moduleName)
+func (s *Store) AddHostFunction(m *ModuleInstance, hf *GoFunc) (*FunctionInstance, error) {
 	typeInstance, err := s.getTypeInstance(hf.functionType)
 	if err != nil {
 		return nil, err
 	}
 
 	f := &FunctionInstance{
-		Name:           fmt.Sprintf("%s.%s", moduleName, hf.wasmFunctionName),
+		Name:           fmt.Sprintf("%s.%s", m.Name, hf.wasmFunctionName),
 		HostFunction:   hf.goFunc,
 		FunctionKind:   hf.functionKind,
 		FunctionType:   typeInstance,
@@ -883,33 +910,31 @@ func (s *Store) AddHostFunction(moduleName string, hf *GoFunc) (*FunctionInstanc
 	return f, nil
 }
 
-func (s *Store) AddGlobal(moduleName, name string, value uint64, valueType ValueType, mutable bool) error {
+func (s *Store) AddGlobal(m *ModuleInstance, name string, value uint64, valueType ValueType, mutable bool) error {
 	g := &GlobalInstance{
 		Val:  value,
 		Type: &GlobalType{Mutable: mutable, ValType: valueType},
 	}
 	s.Globals = append(s.Globals, g)
 
-	m := s.getModuleInstance(moduleName)
 	return m.addExport(name, &ExportInstance{Kind: ExportKindGlobal, Global: g})
 }
 
-func (s *Store) AddTableInstance(moduleName, name string, min uint32, max *uint32) error {
+func (s *Store) AddTableInstance(m *ModuleInstance, name string, min uint32, max *uint32) error {
 	t := newTableInstance(min, max)
 	s.Tables = append(s.Tables, t)
 
-	m := s.getModuleInstance(moduleName)
 	return m.addExport(name, &ExportInstance{Kind: ExportKindTable, Table: t})
 }
 
-func (s *Store) AddMemoryInstance(moduleName, name string, min uint32, max *uint32) error {
+func (s *Store) AddMemoryInstance(m *ModuleInstance, name string, min uint32, max *uint32) error {
 	memory := &MemoryInstance{
 		Buffer: make([]byte, memoryPagesToBytesNum(min)),
 		Min:    min,
 		Max:    max,
 	}
 	s.Memories = append(s.Memories, memory)
-	m := s.getModuleInstance(moduleName)
+
 	return m.addExport(name, &ExportInstance{Kind: ExportKindMemory, Memory: memory})
 }
 
@@ -927,16 +952,6 @@ func (s *Store) getTypeInstance(t *FunctionType) (*TypeInstance, error) {
 	return &TypeInstance{Type: t, TypeID: id}, nil
 }
 
-// getModuleInstance returns an existing ModuleInstance if exists, or assigns a new one.
-func (s *Store) getModuleInstance(name string) *ModuleInstance {
-	m, ok := s.ModuleInstances[name]
-	if !ok {
-		m = &ModuleInstance{Name: name, Exports: map[string]*ExportInstance{}}
-		s.ModuleInstances[name] = m
-	}
-	return m
-}
-
 func newTableInstance(min uint32, max *uint32) *TableInstance {
 	tableInst := &TableInstance{
 		Table:    make([]TableElement, min),
diff --git a/wasi.go b/wasi.go
index b21e0753..570946f0 100644
--- a/wasi.go
+++ b/wasi.go
@@ -98,15 +98,15 @@ func StartWASICommand(store wasm.Store, module *Module) (wasm.ModuleExports, err
 		return nil, err
 	}
 
-	ret, err := InstantiateModule(store, module)
+	instantiated, err := internal.Instantiate(module.wasm, module.name)
 	if err != nil {
 		return nil, err
 	}
+	ctx := instantiated.Context
 
-	exports := internal.ModuleContexts[module.name]
-	start := exports.Function(internalwasi.FunctionStart)
-	if _, err = start(exports.Context()); err != nil {
+	start := ctx.Function(internalwasi.FunctionStart)
+	if _, err = start(ctx.Context()); err != nil {
 		return nil, fmt.Errorf("module[%s] function[%s] failed: %w", module.name, internalwasi.FunctionStart, err)
 	}
-	return ret, nil
+	return instantiated, nil
 }
diff --git a/wasm/wasm.go b/wasm/wasm.go
index fb04d8c9..7fa3f043 100644
--- a/wasm/wasm.go
+++ b/wasm/wasm.go
@@ -8,17 +8,24 @@ import (
 
 // Store allows access to instantiated modules and host functions
 type Store interface {
-	// ModuleExports returns exports from an instantiated module or false if it wasn't.
-	ModuleExports(moduleName string) (ModuleExports, bool)
+	// ModuleExports returns exports from an instantiated module or nil if there aren't any.
+	ModuleExports(moduleName string) ModuleExports
 
-	// HostExports returns exported host functions for the moduleName or false there are none.
-	HostExports(moduleName string) (HostExports, bool)
+	// HostExports returns exported host functions for the moduleName or nil if there aren't any.
+	HostExports(moduleName string) HostExports
 }
 
 // ModuleExports return functions exported in a module, post-instantiation.
 //
 // Note: This is an interface for decoupling, not third-party implementations. All implementations are in wazero.
 type ModuleExports interface {
+	// Memory returns a memory exported from this module or nil if it wasn't.
+	//
+	// Note: WASI modules require exporting a Memory named "memory". This means that a module successfully initialized
+	// as a WASI Command or Reactor will never return nil for this name.
+	// See https://github.com/WebAssembly/WASI/blob/snapshot-01/design/application-abi.md#current-unstable-abi
+	Memory(name string) Memory
+
 	// Function returns a function exported from this module or nil if it wasn't.
 	Function(name string) Function
 }
@@ -93,11 +100,13 @@ type ModuleContext interface {
 // Note: This includes all value types available in WebAssembly 1.0 (MVP) and all are encoded little-endian.
 // See https://www.w3.org/TR/wasm-core-1/#storage%E2%91%A0
 type Memory interface {
-	// Len returns the size in bytes available. Ex. If the underlying memory has 1 page: 65536
+	// Size returns the size in bytes available. Ex. If the underlying memory has 1 page: 65536
 	//
 	// Note: this will not grow during a host function call, even if the underlying memory can.  Ex. If the underlying
 	// memory has min 0 and max 2 pages, this returns zero.
-	Len() uint32
+	//
+	// See https://www.w3.org/TR/wasm-core-1/#-hrefsyntax-instr-memorymathsfmemorysize%E2%91%A0
+	Size() uint32
 
 	// ReadUint32Le reads a uint32 in little-endian encoding from the underlying buffer at the offset in or returns
 	// false if out of range.


Test patch:
diff --git a/internal/wasi/wasi_test.go b/internal/wasi/wasi_test.go
index 978644f9..4c8230ca 100644
--- a/internal/wasi/wasi_test.go
+++ b/internal/wasi/wasi_test.go
@@ -789,13 +789,17 @@ func instantiateWasmStore(t *testing.T, wasiFunction, wasiImport, moduleName str
 	fn, err := wasm.NewGoFunc(wasiFunction, goFunc)
 	require.NoError(t, err)
 
-	wasiFn, err := store.AddHostFunction(wasi.ModuleSnapshotPreview1, fn)
+	// Add the host module
+	hostModule := &wasm.ModuleInstance{Name: wasi.ModuleSnapshotPreview1, Exports: map[string]*wasm.ExportInstance{}}
+	store.ModuleInstances[hostModule.Name] = hostModule
+
+	wasiFn, err := store.AddHostFunction(hostModule, fn)
 	require.NoError(t, err)
 
-	ctx, err := store.Instantiate(mod, moduleName)
+	instantiated, err := store.Instantiate(mod, moduleName)
 	require.NoError(t, err)
 
-	return store, ctx, wasiFn
+	return store, instantiated.Context, wasiFn
 }
 
 // maskMemory overwrites the first memory in the store with '?', so tests can see what's written.
diff --git a/internal/wasm/host_test.go b/internal/wasm/host_test.go
index aea97d5d..940eeaf4 100644
--- a/internal/wasm/host_test.go
+++ b/internal/wasm/host_test.go
@@ -1,7 +1,6 @@
 package internalwasm
 
 import (
-	"context"
 	"math"
 	"testing"
 
@@ -25,19 +24,19 @@ func TestMemoryInstance_HasLen(t *testing.T) {
 		},
 		{
 			name:        "maximum valid sizeInBytes",
-			offset:      memory.Len() - 8,
+			offset:      memory.Size() - 8,
 			sizeInBytes: 8,
 			expected:    true,
 		},
 		{
 			name:        "sizeInBytes exceeds the valid size by 1",
 			offset:      100, // arbitrary valid offset
-			sizeInBytes: uint64(memory.Len() - 99),
+			sizeInBytes: uint64(memory.Size() - 99),
 			expected:    false,
 		},
 		{
 			name:        "offset exceeds the memory size",
-			offset:      memory.Len(),
+			offset:      memory.Size(),
 			sizeInBytes: 1, // arbitrary size
 			expected:    false,
 		},
@@ -47,7 +46,7 @@ func TestMemoryInstance_HasLen(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			require.Equal(t, tc.expected, memory.hasLen(tc.offset, uint32(tc.sizeInBytes)))
+			require.Equal(t, tc.expected, memory.hasSize(tc.offset, uint32(tc.sizeInBytes)))
 		})
 	}
 }
@@ -274,15 +273,15 @@ func TestMemoryInstance_WriteUint32Le(t *testing.T) {
 		},
 		{
 			name:          "maximum boundary valid offset",
-			offset:        memory.Len() - 4, // 4 is the size of uint32
-			v:             1,                // arbitrary valid v
+			offset:        memory.Size() - 4, // 4 is the size of uint32
+			v:             1,                 // arbitrary valid v
 			expectedOk:    true,
 			expectedBytes: []byte{0x1, 0x00, 0x00, 0x00},
 		},
 		{
 			name:          "offset exceeds the maximum valid offset by 1",
-			offset:        memory.Len() - 4 + 1, // 4 is the size of uint32
-			v:             1,                    // arbitrary valid v
+			offset:        memory.Size() - 4 + 1, // 4 is the size of uint32
+			v:             1,                     // arbitrary valid v
 			expectedBytes: []byte{0xff, 0xff, 0xff, 0xff},
 		},
 	}
@@ -324,15 +323,15 @@ func TestMemoryInstance_WriteUint64Le(t *testing.T) {
 		},
 		{
 			name:          "maximum boundary valid offset",
-			offset:        memory.Len() - 8, // 8 is the size of uint64
-			v:             1,                // arbitrary valid v
+			offset:        memory.Size() - 8, // 8 is the size of uint64
+			v:             1,                 // arbitrary valid v
 			expectedOk:    true,
 			expectedBytes: []byte{0x1, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
 		},
 		{
 			name:       "offset exceeds the maximum valid offset by 1",
-			offset:     memory.Len() - 8 + 1, // 8 is the size of uint64
-			v:          1,                    // arbitrary valid v
+			offset:     memory.Size() - 8 + 1, // 8 is the size of uint64
+			v:          1,                     // arbitrary valid v
 			expectedOk: false,
 		},
 	}
@@ -375,15 +374,15 @@ func TestMemoryInstance_WriteFloat32Le(t *testing.T) {
 		},
 		{
 			name:          "maximum boundary valid offset",
-			offset:        memory.Len() - 4, // 4 is the size of float32
-			v:             0.1,              // arbitrary valid v
+			offset:        memory.Size() - 4, // 4 is the size of float32
+			v:             0.1,               // arbitrary valid v
 			expectedOk:    true,
 			expectedBytes: []byte{0xcd, 0xcc, 0xcc, 0x3d},
 		},
 		{
 			name:          "offset exceeds the maximum valid offset by 1",
-			offset:        memory.Len() - 4 + 1, // 4 is the size of float32
-			v:             math.MaxFloat32,      // arbitrary valid v
+			offset:        memory.Size() - 4 + 1, // 4 is the size of float32
+			v:             math.MaxFloat32,       // arbitrary valid v
 			expectedBytes: []byte{0xff, 0xff, 0xff, 0xff},
 		},
 	}
@@ -425,15 +424,15 @@ func TestMemoryInstance_WriteFloat64Le(t *testing.T) {
 		},
 		{
 			name:          "maximum boundary valid offset",
-			offset:        memory.Len() - 8, // 8 is the size of float64
-			v:             math.MaxFloat64,  // arbitrary valid v
+			offset:        memory.Size() - 8, // 8 is the size of float64
+			v:             math.MaxFloat64,   // arbitrary valid v
 			expectedOk:    true,
 			expectedBytes: []byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xef, 0x7f},
 		},
 		{
 			name:       "offset exceeds the maximum valid offset by 1",
-			offset:     memory.Len() - 8 + 1, // 8 is the size of float64
-			v:          math.MaxFloat64,      // arbitrary valid v
+			offset:     memory.Size() - 8 + 1, // 8 is the size of float64
+			v:          math.MaxFloat64,       // arbitrary valid v
 			expectedOk: false,
 		},
 	}
@@ -449,62 +448,3 @@ func TestMemoryInstance_WriteFloat64Le(t *testing.T) {
 		})
 	}
 }
-
-func TestFunction_Call(t *testing.T) {
-	name := "test"
-	fn := "fn"
-	engine := &nopEngine{}
-	s := NewStore(context.Background(), engine)
-	m := &ModuleInstance{
-		Name: name,
-		Exports: map[string]*ExportInstance{
-			fn: {
-				Kind: ExportKindFunc,
-				Function: &FunctionInstance{
-					FunctionKind: FunctionKindWasm,
-					FunctionType: &TypeInstance{
-						Type: &FunctionType{
-							Params:  []ValueType{},
-							Results: []ValueType{},
-						},
-					},
-				},
-			},
-		},
-	}
-	ctx := NewModuleContext(context.Background(), s.Engine, m)
-	s.ModuleInstances[name] = m
-	s.ModuleContexts[name] = ctx
-
-	type testKey struct{}
-	ctxVal := context.WithValue(context.Background(), testKey{}, "test")
-
-	tests := []struct {
-		name      string
-		ctx       context.Context
-		actualCtx context.Context
-	}{
-		{
-			name:      "nil context",
-			ctx:       nil,
-			actualCtx: context.Background(),
-		},
-		{
-			name:      "background context",
-			ctx:       context.Background(),
-			actualCtx: context.Background(),
-		},
-		{
-			name:      "context with value",
-			ctx:       ctxVal,
-			actualCtx: ctxVal,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			_, err := (&function{ctx, m.Exports[fn].Function}).Call(tt.ctx)
-			require.NoError(t, err)
-			require.Equal(t, tt.actualCtx, engine.ctx.ctx)
-		})
-	}
-}
diff --git a/internal/wasm/interpreter/interpreter_test.go b/internal/wasm/interpreter/interpreter_test.go
index 45b080d9..75bd6fed 100644
--- a/internal/wasm/interpreter/interpreter_test.go
+++ b/internal/wasm/interpreter/interpreter_test.go
@@ -50,7 +50,7 @@ func TestInterpreter_CallHostFunc(t *testing.T) {
 		hostFn := reflect.ValueOf(func(ctx publicwasm.ModuleContext) {
 			ctxMemory = ctx.Memory()
 		})
-		module := &wasm.ModuleInstance{Memory: memory}
+		module := &wasm.ModuleInstance{MemoryInstance: memory}
 		it := interpreter{functions: map[wasm.FunctionAddress]*interpreterFunction{
 			0: {hostFn: &hostFn, funcInstance: &wasm.FunctionInstance{
 				FunctionKind: wasm.FunctionKindGoModuleContext,
diff --git a/internal/wasm/jit/engine_test.go b/internal/wasm/jit/engine_test.go
index 8a042b85..91db878b 100644
--- a/internal/wasm/jit/engine_test.go
+++ b/internal/wasm/jit/engine_test.go
@@ -19,7 +19,7 @@ func TestVerifyOffsetValue(t *testing.T) {
 	require.Equal(t, int(unsafe.Offsetof(eng.callFrameStackElementZeroAddress)), engineGlobalContextCallFrameStackElement0AddressOffset)
 	require.Equal(t, int(unsafe.Offsetof(eng.callFrameStackLen)), engineGlobalContextCallFrameStackLenOffset)
 	require.Equal(t, int(unsafe.Offsetof(eng.callFrameStackPointer)), engineGlobalContextCallFrameStackPointerOffset)
-	require.Equal(t, int(unsafe.Offsetof(eng.previousCallFrameStackPointer)), engineGlobalContextPreviouscallFrameStackPointer)
+	require.Equal(t, int(unsafe.Offsetof(eng.previousCallFrameStackPointer)), engineGlobalContextPreviousCallFrameStackPointer)
 	require.Equal(t, int(unsafe.Offsetof(eng.compiledFunctionsElement0Address)), engineGlobalContextCompiledFunctionsElement0AddressOffset)
 
 	// Offsets for engine.moduleContext.
@@ -61,7 +61,7 @@ func TestVerifyOffsetValue(t *testing.T) {
 	// Offsets for wasm.ModuleInstance.
 	var moduleInstance wasm.ModuleInstance
 	require.Equal(t, int(unsafe.Offsetof(moduleInstance.Globals)), moduleInstanceGlobalsOffset)
-	require.Equal(t, int(unsafe.Offsetof(moduleInstance.Memory)), moduleInstanceMemoryOffset)
+	require.Equal(t, int(unsafe.Offsetof(moduleInstance.MemoryInstance)), moduleInstanceMemoryOffset)
 	require.Equal(t, int(unsafe.Offsetof(moduleInstance.Tables)), moduleInstanceTablesOffset)
 
 	// Offsets for wasm.TableInstance.
diff --git a/internal/wasm/jit/jit_amd64_test.go b/internal/wasm/jit/jit_amd64_test.go
index dfd99c57..086558ae 100644
--- a/internal/wasm/jit/jit_amd64_test.go
+++ b/internal/wasm/jit/jit_amd64_test.go
@@ -215,9 +215,9 @@ func TestAmd64Compiler_initializeModuleContext(t *testing.T) {
 		{
 			name: "no nil",
 			moduleInstance: &wasm.ModuleInstance{
-				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				Tables:  []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
-				Globals: []*wasm.GlobalInstance{{Val: 100}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Tables:         []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
+				Globals:        []*wasm.GlobalInstance{{Val: 100}},
 			},
 		},
 		{
@@ -230,47 +230,47 @@ func TestAmd64Compiler_initializeModuleContext(t *testing.T) {
 		{
 			name: "memory zero length",
 			moduleInstance: &wasm.ModuleInstance{
-				Tables:  []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
-				Globals: []*wasm.GlobalInstance{{Val: 100}},
-				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 0)},
+				Tables:         []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
+				Globals:        []*wasm.GlobalInstance{{Val: 100}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 0)},
 			},
 		},
 		{
 			name: "table length zero",
 			moduleInstance: &wasm.ModuleInstance{
-				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				Tables:  []*wasm.TableInstance{{Table: nil}},
-				Globals: []*wasm.GlobalInstance{{Val: 100}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Tables:         []*wasm.TableInstance{{Table: nil}},
+				Globals:        []*wasm.GlobalInstance{{Val: 100}},
 			},
 		},
 		{
 			name: "table length zero part2",
 			moduleInstance: &wasm.ModuleInstance{
-				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				Tables:  []*wasm.TableInstance{{Table: make([]wasm.TableElement, 0)}},
-				Globals: []*wasm.GlobalInstance{{Val: 100}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Tables:         []*wasm.TableInstance{{Table: make([]wasm.TableElement, 0)}},
+				Globals:        []*wasm.GlobalInstance{{Val: 100}},
 			},
 		},
 		{
 			name: "table nil",
 			moduleInstance: &wasm.ModuleInstance{
-				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				Tables:  []*wasm.TableInstance{},
-				Globals: []*wasm.GlobalInstance{{Val: 100}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Tables:         []*wasm.TableInstance{},
+				Globals:        []*wasm.GlobalInstance{{Val: 100}},
 			},
 		},
 		{
 			name: "table nil part2",
 			moduleInstance: &wasm.ModuleInstance{
-				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				Globals: []*wasm.GlobalInstance{{Val: 100}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Globals:        []*wasm.GlobalInstance{{Val: 100}},
 			},
 		},
 		{
 			name: "globals nil",
 			moduleInstance: &wasm.ModuleInstance{
-				Memory: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				Tables: []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Tables:         []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
 			},
 		},
 	} {
@@ -306,8 +306,8 @@ func TestAmd64Compiler_initializeModuleContext(t *testing.T) {
 			bufSliceHeader := (*reflect.SliceHeader)(unsafe.Pointer(&tc.moduleInstance.Globals))
 			require.Equal(t, bufSliceHeader.Data, engine.moduleContext.globalElement0Address)
 
-			if tc.moduleInstance.Memory != nil {
-				bufSliceHeader := (*reflect.SliceHeader)(unsafe.Pointer(&tc.moduleInstance.Memory.Buffer))
+			if tc.moduleInstance.MemoryInstance != nil {
+				bufSliceHeader := (*reflect.SliceHeader)(unsafe.Pointer(&tc.moduleInstance.MemoryInstance.Buffer))
 				require.Equal(t, uint64(bufSliceHeader.Len), engine.moduleContext.memorySliceLen)
 				require.Equal(t, bufSliceHeader.Data, engine.moduleContext.memoryElement0Address)
 			}
@@ -5972,7 +5972,7 @@ func TestAmd64Compiler_callFunction(t *testing.T) {
 					// Each function takes one argument, adds the value with 100 + i and returns the result.
 					addTargetValue := uint32(100 + i)
 					moduleInstance := &wasm.ModuleInstance{
-						Memory: &wasm.MemoryInstance{Buffer: make([]byte, 1024)},
+						MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 1024)},
 					}
 					moduleInstanceToExpectedValueInMemory[moduleInstance] = addTargetValue
 
@@ -6066,7 +6066,7 @@ func TestAmd64Compiler_callFunction(t *testing.T) {
 
 				// Also, in the middle of function call, we write the added value into each memory instance.
 				for mod, expValue := range moduleInstanceToExpectedValueInMemory {
-					require.Equal(t, expValue, binary.LittleEndian.Uint32(mod.Memory.Buffer[0:]))
+					require.Equal(t, expValue, binary.LittleEndian.Uint32(mod.MemoryInstance.Buffer[0:]))
 				}
 			})
 		})
diff --git a/internal/wasm/jit/jit_arm64_test.go b/internal/wasm/jit/jit_arm64_test.go
index bb34d0f1..e0b7a05f 100644
--- a/internal/wasm/jit/jit_arm64_test.go
+++ b/internal/wasm/jit/jit_arm64_test.go
@@ -2107,16 +2107,16 @@ func TestArm64Compiler_compileModuleContextInitialization(t *testing.T) {
 		{
 			name: "no nil",
 			moduleInstance: &wasm.ModuleInstance{
-				Globals: []*wasm.GlobalInstance{{Val: 100}},
-				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				Tables:  []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
+				Globals:        []*wasm.GlobalInstance{{Val: 100}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Tables:         []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
 			},
 		},
 		{
 			name: "globals nil",
 			moduleInstance: &wasm.ModuleInstance{
-				Memory: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				Tables: []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Tables:         []*wasm.TableInstance{{Table: make([]wasm.TableElement, 20)}},
 			},
 		},
 		{
@@ -2129,25 +2129,25 @@ func TestArm64Compiler_compileModuleContextInitialization(t *testing.T) {
 		{
 			name: "table nil",
 			moduleInstance: &wasm.ModuleInstance{
-				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				Tables:  []*wasm.TableInstance{{Table: nil}},
-				Globals: []*wasm.GlobalInstance{{Val: 100}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Tables:         []*wasm.TableInstance{{Table: nil}},
+				Globals:        []*wasm.GlobalInstance{{Val: 100}},
 			},
 		},
 		{
 			name: "table empty",
 			moduleInstance: &wasm.ModuleInstance{
-				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 10)},
-				Tables:  []*wasm.TableInstance{{Table: make([]wasm.TableElement, 0)}},
-				Globals: []*wasm.GlobalInstance{{Val: 100}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 10)},
+				Tables:         []*wasm.TableInstance{{Table: make([]wasm.TableElement, 0)}},
+				Globals:        []*wasm.GlobalInstance{{Val: 100}},
 			},
 		},
 		{
 			name: "memory zero length",
 			moduleInstance: &wasm.ModuleInstance{
-				Globals: []*wasm.GlobalInstance{{Val: 100}},
-				Tables:  []*wasm.TableInstance{{Table: make([]wasm.TableElement, 0)}},
-				Memory:  &wasm.MemoryInstance{Buffer: make([]byte, 0)},
+				Globals:        []*wasm.GlobalInstance{{Val: 100}},
+				Tables:         []*wasm.TableInstance{{Table: make([]wasm.TableElement, 0)}},
+				MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, 0)},
 			},
 		},
 		{
@@ -2186,8 +2186,8 @@ func TestArm64Compiler_compileModuleContextInitialization(t *testing.T) {
 			bufSliceHeader := (*reflect.SliceHeader)(unsafe.Pointer(&tc.moduleInstance.Globals))
 			require.Equal(t, bufSliceHeader.Data, engine.moduleContext.globalElement0Address)
 
-			if tc.moduleInstance.Memory != nil {
-				bufSliceHeader := (*reflect.SliceHeader)(unsafe.Pointer(&tc.moduleInstance.Memory.Buffer))
+			if tc.moduleInstance.MemoryInstance != nil {
+				bufSliceHeader := (*reflect.SliceHeader)(unsafe.Pointer(&tc.moduleInstance.MemoryInstance.Buffer))
 				require.Equal(t, uint64(bufSliceHeader.Len), engine.moduleContext.memorySliceLen)
 				require.Equal(t, bufSliceHeader.Data, engine.moduleContext.memoryElement0Address)
 			}
diff --git a/internal/wasm/jit/jit_test.go b/internal/wasm/jit/jit_test.go
index 562fb47c..f80fe2f9 100644
--- a/internal/wasm/jit/jit_test.go
+++ b/internal/wasm/jit/jit_test.go
@@ -47,7 +47,7 @@ func (j *jitEnv) stackTopAsFloat64() float64 {
 }
 
 func (j *jitEnv) memory() []byte {
-	return j.moduleInstance.Memory.Buffer
+	return j.moduleInstance.MemoryInstance.Buffer
 }
 
 func (j *jitEnv) stack() []uint64 {
@@ -133,9 +133,9 @@ func newJITEnvironment() *jitEnv {
 	return &jitEnv{
 		eng: newEngine(),
 		moduleInstance: &wasm.ModuleInstance{
-			Memory:  &wasm.MemoryInstance{Buffer: make([]byte, wasm.MemoryPageSize*defaultMemoryPageNumInTest)},
-			Tables:  []*wasm.TableInstance{{}},
-			Globals: []*wasm.GlobalInstance{},
+			MemoryInstance: &wasm.MemoryInstance{Buffer: make([]byte, wasm.MemoryPageSize*defaultMemoryPageNumInTest)},
+			Tables:         []*wasm.TableInstance{{}},
+			Globals:        []*wasm.GlobalInstance{},
 		},
 	}
 }
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 38ca7f11..cc4d9688 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -15,26 +15,91 @@ import (
 	"github.com/tetratelabs/wazero/wasm"
 )
 
-func TestStore_GetModuleInstance(t *testing.T) {
-	name := "test"
+func TestModuleInstance_Memory(t *testing.T) {
+	tests := []struct {
+		name        string
+		input       *Module
+		expected    bool
+		expectedLen uint32
+	}{
+		{
+			name:  "no memory",
+			input: &Module{},
+		},
+		{
+			name:  "memory not exported",
+			input: &Module{MemorySection: []*MemoryType{{1, nil}}},
+		},
+		{
+			name:  "memory not exported, one page",
+			input: &Module{MemorySection: []*MemoryType{{1, nil}}},
+		},
+		{
+			name: "memory exported, different name",
+			input: &Module{
+				MemorySection: []*MemoryType{{1, nil}},
+				ExportSection: map[string]*Export{"momory": {Kind: ExportKindMemory, Name: "momory", Index: 0}},
+			},
+		},
+		{
+			name: "memory exported, but zero length",
+			input: &Module{
+				MemorySection: []*MemoryType{{0, nil}},
+				ExportSection: map[string]*Export{"memory": {Kind: ExportKindMemory, Name: "memory", Index: 0}},
+			},
+			expected: true,
+		},
+		{
+			name: "memory exported, one page",
+			input: &Module{
+				MemorySection: []*MemoryType{{1, nil}},
+				ExportSection: map[string]*Export{"memory": {Kind: ExportKindMemory, Name: "memory", Index: 0}},
+			},
+			expected:    true,
+			expectedLen: 65536,
+		},
+		{
+			name: "memory exported, two pages",
+			input: &Module{
+				MemorySection: []*MemoryType{{2, nil}},
+				ExportSection: map[string]*Export{"memory": {Kind: ExportKindMemory, Name: "memory", Index: 0}},
+			},
+			expected:    true,
+			expectedLen: 65536 * 2,
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
 
-	s := NewStore(context.Background(), nopEngineInstance)
+		t.Run(tc.name, func(t *testing.T) {
+			s := NewStore(context.Background(), &catchContext{})
 
-	m1 := s.getModuleInstance(name)
-	require.Equal(t, m1, s.ModuleInstances[name])
-	require.NotNil(t, m1.Exports)
+			instance, err := s.Instantiate(tc.input, "test")
+			require.NoError(t, err)
 
-	m2 := s.getModuleInstance(name)
-	require.Equal(t, m1, m2)
+			mem := instance.Memory("memory")
+			if tc.expected {
+				require.Equal(t, tc.expectedLen, mem.Size())
+			} else {
+				require.Nil(t, mem)
+			}
+		})
+	}
 }
 
 func TestStore_AddHostFunction(t *testing.T) {
-	s := NewStore(context.Background(), nopEngineInstance)
+	s := NewStore(context.Background(), &catchContext{})
 
 	hf, err := NewGoFunc("fn", func(wasm.ModuleContext) {
 	})
 	require.NoError(t, err)
-	_, err = s.AddHostFunction("test", hf)
+
+	// Add the host module
+	hostModule := &ModuleInstance{Name: "test", Exports: make(map[string]*ExportInstance, 1)}
+	s.ModuleInstances[hostModule.Name] = hostModule
+
+	_, err = s.AddHostFunction(hostModule, hf)
 	require.NoError(t, err)
 
 	// The function was added to the store, prefixed by the owning module name
@@ -43,27 +108,30 @@ func TestStore_AddHostFunction(t *testing.T) {
 	require.Equal(t, "test.fn", fn.Name)
 
 	// The function was exported in the module
-	m := s.getModuleInstance("test")
-	require.Equal(t, 1, len(m.Exports))
-	exp, ok := m.Exports["fn"]
+	require.Equal(t, 1, len(hostModule.Exports))
+	exp, ok := hostModule.Exports["fn"]
 	require.True(t, ok)
 
 	// Trying to register it again should fail
-	_, err = s.AddHostFunction("test", hf)
+	_, err = s.AddHostFunction(hostModule, hf)
 	require.EqualError(t, err, `"fn" is already exported in module "test"`)
 
 	// Any side effects should be reverted
 	require.Equal(t, []*FunctionInstance{fn}, s.Functions)
-	require.Equal(t, map[string]*ExportInstance{"fn": exp}, m.Exports)
+	require.Equal(t, map[string]*ExportInstance{"fn": exp}, hostModule.Exports)
 }
 
 func TestStore_ExportImportedHostFunction(t *testing.T) {
-	s := NewStore(context.Background(), nopEngineInstance)
+	s := NewStore(context.Background(), &catchContext{})
 
 	hf, err := NewGoFunc("host_fn", func(wasm.ModuleContext) {
 	})
 	require.NoError(t, err)
-	_, err = s.AddHostFunction("", hf)
+
+	// Add the host module
+	hostModule := &ModuleInstance{Name: "", Exports: make(map[string]*ExportInstance, 1)}
+	s.ModuleInstances[hostModule.Name] = hostModule
+	_, err = s.AddHostFunction(hostModule, hf)
 	require.NoError(t, err)
 
 	t.Run("ModuleInstance is the importing module", func(t *testing.T) {
@@ -85,11 +153,78 @@ func TestStore_ExportImportedHostFunction(t *testing.T) {
 	})
 }
 
+func TestFunctionInstance_call(t *testing.T) {
+	type key string
+	storeCtx := context.WithValue(context.Background(), key("wa"), "zero")
+
+	notStoreCtx := context.WithValue(context.Background(), key("wazer"), "o")
+
+	tests := []struct {
+		name     string
+		ctx      context.Context
+		expected context.Context
+	}{
+		{
+			name:     "nil defaults to store context",
+			ctx:      nil,
+			expected: storeCtx,
+		},
+		{
+			name:     "set overrides store context",
+			ctx:      notStoreCtx,
+			expected: notStoreCtx,
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			engine := &catchContext{}
+			store := NewStore(storeCtx, engine)
+
+			// Define a fake host function
+			functionName := "fn"
+			hostFn := func(ctx wasm.ModuleContext) {
+			}
+			fn, err := NewGoFunc(functionName, hostFn)
+			require.NoError(t, err)
+
+			// Add the host module
+			hostModule := &ModuleInstance{Name: "host", Exports: map[string]*ExportInstance{}}
+			store.ModuleInstances[hostModule.Name] = hostModule
+			_, err = store.AddHostFunction(hostModule, fn)
+			require.NoError(t, err)
+
+			// Make a module to import the function
+			instantiated, err := store.Instantiate(&Module{
+				TypeSection: []*FunctionType{{}},
+				ImportSection: []*Import{{
+					Kind:     ImportKindFunc,
+					Module:   hostModule.Name,
+					Name:     functionName,
+					DescFunc: 0,
+				}},
+				MemorySection: []*MemoryType{{1, nil}},
+				ExportSection: map[string]*Export{functionName: {Kind: ExportKindFunc, Name: functionName, Index: 0}},
+			}, "test")
+			require.NoError(t, err)
+
+			// This fails if the function wasn't invoked, or had an unexpected context.
+			_, err = instantiated.Function(functionName)(tc.ctx)
+			require.NoError(t, err)
+			if tc.expected == nil {
+				require.Nil(t, engine.ctx)
+			} else {
+				require.Equal(t, tc.expected, engine.ctx.Context())
+			}
+		})
+	}
+}
+
 func TestStore_BuildFunctionInstances_FunctionNames(t *testing.T) {
 	name := "test"
-
-	s := NewStore(context.Background(), nopEngineInstance)
-	mi := s.getModuleInstance(name)
+	s := NewStore(context.Background(), &catchContext{})
 
 	zero := Index(0)
 	nopCode := &Code{nil, []byte{OpcodeEnd}}
@@ -106,6 +241,10 @@ func TestStore_BuildFunctionInstances_FunctionNames(t *testing.T) {
 		CodeSection: []*Code{nopCode, nopCode, nopCode, nopCode, nopCode},
 	}
 
+	// Make a fake module for test
+	mi := &ModuleInstance{Name: name}
+	s.ModuleInstances[mi.Name] = mi
+
 	_, err := s.buildFunctionInstances(m, mi)
 	require.NoError(t, err)
 
@@ -118,31 +257,29 @@ func TestStore_BuildFunctionInstances_FunctionNames(t *testing.T) {
 	require.Equal(t, []string{"unknown", "two", "unknown", "four", "five"}, names)
 }
 
-var nopEngineInstance Engine = &nopEngine{}
-
-type nopEngine struct {
+type catchContext struct {
 	ctx *ModuleContext
 }
 
-func (e *nopEngine) Call(ctx *ModuleContext, _ *FunctionInstance, _ ...uint64) (results []uint64, err error) {
+func (e *catchContext) Call(ctx *ModuleContext, _ *FunctionInstance, _ ...uint64) (results []uint64, err error) {
 	e.ctx = ctx
-	return nil, nil
+	return
 }
 
-func (e *nopEngine) Compile(_ *FunctionInstance) error {
+func (e *catchContext) Compile(_ *FunctionInstance) error {
 	return nil
 }
 
 func TestStore_addHostFunction(t *testing.T) {
 	t.Run("too many functions", func(t *testing.T) {
-		s := NewStore(context.Background(), nopEngineInstance)
+		s := NewStore(context.Background(), &catchContext{})
 		const max = 10
 		s.maximumFunctionAddress = max
 		err := s.addFunctionInstance(&FunctionInstance{Address: max + 1})
 		require.Error(t, err)
 	})
 	t.Run("ok", func(t *testing.T) {
-		s := NewStore(context.Background(), nopEngineInstance)
+		s := NewStore(context.Background(), &catchContext{})
 		for i := 0; i < 10; i++ {
 			f := &FunctionInstance{FunctionKind: FunctionKindGoNoContext}
 			require.Len(t, s.Functions, i)
@@ -158,7 +295,7 @@ func TestStore_addHostFunction(t *testing.T) {
 
 func TestStore_getTypeInstance(t *testing.T) {
 	t.Run("too many functions", func(t *testing.T) {
-		s := NewStore(context.Background(), nopEngineInstance)
+		s := NewStore(context.Background(), &catchContext{})
 		const max = 10
 		s.maximumFunctionTypes = max
 		s.TypeIDs = make(map[string]FunctionTypeID)
@@ -177,7 +314,7 @@ func TestStore_getTypeInstance(t *testing.T) {
 		} {
 			tc := tc
 			t.Run(tc.String(), func(t *testing.T) {
-				s := NewStore(context.Background(), nopEngineInstance)
+				s := NewStore(context.Background(), &catchContext{})
 				actual, err := s.getTypeInstance(tc)
 				require.NoError(t, err)
 
@@ -193,7 +330,7 @@ func TestStore_getTypeInstance(t *testing.T) {
 func TestStore_buildGlobalInstances(t *testing.T) {
 	t.Run("too many globals", func(t *testing.T) {
 		// Setup a store to have the reasonably low max on globals for testing.
-		s := NewStore(context.Background(), nopEngineInstance)
+		s := NewStore(context.Background(), &catchContext{})
 		const max = 10
 		s.maximumGlobals = max
 
@@ -202,7 +339,7 @@ func TestStore_buildGlobalInstances(t *testing.T) {
 		require.Error(t, err)
 	})
 	t.Run("invalid constant expression", func(t *testing.T) {
-		s := NewStore(context.Background(), nopEngineInstance)
+		s := NewStore(context.Background(), &catchContext{})
 
 		// Empty constant expression is invalid.
 		m := &Module{GlobalSection: []*Global{{Init: &ConstantExpression{}}}}
@@ -211,7 +348,7 @@ func TestStore_buildGlobalInstances(t *testing.T) {
 	})
 
 	t.Run("global type mismatch", func(t *testing.T) {
-		s := NewStore(context.Background(), nopEngineInstance)
+		s := NewStore(context.Background(), &catchContext{})
 		m := &Module{GlobalSection: []*Global{{
 			// Global with i32.const initial value, but with type specified as f64 must be error.
 			Init: &ConstantExpression{Opcode: OpcodeI32Const, Data: []byte{0}},
@@ -230,7 +367,7 @@ func TestStore_buildGlobalInstances(t *testing.T) {
 
 		m := &Module{GlobalSection: []*Global{global}}
 
-		s := NewStore(context.Background(), nopEngineInstance)
+		s := NewStore(context.Background(), &catchContext{})
 		target := &ModuleInstance{}
 		_, err = s.buildGlobalInstances(m, target)
 		require.NoError(t, err)
diff --git a/store_test.go b/store_test.go
index 6cfb1005..556d1c65 100644
--- a/store_test.go
+++ b/store_test.go
@@ -11,6 +11,47 @@ import (
 	"github.com/tetratelabs/wazero/wasm"
 )
 
+// TestModuleExports_Memory only covers a couple cases to avoid duplication of internal/wasm/store_test.go
+func TestModuleExports_Memory(t *testing.T) {
+	tests := []struct {
+		name, wat   string
+		expected    bool
+		expectedLen uint32
+	}{
+		{
+			name: "no memory",
+			wat:  `(module)`,
+		},
+		{
+			name:        "memory exported, one page",
+			wat:         `(module (memory $mem 1) (export "memory" (memory $mem)))`,
+			expected:    true,
+			expectedLen: 65536,
+		},
+	}
+
+	for _, tt := range tests {
+
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			mod, err := DecodeModuleText([]byte(tc.wat))
+			require.NoError(t, err)
+
+			// Instantiate the module and get the export of the above hostFn
+			exports, err := InstantiateModule(NewStore(), mod)
+			require.NoError(t, err)
+
+			mem := exports.Memory("memory")
+			if tc.expected {
+				require.Equal(t, tc.expectedLen, mem.Size())
+			} else {
+				require.Nil(t, mem)
+			}
+		})
+	}
+}
+
 func TestFunction_Context(t *testing.T) {
 	type key string
 	storeCtx := context.WithValue(context.Background(), key("wa"), "zero")
diff --git a/tests/spectest/spec_test.go b/tests/spectest/spec_test.go
index 45aac5e8..b12c1464 100644
--- a/tests/spectest/spec_test.go
+++ b/tests/spectest/spec_test.go
@@ -186,6 +186,10 @@ func (c command) expectedError() (err error) {
 }
 
 func addSpectestModule(t *testing.T, store *wasm.Store) {
+	// Add the host module
+	spectest := &wasm.ModuleInstance{Name: "spectest", Exports: map[string]*wasm.ExportInstance{}}
+	store.ModuleInstances[spectest.Name] = spectest
+
 	var printV = func() {}
 	var printI32 = func(uint32) {}
 	var printF32 = func(float32) {}
@@ -205,7 +209,7 @@ func addSpectestModule(t *testing.T, store *wasm.Store) {
 	} {
 		fn, err := wasm.NewGoFunc(n, v)
 		require.NoError(t, err)
-		_, err = store.AddHostFunction("spectest", fn)
+		_, err = store.AddHostFunction(spectest, fn)
 		require.NoError(t, err, "AddHostFunction(%s)", n)
 	}
 
@@ -219,14 +223,14 @@ func addSpectestModule(t *testing.T, store *wasm.Store) {
 		{name: "global_f32", valueType: wasm.ValueTypeF32, value: uint64(uint32(0x44268000))},
 		{name: "global_f64", valueType: wasm.ValueTypeF64, value: uint64(0x4084d00000000000)},
 	} {
-		require.NoError(t, store.AddGlobal("spectest", g.name, g.value, g.valueType, false), "AddGlobal(%s)", g.name)
+		require.NoError(t, store.AddGlobal(spectest, g.name, g.value, g.valueType, false), "AddGlobal(%s)", g.name)
 	}
 
 	tableLimitMax := uint32(20)
-	require.NoError(t, store.AddTableInstance("spectest", "table", 10, &tableLimitMax))
+	require.NoError(t, store.AddTableInstance(spectest, "table", 10, &tableLimitMax))
 
 	memoryLimitMax := uint32(2)
-	require.NoError(t, store.AddMemoryInstance("spectest", "memory", 1, &memoryLimitMax))
+	require.NoError(t, store.AddMemoryInstance(spectest, "memory", 1, &memoryLimitMax))
 }
 
 func TestJIT(t *testing.T) {
@@ -241,7 +245,6 @@ func TestInterpreter(t *testing.T) {
 }
 
 func runTest(t *testing.T, newEngine func() wasm.Engine) {
-	ctx := context.Background()
 	files, err := testcases.ReadDir("testdata")
 	require.NoError(t, err)
 
@@ -307,7 +310,7 @@ func runTest(t *testing.T, newEngine func() wasm.Engine) {
 							if c.Action.Module != "" {
 								msg += " in module " + c.Action.Module
 							}
-							vals, types, err := callFunction(store, ctx, moduleName, c.Action.Field, args...)
+							vals, types, err := callFunction(store, moduleName, c.Action.Field, args...)
 							require.NoError(t, err, msg)
 							require.Equal(t, len(exps), len(vals), msg)
 							require.Equal(t, len(exps), len(types), msg)
@@ -367,7 +370,7 @@ func runTest(t *testing.T, newEngine func() wasm.Engine) {
 							if c.Action.Module != "" {
 								msg += " in module " + c.Action.Module
 							}
-							_, _, err := callFunction(store, ctx, moduleName, c.Action.Field, args...)
+							_, _, err := callFunction(store, moduleName, c.Action.Field, args...)
 							require.ErrorIs(t, err, c.expectedError(), msg)
 						default:
 							t.Fatalf("unsupported action type type: %v", c)
@@ -396,7 +399,7 @@ func runTest(t *testing.T, newEngine func() wasm.Engine) {
 							if c.Action.Module != "" {
 								msg += " in module " + c.Action.Module
 							}
-							_, _, err := callFunction(store, ctx, moduleName, c.Action.Field, args...)
+							_, _, err := callFunction(store, moduleName, c.Action.Field, args...)
 							require.ErrorIs(t, err, wasm.ErrRuntimeCallStackOverflow, msg)
 						default:
 							t.Fatalf("unsupported action type type: %v", c)
@@ -460,9 +463,10 @@ func requireValueEq(t *testing.T, actual, expected uint64, valType wasm.ValueTyp
 
 // callFunction is inlined here as the spectest needs to validate the signature was correct
 // TODO: This is likely already covered with unit tests!
-func callFunction(s *wasm.Store, ctx context.Context, moduleName, funcName string, params ...uint64) (results []uint64, resultTypes []wasm.ValueType, err error) {
+func callFunction(s *wasm.Store, moduleName, funcName string, params ...uint64) (results []uint64, resultTypes []wasm.ValueType, err error) {
+	ctx := s.ModuleContexts[moduleName]
 	var exp *wasm.ExportInstance
-	if exp, err = s.ModuleContexts[moduleName].Module.GetExport(funcName, wasm.ExportKindFunc); err != nil {
+	if exp, err = ctx.Module.GetExport(funcName, wasm.ExportKindFunc); err != nil {
 		return
 	}
 
@@ -472,8 +476,7 @@ func callFunction(s *wasm.Store, ctx context.Context, moduleName, funcName strin
 		return
 	}
 
-	hostCtx := s.ModuleContexts[moduleName].WithContext(ctx)
-	results, err = s.Engine.Call(hostCtx, f, params...)
+	results, err = s.Engine.Call(ctx, f, params...)
 	resultTypes = f.FunctionType.Type.Results
 	return
 }