Adds ModuleBuilder.ExportMemory and ExportMemoryWithMax (#440)

This adds functions to configure memory with ModuleBuilder. This uses
two functions, ExportMemory and ExportMemoryWithMax, as working with
uint32 pointers is awkward.

Signed-off-by: Adrian Cole <adrian@tetrate.io>

builder.go

@@ -1,6 +1,8 @@
 package wazero
 
 import (
+	"fmt"
+
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/internal/wasm"
 )
@@ -67,6 +69,22 @@ type ModuleBuilder interface {
 	// ExportFunctions is a convenience that calls ExportFunction for each key/value in the provided map.
 	ExportFunctions(nameToGoFunc map[string]interface{}) ModuleBuilder
 
+	// ExportMemory adds linear memory, which a WebAssembly module can import and become available via api.Memory.
+	//
+	// * name - the name to export. Ex "memory" for wasi.ModuleSnapshotPreview1
+	// * minPages - the possibly zero initial size in pages (65536 bytes per page).
+	//
+	// Note: This is allowed to grow to RuntimeConfig.WithMemoryMaxPages (4GiB). To bound it, use ExportMemoryWithMax.
+	// Note: If a memory is already exported with the same name, this overwrites it.
+	// Note: Version 1.0 (20191205) of the WebAssembly spec allows at most one memory per module.
+	// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#memory-section%E2%91%A0
+	ExportMemory(name string, minPages uint32) ModuleBuilder
+
+	// ExportMemoryWithMax is like ExportMemory, but can prevent overuse of memory.
+	//
+	// Note: maxPages must be at least minPages and no larger than RuntimeConfig.WithMemoryMaxPages
+	ExportMemoryWithMax(name string, minPages, maxPages uint32) ModuleBuilder
+
 	// Build returns a module to instantiate, or returns an error if any of the configuration is invalid.
 	Build() (*CompiledCode, error)
 
@@ -81,6 +99,7 @@ type moduleBuilder struct {
 	r            *runtime
 	moduleName   string
 	nameToGoFunc map[string]interface{}
+	nameToMemory map[string]*wasm.Memory
 }
 
 // NewModuleBuilder implements Runtime.NewModuleBuilder
@@ -89,6 +108,7 @@ func (r *runtime) NewModuleBuilder(moduleName string) ModuleBuilder {
 		r:            r,
 		moduleName:   moduleName,
 		nameToGoFunc: map[string]interface{}{},
+		nameToMemory: map[string]*wasm.Memory{},
 	}
 }
 
@@ -106,10 +126,31 @@ func (b *moduleBuilder) ExportFunctions(nameToGoFunc map[string]interface{}) Mod
 	return b
 }
 
