Reduces exports and prepares ReleaseModuleInstance for testing (#327)

This removes more exports and adds a toehold test for
`ReleaseModuleInstance` so that it can later be completed for both
wasm and host modules.

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

internal/wasi/wasi_test.go

@@ -54,14 +54,13 @@ func TestSnapshotPreview1_ArgsGet(t *testing.T) {
 		'?', // stopped after encoding
 	}
 
-	mod, fn := instantiateModule(t, FunctionArgsGet, ImportArgsGet, moduleName, args)
-	mem := mod.Memory("memory")
+	mem, fn := instantiateModule(t, FunctionArgsGet, ImportArgsGet, moduleName, args)
 
 	t.Run("SnapshotPreview1.ArgsGet", func(t *testing.T) {
 		maskMemory(t, mem, len(expectedMemory))
 
 		// Invoke ArgsGet directly and check the memory side effects.
-		errno := NewAPI(args).ArgsGet(mod.Instance.Ctx, argv, argvBuf)
+		errno := NewAPI(args).ArgsGet(ctx(mem), argv, argvBuf)
 		require.Equal(t, wasi.ErrnoSuccess, errno)
 
 		actual, ok := mem.Read(0, uint32(len(expectedMemory)))
@@ -82,12 +81,17 @@ func TestSnapshotPreview1_ArgsGet(t *testing.T) {
 	})
 }
 
+func ctx(mem publicwasm.Memory) *wasm.ModuleContext {
+	ctx := (&wasm.ModuleContext{}).WithMemory(mem.(*wasm.MemoryInstance))
+	return ctx
+}
+
 func TestSnapshotPreview1_ArgsGet_Errors(t *testing.T) {
 	args, err := Args("a", "bc")
 	require.NoError(t, err)
-	mod, fn := instantiateModule(t, FunctionArgsGet, ImportArgsGet, moduleName, args)
+	mem, fn := instantiateModule(t, FunctionArgsGet, ImportArgsGet, moduleName, args)
 
-	memorySize := mod.Memory("memory").Size()
+	memorySize := mem.Size()
 	validAddress := uint32(0) // arbitrary valid address as arguments to args_get. We chose 0 here.
 
 	tests := []struct {
@@ -143,15 +147,13 @@ func TestSnapshotPreview1_ArgsSizesGet(t *testing.T) {
 		'?', // stopped after encoding
 	}
 
-	mod, fn := instantiateModule(t, FunctionArgsSizesGet, ImportArgsSizesGet, moduleName, args)
-
-	mem := mod.Memory("memory")
+	mem, fn := instantiateModule(t, FunctionArgsSizesGet, ImportArgsSizesGet, moduleName, args)
 
 	t.Run("SnapshotPreview1.ArgsSizesGet", func(t *testing.T) {
 		maskMemory(t, mem, len(expectedMemory))
 
 		// Invoke ArgsSizesGet directly and check the memory side effects.
-		errno := NewAPI(args).ArgsSizesGet(mod.Instance.Ctx, resultArgc, resultArgvBufSize)
+		errno := NewAPI(args).ArgsSizesGet(ctx(mem), resultArgc, resultArgvBufSize)
 		require.Equal(t, wasi.ErrnoSuccess, errno)
 
 		actual, ok := mem.Read(0, uint32(len(expectedMemory)))
@@ -176,9 +178,9 @@ func TestSnapshotPreview1_ArgsSizesGet_Errors(t *testing.T) {
 	args, err := Args("a", "bc")
 	require.NoError(t, err)
 
-	mod, fn := instantiateModule(t, FunctionArgsSizesGet, ImportArgsSizesGet, moduleName, args)
+	mem, fn := instantiateModule(t, FunctionArgsSizesGet, ImportArgsSizesGet, moduleName, args)
 
-	memorySize := mod.Memory("memory").Size()
+	memorySize := mem.Size()
 	validAddress := uint32(0) // arbitrary valid address as arguments to args_sizes_get. We chose 0 here.
 
 	tests := []struct {
@@ -270,14 +272,13 @@ func TestSnapshotPreview1_EnvironGet(t *testing.T) {
 		'?', // stopped after encoding
 	}
 
-	mod, fn := instantiateModule(t, FunctionEnvironGet, ImportEnvironGet, moduleName, envOpt)
-	mem := mod.Memory("memory")
+	mem, fn := instantiateModule(t, FunctionEnvironGet, ImportEnvironGet, moduleName, envOpt)
 
 	t.Run("SnapshotPreview1.EnvironGet", func(t *testing.T) {
 		maskMemory(t, mem, len(expectedMemory))
 
 		// Invoke EnvironGet directly and check the memory side effects.
-		errno := NewAPI(envOpt).EnvironGet(mod.Instance.Ctx, resultEnviron, resultEnvironBuf)
+		errno := NewAPI(envOpt).EnvironGet(ctx(mem), resultEnviron, resultEnvironBuf)
 		require.Equal(t, wasi.ErrnoSuccess, errno)
 
 		actual, ok := mem.Read(0, uint32(len(expectedMemory)))
@@ -302,9 +303,9 @@ func TestSnapshotPreview1_EnvironGet_Errors(t *testing.T) {
 	envOpt, err := Environ("a=bc", "b=cd")
 	require.NoError(t, err)
 
-	mod, fn := instantiateModule(t, FunctionEnvironGet, ImportEnvironGet, moduleName, envOpt)
+	mem, fn := instantiateModule(t, FunctionEnvironGet, ImportEnvironGet, moduleName, envOpt)
 
-	memorySize := mod.Memory("memory").Size()
+	memorySize := mem.Size()
 	validAddress := uint32(0) // arbitrary valid address as arguments to environ_get. We chose 0 here.
 
 	tests := []struct {
@@ -360,14 +361,13 @@ func TestSnapshotPreview1_EnvironSizesGet(t *testing.T) {
 		'?', // stopped after encoding
 	}
 
-	mod, fn := instantiateModule(t, FunctionEnvironSizesGet, ImportEnvironSizesGet, moduleName, envOpt)
-	mem := mod.Memory("memory")
+	mem, fn := instantiateModule(t, FunctionEnvironSizesGet, ImportEnvironSizesGet, moduleName, envOpt)
 
 	t.Run("SnapshotPreview1.EnvironSizesGet", func(t *testing.T) {
 		maskMemory(t, mem, len(expectedMemory))
 
 		// Invoke EnvironSizesGet directly and check the memory side effects.
-		errno := NewAPI(envOpt).EnvironSizesGet(mod.Instance.Ctx, resultEnvironc, resultEnvironBufSize)
+		errno := NewAPI(envOpt).EnvironSizesGet(ctx(mem), resultEnvironc, resultEnvironBufSize)
 		require.Equal(t, wasi.ErrnoSuccess, errno)
 
 		actual, ok := mem.Read(0, uint32(len(expectedMemory)))
@@ -392,9 +392,9 @@ func TestSnapshotPreview1_EnvironSizesGet_Errors(t *testing.T) {
 	envOpt, err := Environ("a=b", "b=cd")
 	require.NoError(t, err)
 
-	mod, fn := instantiateModule(t, FunctionEnvironSizesGet, ImportEnvironSizesGet, moduleName, envOpt)
+	mem, fn := instantiateModule(t, FunctionEnvironSizesGet, ImportEnvironSizesGet, moduleName, envOpt)
 
-	memorySize := mod.Memory("memory").Size()
+	memorySize := mem.Size()
 	validAddress := uint32(0) // arbitrary valid address as arguments to environ_sizes_get. We chose 0 here.
 
 	tests := []struct {
@@ -437,10 +437,10 @@ func TestSnapshotPreview1_EnvironSizesGet_Errors(t *testing.T) {
 
 // TestSnapshotPreview1_ClockResGet only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_ClockResGet(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionClockResGet, ImportClockResGet, moduleName)
+	mem, fn := instantiateModule(t, FunctionClockResGet, ImportClockResGet, moduleName)
 
 	t.Run("SnapshotPreview1.ClockResGet", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().ClockResGet(mod.Instance.Ctx, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().ClockResGet(ctx(mem), 0, 0))
 	})
 
 	t.Run(FunctionClockResGet, func(t *testing.T) {
@@ -462,14 +462,13 @@ func TestSnapshotPreview1_ClockTimeGet(t *testing.T) {
 	clockOpt := func(api *wasiAPI) {
 		api.timeNowUnixNano = func() uint64 { return epochNanos }
 	}
-	mod, fn := instantiateModule(t, FunctionClockTimeGet, ImportClockTimeGet, moduleName, clockOpt)
-	mem := mod.Memory("memory")
+	mem, fn := instantiateModule(t, FunctionClockTimeGet, ImportClockTimeGet, moduleName, clockOpt)
 
 	t.Run("SnapshotPreview1.ClockTimeGet", func(t *testing.T) {
 		maskMemory(t, mem, len(expectedMemory))
 
 		// invoke ClockTimeGet directly and check the memory side effects!
-		errno := NewAPI(clockOpt).ClockTimeGet(mod.Instance.Ctx, 0 /* TODO: id */, 0 /* TODO: precision */, resultTimestamp)
+		errno := NewAPI(clockOpt).ClockTimeGet(ctx(mem), 0 /* TODO: id */, 0 /* TODO: precision */, resultTimestamp)
 		require.Equal(t, wasi.ErrnoSuccess, errno)
 
 		actual, ok := mem.Read(0, uint32(len(expectedMemory)))
@@ -493,11 +492,11 @@ func TestSnapshotPreview1_ClockTimeGet(t *testing.T) {
 func TestSnapshotPreview1_ClockTimeGet_Errors(t *testing.T) {
 	epochNanos := uint64(1640995200000000000) // midnight UTC 2022-01-01
 
-	mod, fn := instantiateModule(t, FunctionClockTimeGet, ImportClockTimeGet, moduleName, func(api *wasiAPI) {
+	mem, fn := instantiateModule(t, FunctionClockTimeGet, ImportClockTimeGet, moduleName, func(api *wasiAPI) {
 		api.timeNowUnixNano = func() uint64 { return epochNanos }
 	})
 
-	memorySize := mod.Memory("memory").Size()
+	memorySize := mem.Size()
 
 	tests := []struct {
 		name            string
@@ -528,10 +527,10 @@ func TestSnapshotPreview1_ClockTimeGet_Errors(t *testing.T) {
 
 // TestSnapshotPreview1_FdAdvise only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdAdvise(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdAdvise, ImportFdAdvise, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdAdvise, ImportFdAdvise, moduleName)
 
 	t.Run("SnapshotPreview1.FdAdvise", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdAdvise(mod.Instance.Ctx, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdAdvise(ctx(mem), 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionFdAdvise, func(t *testing.T) {
@@ -543,10 +542,10 @@ func TestSnapshotPreview1_FdAdvise(t *testing.T) {
 
 // TestSnapshotPreview1_FdAllocate only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdAllocate(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdAllocate, ImportFdAllocate, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdAllocate, ImportFdAllocate, moduleName)
 
 	t.Run("SnapshotPreview1.FdAllocate", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdAllocate(mod.Instance.Ctx, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdAllocate(ctx(mem), 0, 0, 0))
 	})
 
 	t.Run(FunctionFdAllocate, func(t *testing.T) {
@@ -560,9 +559,9 @@ func TestSnapshotPreview1_FdClose(t *testing.T) {
 	fdToClose := uint32(3) // arbitrary fd
 	fdToKeep := uint32(4)  // another arbitrary fd
 
-	setupFD := func() (*wasm.PublicModule, publicwasm.Function, *wasiAPI) {
+	setupFD := func() (publicwasm.Memory, publicwasm.Function, *wasiAPI) {
 		var api *wasiAPI
-		mod, fn := instantiateModule(t, FunctionFdClose, ImportFdClose, moduleName, func(a *wasiAPI) {
+		mem, fn := instantiateModule(t, FunctionFdClose, ImportFdClose, moduleName, func(a *wasiAPI) {
 			memFs := &MemFS{}
 			a.opened = map[uint32]fileEntry{
 				fdToClose: {
@@ -576,13 +575,13 @@ func TestSnapshotPreview1_FdClose(t *testing.T) {
 			}
 			api = a // for later tests
 		})
-		return mod, fn, api
+		return mem, fn, api
 	}
 
 	t.Run("SnapshotPreview1.FdClose", func(t *testing.T) {
-		mod, _, api := setupFD()
+		mem, _, api := setupFD()
 
-		errno := api.FdClose(mod.Instance.Ctx, fdToClose)
+		errno := api.FdClose(ctx(mem), fdToClose)
 		require.Equal(t, wasi.ErrnoSuccess, errno)
 		require.NotContains(t, api.opened, fdToClose) // Fd is closed and removed from the opened FDs.
 		require.Contains(t, api.opened, fdToKeep)
@@ -597,18 +596,19 @@ func TestSnapshotPreview1_FdClose(t *testing.T) {
 		require.Contains(t, api.opened, fdToKeep)
 	})
 	t.Run("ErrnoBadF for an invalid FD", func(t *testing.T) {
-		mod, _, api := setupFD()
-		errno := api.FdClose(mod.Instance.Ctx, 42) // 42 is an arbitrary invalid FD
+		mem, _, api := setupFD()
+
+		errno := api.FdClose(ctx(mem), 42) // 42 is an arbitrary invalid FD
 		require.Equal(t, wasi.ErrnoBadf, errno)
 	})
 }
 
 // TestSnapshotPreview1_FdDatasync only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdDatasync(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdDatasync, ImportFdDatasync, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdDatasync, ImportFdDatasync, moduleName)
 
 	t.Run("SnapshotPreview1.FdDatasync", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdDatasync(mod.Instance.Ctx, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdDatasync(ctx(mem), 0))
 	})
 
 	t.Run(FunctionFdDatasync, func(t *testing.T) {
@@ -622,10 +622,10 @@ func TestSnapshotPreview1_FdDatasync(t *testing.T) {
 
 // TestSnapshotPreview1_FdFdstatSetFlags only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdFdstatSetFlags(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdFdstatSetFlags, ImportFdFdstatSetFlags, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdFdstatSetFlags, ImportFdFdstatSetFlags, moduleName)
 
 	t.Run("SnapshotPreview1.FdFdstatSetFlags", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdFdstatSetFlags(mod.Instance.Ctx, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdFdstatSetFlags(ctx(mem), 0, 0))
 	})
 
 	t.Run(FunctionFdFdstatSetFlags, func(t *testing.T) {
@@ -637,10 +637,10 @@ func TestSnapshotPreview1_FdFdstatSetFlags(t *testing.T) {
 
 // TestSnapshotPreview1_FdFdstatSetRights only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdFdstatSetRights(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdFdstatSetRights, ImportFdFdstatSetRights, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdFdstatSetRights, ImportFdFdstatSetRights, moduleName)
 
 	t.Run("SnapshotPreview1.FdFdstatSetRights", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdFdstatSetRights(mod.Instance.Ctx, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdFdstatSetRights(ctx(mem), 0, 0, 0))
 	})
 
 	t.Run(FunctionFdFdstatSetRights, func(t *testing.T) {
@@ -652,10 +652,10 @@ func TestSnapshotPreview1_FdFdstatSetRights(t *testing.T) {
 
 // TestSnapshotPreview1_FdFilestatGet only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdFilestatGet(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdFilestatGet, ImportFdFilestatGet, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdFilestatGet, ImportFdFilestatGet, moduleName)
 
 	t.Run("SnapshotPreview1.FdFilestatGet", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdFilestatGet(mod.Instance.Ctx, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdFilestatGet(ctx(mem), 0, 0))
 	})
 
 	t.Run(FunctionFdFilestatGet, func(t *testing.T) {
@@ -667,10 +667,10 @@ func TestSnapshotPreview1_FdFilestatGet(t *testing.T) {
 
 // TestSnapshotPreview1_FdFilestatSetSize only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdFilestatSetSize(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdFilestatSetSize, ImportFdFilestatSetSize, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdFilestatSetSize, ImportFdFilestatSetSize, moduleName)
 
 	t.Run("SnapshotPreview1.FdFilestatSetSize", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdFilestatSetSize(mod.Instance.Ctx, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdFilestatSetSize(ctx(mem), 0, 0))
 	})
 
 	t.Run(FunctionFdFilestatSetSize, func(t *testing.T) {
@@ -682,10 +682,10 @@ func TestSnapshotPreview1_FdFilestatSetSize(t *testing.T) {
 
 // TestSnapshotPreview1_FdFilestatSetTimes only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdFilestatSetTimes(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdFilestatSetTimes, ImportFdFilestatSetTimes, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdFilestatSetTimes, ImportFdFilestatSetTimes, moduleName)
 
 	t.Run("SnapshotPreview1.FdFilestatSetTimes", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdFilestatSetTimes(mod.Instance.Ctx, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdFilestatSetTimes(ctx(mem), 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionFdFilestatSetTimes, func(t *testing.T) {
@@ -697,10 +697,10 @@ func TestSnapshotPreview1_FdFilestatSetTimes(t *testing.T) {
 
 // TestSnapshotPreview1_FdPread only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdPread(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdPread, ImportFdPread, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdPread, ImportFdPread, moduleName)
 
 	t.Run("SnapshotPreview1.FdPread", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdPread(mod.Instance.Ctx, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdPread(ctx(mem), 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionFdPread, func(t *testing.T) {
@@ -715,14 +715,13 @@ func TestSnapshotPreview1_FdPread(t *testing.T) {
 func TestSnapshotPreview1_FdPrestatDirName(t *testing.T) {
 	fd := uint32(3) // arbitrary fd after 0, 1, and 2, that are stdin/out/err
 	var api *wasiAPI
-	mod, fn := instantiateModule(t, FunctionFdPrestatDirName, ImportFdPrestatDirName, moduleName, func(a *wasiAPI) {
+	mem, fn := instantiateModule(t, FunctionFdPrestatDirName, ImportFdPrestatDirName, moduleName, func(a *wasiAPI) {
 		a.opened[fd] = fileEntry{
 			path:    "/tmp",
 			fileSys: &MemFS{},
 		}
 		api = a // for later tests
 	})
-	mem := mod.Memory("memory")
 
 	path := uint32(1)    // arbitrary offset
 	pathLen := uint32(3) // shorter than len("/tmp") to test the path is written for the length of pathLen
@@ -735,7 +734,7 @@ func TestSnapshotPreview1_FdPrestatDirName(t *testing.T) {
 	t.Run("SnapshotPreview1.FdPrestatDirName", func(t *testing.T) {
 		maskMemory(t, mem, len(expectedMemory))
 
-		errno := api.FdPrestatDirName(mod.Instance.Ctx, fd, path, pathLen)
+		errno := api.FdPrestatDirName(ctx(mem), fd, path, pathLen)
 		require.Equal(t, wasi.ErrnoSuccess, errno)
 
 		actual, ok := mem.Read(0, uint32(len(expectedMemory)))
@@ -759,9 +758,7 @@ func TestSnapshotPreview1_FdPrestatDirName(t *testing.T) {
 func TestSnapshotPreview1_FdPrestatDirName_Errors(t *testing.T) {
 	dirName := "/tmp"
 	opt := Preopen(dirName, &MemFS{})
-	mod, fn := instantiateModule(t, FunctionFdPrestatDirName, ImportFdPrestatDirName, moduleName, opt)
-
-	mem := mod.Memory("memory")
+	mem, fn := instantiateModule(t, FunctionFdPrestatDirName, ImportFdPrestatDirName, moduleName, opt)
 
 	memorySize := mem.Size()
 	validAddress := uint32(0) // Arbitrary valid address as arguments to fd_prestat_dir_name. We chose 0 here.
@@ -817,10 +814,10 @@ func TestSnapshotPreview1_FdPrestatDirName_Errors(t *testing.T) {
 
 // TestSnapshotPreview1_FdPwrite only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdPwrite(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdPwrite, ImportFdPwrite, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdPwrite, ImportFdPwrite, moduleName)
 
 	t.Run("SnapshotPreview1.FdPwrite", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdPwrite(mod.Instance.Ctx, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdPwrite(ctx(mem), 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionFdPwrite, func(t *testing.T) {
@@ -854,12 +851,10 @@ func TestSnapshotPreview1_FdRead(t *testing.T) {
 	)
 
 	var api *wasiAPI
-	mod, fn := instantiateModule(t, FunctionFdRead, ImportFdRead, moduleName, func(a *wasiAPI) {
+	mem, fn := instantiateModule(t, FunctionFdRead, ImportFdRead, moduleName, func(a *wasiAPI) {
 		api = a // for later tests
 	})
 
-	mem := mod.Memory("memory")
-
 	// TestSnapshotPreview1_FdRead uses a matrix because setting up test files is complicated and has to be clean each time.
 	type fdReadFn func(ctx publicwasm.ModuleContext, fd, iovs, iovsCount, resultSize uint32) wasi.Errno
 	tests := []struct {
@@ -893,7 +888,7 @@ func TestSnapshotPreview1_FdRead(t *testing.T) {
 			ok := mem.Write(0, initialMemory)
 			require.True(t, ok)
 
-			errno := tc.fdRead()(mod.Instance.Ctx, fd, iovs, iovsCount, resultSize)
+			errno := tc.fdRead()(ctx(mem), fd, iovs, iovsCount, resultSize)
 			require.Equal(t, wasi.ErrnoSuccess, errno)
 
 			actual, ok := mem.Read(0, uint32(len(expectedMemory)))
@@ -906,7 +901,7 @@ func TestSnapshotPreview1_FdRead(t *testing.T) {
 func TestSnapshotPreview1_FdRead_Errors(t *testing.T) {
 	validFD := uint64(3)                                // arbitrary valid fd after 0, 1, and 2, that are stdin/out/err
 	file, memFS := createFile(t, "test_path", []byte{}) // file with empty contents
-	mod, fn := instantiateModule(t, FunctionFdRead, ImportFdRead, moduleName, func(a *wasiAPI) {
+	mem, fn := instantiateModule(t, FunctionFdRead, ImportFdRead, moduleName, func(a *wasiAPI) {
 		a.opened[uint32(validFD)] = fileEntry{
 			path:    "test_path",
 			fileSys: memFS,
@@ -914,8 +909,6 @@ func TestSnapshotPreview1_FdRead_Errors(t *testing.T) {
 		}
 	})
 
-	mem := mod.Memory("memory")
-
 	tests := []struct {
 		name                            string
 		fd, iovs, iovsCount, resultSize uint64
@@ -999,10 +992,10 @@ func TestSnapshotPreview1_FdRead_Errors(t *testing.T) {
 
 // TestSnapshotPreview1_FdReaddir only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdReaddir(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdReaddir, ImportFdReaddir, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdReaddir, ImportFdReaddir, moduleName)
 
 	t.Run("SnapshotPreview1.FdReaddir", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdReaddir(mod.Instance.Ctx, 0, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdReaddir(ctx(mem), 0, 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionFdReaddir, func(t *testing.T) {
@@ -1014,10 +1007,10 @@ func TestSnapshotPreview1_FdReaddir(t *testing.T) {
 
 // TestSnapshotPreview1_FdRenumber only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdRenumber(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdRenumber, ImportFdRenumber, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdRenumber, ImportFdRenumber, moduleName)
 
 	t.Run("SnapshotPreview1.FdRenumber", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdRenumber(mod.Instance.Ctx, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdRenumber(ctx(mem), 0, 0))
 	})
 
 	t.Run(FunctionFdRenumber, func(t *testing.T) {
@@ -1029,10 +1022,10 @@ func TestSnapshotPreview1_FdRenumber(t *testing.T) {
 
 // TestSnapshotPreview1_FdSeek only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdSeek(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdSeek, ImportFdSeek, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdSeek, ImportFdSeek, moduleName)
 
 	t.Run("SnapshotPreview1.FdSeek", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdSeek(mod.Instance.Ctx, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdSeek(ctx(mem), 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionFdSeek, func(t *testing.T) {
@@ -1044,10 +1037,10 @@ func TestSnapshotPreview1_FdSeek(t *testing.T) {
 
 // TestSnapshotPreview1_FdSync only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdSync(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdSync, ImportFdSync, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdSync, ImportFdSync, moduleName)
 
 	t.Run("SnapshotPreview1.FdSync", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdSync(mod.Instance.Ctx, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdSync(ctx(mem), 0))
 	})
 
 	t.Run(FunctionFdSync, func(t *testing.T) {
@@ -1059,10 +1052,10 @@ func TestSnapshotPreview1_FdSync(t *testing.T) {
 
 // TestSnapshotPreview1_FdTell only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdTell(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionFdTell, ImportFdTell, moduleName)
+	mem, fn := instantiateModule(t, FunctionFdTell, ImportFdTell, moduleName)
 
 	t.Run("SnapshotPreview1.FdTell", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdTell(mod.Instance.Ctx, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().FdTell(ctx(mem), 0, 0))
 	})
 
 	t.Run(FunctionFdTell, func(t *testing.T) {
@@ -1096,12 +1089,10 @@ func TestSnapshotPreview1_FdWrite(t *testing.T) {
 	)
 
 	var api *wasiAPI
-	mod, fn := instantiateModule(t, FunctionFdWrite, ImportFdWrite, moduleName, func(a *wasiAPI) {
+	mem, fn := instantiateModule(t, FunctionFdWrite, ImportFdWrite, moduleName, func(a *wasiAPI) {
 		api = a // for later tests
 	})
 
-	mem := mod.Memory("memory")
-
 	// TestSnapshotPreview1_FdWrite uses a matrix because setting up test files is complicated and has to be clean each time.
 	type fdWriteFn func(ctx publicwasm.ModuleContext, fd, iovs, iovsCount, resultSize uint32) wasi.Errno
 	tests := []struct {
@@ -1134,7 +1125,7 @@ func TestSnapshotPreview1_FdWrite(t *testing.T) {
 			ok := mem.Write(0, initialMemory)
 			require.True(t, ok)
 
-			errno := tc.fdWrite()(mod.Instance.Ctx, fd, iovs, iovsCount, resultSize)
+			errno := tc.fdWrite()(ctx(mem), fd, iovs, iovsCount, resultSize)
 			require.Equal(t, wasi.ErrnoSuccess, errno)
 
 			actual, ok := mem.Read(0, uint32(len(expectedMemory)))
@@ -1148,7 +1139,7 @@ func TestSnapshotPreview1_FdWrite(t *testing.T) {
 func TestSnapshotPreview1_FdWrite_Errors(t *testing.T) {
 	validFD := uint64(3)                                // arbitrary valid fd after 0, 1, and 2, that are stdin/out/err
 	file, memFS := createFile(t, "test_path", []byte{}) // file with empty contents
-	mod, fn := instantiateModule(t, FunctionFdWrite, ImportFdWrite, moduleName, func(a *wasiAPI) {
+	mem, fn := instantiateModule(t, FunctionFdWrite, ImportFdWrite, moduleName, func(a *wasiAPI) {
 		a.opened[uint32(validFD)] = fileEntry{
 			path:    "test_path",
 			fileSys: memFS,
@@ -1156,8 +1147,6 @@ func TestSnapshotPreview1_FdWrite_Errors(t *testing.T) {
 		}
 	})
 
-	mem := mod.Memory("memory")
-
 	tests := []struct {
 		name                            string
 		fd, iovs, iovsCount, resultSize uint64
@@ -1253,10 +1242,10 @@ func createFile(t *testing.T, path string, contents []byte) (*memFile, *MemFS) {
 
 // TestSnapshotPreview1_PathCreateDirectory only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathCreateDirectory(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionPathCreateDirectory, ImportPathCreateDirectory, moduleName)
+	mem, fn := instantiateModule(t, FunctionPathCreateDirectory, ImportPathCreateDirectory, moduleName)
 
 	t.Run("SnapshotPreview1.PathCreateDirectory", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathCreateDirectory(mod.Instance.Ctx, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathCreateDirectory(ctx(mem), 0, 0, 0))
 	})
 
 	t.Run(FunctionPathCreateDirectory, func(t *testing.T) {
@@ -1268,10 +1257,10 @@ func TestSnapshotPreview1_PathCreateDirectory(t *testing.T) {
 
 // TestSnapshotPreview1_PathFilestatGet only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathFilestatGet(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionPathFilestatGet, ImportPathFilestatGet, moduleName)
+	mem, fn := instantiateModule(t, FunctionPathFilestatGet, ImportPathFilestatGet, moduleName)
 
 	t.Run("SnapshotPreview1.PathFilestatGet", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathFilestatGet(mod.Instance.Ctx, 0, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathFilestatGet(ctx(mem), 0, 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionPathFilestatGet, func(t *testing.T) {
@@ -1283,10 +1272,10 @@ func TestSnapshotPreview1_PathFilestatGet(t *testing.T) {
 
 // TestSnapshotPreview1_PathFilestatSetTimes only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathFilestatSetTimes(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionPathFilestatSetTimes, ImportPathFilestatSetTimes, moduleName)
+	mem, fn := instantiateModule(t, FunctionPathFilestatSetTimes, ImportPathFilestatSetTimes, moduleName)
 
 	t.Run("SnapshotPreview1.PathFilestatSetTimes", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathFilestatSetTimes(mod.Instance.Ctx, 0, 0, 0, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathFilestatSetTimes(ctx(mem), 0, 0, 0, 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionPathFilestatSetTimes, func(t *testing.T) {
@@ -1298,10 +1287,10 @@ func TestSnapshotPreview1_PathFilestatSetTimes(t *testing.T) {
 
 // TestSnapshotPreview1_PathLink only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathLink(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionPathLink, ImportPathLink, moduleName)
+	mem, fn := instantiateModule(t, FunctionPathLink, ImportPathLink, moduleName)
 
 	t.Run("SnapshotPreview1.PathLink", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathLink(mod.Instance.Ctx, 0, 0, 0, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathLink(ctx(mem), 0, 0, 0, 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionPathLink, func(t *testing.T) {
@@ -1315,10 +1304,10 @@ func TestSnapshotPreview1_PathLink(t *testing.T) {
 
 // TestSnapshotPreview1_PathReadlink only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathReadlink(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionPathReadlink, ImportPathReadlink, moduleName)
+	mem, fn := instantiateModule(t, FunctionPathReadlink, ImportPathReadlink, moduleName)
 
 	t.Run("SnapshotPreview1.PathLink", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathReadlink(mod.Instance.Ctx, 0, 0, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathReadlink(ctx(mem), 0, 0, 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionPathReadlink, func(t *testing.T) {
@@ -1330,10 +1319,10 @@ func TestSnapshotPreview1_PathReadlink(t *testing.T) {
 
 // TestSnapshotPreview1_PathRemoveDirectory only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathRemoveDirectory(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionPathRemoveDirectory, ImportPathRemoveDirectory, moduleName)
+	mem, fn := instantiateModule(t, FunctionPathRemoveDirectory, ImportPathRemoveDirectory, moduleName)
 
 	t.Run("SnapshotPreview1.PathRemoveDirectory", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathRemoveDirectory(mod.Instance.Ctx, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathRemoveDirectory(ctx(mem), 0, 0, 0))
 	})
 
 	t.Run(FunctionPathRemoveDirectory, func(t *testing.T) {
@@ -1345,10 +1334,10 @@ func TestSnapshotPreview1_PathRemoveDirectory(t *testing.T) {
 
 // TestSnapshotPreview1_PathRename only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathRename(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionPathRename, ImportPathRename, moduleName)
+	mem, fn := instantiateModule(t, FunctionPathRename, ImportPathRename, moduleName)
 
 	t.Run("SnapshotPreview1.PathRename", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathRename(mod.Instance.Ctx, 0, 0, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathRename(ctx(mem), 0, 0, 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionPathRename, func(t *testing.T) {
@@ -1360,10 +1349,10 @@ func TestSnapshotPreview1_PathRename(t *testing.T) {
 
 // TestSnapshotPreview1_PathSymlink only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathSymlink(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionPathSymlink, ImportPathSymlink, moduleName)
+	mem, fn := instantiateModule(t, FunctionPathSymlink, ImportPathSymlink, moduleName)
 
 	t.Run("SnapshotPreview1.PathSymlink", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathSymlink(mod.Instance.Ctx, 0, 0, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathSymlink(ctx(mem), 0, 0, 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionPathSymlink, func(t *testing.T) {
@@ -1375,10 +1364,10 @@ func TestSnapshotPreview1_PathSymlink(t *testing.T) {
 
 // TestSnapshotPreview1_PathUnlinkFile only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathUnlinkFile(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionPathUnlinkFile, ImportPathUnlinkFile, moduleName)
+	mem, fn := instantiateModule(t, FunctionPathUnlinkFile, ImportPathUnlinkFile, moduleName)
 
 	t.Run("SnapshotPreview1.PathUnlinkFile", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathUnlinkFile(mod.Instance.Ctx, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().PathUnlinkFile(ctx(mem), 0, 0, 0))
 	})
 
 	t.Run(FunctionPathUnlinkFile, func(t *testing.T) {
@@ -1390,10 +1379,10 @@ func TestSnapshotPreview1_PathUnlinkFile(t *testing.T) {
 
 // TestSnapshotPreview1_PollOneoff only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PollOneoff(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionPollOneoff, ImportPollOneoff, moduleName)
+	mem, fn := instantiateModule(t, FunctionPollOneoff, ImportPollOneoff, moduleName)
 
 	t.Run("SnapshotPreview1.PollOneoff", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().PollOneoff(mod.Instance.Ctx, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().PollOneoff(ctx(mem), 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionPollOneoff, func(t *testing.T) {
@@ -1436,10 +1425,10 @@ func TestSnapshotPreview1_ProcExit(t *testing.T) {
 
 // TestSnapshotPreview1_ProcRaise only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_ProcRaise(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionProcRaise, ImportProcRaise, moduleName)
+	mem, fn := instantiateModule(t, FunctionProcRaise, ImportProcRaise, moduleName)
 
 	t.Run("SnapshotPreview1.ProcRaise", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().ProcRaise(mod.Instance.Ctx, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().ProcRaise(ctx(mem), 0))
 	})
 
 	t.Run(FunctionProcRaise, func(t *testing.T) {
@@ -1451,10 +1440,10 @@ func TestSnapshotPreview1_ProcRaise(t *testing.T) {
 
 // TestSnapshotPreview1_SchedYield only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_SchedYield(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionSchedYield, ImportSchedYield, moduleName)
+	mem, fn := instantiateModule(t, FunctionSchedYield, ImportSchedYield, moduleName)
 
 	t.Run("SnapshotPreview1.SchedYield", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().SchedYield(mod.Instance.Ctx))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().SchedYield(ctx(mem)))
 	})
 
 	t.Run(FunctionSchedYield, func(t *testing.T) {
@@ -1485,13 +1474,12 @@ func TestSnapshotPreview1_RandomGet(t *testing.T) {
 		}
 	}
 
-	mod, fn := instantiateModule(t, FunctionRandomGet, ImportRandomGet, moduleName, randOpt)
-	mem := mod.Memory("memory")
+	mem, fn := instantiateModule(t, FunctionRandomGet, ImportRandomGet, moduleName, randOpt)
 	t.Run("SnapshotPreview1.RandomGet", func(t *testing.T) {
 		maskMemory(t, mem, len(expectedMemory))
 
 		// Invoke RandomGet directly and check the memory side effects!
-		errno := NewAPI(randOpt).RandomGet(mod.Instance.Ctx, offset, length)
+		errno := NewAPI(randOpt).RandomGet(ctx(mem), offset, length)
 		require.Equal(t, wasi.ErrnoSuccess, errno)
 
 		actual, ok := mem.Read(0, offset+length+1)
@@ -1515,8 +1503,8 @@ func TestSnapshotPreview1_RandomGet(t *testing.T) {
 func TestSnapshotPreview1_RandomGet_Errors(t *testing.T) {
 	validAddress := uint32(0) // arbitrary valid address
 
-	mod, fn := instantiateModule(t, FunctionRandomGet, ImportRandomGet, moduleName)
-	memorySize := mod.Memory("memory").Size()
+	mem, fn := instantiateModule(t, FunctionRandomGet, ImportRandomGet, moduleName)
+	memorySize := mem.Size()
 
 	tests := []struct {
 		name   string
@@ -1561,10 +1549,10 @@ func TestSnapshotPreview1_RandomGet_SourceError(t *testing.T) {
 
 // TestSnapshotPreview1_SockRecv only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_SockRecv(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionSockRecv, ImportSockRecv, moduleName)
+	mem, fn := instantiateModule(t, FunctionSockRecv, ImportSockRecv, moduleName)
 
 	t.Run("SnapshotPreview1.SockRecv", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().SockRecv(mod.Instance.Ctx, 0, 0, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().SockRecv(ctx(mem), 0, 0, 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionSockRecv, func(t *testing.T) {
@@ -1576,10 +1564,10 @@ func TestSnapshotPreview1_SockRecv(t *testing.T) {
 
 // TestSnapshotPreview1_SockSend only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_SockSend(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionSockSend, ImportSockSend, moduleName)
+	mem, fn := instantiateModule(t, FunctionSockSend, ImportSockSend, moduleName)
 
 	t.Run("SnapshotPreview1.SockSend", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().SockSend(mod.Instance.Ctx, 0, 0, 0, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().SockSend(ctx(mem), 0, 0, 0, 0, 0))
 	})
 
 	t.Run(FunctionSockSend, func(t *testing.T) {
@@ -1591,10 +1579,10 @@ func TestSnapshotPreview1_SockSend(t *testing.T) {
 
 // TestSnapshotPreview1_SockShutdown only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_SockShutdown(t *testing.T) {
-	mod, fn := instantiateModule(t, FunctionSockShutdown, ImportSockShutdown, moduleName)
+	mem, fn := instantiateModule(t, FunctionSockShutdown, ImportSockShutdown, moduleName)
 
 	t.Run("SnapshotPreview1.SockShutdown", func(t *testing.T) {
-		require.Equal(t, wasi.ErrnoNosys, NewAPI().SockShutdown(mod.Instance.Ctx, 0, 0))
+		require.Equal(t, wasi.ErrnoNosys, NewAPI().SockShutdown(ctx(mem), 0, 0))
 	})
 
 	t.Run(FunctionSockShutdown, func(t *testing.T) {
@@ -1606,9 +1594,9 @@ func TestSnapshotPreview1_SockShutdown(t *testing.T) {
 
 const testMemoryPageSize = 1
 
-func instantiateModule(t *testing.T, wasiFunction, wasiImport, moduleName string, opts ...Option) (*wasm.PublicModule, publicwasm.Function) {
+func instantiateModule(t *testing.T, wasiFunction, wasiImport, moduleName string, opts ...Option) (publicwasm.Memory, publicwasm.Function) {
 	enabledFeatures := wasm.Features20191205
-	mod, err := text.DecodeModule([]byte(fmt.Sprintf(`(module
+	mem, err := text.DecodeModule([]byte(fmt.Sprintf(`(module
   %[2]s
   (memory 1)  ;; just an arbitrary size big enough for tests
   (export "memory" (memory 0))
@@ -1616,18 +1604,18 @@ func instantiateModule(t *testing.T, wasiFunction, wasiImport, moduleName string
 )`, wasiFunction, wasiImport)), enabledFeatures)
 	require.NoError(t, err)
 
+	// The package `wazero` has a simpler interface for adding host modules, but we can't use that as it would create an
+	// import cycle. Instead, we export Store.NewHostModule and use it here.
 	store := wasm.NewStore(context.Background(), interpreter.NewEngine(), enabledFeatures)
-
-	snapshotPreview1Functions := SnapshotPreview1Functions(opts...)
-	_, err = store.NewHostModule(wasi.ModuleSnapshotPreview1, snapshotPreview1Functions)
+	_, err = store.NewHostModule(wasi.ModuleSnapshotPreview1, SnapshotPreview1Functions(opts...))
 	require.NoError(t, err)
 
-	instantiated, err := store.Instantiate(mod, moduleName)
+	instantiated, err := store.Instantiate(mem, moduleName)
 	require.NoError(t, err)
 
 	fn := instantiated.Function(wasiFunction)
 	require.NotNil(t, fn)
-	return instantiated, fn
+	return instantiated.Memory("memory"), fn
 }
 
 // maskMemory sets the first memory in the store to '?' * size, so tests can see what's written.

internal/wasm/global.go

@@ -124,7 +124,7 @@ func (g globalF64) String() string {
 
 // Global implements wasm.Module Global
 func (m *PublicModule) Global(name string) publicwasm.Global {
-	exp, err := m.Instance.getExport(name, ExternTypeGlobal)
+	exp, err := m.instance.getExport(name, ExternTypeGlobal)
 	if err != nil {
 		return nil
 	}

internal/wasm/host.go

@@ -91,23 +91,24 @@ func (f *exportedFunction) Call(ctx context.Context, params ...uint64) ([]uint64
 }
 
 // NewHostModule is defined internally for use in WASI tests and to keep the code size in the root directory small.
-func (s *Store) NewHostModule(moduleName string, nameToGoFunc map[string]interface{}) (publicwasm.HostModule, error) {
+func (s *Store) NewHostModule(moduleName string, nameToGoFunc map[string]interface{}) (*HostModule, error) {
 	if err := s.requireModuleUnused(moduleName); err != nil {
 		return nil, err
 	}
 
 	exportCount := len(nameToGoFunc)
 
-	ret := &HostModule{NameToFunctionInstance: make(map[string]*FunctionInstance, exportCount)}
+	ret := &HostModule{name: moduleName, NameToFunctionInstance: make(map[string]*FunctionInstance, exportCount)}
 	hostModule := &ModuleInstance{
-		Name: moduleName, Exports: make(map[string]*ExportInstance, exportCount),
+		Name:       moduleName,
+		Exports:    make(map[string]*ExportInstance, exportCount),
 		hostModule: ret,
 	}
 	s.moduleInstances[moduleName] = hostModule
 	for name, goFunc := range nameToGoFunc {
 		if hf, err := NewGoFunc(name, goFunc); err != nil {
 			return nil, err
-		} else if function, err := s.AddHostFunction(hostModule, hf); err != nil {
+		} else if function, err := s.addHostFunction(hostModule, hf); err != nil {
 			return nil, err
 		} else {
 			ret.NameToFunctionInstance[name] = function
@@ -120,7 +121,7 @@ func (s *Store) NewHostModule(moduleName string, nameToGoFunc map[string]interfa
 // 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.
-func (s *Store) AddHostFunction(m *ModuleInstance, hf *GoFunc) (*FunctionInstance, error) {
+func (s *Store) addHostFunction(m *ModuleInstance, hf *GoFunc) (*FunctionInstance, error) {
 	typeInstance, err := s.getTypeInstance(hf.functionType)
 	if err != nil {
 		return nil, err
@@ -144,14 +145,7 @@ func (s *Store) AddHostFunction(m *ModuleInstance, hf *GoFunc) (*FunctionInstanc
 		return nil, fmt.Errorf("failed to compile %s: %v", f.Name, err)
 	}
 
-	if err = m.addExport(hf.wasmFunctionName, &ExportInstance{Type: ExternTypeFunc, Function: f}); err != nil {
-		// On failure, we must release the function instance.
-		if err := s.releaseFunctionInstances(f); err != nil {
-			return nil, err
-		}
-		return nil, err
-	}
-
+	m.Exports[hf.wasmFunctionName] = &ExportInstance{Type: ExternTypeFunc, Function: f}
 	m.Functions = append(m.Functions, f)
 	return f, nil
 }
@@ -165,9 +159,16 @@ func (s *Store) requireModuleUnused(moduleName string) error {
 
 // HostModule implements wasm.HostModule
 type HostModule struct {
+	// name is for String and Store.ReleaseModuleInstance
+	name                   string
 	NameToFunctionInstance map[string]*FunctionInstance
 }
 
+// String implements fmt.Stringer
+func (m *HostModule) String() string {
+	return fmt.Sprintf("HostModule[%s]", m.name)
+}
+
 // ParamTypes implements wasm.HostFunction ParamTypes
 func (f *FunctionInstance) ParamTypes() []publicwasm.ValueType {
 	return f.FunctionType.Type.Params
@@ -188,6 +189,6 @@ func (f *FunctionInstance) Call(ctx publicwasm.ModuleContext, params ...uint64)
 }
 
 // Function implements wasm.HostModule Function
-func (g *HostModule) Function(name string) publicwasm.HostFunction {
-	return g.NameToFunctionInstance[name]
+func (m *HostModule) Function(name string) publicwasm.HostFunction {
+	return m.NameToFunctionInstance[name]
 }

internal/wasm/host_test.go

@@ -5,6 +5,8 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
+	"github.com/tetratelabs/wazero/wasm"
 )
 
 func TestMemoryInstance_HasLen(t *testing.T) {
@@ -448,3 +450,38 @@ func TestMemoryInstance_WriteFloat64Le(t *testing.T) {
 		})
 	}
 }
+
+func TestStore_NewHostModule(t *testing.T) {
+	s := newStore()
+
+	// Add the host module
+	_, err := s.NewHostModule("test", map[string]interface{}{"fn": func(wasm.ModuleContext) {}})
+	require.NoError(t, err)
+
+	// Ensure it was added to module instances
+	hm := s.moduleInstances["test"]
+	require.NotNil(t, hm)
+
+	// The function was added to the store, prefixed by the owning module name
+	require.Equal(t, 1, len(s.functions))
+	fn := s.functions[0]
+	require.Equal(t, "test.fn", fn.Name)
+
+	// The function was exported in the module
+	require.Equal(t, 1, len(hm.Exports))
+	_, ok := hm.Exports["fn"]
+	require.True(t, ok)
+
+	// Trying to register it again should fail
+	_, err = s.NewHostModule("test", map[string]interface{}{"host_fn": func(wasm.ModuleContext) {}})
+	require.EqualError(t, err, "module test has already been instantiated")
+}
+
+func TestHostModule_String(t *testing.T) {
+	s := newStore()
+
+	// Ensure paths that can create the host module can see the name.
+	hm, err := s.NewHostModule("host", map[string]interface{}{"host_fn": func(wasm.ModuleContext) {}})
+	require.NoError(t, err)
+	require.Equal(t, "HostModule[host]", hm.String())
+}

internal/wasm/store.go

@@ -46,28 +46,29 @@ type (
 		// maximumFunctionIndex represents the limit on the number of function addresses (= function instances) in a store.
 		// Note: this is fixed to 2^27 but have this a field for testability.
 		maximumFunctionIndex FunctionIndex
+
 		//  maximumFunctionTypes represents the limit on the number of function types in a store.
 		// Note: this is fixed to 2^27 but have this a field for testability.
 		maximumFunctionTypes int
 
 		// releasedFunctionIndex holds reusable FunctionIndexes. An index is added when
-		// an function instance is released in releaseFunctionInstances, and is popped when
-		// an new instance is added in store.addFunctionInstances.
+		// a function instance is released in releaseFunctionInstances, and is popped when
+		// a new instance is added in addFunctionInstances.
 		releasedFunctionIndex []FunctionIndex
 
 		// releasedMemoryIndex holds reusable memoryIndexes. An index is added when
 		// an memory instance is released in releaseMemoryInstance, and is popped when
-		// an new instance is added in store.addMemoryInstance.
+		// a new instance is added in addMemoryInstance.
 		releasedMemoryIndex []memoryIndex
 
 		// releasedTableIndex holds reusable tableIndexes. An index is added when
 		// an table instance is released in releaseTableInstance, and is popped when
-		// an new instance is added in store.addTableInstance.
+		// a new instance is added in addTableInstance.
 		releasedTableIndex []tableIndex
 
 		// releasedGlobalIndex holds reusable globalIndexes. An index is added when
 		// an global instance is released in releaseGlobalInstances, and is popped when
-		// an new instance is added in store.addGlobalInstances.
+		// a new instance is added in addGlobalInstances.
 		releasedGlobalIndex []globalIndex
 
 		// The followings fields match the definition of Store in the specification.
@@ -112,7 +113,7 @@ type (
 		// hostModule holds HostModule if this is a "host module" which is created in store.NewHostModule.
 		hostModule *HostModule
 
-		// TODO per https://github.com/tetratelabs/wazero/issues/293
+		// TODO: https://github.com/tetratelabs/wazero/issues/293
 		refCount      int
 		moduleImports map[*ModuleInstance]struct{}
 	}
@@ -455,7 +456,7 @@ func (s *Store) Instantiate(module *Module, name string) (*PublicModule, error)
 	instance.applyElements(module.ElementSection)
 	instance.applyData(module.DataSection)
 
-	// Persist the instances other than functions (which we already persisted before compilation).
+	// Persist the instances other tha functions (which we already persisted before compilation).
 	s.addGlobalInstances(globals...)
 	s.addTableInstance(table)
 	s.addMemoryInstance(instance.MemoryInstance)
@@ -477,21 +478,28 @@ func (s *Store) Instantiate(module *Module, name string) (*PublicModule, error)
 			return nil, fmt.Errorf("module[%s] start function failed: %w", name, err)
 		}
 	}
-	return &PublicModule{s, instance}, nil
+	return &PublicModule{s: s, instance: instance}, nil
 }
 
-func (s *Store) ReleaseModuleInstance(instance *ModuleInstance) error {
+// ReleaseModuleInstance deallocates resources if a module with the given name exists.
+func (s *Store) ReleaseModuleInstance(moduleName string) error {
+	// TODO: none of this is goroutine safe
+	instance := s.moduleInstances[moduleName]
+	if instance == nil {
+		return nil // already released
+	}
+
 	instance.refCount--
 	if instance.refCount > 0 {
 		// This case other modules are importing this module instance and still alive.
-		return nil
+		return fmt.Errorf("%d modules import this and need to be closed first", instance.refCount)
 	}
 
 	// TODO: check outstanding calls and wait until they exit.
 
 	// Recursively release the imported instances.
 	for mod := range instance.moduleImports {
-		if err := s.ReleaseModuleInstance(mod); err != nil {
+		if err := s.ReleaseModuleInstance(mod.Name); err != nil {
 			return fmt.Errorf("unable to release imported module [%s]: %w", mod.Name, err)
 		}
 	}
@@ -499,8 +507,15 @@ func (s *Store) ReleaseModuleInstance(instance *ModuleInstance) error {
 	if err := s.releaseFunctionInstances(instance.Functions...); err != nil {
 		return fmt.Errorf("unable to release function instance: %w", err)
 	}
-	s.releaseMemoryInstance(instance.MemoryInstance)
-	s.releaseTableInstance(instance.TableInstance)
+
+	if instance.MemoryInstance != nil {
+		s.releaseMemoryInstance(instance.MemoryInstance)
+	}
+
+	if instance.TableInstance != nil {
+		s.releaseTableInstance(instance.TableInstance)
+	}
+
 	s.releaseGlobalInstances(instance.Globals...)
 
 	// Explicitly assign nil so that we ensure this moduleInstance no longer holds reference to instances.
@@ -521,7 +536,7 @@ func (s *Store) releaseFunctionInstances(fs ...*FunctionInstance) error {
 			return err
 		}
 
-		// Release refernce to the function instance.
+		// Release reference to the function instance.
 		s.functions[f.Index] = nil
 
 		// Append the address so that we can reuse it in order to avoid index space explosion.
@@ -573,7 +588,7 @@ func (s *Store) addGlobalInstances(gs ...*GlobalInstance) {
 }
 
 func (s *Store) releaseTableInstance(t *TableInstance) {
-	// Release refernce to the table instance.
+	// Release reference to the table instance.
 	s.tables[t.index] = nil
 
 	// Append the index so that we can reuse it in order to avoid index space explosion.
@@ -599,7 +614,7 @@ func (s *Store) addTableInstance(t *TableInstance) {
 }
 
 func (s *Store) releaseMemoryInstance(m *MemoryInstance) {
-	// Release refernce to the memory instance.
+	// Release reference to the memory instance.
 	s.memories[m.index] = nil
 
 	// Append the index so that we can reuse it in order to avoid index space explosion.
@@ -629,29 +644,33 @@ func (s *Store) Module(moduleName string) publicwasm.Module {
 	if m, ok := s.moduleInstances[moduleName]; !ok {
 		return nil
 	} else {
-		return &PublicModule{s, m}
+		return &PublicModule{s: s, instance: m}
 	}
 }
 
 // PublicModule implements wasm.Module
 type PublicModule struct {
-	s *Store
-	// Context is exported for /wasi.go
-	Instance *ModuleInstance
+	s        *Store
+	instance *ModuleInstance
+}
+
+// String implements fmt.Stringer
+func (m *PublicModule) String() string {
+	return fmt.Sprintf("Module[%s]", m.instance.Name)
 }
 
 // Function implements wasm.Module Function
 func (m *PublicModule) Function(name string) publicwasm.Function {
-	exp, err := m.Instance.getExport(name, ExternTypeFunc)
+	exp, err := m.instance.getExport(name, ExternTypeFunc)
 	if err != nil {
 		return nil
 	}
-	return &exportedFunction{module: m.Instance.Ctx, function: exp.Function}
+	return &exportedFunction{module: m.instance.Ctx, function: exp.Function}
 }
 
 // Memory implements wasm.Module Memory
 func (m *PublicModule) Memory(name string) publicwasm.Memory {
-	exp, err := m.Instance.getExport(name, ExternTypeMemory)
+	exp, err := m.instance.getExport(name, ExternTypeMemory)
 	if err != nil {
 		return nil
 	}

internal/wasm/store_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/binary"
 	"math"
-	"os"
 	"strconv"
 	"testing"
 
@@ -86,48 +85,49 @@ func TestModuleInstance_Memory(t *testing.T) {
 	}
 }
 
-func TestStore_AddHostFunction(t *testing.T) {
+func TestPublicModule_String(t *testing.T) {
 	s := newStore()
 
-	hf, err := NewGoFunc("fn", func(wasm.ModuleContext) {})
+	// Ensure paths that can create the host module can see the name.
+	m, err := s.Instantiate(&Module{}, "module")
 	require.NoError(t, err)
+	require.Equal(t, "Module[module]", m.String())
+	require.Equal(t, "Module[module]", s.Module(m.instance.Name).String())
+}
+
+func TestStore_ReleaseModule(t *testing.T) {
+	s := newStore()
 
 	// Add the host module
-	hostModule := &ModuleInstance{Name: "test", Exports: make(map[string]*ExportInstance, 1)}
-	s.moduleInstances[hostModule.Name] = hostModule
-
-	_, err = s.AddHostFunction(hostModule, hf)
+	hm, err := s.NewHostModule("", map[string]interface{}{"host_fn": func(wasm.ModuleContext) {}})
 	require.NoError(t, err)
 
-	// The function was added to the store, prefixed by the owning module name
-	require.Equal(t, 1, len(s.functions))
-	fn := s.functions[0]
-	require.Equal(t, "test.fn", fn.Name)
+	t.Run("Module imports HostModule", func(t *testing.T) {
+		name := "test"
+		_, err = s.Instantiate(&Module{
+			TypeSection:   []*FunctionType{{}},
+			ImportSection: []*Import{{Type: ExternTypeFunc, Name: "host_fn", DescFunc: 0}},
+			MemorySection: []*MemoryType{{1, nil}},
+		}, name)
+		require.NoError(t, err)
 
-	// The function was exported in the module
-	require.Equal(t, 1, len(hostModule.Exports))
-	exp, ok := hostModule.Exports["fn"]
-	require.True(t, ok)
+		// TODO: We shouldn't be able to release the host module as it is in use!
+		require.NoError(t, s.ReleaseModuleInstance(hm.name))
 
-	// Trying to register it again should fail
-	_, err = s.AddHostFunction(hostModule, hf)
-	require.EqualError(t, err, `"fn" is already exported in module "test"`)
+		// Can release the importing module
+		require.NoError(t, s.ReleaseModuleInstance(name))
+		require.Nil(t, s.moduleInstances[name])
 
-	// Any side effects should be reverted
-	require.Equal(t, []*FunctionInstance{fn, nil}, s.functions)
-	require.Equal(t, map[string]*ExportInstance{"fn": exp}, hostModule.Exports)
+		// Can re-release the importing module
+		require.NoError(t, s.ReleaseModuleInstance(name))
+	})
 }
 
 func TestStore_ExportImportedHostFunction(t *testing.T) {
 	s := newStore()
 
-	hf, err := NewGoFunc("host_fn", func(wasm.ModuleContext) {})
-	require.NoError(t, err)
-
 	// Add the host module
-	hostModule := &ModuleInstance{Name: "", Exports: make(map[string]*ExportInstance, 1)}
-	s.moduleInstances[hostModule.Name] = hostModule
-	_, err = s.AddHostFunction(hostModule, hf)
+	_, err := s.NewHostModule("", map[string]interface{}{"host_fn": func(wasm.ModuleContext) {}})
 	require.NoError(t, err)
 
 	t.Run("ModuleInstance is the importing module", func(t *testing.T) {
@@ -144,7 +144,6 @@ func TestStore_ExportImportedHostFunction(t *testing.T) {
 
 		ei, err := mod.getExport("host.fn", ExternTypeFunc)
 		require.NoError(t, err)
-		os.Environ()
 		// We expect the host function to be called in context of the importing module.
 		// Otherwise, it would be the pseudo-module of the host, which only includes types and function definitions.
 		// Notably, this ensures the host function call context has the correct memory (from the importing module).
@@ -182,17 +181,11 @@ func TestFunctionInstance_Call(t *testing.T) {
 			engine := &catchContext{}
 			store := NewStore(storeCtx, engine, Features20191205)
 
-			// 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)
+			functionName := "fn"
+			hm, err := store.NewHostModule("host",
+				map[string]interface{}{functionName: func(wasm.ModuleContext) {}},
+			)
 			require.NoError(t, err)
 
 			// Make a module to import the function
@@ -200,7 +193,7 @@ func TestFunctionInstance_Call(t *testing.T) {
 				TypeSection: []*FunctionType{{}},
 				ImportSection: []*Import{{
 					Type:     ExternTypeFunc,
-					Module:   hostModule.Name,
+					Module:   hm.name,
 					Name:     functionName,
 					DescFunc: 0,
 				}},

wasi.go

@@ -127,7 +127,7 @@ func StartWASICommand(r Runtime, module *DecodedModule) (wasm.Module, error) {
 	}
 
 	start := mod.Function(internalwasi.FunctionStart)
-	if _, err = start.Call(mod.Instance.Ctx.Context()); err != nil {
+	if _, err = start.Call(internal.ctx); err != nil {
 		return nil, fmt.Errorf("module[%s] function[%s] failed: %w", module.name, internalwasi.FunctionStart, err)
 	}
 	return mod, nil

wasm.go

@@ -2,6 +2,7 @@ package wazero
 
 import (
 	"bytes"
+	"context"
 	"errors"
 
 	internalwasm "github.com/tetratelabs/wazero/internal/wasm"
@@ -66,6 +67,7 @@ func NewRuntime() Runtime {
 // NewRuntimeWithConfig returns a runtime with the given configuration.
 func NewRuntimeWithConfig(config *RuntimeConfig) Runtime {
 	return &runtime{
+		ctx:             config.ctx,
 		store:           internalwasm.NewStore(config.ctx, config.engine, config.enabledFeatures),
 		enabledFeatures: config.enabledFeatures,
 	}
@@ -73,6 +75,7 @@ func NewRuntimeWithConfig(config *RuntimeConfig) Runtime {
 
 // runtime allows decoupling of public interfaces from internal representation.
 type runtime struct {
+	ctx             context.Context
 	store           *internalwasm.Store
 	enabledFeatures internalwasm.Features
 }

wasm/wasm.go

@@ -69,6 +69,8 @@ type Store interface {
 // Note: This is an interface for decoupling, not third-party implementations. All implementations are in wazero.
 // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#external-types%E2%91%A0
 type Module interface {
+	fmt.Stringer
+
 	// Function returns a function exported from this module or nil if it wasn't.
 	Function(name string) Function
 
@@ -157,6 +159,8 @@ type MutableGlobal interface {
 // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#syntax-hostfunc
 // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#external-types%E2%91%A0
 type HostModule interface {
+	fmt.Stringer
+
 	// Function returns a host function exported under this module name or nil if it wasn't.
 	Function(name string) HostFunction
 }