Adds ModuleConfig.WithImport and WithImportModule (#444)

Before, complicated wasm could be hard to implement, particularly as it
might have cyclic imports. This change allows users to re-map imports to
untangle any cycles or to break up monolithic modules like "env".

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

config.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"io/fs"
 	"math"
+	"strings"
 
 	"github.com/tetratelabs/wazero/internal/wasm"
 	"github.com/tetratelabs/wazero/internal/wasm/interpreter"
@@ -143,6 +144,11 @@ type ModuleConfig struct {
 	preopens map[uint32]*wasm.FileEntry
 	// preopenPaths allow overwriting of existing paths.
 	preopenPaths map[string]uint32
+	// replacedImports holds the latest state of WithImport
+	// Note: Key is NUL delimited as import module and name can both include any UTF-8 characters.
+	replacedImports map[string][2]string
+	// replacedImportModules holds the latest state of WithImportModule
+	replacedImportModules map[string]string
 }
 
 func NewModuleConfig() *ModuleConfig {
@@ -170,6 +176,60 @@ func (c *ModuleConfig) WithName(name string) *ModuleConfig {
 	return c
 }
 
+// WithImport replaces a specific import module and name with a new one. This allows you to break up a monolithic
+// module imports, such as "env". This can also help reduce cyclic dependencies.
+//
+// For example, if a module was compiled with one module owning all imports:
+//	(import "js" "tbl" (table $tbl 4 funcref))
+//	(import "js" "increment" (func $increment (result i32)))
+//	(import "js" "decrement" (func $decrement (result i32)))
+//	(import "js" "wasm_increment" (func $wasm_increment (result i32)))
+//	(import "js" "wasm_decrement" (func $wasm_decrement (result i32)))
+//
+// Use this function to import "increment" and "decrement" from the module "go" and other imports from "wasm":
+//	config.WithImportModule("js", "wasm")
+//	config.WithImport("wasm", "increment", "go", "increment")
+//	config.WithImport("wasm", "decrement", "go", "decrement")
+//
+// Upon instantiation, imports resolve as if they were compiled like so:
+//	(import "wasm" "tbl" (table $tbl 4 funcref))
+//	(import "go" "increment" (func $increment (result i32)))
+//	(import "go" "decrement" (func $decrement (result i32)))
+//	(import "wasm" "wasm_increment" (func $wasm_increment (result i32)))
+//	(import "wasm" "wasm_decrement" (func $wasm_decrement (result i32)))
+//
+// Note: Any WithImport instructions happen in order, after any WithImportModule instructions.
+func (c *ModuleConfig) WithImport(oldModule, oldName, newModule, newName string) *ModuleConfig {
+	if c.replacedImports == nil {
+		c.replacedImports = map[string][2]string{}
+	}
+	var builder strings.Builder
+	builder.WriteString(oldModule)
+	builder.WriteByte(0) // delimit with NUL as module and name can be any UTF-8 characters.
+	builder.WriteString(oldName)
+	c.replacedImports[builder.String()] = [2]string{newModule, newName}
+	return c
+}
+
+// WithImportModule replaces every import with oldModule with newModule. This is helpful for modules who have
+// transitioned to a stable status since the underlying wasm was compiled.
+//
+// For example, if a module was compiled like below, with an old module for WASI:
+//	(import "wasi_unstable" "args_get" (func (param i32, i32) (result i32)))
+//
+// Use this function to update it to the current version:
+//	config.WithImportModule("wasi_unstable", wasi.ModuleSnapshotPreview1)
+//
+// See WithImport for a comprehensive example.
+// Note: Any WithImportModule instructions happen in order, before any WithImport instructions.
+func (c *ModuleConfig) WithImportModule(oldModule, newModule string) *ModuleConfig {
+	if c.replacedImportModules == nil {
+		c.replacedImportModules = map[string]string{}
+	}
+	c.replacedImportModules[oldModule] = newModule
+	return c
+}
+
 // WithStartFunctions configures the functions to call after the module is instantiated. Defaults to "_start".
 //
 // Note: If any function doesn't exist, it is skipped. However, all functions that do exist are called in order.
@@ -352,3 +412,53 @@ func (c *ModuleConfig) toSysContext() (sys *wasm.SysContext, err error) {
 
 	return wasm.NewSysContext(math.MaxUint32, c.args, environ, c.stdin, c.stdout, c.stderr, preopens)
 }
+
+func (c *ModuleConfig) replaceImports(module *wasm.Module) *wasm.Module {
+	if (c.replacedImportModules == nil && c.replacedImports == nil) || module.ImportSection == nil {
+		return module
+	}
+
+	changed := false
+
+	ret := *module // shallow copy
+	replacedImports := make([]*wasm.Import, len(module.ImportSection))
+	copy(replacedImports, module.ImportSection)
+
+	// First, replace any import.Module
+	for oldModule, newModule := range c.replacedImportModules {
+		for i, imp := range replacedImports {
+			if imp.Module == oldModule {
+				changed = true
+				cp := *imp // shallow copy
+				cp.Module = newModule
+				replacedImports[i] = &cp
+			} else {
+				replacedImports[i] = imp
+			}
+		}
+	}
+
+	// Now, replace any import.Module+import.Name
+	for oldImport, newImport := range c.replacedImports {
+		for i, imp := range replacedImports {
+			nulIdx := strings.IndexByte(oldImport, 0)
+			oldModule := oldImport[0:nulIdx]
+			oldName := oldImport[nulIdx+1:]
+			if imp.Module == oldModule && imp.Name == oldName {
+				changed = true
+				cp := *imp // shallow copy
+				cp.Module = newImport[0]
+				cp.Name = newImport[1]
+				replacedImports[i] = &cp
+			} else {
+				replacedImports[i] = imp
+			}
+		}
+	}
+
+	if !changed {
+		return module
+	}
+	ret.ImportSection = replacedImports
+	return &ret
+}

config_test.go

@@ -124,6 +124,386 @@ func TestRuntimeConfig_FeatureToggle(t *testing.T) {
 	}
 }
 
+func TestModuleConfig(t *testing.T) {
+	tests := []struct {
+		name     string
+		with     func(*ModuleConfig) *ModuleConfig
+		expected *ModuleConfig
+	}{
+		{
+			name: "WithName",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithName("wazero")
+			},
+			expected: &ModuleConfig{
+				name: "wazero",
+			},
+		},
+		{
+			name: "WithName - empty",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithName("")
+			},
+			expected: &ModuleConfig{},
+		},
+		{
+			name: "WithImport",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImport("env", "abort", "assemblyscript", "abort")
+			},
+			expected: &ModuleConfig{
+				replacedImports: map[string][2]string{"env\000abort": {"assemblyscript", "abort"}},
+			},
+		},
+		{
+			name: "WithImport - empty to non-empty - module",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImport("", "abort", "assemblyscript", "abort")
+			},
+			expected: &ModuleConfig{
+				replacedImports: map[string][2]string{"\000abort": {"assemblyscript", "abort"}},
+			},
+		},
+		{
+			name: "WithImport - non-empty to empty - module",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImport("env", "abort", "", "abort")
+			},
+			expected: &ModuleConfig{
+				replacedImports: map[string][2]string{"env\000abort": {"", "abort"}},
+			},
+		},
+		{
+			name: "WithImport - empty to non-empty - name",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImport("env", "", "assemblyscript", "abort")
+			},
+			expected: &ModuleConfig{
+				replacedImports: map[string][2]string{"env\000": {"assemblyscript", "abort"}},
+			},
+		},
+		{
+			name: "WithImport - non-empty to empty - name",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImport("env", "abort", "assemblyscript", "")
+			},
+			expected: &ModuleConfig{
+				replacedImports: map[string][2]string{"env\000abort": {"assemblyscript", ""}},
+			},
+		},
+		{
+			name: "WithImport - override",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImport("env", "abort", "assemblyscript", "abort").
+					WithImport("env", "abort", "go", "exit")
+			},
+			expected: &ModuleConfig{
+				replacedImports: map[string][2]string{"env\000abort": {"go", "exit"}},
+			},
+		},
+		{
+			name: "WithImport - twice",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImport("env", "abort", "assemblyscript", "abort").
+					WithImport("wasi_unstable", "proc_exit", "wasi_snapshot_preview1", "proc_exit")
+			},
+			expected: &ModuleConfig{
+				replacedImports: map[string][2]string{
+					"env\000abort":               {"assemblyscript", "abort"},
+					"wasi_unstable\000proc_exit": {"wasi_snapshot_preview1", "proc_exit"},
+				},
+			},
+		},
+		{
+			name: "WithImportModule",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImportModule("env", "assemblyscript")
+			},
+			expected: &ModuleConfig{
+				replacedImportModules: map[string]string{"env": "assemblyscript"},
+			},
+		},
+		{
+			name: "WithImportModule - empty to non-empty",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImportModule("", "assemblyscript")
+			},
+			expected: &ModuleConfig{
+				replacedImportModules: map[string]string{"": "assemblyscript"},
+			},
+		},
+		{
+			name: "WithImportModule - non-empty to empty",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImportModule("env", "")
+			},
+			expected: &ModuleConfig{
+				replacedImportModules: map[string]string{"env": ""},
+			},
+		},
+		{
+			name: "WithImportModule - override",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImportModule("env", "assemblyscript").
+					WithImportModule("env", "go")
+			},
+			expected: &ModuleConfig{
+				replacedImportModules: map[string]string{"env": "go"},
+			},
+		},
+		{
+			name: "WithImportModule - twice",
+			with: func(c *ModuleConfig) *ModuleConfig {
+				return c.WithImportModule("env", "go").
+					WithImportModule("wasi_unstable", "wasi_snapshot_preview1")
+			},
+			expected: &ModuleConfig{
+				replacedImportModules: map[string]string{
+					"env":           "go",
+					"wasi_unstable": "wasi_snapshot_preview1",
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			input := &ModuleConfig{}
+			rc := tc.with(input)
+			require.Equal(t, tc.expected, rc)
+		})
+	}
+}
+
+func TestModuleConfig_replaceImports(t *testing.T) {
+	tests := []struct {
+		name       string
+		config     *ModuleConfig
+		input      *wasm.Module
+		expected   *wasm.Module
+		expectSame bool
+	}{
+		{
+			name:       "no config, no imports",
+			config:     &ModuleConfig{},
+			input:      &wasm.Module{},
+			expected:   &wasm.Module{},
+			expectSame: true,
+		},
+		{
+			name:   "no config",
+			config: &ModuleConfig{},
+			input: &wasm.Module{
+				ImportSection: []*wasm.Import{
+					{
+						Module: "wasi_snapshot_preview1", Name: "args_sizes_get",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "wasi_snapshot_preview1", Name: "fd_write",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 2,
+					},
+				},
+			},
+			expectSame: true,
+		},
+		{
+			name: "replacedImportModules",
+			config: &ModuleConfig{
+				replacedImportModules: map[string]string{"wasi_unstable": "wasi_snapshot_preview1"},
+			},
+			input: &wasm.Module{
+				ImportSection: []*wasm.Import{
+					{
+						Module: "wasi_unstable", Name: "args_sizes_get",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "wasi_unstable", Name: "fd_write",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 2,
+					},
+				},
+			},
+			expected: &wasm.Module{
+				ImportSection: []*wasm.Import{
+					{
+						Module: "wasi_snapshot_preview1", Name: "args_sizes_get",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "wasi_snapshot_preview1", Name: "fd_write",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 2,
+					},
+				},
+			},
+		},
+		{
+			name: "replacedImportModules doesn't match",
+			config: &ModuleConfig{
+				replacedImportModules: map[string]string{"env": ""},
+			},
+			input: &wasm.Module{
+				ImportSection: []*wasm.Import{
+					{
+						Module: "wasi_snapshot_preview1", Name: "args_sizes_get",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "wasi_snapshot_preview1", Name: "fd_write",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 2,
+					},
+				},
+			},
+			expectSame: true,
+		},
+		{
+			name: "replacedImports",
+			config: &ModuleConfig{
+				replacedImports: map[string][2]string{"env\000abort": {"assemblyscript", "abort"}},
+			},
+			input: &wasm.Module{
+				ImportSection: []*wasm.Import{
+					{
+						Module: "env", Name: "abort",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "env", Name: "seed",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 2,
+					},
+				},
+			},
+			expected: &wasm.Module{
+				ImportSection: []*wasm.Import{
+					{
+						Module: "assemblyscript", Name: "abort",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "env", Name: "seed",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 2,
+					},
+				},
+			},
+		},
+		{
+			name: "replacedImports don't match",
+			config: &ModuleConfig{
+				replacedImports: map[string][2]string{"env\000abort": {"assemblyscript", "abort"}},
+			},
+			input: &wasm.Module{
+				ImportSection: []*wasm.Import{
+					{
+						Module: "wasi_snapshot_preview1", Name: "args_sizes_get",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "wasi_snapshot_preview1", Name: "fd_write",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 2,
+					},
+				},
+			},
+			expectSame: true,
+		},
+		{
+			name: "replacedImportModules and replacedImports",
+			config: &ModuleConfig{
+				replacedImportModules: map[string]string{"js": "wasm"},
+				replacedImports: map[string][2]string{
+					"wasm\000increment": {"go", "increment"},
+					"wasm\000decrement": {"go", "decrement"},
+				},
+			},
+			input: &wasm.Module{
+				ImportSection: []*wasm.Import{
+					{
+						Module: "js", Name: "tbl",
+						Type:      wasm.ExternTypeTable,
+						DescTable: &wasm.Table{Min: 4},
+					},
+					{
+						Module: "js", Name: "increment",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "js", Name: "decrement",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "js", Name: "wasm_increment",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "js", Name: "wasm_increment",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+				},
+			},
+			expected: &wasm.Module{
+				ImportSection: []*wasm.Import{
+					{
+						Module: "wasm", Name: "tbl",
+						Type:      wasm.ExternTypeTable,
+						DescTable: &wasm.Table{Min: 4},
+					},
+					{
+						Module: "go", Name: "increment",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "go", Name: "decrement",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "wasm", Name: "wasm_increment",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+					{
+						Module: "wasm", Name: "wasm_increment",
+						Type:     wasm.ExternTypeFunc,
+						DescFunc: 0,
+					},
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			actual := tc.config.replaceImports(tc.input)
+			if tc.expectSame {
+				require.Same(t, tc.input, actual)
+			} else {
+				require.NotSame(t, tc.input, actual)
+				require.Equal(t, tc.expected, actual)
+			}
+		})
+	}
+}
+
 func TestModuleConfig_toSysContext(t *testing.T) {
 	testFS := fstest.MapFS{}
 	testFS2 := fstest.MapFS{}

examples/file_system_test.go

@@ -42,19 +42,15 @@ func Test_Cat(t *testing.T) {
 	// Combine the above into our baseline config, overriding defaults (which discard stdout and have no file system).
 	config := wazero.NewModuleConfig().WithStdout(stdoutBuf).WithFS(rooted)
 
-	// Compile the `cat` module.
-	code, err := r.CompileModule(catWasm)
-	require.NoError(t, err)
-
 	// Instantiate WASI, which implements system I/O such as console output.
 	wm, err := wasi.InstantiateSnapshotPreview1(r)
 	require.NoError(t, err)
 	defer wm.Close()
 
-	// InstantiateModuleWithConfig by default runs the "_start" function which is what TinyGo compiles "main" to.
+	// InstantiateModuleFromCodeWithConfig runs the "_start" function which is what TinyGo compiles "main" to.
 	// * Set the program name (arg[0]) to "cat" and add args to write "cat.go" to stdout twice.
 	// * We use both "/cat.go" and "./cat.go" because WithFS by default maps the workdir "." to "/".
-	cat, err := r.InstantiateModuleWithConfig(code, config.WithArgs("cat", "/cat.go", "./cat.go"))
+	cat, err := r.InstantiateModuleFromCodeWithConfig(catWasm, config.WithArgs("cat", "/cat.go", "./cat.go"))
 	require.NoError(t, err)
 	defer cat.Close()
 

examples/host_func_test.go

@@ -60,10 +60,6 @@ func Test_hostFunc(t *testing.T) {
 	_, err := r.NewModuleBuilder("env").ExportFunction("get_random_bytes", getRandomBytes).Instantiate()
 	require.NoError(t, err)
 
-	// Compile the `hostFunc` module.
-	code, err := r.CompileModule(hostFuncWasm)
-	require.NoError(t, err)
-
 	// Configure stdout (console) to write to a buffer.
 	stdout := bytes.NewBuffer(nil)
 	config := wazero.NewModuleConfig().WithStdout(stdout)
@@ -73,8 +69,8 @@ func Test_hostFunc(t *testing.T) {
 	require.NoError(t, err)
 	defer wm.Close()
 
-	// InstantiateModuleWithConfig runs the "_start" function which is what TinyGo compiles "main" to.
-	module, err := r.InstantiateModuleWithConfig(code, config)
+	// InstantiateModuleFromCodeWithConfig runs the "_start" function which is what TinyGo compiles "main" to.
+	module, err := r.InstantiateModuleFromCodeWithConfig(hostFuncWasm, config)
 	require.NoError(t, err)
 	defer module.Close()
 

examples/replace_test.go

@@ -0,0 +1,42 @@
+package examples
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/api"
+)
+
+// Test_Replace shows how you can replace a module import when it doesn't match instantiated modules.
+func Test_Replace(t *testing.T) {
+	r := wazero.NewRuntime()
+
+	// Instantiate a function that closes the module under "assemblyscript.abort".
+	host, err := r.NewModuleBuilder("assemblyscript").
+		ExportFunction("abort", func(m api.Module, messageOffset, fileNameOffset, line, col uint32) {
+			_ = m.CloseWithExitCode(255)
+		}).Instantiate()
+	require.NoError(t, err)
+	defer host.Close()
+
+	// Compile code that needs the function "env.abort".
+	code, err := r.CompileModule([]byte(`(module
+	(import "env" "abort" (func $~lib/builtins/abort (param i32 i32 i32 i32)))
+
+	(export "abort" (func 0)) ;; exports the import for testing
+)`))
+	require.NoError(t, err)
+
+	// Instantiate the module, replacing the import "env.abort" with "assemblyscript.abort".
+	mod, err := r.InstantiateModuleWithConfig(code, wazero.NewModuleConfig().
+		WithName(t.Name()).
+		WithImport("env", "abort", "assemblyscript", "abort"))
+	require.NoError(t, err)
+	defer mod.Close()
+
+	// Since the above worked, the exported function closes the module.
+	_, err = mod.ExportedFunction("abort").Call(nil, 0, 0, 0, 0)
+	require.EqualError(t, err, `module "Test_Replace" closed with exit_code(255)`)
+}

examples/stdio_test.go

@@ -19,10 +19,6 @@ var stdioWasm []byte
 func Test_stdio(t *testing.T) {
 	r := wazero.NewRuntime()
 
-	// Compile the `stdioWasm` module.
-	code, err := r.CompileModule(stdioWasm)
-	require.NoError(t, err)
-
 	// Configure standard I/O (ex stdout) to write to buffers instead of no-op.
 	stdinBuf := bytes.NewBuffer([]byte("WASI\n"))
 	stdoutBuf := bytes.NewBuffer(nil)
@@ -34,8 +30,8 @@ func Test_stdio(t *testing.T) {
 	require.NoError(t, err)
 	defer wm.Close()
 
-	// InstantiateModuleWithConfig runs the "_start" function which is what TinyGo compiles "main" to.
-	module, err := r.InstantiateModuleWithConfig(code, config)
+	// InstantiateModuleFromCodeWithConfig runs the "_start" function which is what TinyGo compiles "main" to.
+	module, err := r.InstantiateModuleFromCodeWithConfig(stdioWasm, config)
 	require.NoError(t, err)
 	defer module.Close()
 

wasm.go

@@ -57,6 +57,16 @@ type Runtime interface {
 	// source multiple times, use CompileModule as InstantiateModule avoids redundant decoding and/or compilation.
 	InstantiateModuleFromCode(source []byte) (api.Module, error)
 
+	// InstantiateModuleFromCodeWithConfig is a convenience function that chains CompileModule to
+	// InstantiateModuleWithConfig.
+	//
+	// Ex. To only change the module name:
+	//	wasm, _ := wazero.NewRuntime().InstantiateModuleFromCodeWithConfig(source, wazero.NewModuleConfig().
+	//		WithName("wasm")
+	//	)
+	//	defer wasm.Close()
+	InstantiateModuleFromCodeWithConfig(source []byte, config *ModuleConfig) (api.Module, error)
+
 	// InstantiateModule instantiates the module namespace or errs if the configuration was invalid.
 	//
 	// Ex.
@@ -164,6 +174,15 @@ func (r *runtime) InstantiateModuleFromCode(source []byte) (api.Module, error) {
 	}
 }
 
+// InstantiateModuleFromCodeWithConfig implements Runtime.InstantiateModuleFromCodeWithConfig
+func (r *runtime) InstantiateModuleFromCodeWithConfig(source []byte, config *ModuleConfig) (api.Module, error) {
+	if code, err := r.CompileModule(source); err != nil {
+		return nil, err
+	} else {
+		return r.InstantiateModuleWithConfig(code, config)
+	}
+}
+
 // InstantiateModule implements Runtime.InstantiateModule
 func (r *runtime) InstantiateModule(code *CompiledCode) (mod api.Module, err error) {
 	return r.InstantiateModuleWithConfig(code, NewModuleConfig())
@@ -181,7 +200,9 @@ func (r *runtime) InstantiateModuleWithConfig(code *CompiledCode, config *Module
 		name = code.module.NameSection.ModuleName
 	}
 
-	mod, err = r.store.Instantiate(r.ctx, code.module, name, sys)
+	module := config.replaceImports(code.module)
+
+	mod, err = r.store.Instantiate(r.ctx, module, name, sys)
 	if err != nil {
 		return
 	}

wasm_test.go

@@ -288,10 +288,7 @@ func TestFunction_Context(t *testing.T) {
 			defer closer() // nolint
 
 			// Instantiate the module and get the export of the above hostFn
-			code, err := r.CompileModule(source)
-			require.NoError(t, err)
-
-			module, err := r.InstantiateModuleWithConfig(code, NewModuleConfig().WithName(t.Name()))
+			module, err := r.InstantiateModuleFromCodeWithConfig(source, NewModuleConfig().WithName(t.Name()))
 			require.NoError(t, err)
 			defer module.Close()