+// ExportMemory implements ModuleBuilder.ExportMemory
+func (b *moduleBuilder) ExportMemory(name string, minPages uint32) ModuleBuilder {
+	b.nameToMemory[name] = &wasm.Memory{Min: minPages, Max: b.r.memoryMaxPages}
+	return b
+}
+
+// ExportMemoryWithMax implements ModuleBuilder.ExportMemoryWithMax
+func (b *moduleBuilder) ExportMemoryWithMax(name string, minPages, maxPages uint32) ModuleBuilder {
+	b.nameToMemory[name] = &wasm.Memory{Min: minPages, Max: maxPages}
+	return b
+}
+
 // Build implements ModuleBuilder.Build
 func (b *moduleBuilder) Build() (*CompiledCode, error) {
+	// Verify the maximum limit here, so we don't have to pass it to wasm.NewHostModule
+	maxLimit := b.r.memoryMaxPages
+	for name, mem := range b.nameToMemory {
+		if mem.Max > maxLimit {
+			max := mem.Max
+			return nil, fmt.Errorf("memory[%s] max %d pages (%s) outside range of %d pages (%s)", name, max, wasm.PagesToUnitOfBytes(max), maxLimit, wasm.PagesToUnitOfBytes(maxLimit))
+		}
+	}
+
 	// TODO: we can use r.enabledFeatures to fail early on things like mutable globals
-	if module, err := wasm.NewHostModule(b.moduleName, b.nameToGoFunc); err != nil {
+	if module, err := wasm.NewHostModule(b.moduleName, b.nameToGoFunc, b.nameToMemory); err != nil {
 		return nil, err
 	} else {
 		return &CompiledCode{module: module}, nil

builder_test.go

@@ -1,6 +1,7 @@
 package wazero
 
 import (
+	"math"
 	"reflect"
 	"testing"
 
@@ -42,6 +43,18 @@ func TestNewModuleBuilder_Build(t *testing.T) {
 			},
 			expected: &wasm.Module{NameSection: &wasm.NameSection{ModuleName: "env"}},
 		},
+		{
+			name: "ExportMemory",
+			input: func(r Runtime) ModuleBuilder {
+				return r.NewModuleBuilder("").ExportMemory("memory", 1)
+			},
+			expected: &wasm.Module{
+				MemorySection: &wasm.Memory{Min: 1, Max: wasm.MemoryMaxPages},
+				ExportSection: map[string]*wasm.Export{
+					"memory": {Name: "memory", Type: wasm.ExternTypeMemory, Index: 0},
+				},
+			},
+		},
 		{
 			name: "ExportFunction",
 			input: func(r Runtime) ModuleBuilder {
@@ -164,8 +177,41 @@ func TestNewModuleBuilder_Build(t *testing.T) {
 	}
 }
 
-// TestNewModuleBuilder_InstantiateModule ensures Runtime.InstantiateModule is called on success.
-func TestNewModuleBuilder_InstantiateModule(t *testing.T) {
+// TestNewModuleBuilder_Build_Errors only covers a few scenarios to avoid duplicating tests in internal/wasm/host_test.go
+func TestNewModuleBuilder_Build_Errors(t *testing.T) {
+	tests := []struct {
+		name        string
+		input       func(Runtime) ModuleBuilder
+		expectedErr string
+	}{
+		{
+			name: "memory max > limit",
+			input: func(r Runtime) ModuleBuilder {
+				return r.NewModuleBuilder("").ExportMemory("memory", math.MaxUint32)
+			},
+			expectedErr: "memory[memory] min 4294967295 pages (3 Ti) > max 65536 pages (4 Gi)",
+		},
+		{
+			name: "memory min > limit",
+			input: func(r Runtime) ModuleBuilder {
+				return r.NewModuleBuilder("").ExportMemoryWithMax("memory", 1, math.MaxUint32)
+			},
+			expectedErr: "memory[memory] max 4294967295 pages (3 Ti) outside range of 65536 pages (4 Gi)",
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			_, e := tc.input(NewRuntime()).Build()
+			require.EqualError(t, e, tc.expectedErr)
+		})
+	}
+}
+
+// TestNewModuleBuilder_Instantiate ensures Runtime.InstantiateModule is called on success.
+func TestNewModuleBuilder_Instantiate(t *testing.T) {
 	r := NewRuntime()
 	m, err := r.NewModuleBuilder("env").Instantiate()
 	require.NoError(t, err)
@@ -174,8 +220,8 @@ func TestNewModuleBuilder_InstantiateModule(t *testing.T) {
 	require.Equal(t, r.(*runtime).store.Module("env"), m)
 }
 
-// TestNewModuleBuilder_InstantiateModule_Errors ensures errors propagate from Runtime.InstantiateModule
-func TestNewModuleBuilder_InstantiateModule_Errors(t *testing.T) {
+// TestNewModuleBuilder_Instantiate_Errors ensures errors propagate from Runtime.InstantiateModule
+func TestNewModuleBuilder_Instantiate_Errors(t *testing.T) {
 	r := NewRuntime()
 	_, err := r.NewModuleBuilder("env").Instantiate()
 	require.NoError(t, err)

internal/wasm/host.go

@@ -4,40 +4,62 @@ import (
 	"fmt"
 	"reflect"
 	"sort"
+	"strings"
 
 	"github.com/tetratelabs/wazero/internal/wasmdebug"
 )
 
 // NewHostModule is defined internally for use in WASI tests and to keep the code size in the root directory small.
-func NewHostModule(moduleName string, nameToGoFunc map[string]interface{}) (*Module, error) {
-	hostFunctionCount := uint32(len(nameToGoFunc))
-	if hostFunctionCount == 0 {
+func NewHostModule(moduleName string, nameToGoFunc map[string]interface{}, nameToMemory map[string]*Memory) (m *Module, err error) {
 	if moduleName != "" {
-			return &Module{NameSection: &NameSection{ModuleName: moduleName}}, nil
+		m = &Module{NameSection: &NameSection{ModuleName: moduleName}}
 	} else {
-			return &Module{}, nil
+		m = &Module{}
+	}
+
+	funcCount := uint32(len(nameToGoFunc))
+	memoryCount := uint32(len(nameToMemory))
+	exportCount := funcCount + memoryCount
+	if exportCount > 0 {
+		m.ExportSection = make(map[string]*Export, exportCount)
+	}
+
+	if funcCount > 0 {
+		if err = addFuncs(m, nameToGoFunc); err != nil {
+			return
 		}
 	}
 
-	m := &Module{
-		NameSection:         &NameSection{ModuleName: moduleName, FunctionNames: make([]*NameAssoc, 0, hostFunctionCount)},
-		HostFunctionSection: make([]*reflect.Value, 0, hostFunctionCount),
-		ExportSection:       make(map[string]*Export, hostFunctionCount),
+	if memoryCount > 0 {
+		if err = addMemory(m, nameToMemory); err != nil {
+			return
 		}
+	}
+	return
+}
 
-	// Ensure insertion order is consistent
-	names := make([]string, 0, hostFunctionCount)
+func addFuncs(m *Module, nameToGoFunc map[string]interface{}) error {
+	funcCount := uint32(len(nameToGoFunc))
+	funcNames := make([]string, 0, funcCount)
+	if m.NameSection == nil {
+		m.NameSection = &NameSection{}
+	}
+	m.NameSection.FunctionNames = make([]*NameAssoc, 0, funcCount)
+	m.FunctionSection = make([]Index, 0, funcCount)
+	m.HostFunctionSection = make([]*reflect.Value, 0, funcCount)
+
+	// Sort names for consistent iteration
 	for k := range nameToGoFunc {
-		names = append(names, k)
+		funcNames = append(funcNames, k)
 	}
-	sort.Strings(names)
+	sort.Strings(funcNames)
 
-	for idx := Index(0); idx < hostFunctionCount; idx++ {
-		name := names[idx]
+	for idx := Index(0); idx < funcCount; idx++ {
+		name := funcNames[idx]
 		fn := reflect.ValueOf(nameToGoFunc[name])
 		_, functionType, _, err := getFunctionType(&fn, false)
 		if err != nil {
-			return nil, fmt.Errorf("func[%s] %w", name, err)
+			return fmt.Errorf("func[%s] %w", name, err)
 		}
 
 		m.FunctionSection = append(m.FunctionSection, m.maybeAddType(functionType))
@@ -45,7 +67,38 @@ func NewHostModule(moduleName string, nameToGoFunc map[string]interface{}) (*Mod
 		m.ExportSection[name] = &Export{Type: ExternTypeFunc, Name: name, Index: idx}
 		m.NameSection.FunctionNames = append(m.NameSection.FunctionNames, &NameAssoc{Index: idx, Name: name})
 	}
-	return m, nil
+	return nil
+}
+
+func addMemory(m *Module, nameToMemory map[string]*Memory) error {
+	memoryCount := uint32(len(nameToMemory))
+
+	// Only one memory can be defined or imported
+	if memoryCount > 1 {
+		memoryNames := make([]string, 0, memoryCount)
+		for k := range nameToMemory {
+			memoryNames = append(memoryNames, k)
+		}
+		sort.Strings(memoryNames) // For consistent error messages
+		return fmt.Errorf("only one memory is allowed, but configured: %s", strings.Join(memoryNames, ", "))
+	}
+
+	// Find the memory name to export.
+	var name string
+	for k, v := range nameToMemory {
+		name = k
+		if v.Min > v.Max {
+			return fmt.Errorf("memory[%s] min %d pages (%s) > max %d pages (%s)", name, v.Min, PagesToUnitOfBytes(v.Min), v.Max, PagesToUnitOfBytes(v.Max))
+		}
+		m.MemorySection = v
+	}
+
+	if e, ok := m.ExportSection[name]; ok { // Exports cannot collide on names, regardless of type.
+		return fmt.Errorf("memory[%s] exports the same name as a %s", name, ExternTypeName(e.Type))
+	}
+
+	m.ExportSection[name] = &Export{Type: ExternTypeMemory, Name: name, Index: 0}
+	return nil
 }
 
 func (m *Module) maybeAddType(ft *FunctionType) Index {
@@ -60,36 +113,6 @@ func (m *Module) maybeAddType(ft *FunctionType) Index {
 	return result
 }
 
-func (m *Module) validateHostFunctions() error {
-	functionCount := m.SectionElementCount(SectionIDFunction)
-	hostFunctionCount := m.SectionElementCount(SectionIDHostFunction)
-	if functionCount == 0 && hostFunctionCount == 0 {
-		return nil
-	}
-
-	typeCount := m.SectionElementCount(SectionIDType)
-	if hostFunctionCount != functionCount {
-		return fmt.Errorf("host function count (%d) != function count (%d)", hostFunctionCount, functionCount)
-	}
-
-	for idx, typeIndex := range m.FunctionSection {
-		_, ft, _, err := getFunctionType(m.HostFunctionSection[idx], false)
-		if err != nil {
-			return fmt.Errorf("%s is not a valid go func: %w", m.funcDesc(SectionIDHostFunction, Index(idx)), err)
-		}
-
-		if typeIndex >= typeCount {
-			return fmt.Errorf("%s type section index out of range: %d", m.funcDesc(SectionIDHostFunction, Index(idx)), typeIndex)
-		}
-
-		t := m.TypeSection[typeIndex]
-		if !t.EqualsSignature(ft.Params, ft.Results) {
-			return fmt.Errorf("%s signature doesn't match type section: %s != %s", m.funcDesc(SectionIDHostFunction, Index(idx)), ft, t)
-		}
-	}
-	return nil
-}
-
 func (m *Module) buildHostFunctions(moduleName string) (functions []*FunctionInstance) {
 	// ModuleBuilder has no imports, which means the FunctionSection index is the same as the position in the function
 	// index namespace. Also, it ensures every function has a name. That's why there is less error checking here.

internal/wasm/host_test.go

@@ -13,6 +13,10 @@ import (
 type wasiAPI struct {
 }
 
+func ArgsSizesGet(ctx api.Module, resultArgc, resultArgvBufSize uint32) uint32 {
+	return 0
+}
+
 func (a *wasiAPI) ArgsSizesGet(ctx api.Module, resultArgc, resultArgvBufSize uint32) uint32 {
 	return 0
 }
@@ -32,7 +36,8 @@ func TestNewHostModule(t *testing.T) {
 
 	tests := []struct {
 		name, moduleName string
-		goFuncs          map[string]interface{}
+		nameToGoFunc     map[string]interface{}
+		nameToMemory     map[string]*Memory
 		expected         *Module
 	}{
 		{
@@ -44,10 +49,20 @@ func TestNewHostModule(t *testing.T) {
 			moduleName: "test",
 			expected:   &Module{NameSection: &NameSection{ModuleName: "test"}},
 		},
+		{
+			name:         "memory",
+			nameToMemory: map[string]*Memory{"memory": {1, 2}},
+			expected: &Module{
+				MemorySection: &Memory{Min: 1, Max: 2},
+				ExportSection: map[string]*Export{
+					"memory": {Name: "memory", Type: ExternTypeMemory, Index: 0},
+				},
+			},
+		},
 		{
 			name:       "two struct funcs",
 			moduleName: "wasi_snapshot_preview1",
-			goFuncs: map[string]interface{}{
+			nameToGoFunc: map[string]interface{}{
 				functionArgsSizesGet: a.ArgsSizesGet,
 				functionFdWrite:      a.FdWrite,
 			},
@@ -71,13 +86,41 @@ func TestNewHostModule(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:       "one of each",
+			moduleName: "env",
+			nameToGoFunc: map[string]interface{}{
+				functionArgsSizesGet: a.ArgsSizesGet,
+			},
+			nameToMemory: map[string]*Memory{
+				"memory": {1, 1},
+			},
+			expected: &Module{
+				TypeSection: []*FunctionType{
+					{Params: []ValueType{i32, i32}, Results: []ValueType{i32}},
+				},
+				FunctionSection:     []Index{0},
+				HostFunctionSection: []*reflect.Value{&fnArgsSizesGet},
+				ExportSection: map[string]*Export{
+					"args_sizes_get": {Name: "args_sizes_get", Type: ExternTypeFunc, Index: 0},
+					"memory":         {Name: "memory", Type: ExternTypeMemory, Index: 0},
+				},
+				MemorySection: &Memory{Min: 1, Max: 1},
+				NameSection: &NameSection{
+					ModuleName: "env",
+					FunctionNames: NameMap{
+						{Index: 0, Name: "args_sizes_get"},
+					},
+				},
+			},
+		},
 	}
 
 	for _, tt := range tests {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			m, e := NewHostModule(tc.moduleName, tc.goFuncs)
+			m, e := NewHostModule(tc.moduleName, tc.nameToGoFunc, tc.nameToMemory)
 			require.NoError(t, e)
 			requireHostModuleEquals(t, tc.expected, m)
 		})
@@ -107,110 +150,41 @@ func requireHostModuleEquals(t *testing.T, expected, actual *Module) {
 }
 
 func TestNewHostModule_Errors(t *testing.T) {
-	t.Run("Adds export name to error message", func(t *testing.T) {
-		_, err := NewHostModule("test", map[string]interface{}{"fn": "hello"})
-		require.EqualError(t, err, "func[fn] kind != func: string")
-	})
-}
-
-func TestModule_validateHostFunctions(t *testing.T) {
-	notFn := reflect.ValueOf(t)
-	fn := reflect.ValueOf(func(api.Module) {})
-
-	t.Run("ok", func(t *testing.T) {
-		m := Module{
-			TypeSection:         []*FunctionType{{}},
-			FunctionSection:     []uint32{0},
-			HostFunctionSection: []*reflect.Value{&fn},
-		}
-		err := m.validateHostFunctions()
-		require.NoError(t, err)
-	})
-	t.Run("function, but no host function", func(t *testing.T) {
-		m := Module{
-			TypeSection:         []*FunctionType{{}},
-			FunctionSection:     []Index{0},
-			HostFunctionSection: nil,
-		}
-		err := m.validateHostFunctions()
-		require.Error(t, err)
-		require.EqualError(t, err, "host function count (0) != function count (1)")
-	})
-	t.Run("function out of range of host functions", func(t *testing.T) {
-		m := Module{
-			TypeSection:         []*FunctionType{{}},
-			FunctionSection:     []Index{1},
-			HostFunctionSection: []*reflect.Value{&fn},
-		}
-		err := m.validateHostFunctions()
-		require.Error(t, err)
-		require.EqualError(t, err, "host_function[0] type section index out of range: 1")
-	})
-	t.Run("mismatch params", func(t *testing.T) {
-		m := Module{
-			TypeSection:         []*FunctionType{{Params: []ValueType{ValueTypeF32}}},
-			FunctionSection:     []Index{0},
-			HostFunctionSection: []*reflect.Value{&fn},
-		}
-		err := m.validateHostFunctions()
-		require.Error(t, err)
-		require.EqualError(t, err, "host_function[0] signature doesn't match type section: v_v != f32_v")
-	})
-	t.Run("mismatch results", func(t *testing.T) {
-		m := Module{
-			TypeSection:         []*FunctionType{{Results: []ValueType{ValueTypeF32}}},
-			FunctionSection:     []Index{0},
-			HostFunctionSection: []*reflect.Value{&fn},
-		}
-		err := m.validateHostFunctions()
-		require.Error(t, err)
-		require.EqualError(t, err, "host_function[0] signature doesn't match type section: v_v != v_f32")
-	})
-	t.Run("not a function", func(t *testing.T) {
-		m := Module{
-			TypeSection:         []*FunctionType{{}},
-			FunctionSection:     []Index{0},
-			HostFunctionSection: []*reflect.Value{&notFn},
-		}
-		err := m.validateHostFunctions()
-		require.Error(t, err)
-		require.EqualError(t, err, "host_function[0] is not a valid go func: kind != func: ptr")
-	})
-	t.Run("not a function - exported", func(t *testing.T) {
-		m := Module{
-			TypeSection:         []*FunctionType{{}},
-			FunctionSection:     []Index{0},
-			HostFunctionSection: []*reflect.Value{&notFn},
-			ExportSection:       map[string]*Export{"f1": {Name: "f1", Type: ExternTypeFunc, Index: 0}},
-		}
-		err := m.validateHostFunctions()
-		require.Error(t, err)
-		require.EqualError(t, err, `host_function[0] export["f1"] is not a valid go func: kind != func: ptr`)
-	})
-	t.Run("not a function  - exported after import", func(t *testing.T) {
-		m := Module{
-			TypeSection:         []*FunctionType{{}},
-			ImportSection:       []*Import{{Type: ExternTypeFunc}},
-			FunctionSection:     []Index{1},
-			HostFunctionSection: []*reflect.Value{&notFn},
-			ExportSection:       map[string]*Export{"f1": {Name: "f1", Type: ExternTypeFunc, Index: 1}},
-		}
-		err := m.validateHostFunctions()
-		require.Error(t, err)
-		require.EqualError(t, err, `host_function[0] export["f1"] is not a valid go func: kind != func: ptr`)
-	})
-	t.Run("not a function - exported twice", func(t *testing.T) {
-		m := Module{
-			TypeSection:         []*FunctionType{{}},
-			FunctionSection:     []Index{0},
-			HostFunctionSection: []*reflect.Value{&notFn},
-			ExportSection: map[string]*Export{
-				"f1": {Name: "f1", Type: ExternTypeFunc, Index: 0},
-				"f2": {Name: "f2", Type: ExternTypeFunc, Index: 0},
+	tests := []struct {
+		name, moduleName string
+		nameToGoFunc     map[string]interface{}
+		nameToMemory     map[string]*Memory
+		expectedErr      string
+	}{
+		{
+			name:         "not a function",
+			nameToGoFunc: map[string]interface{}{"fn": t},
+			expectedErr:  "func[fn] kind != func: ptr",
+		},
+		{
+			name:         "memory collides on func name",
+			nameToGoFunc: map[string]interface{}{"fn": ArgsSizesGet},
+			nameToMemory: map[string]*Memory{"fn": {1, 1}},
+			expectedErr:  "memory[fn] exports the same name as a func",
+		},
+		{
+			name:         "multiple memories",
+			nameToMemory: map[string]*Memory{"memory": {1, 1}, "mem": {2, 2}},
+			expectedErr:  "only one memory is allowed, but configured: mem, memory",
+		},
+		{
+			name:         "memory max < min",
+			nameToMemory: map[string]*Memory{"memory": {1, 0}},
+			expectedErr:  "memory[memory] min 1 pages (64 Ki) > max 0 pages (0 Ki)",
 		},
 	}
-		err := m.validateHostFunctions()
-		require.Error(t, err)
-		require.EqualError(t, err, `host_function[0] export["f1","f2"] is not a valid go func: kind != func: ptr`)
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			_, e := NewHostModule(tc.moduleName, tc.nameToGoFunc, tc.nameToMemory)
+			require.EqualError(t, e, tc.expectedErr)
 		})
+	}
 }

internal/wasm/jit/engine_test.go

@@ -408,7 +408,7 @@ func TestJIT_SliceAllocatedOnHeap(t *testing.T) {
 		// Trigger relocation of goroutine stack because at this point we have the majority of
 		// goroutine stack unused after recursive call.
 		runtime.GC()
-	}})
+	}}, map[string]*wasm.Memory{})
 	require.NoError(t, err)
 
 	_, err = store.Instantiate(context.Background(), hm, hostModuleName, nil)

internal/wasm/module.go

@@ -231,11 +231,7 @@ func (m *Module) Validate(enabledFeatures Features) error {
 		if err = m.validateFunctions(enabledFeatures, functions, globals, memory, table, MaximumFunctionIndex); err != nil {
 			return err
 		}
-	} else {
-		if err = m.validateHostFunctions(); err != nil {
-			return err
-		}
-	}
+	} // No need to validate host functions as NewHostModule validates
 
 	if _, err = m.validateTable(); err != nil {
 		return err

internal/wasm/store_test.go

@@ -90,7 +90,7 @@ func TestModuleInstance_Memory(t *testing.T) {
 
 func TestStore_Instantiate(t *testing.T) {
 	s := newStore()
-	m, err := NewHostModule("", map[string]interface{}{"fn": func(api.Module) {}})
+	m, err := NewHostModule("", map[string]interface{}{"fn": func(api.Module) {}}, map[string]*Memory{})
 	require.NoError(t, err)
 
 	type key string
@@ -120,7 +120,7 @@ func TestStore_CloseModule(t *testing.T) {
 		{
 			name: "Module imports HostModule",
 			initializer: func(t *testing.T, s *Store) {
-				m, err := NewHostModule(importedModuleName, map[string]interface{}{"fn": func(api.Module) {}})
+				m, err := NewHostModule(importedModuleName, map[string]interface{}{"fn": func(api.Module) {}}, map[string]*Memory{})
 				require.NoError(t, err)
 				_, err = s.Instantiate(context.Background(), m, importedModuleName, nil)
 				require.NoError(t, err)
@@ -177,7 +177,7 @@ func TestStore_CloseModule(t *testing.T) {
 func TestStore_hammer(t *testing.T) {
 	const importedModuleName = "imported"
 
-	m, err := NewHostModule(importedModuleName, map[string]interface{}{"fn": func(api.Module) {}})
+	m, err := NewHostModule(importedModuleName, map[string]interface{}{"fn": func(api.Module) {}}, map[string]*Memory{})
 	require.NoError(t, err)
 
 	s := newStore()
@@ -227,7 +227,7 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 	const importedModuleName = "imported"
 	const importingModuleName = "test"
 
-	m, err := NewHostModule(importedModuleName, map[string]interface{}{"fn": func(api.Module) {}})
+	m, err := NewHostModule(importedModuleName, map[string]interface{}{"fn": func(api.Module) {}}, map[string]*Memory{})
 	require.NoError(t, err)
 
 	t.Run("Fails if module name already in use", func(t *testing.T) {
@@ -312,7 +312,7 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 }
 
 func TestModuleContext_ExportedFunction(t *testing.T) {
-	host, err := NewHostModule("host", map[string]interface{}{"host_fn": func(api.Module) {}})
+	host, err := NewHostModule("host", map[string]interface{}{"host_fn": func(api.Module) {}}, map[string]*Memory{})
 	require.NoError(t, err)
 
 	s := newStore()
@@ -347,9 +347,7 @@ func TestFunctionInstance_Call(t *testing.T) {
 	functionName := "fn"
 
 	// This is a fake engine, so we don't capture inside the function body.
-	m, err := NewHostModule("host",
-		map[string]interface{}{functionName: func(api.Module) {}},
-	)
+	m, err := NewHostModule("host", map[string]interface{}{functionName: func(api.Module) {}}, map[string]*Memory{})
 	require.NoError(t, err)
 
 	// Add the host module

wasm_test.go

@@ -121,17 +121,22 @@ func TestRuntime_DecodeModule_Errors(t *testing.T) {
 // TestModule_Memory only covers a couple cases to avoid duplication of internal/wasm/runtime_test.go
 func TestModule_Memory(t *testing.T) {
 	tests := []struct {
-		name, wat   string
+		name        string
+		builder     func(Runtime) ModuleBuilder
 		expected    bool
 		expectedLen uint32
 	}{
 		{
 			name: "no memory",
-			wat:  `(module)`,
+			builder: func(r Runtime) ModuleBuilder {
+				return r.NewModuleBuilder(t.Name())
+			},
 		},
 		{
 			name: "memory exported, one page",
-			wat:         `(module (memory $mem 1) (export "memory" (memory $mem)))`,
+			builder: func(r Runtime) ModuleBuilder {
+				return r.NewModuleBuilder(t.Name()).ExportMemory("memory", 1)
+			},
 			expected:    true,
 			expectedLen: 65536,
 		},
@@ -142,12 +147,10 @@ func TestModule_Memory(t *testing.T) {
 
 		r := NewRuntime()
 		t.Run(tc.name, func(t *testing.T) {
-			code, err := r.CompileModule([]byte(tc.wat))
-			require.NoError(t, err)
-
-			// Instantiate the module and get the export of the above hostFn
-			module, err := r.InstantiateModule(code)
+			// Instantiate the module and get the export of the above memory
+			module, err := tc.builder(r).Instantiate()
 			require.NoError(t, err)
+			defer module.Close()
 
 			mem := module.ExportedMemory("memory")
 			if tc.expected {
@@ -222,6 +225,7 @@ func TestModule_Global(t *testing.T) {
 			// Instantiate the module and get the export of the above global
 			module, err := r.InstantiateModule(&CompiledCode{module: tc.module})
 			require.NoError(t, err)
+			defer module.Close()
 
 			global := module.ExportedGlobal("global")
 			if !tc.expected {
@@ -296,7 +300,7 @@ func TestFunction_Context(t *testing.T) {
 	}
 }
 
-func TestRuntime_NewModule_UsesStoreContext(t *testing.T) {
+func TestRuntime_NewModule_UsesConfiguredContext(t *testing.T) {
 	type key string
 	runtimeCtx := context.WithValue(context.Background(), key("wa"), "zero")
 	config := NewRuntimeConfig().WithContext(runtimeCtx)
@@ -309,8 +313,9 @@ func TestRuntime_NewModule_UsesStoreContext(t *testing.T) {
 		require.Equal(t, runtimeCtx, ctx.Context())
 	}
 
-	_, err := r.NewModuleBuilder("env").ExportFunction("start", start).Instantiate()
+	env, err := r.NewModuleBuilder("env").ExportFunction("start", start).Instantiate()
 	require.NoError(t, err)
+	defer env.Close()
 
 	code, err := r.CompileModule([]byte(`(module $runtime_test.go
 	(import "env" "start" (func $start))
@@ -319,8 +324,10 @@ func TestRuntime_NewModule_UsesStoreContext(t *testing.T) {
 	require.NoError(t, err)
 
 	// Instantiate the module, which calls the start function. This will fail if the context wasn't as intended.
-	_, err = r.InstantiateModule(code)
+	m, err := r.InstantiateModule(code)
 	require.NoError(t, err)
+	defer m.Close()
+
 	require.True(t, calledStart)
 }
 
@@ -378,11 +385,15 @@ func TestInstantiateModuleWithConfig_WithName(t *testing.T) {
 	internal := r.(*runtime).store
 	m1, err := r.InstantiateModuleWithConfig(base, NewModuleConfig().WithName("1"))
 	require.NoError(t, err)
+	defer m1.Close()
+
 	require.Nil(t, internal.Module("0"))
 	require.Equal(t, internal.Module("1"), m1)
 
 	m2, err := r.InstantiateModuleWithConfig(base, NewModuleConfig().WithName("2"))
 	require.NoError(t, err)
+	defer m2.Close()
+
 	require.Nil(t, internal.Module("0"))
 	require.Equal(t, internal.Module("2"), m2)
 }