Repository: tetratelabs/wazero
Pull request #1218 (2023-03-10 02:14:24)

Generates typeIDs at compilation time
Previously we generated the typeIDs for function types unnecessarily at each instantiation time. 
TypeIDs are unique per Runtime/Store, and we don't need to create slice []FunctionTypeID for each instance.
When generating a typeID, we need to take a lock on store in order to keep uniqueness across
instances and this caused the lock contention on the instantiation heavy use-cases.

tldr is that this improves the concurrent instantiation performance!

```
$ benchstat main.txt new.txt
name                                    old time/op    new time/op    delta
Initialization/interpreter-10             29.9µs ± 3%    29.6µs ± 1%     ~     (p=0.151 n=5+5)
Initialization/interpreter-multiple-10    17.1µs ± 7%    13.2µs ± 9%  -23.13%  (p=0.008 n=5+5)
Initialization/compiler-10                35.1µs ± 3%    27.2µs ±47%     ~     (p=0.905 n=4+5)
Initialization/compiler-multiple-10       14.8µs ±15%    13.0µs ± 6%  -12.50%  (p=0.008 n=5+5)
Compilation/with_extern_cache-10           148µs ± 1%     152µs ± 3%     ~     (p=0.056 n=5+5)
Compilation/without_extern_cache-10       2.76ms ± 1%    2.74ms ± 2%     ~     (p=0.151 n=5+5)

name                                    old alloc/op   new alloc/op   delta
Initialization/interpreter-10              137kB ± 0%     137kB ± 0%   -0.04%  (p=0.000 n=5+4)
Initialization/interpreter-multiple-10     137kB ± 0%     137kB ± 0%     ~     (p=0.159 n=4+5)
Initialization/compiler-10                 141kB ± 0%     141kB ± 0%   -0.03%  (p=0.029 n=4+4)
Initialization/compiler-multiple-10        142kB ± 0%     142kB ± 0%   -0.10%  (p=0.008 n=5+5)
Compilation/with_extern_cache-10          55.4kB ± 0%    55.5kB ± 0%   +0.22%  (p=0.008 n=5+5)
Compilation/without_extern_cache-10       1.12MB ± 0%    1.12MB ± 0%     ~     (p=1.000 n=5+5)

name                                    old allocs/op  new allocs/op  delta
Initialization/interpreter-10               52.0 ± 0%      51.0 ± 0%   -1.92%  (p=0.008 n=5+5)
Initialization/interpreter-multiple-10      58.0 ± 0%      57.0 ± 0%   -1.72%  (p=0.029 n=4+4)
Initialization/compiler-10                  42.0 ± 0%      41.0 ± 0%   -2.38%  (p=0.008 n=5+5)
Initialization/compiler-multiple-10         48.0 ± 0%      47.0 ± 0%   -2.08%  (p=0.008 n=5+5)
Compilation/with_extern_cache-10           1.10k ± 0%     1.10k ± 0%   +0.18%  (p=0.008 n=5+5)
Compilation/without_extern_cache-10        11.4k ± 0%     11.4k ± 0%     ~     (p=0.079 n=5+5)
```

mathetake: need to backfill tests and comments
mathetake: hmmm. this results in preventing the reuse of "compiledModule" across multiple Runtime.
mathetake: got an idea
mathetake: got an idea

Patch:
diff --git a/builder.go b/builder.go
index dd3e01f8..6093fdd1 100644
--- a/builder.go
+++ b/builder.go
@@ -328,6 +328,13 @@ func (b *hostModuleBuilder) Compile(ctx context.Context) (CompiledModule, error)
 		return nil, err
 	}
 
+	// typeIDs are static and compile-time known.
+	typeIDs, err := b.r.store.GetFunctionTypeIDs(module.TypeSection)
+	if err != nil {
+		return nil, err
+	}
+	c.typeIDs = typeIDs
+
 	return c, nil
 }
 
diff --git a/config.go b/config.go
index df4fa8f3..e66c0b85 100644
--- a/config.go
+++ b/config.go
@@ -333,6 +333,7 @@ type compiledModule struct {
 	compiledEngine wasm.Engine
 	// closeWithModule prevents leaking compiled code when a module is compiled implicitly.
 	closeWithModule bool
+	typeIDs         []wasm.FunctionTypeID
 }
 
 // Name implements CompiledModule.Name
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index 1f559647..c1bc55a0 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -286,6 +286,7 @@ func (s *Store) Instantiate(
 	module *Module,
 	name string,
 	sys *internalsys.Context,
+	typeIDs []FunctionTypeID,
 ) (*CallContext, error) {
 	// Collect any imported modules to avoid locking the store too long.
 	importedModuleNames := map[string]struct{}{}
@@ -305,7 +306,7 @@ func (s *Store) Instantiate(
 	}
 
 	// Instantiate the module and add it to the store so that other modules can import it.
-	if callCtx, err := s.instantiate(ctx, module, name, sys, importedModules); err != nil {
+	if callCtx, err := s.instantiate(ctx, module, name, sys, importedModules, typeIDs); err != nil {
 		_ = s.deleteModule(name)
 		return nil, err
 	} else {
@@ -325,12 +326,8 @@ func (s *Store) instantiate(
 	name string,
 	sysCtx *internalsys.Context,
 	modules map[string]*ModuleInstance,
+	typeIDs []FunctionTypeID,
 ) (*CallContext, error) {
-	typeIDs, err := s.getFunctionTypeIDs(module.TypeSection)
-	if err != nil {
-		return nil, err
-	}
-
 	importedFunctions, importedGlobals, importedTables, importedMemory, err := resolveImports(module, modules)
 	if err != nil {
 		return nil, err
@@ -562,7 +559,7 @@ func executeConstExpression(importedGlobals []*GlobalInstance, expr *ConstantExp
 	return
 }
 
-func (s *Store) getFunctionTypeIDs(ts []*FunctionType) ([]FunctionTypeID, error) {
+func (s *Store) GetFunctionTypeIDs(ts []*FunctionType) ([]FunctionTypeID, error) {
 	ret := make([]FunctionTypeID, len(ts))
 	for i, t := range ts {
 		inst, err := s.getFunctionTypeID(t)
diff --git a/runtime.go b/runtime.go
index c48f7d73..edd4252f 100644
--- a/runtime.go
+++ b/runtime.go
@@ -203,6 +203,13 @@ func (r *runtime) CompileModule(ctx context.Context, binary []byte) (CompiledMod
 
 	c := &compiledModule{module: internal, compiledEngine: r.store.Engine}
 
+	// typeIDs are static and compile-time known.
+	typeIDs, err := r.store.GetFunctionTypeIDs(internal.TypeSection)
+	if err != nil {
+		return nil, err
+	}
+	c.typeIDs = typeIDs
+
 	listeners, err := buildListeners(ctx, internal)
 	if err != nil {
 		return nil, err
@@ -276,7 +283,7 @@ func (r *runtime) InstantiateModule(
 	}
 
 	// Instantiate the module.
-	mod, err = r.store.Instantiate(ctx, code.module, name, sysCtx)
+	mod, err = r.store.Instantiate(ctx, code.module, name, sysCtx, code.typeIDs)
 	if err != nil {
 		// If there was an error, don't leak the compiled module.
 		if code.closeWithModule {


Test patch:
diff --git a/builder_test.go b/builder_test.go
index 02b9567c..1ca4a6e7 100644
--- a/builder_test.go
+++ b/builder_test.go
@@ -279,6 +279,11 @@ func TestNewHostModuleBuilder_Compile(t *testing.T) {
 
 			require.Equal(t, b.r.store.Engine, m.compiledEngine)
 
+			// TypeIDs must be assigned to compiledModule.
+			expTypeIDs, err := b.r.store.GetFunctionTypeIDs(tc.expected.TypeSection)
+			require.NoError(t, err)
+			require.Equal(t, expTypeIDs, m.typeIDs)
+
 			// Built module must be instantiable by Engine.
 			mod, err := b.r.InstantiateModule(testCtx, m, NewModuleConfig())
 			require.NoError(t, err)
@@ -353,6 +358,8 @@ func requireHostModuleEquals(t *testing.T, expected, actual *wasm.Module) {
 	// `require.Equal(t, expected, actual)` fails reflect pointers don't match, so brute compare:
 	for _, tp := range expected.TypeSection {
 		tp.CacheNumInUint64()
+		// When creating the compiled module, we get the type IDs for types, which results in caching type keys.
+		_ = tp.String()
 	}
 	require.Equal(t, expected.TypeSection, actual.TypeSection)
 	require.Equal(t, expected.ImportSection, actual.ImportSection)
diff --git a/cache_test.go b/cache_test.go
index 2df01a07..fcc6ad28 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -27,6 +27,18 @@ func TestCompilationCache(t *testing.T) {
 		foo, bar := getCacheSharedRuntimes(ctx, t)
 		cacheInst := foo.cache
 
+		// Create a different type id on the bar's store so that we can emulate that bar instantiated the module before facWasm.
+		_, err := bar.store.GetFunctionTypeIDs(
+			// Arbitrary one is fine as long as it is not used in facWasm.
+			[]*wasm.FunctionType{{Params: []wasm.ValueType{
+				wasm.ValueTypeI32, wasm.ValueTypeI32, wasm.ValueTypeI32, wasm.ValueTypeI32, wasm.ValueTypeI32, wasm.ValueTypeI32,
+				wasm.ValueTypeI32, wasm.ValueTypeV128, wasm.ValueTypeI32, wasm.ValueTypeV128, wasm.ValueTypeI32, wasm.ValueTypeI32,
+				wasm.ValueTypeI32, wasm.ValueTypeV128, wasm.ValueTypeI32, wasm.ValueTypeV128, wasm.ValueTypeI32, wasm.ValueTypeI32,
+				wasm.ValueTypeI32, wasm.ValueTypeV128, wasm.ValueTypeI32, wasm.ValueTypeV128, wasm.ValueTypeI32, wasm.ValueTypeI32,
+				wasm.ValueTypeI32, wasm.ValueTypeI32, wasm.ValueTypeI32, wasm.ValueTypeI32, wasm.ValueTypeI32, wasm.ValueTypeI32,
+			}}})
+		require.NoError(t, err)
+
 		// add interpreter first, to ensure compiler support isn't order dependent
 		eng := foo.cache.engs[engineKindInterpreter]
 		if platform.CompilerSupported() {
@@ -41,8 +53,12 @@ func TestCompilationCache(t *testing.T) {
 		barCompiled, err := bar.CompileModule(ctx, facWasm)
 		require.NoError(t, err)
 
-		// Ensures compiled modules are the same.
-		require.Equal(t, compiled, barCompiled)
+		// Ensures compiled modules are the same modulo type IDs, which is unique per store.
+		require.Equal(t, compiled.(*compiledModule).module, barCompiled.(*compiledModule).module)
+		require.Equal(t, compiled.(*compiledModule).closeWithModule, barCompiled.(*compiledModule).closeWithModule)
+		require.Equal(t, compiled.(*compiledModule).compiledEngine, barCompiled.(*compiledModule).compiledEngine)
+		// TypeIDs must be different as we create a different type ID on bar beforehand.
+		require.NotEqual(t, compiled.(*compiledModule).typeIDs, barCompiled.(*compiledModule).typeIDs)
 
 		// Two runtimes are completely separate except the compilation cache,
 		// therefore it should be ok to instantiate the same name module for each of them.
diff --git a/internal/engine/compiler/engine_test.go b/internal/engine/compiler/engine_test.go
index 7d98b842..fcd5dbdb 100644
--- a/internal/engine/compiler/engine_test.go
+++ b/internal/engine/compiler/engine_test.go
@@ -264,7 +264,10 @@ func TestCompiler_SliceAllocatedOnHeap(t *testing.T) {
 	err = s.Engine.CompileModule(testCtx, hm, nil, false)
 	require.NoError(t, err)
 
-	_, err = s.Instantiate(testCtx, hm, hostModuleName, nil)
+	typeIDs, err := s.GetFunctionTypeIDs(hm.TypeSection)
+	require.NoError(t, err)
+
+	_, err = s.Instantiate(testCtx, hm, hostModuleName, nil, typeIDs)
 	require.NoError(t, err)
 
 	const stackCorruption = "value_stack_corruption"
@@ -320,7 +323,10 @@ func TestCompiler_SliceAllocatedOnHeap(t *testing.T) {
 	err = s.Engine.CompileModule(testCtx, m, nil, false)
 	require.NoError(t, err)
 
-	mi, err := s.Instantiate(testCtx, m, t.Name(), nil)
+	typeIDs, err = s.GetFunctionTypeIDs(m.TypeSection)
+	require.NoError(t, err)
+
+	mi, err := s.Instantiate(testCtx, m, t.Name(), nil, typeIDs)
 	require.NoError(t, err)
 
 	for _, fnName := range []string{stackCorruption, callStackCorruption} {
diff --git a/internal/integration_test/spectest/spectest.go b/internal/integration_test/spectest/spectest.go
index 0028dda5..4415e59a 100644
--- a/internal/integration_test/spectest/spectest.go
+++ b/internal/integration_test/spectest/spectest.go
@@ -337,7 +337,10 @@ func addSpectestModule(t *testing.T, ctx context.Context, s *wasm.Store, enabled
 	err = s.Engine.CompileModule(ctx, mod, nil, false)
 	require.NoError(t, err)
 
-	_, err = s.Instantiate(ctx, mod, mod.NameSection.ModuleName, sys.DefaultContext(nil))
+	typeIDs, err := s.GetFunctionTypeIDs(mod.TypeSection)
+	require.NoError(t, err)
+
+	_, err = s.Instantiate(ctx, mod, mod.NameSection.ModuleName, sys.DefaultContext(nil), typeIDs)
 	require.NoError(t, err)
 }
 
@@ -404,7 +407,10 @@ func Run(t *testing.T, testDataFS embed.FS, ctx context.Context, fc filecache.Ca
 						err = s.Engine.CompileModule(ctx, mod, nil, false)
 						require.NoError(t, err, msg)
 
-						_, err = s.Instantiate(ctx, mod, moduleName, nil)
+						typeIDs, err := s.GetFunctionTypeIDs(mod.TypeSection)
+						require.NoError(t, err)
+
+						_, err = s.Instantiate(ctx, mod, moduleName, nil, typeIDs)
 						lastInstantiatedModuleName = moduleName
 						require.NoError(t, err)
 					case "register":
@@ -542,7 +548,10 @@ func Run(t *testing.T, testDataFS embed.FS, ctx context.Context, fc filecache.Ca
 							err = s.Engine.CompileModule(ctx, mod, nil, false)
 							require.NoError(t, err, msg)
 
-							_, err = s.Instantiate(ctx, mod, t.Name(), nil)
+							typeIDs, err := s.GetFunctionTypeIDs(mod.TypeSection)
+							require.NoError(t, err)
+
+							_, err = s.Instantiate(ctx, mod, t.Name(), nil, typeIDs)
 							require.NoError(t, err, msg)
 						} else {
 							requireInstantiationError(t, ctx, s, buf, msg)
@@ -578,7 +587,10 @@ func requireInstantiationError(t *testing.T, ctx context.Context, s *wasm.Store,
 		return
 	}
 
-	_, err = s.Instantiate(ctx, mod, t.Name(), nil)
+	typeIDs, err := s.GetFunctionTypeIDs(mod.TypeSection)
+	require.NoError(t, err)
+
+	_, err = s.Instantiate(ctx, mod, t.Name(), nil, typeIDs)
 	require.Error(t, err, msg)
 }
 
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index 111aa95a..cfe42b05 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -74,7 +74,10 @@ func RunTestEngine_MemoryGrowInRecursiveCall(t *testing.T, et EngineTester) {
 	err = s.Engine.CompileModule(testCtx, hm, nil, false)
 	require.NoError(t, err)
 
-	_, err = s.Instantiate(testCtx, hm, hostModuleName, nil)
+	typeIDs, err := s.GetFunctionTypeIDs(hm.TypeSection)
+	require.NoError(t, err)
+
+	_, err = s.Instantiate(testCtx, hm, hostModuleName, nil, typeIDs)
 	require.NoError(t, err)
 
 	m := &wasm.Module{
@@ -106,7 +109,10 @@ func RunTestEngine_MemoryGrowInRecursiveCall(t *testing.T, et EngineTester) {
 	err = s.Engine.CompileModule(testCtx, m, nil, false)
 	require.NoError(t, err)
 
-	inst, err := s.Instantiate(testCtx, m, t.Name(), nil)
+	typeIDs, err = s.GetFunctionTypeIDs(m.TypeSection)
+	require.NoError(t, err)
+
+	inst, err := s.Instantiate(testCtx, m, t.Name(), nil, typeIDs)
 	require.NoError(t, err)
 
 	growFn = inst.Function(2)
diff --git a/internal/wasm/call_context_test.go b/internal/wasm/call_context_test.go
index 4cec6e84..b0a0e616 100644
--- a/internal/wasm/call_context_test.go
+++ b/internal/wasm/call_context_test.go
@@ -39,7 +39,7 @@ func TestCallContext_String(t *testing.T) {
 
 		t.Run(tc.name, func(t *testing.T) {
 			// Ensure paths that can create the host module can see the name.
-			m, err := s.Instantiate(testCtx, &Module{}, tc.moduleName, nil)
+			m, err := s.Instantiate(testCtx, &Module{}, tc.moduleName, nil, nil)
 			defer m.Close(testCtx) //nolint
 
 			require.NoError(t, err)
@@ -78,7 +78,7 @@ func TestCallContext_Close(t *testing.T) {
 		t.Run(fmt.Sprintf("%s calls ns.CloseWithExitCode(module.name))", tc.name), func(t *testing.T) {
 			for _, ctx := range []context.Context{nil, testCtx} { // Ensure it doesn't crash on nil!
 				moduleName := t.Name()
-				m, err := s.Instantiate(ctx, &Module{}, moduleName, nil)
+				m, err := s.Instantiate(ctx, &Module{}, moduleName, nil, nil)
 				require.NoError(t, err)
 
 				// We use side effects to see if Close called ns.CloseWithExitCode (without repeating store_test.go).
@@ -107,7 +107,7 @@ func TestCallContext_Close(t *testing.T) {
 		_, err := fsCtx.OpenFile(testFS, "/foo", os.O_RDONLY, 0)
 		require.NoError(t, err)
 
-		m, err := s.Instantiate(testCtx, &Module{}, t.Name(), sysCtx)
+		m, err := s.Instantiate(testCtx, &Module{}, t.Name(), sysCtx, nil)
 		require.NoError(t, err)
 
 		// We use side effects to determine if Close in fact called Context.Close (without repeating sys_test.go).
@@ -142,7 +142,7 @@ func TestCallContext_Close(t *testing.T) {
 		_, err := fsCtx.OpenFile(testFS, "/foo", os.O_RDONLY, 0)
 		require.NoError(t, err)
 
-		m, err := s.Instantiate(testCtx, &Module{}, t.Name(), sysCtx)
+		m, err := s.Instantiate(testCtx, &Module{}, t.Name(), sysCtx, nil)
 		require.NoError(t, err)
 
 		require.EqualError(t, m.Close(testCtx), "error closing")
@@ -182,7 +182,7 @@ func TestCallContext_CallDynamic(t *testing.T) {
 		t.Run(fmt.Sprintf("%s calls ns.CloseWithExitCode(module.name))", tc.name), func(t *testing.T) {
 			for _, ctx := range []context.Context{nil, testCtx} { // Ensure it doesn't crash on nil!
 				moduleName := t.Name()
-				m, err := s.Instantiate(ctx, &Module{}, moduleName, nil)
+				m, err := s.Instantiate(ctx, &Module{}, moduleName, nil, nil)
 				require.NoError(t, err)
 
 				// We use side effects to see if Close called ns.CloseWithExitCode (without repeating store_test.go).
@@ -211,7 +211,7 @@ func TestCallContext_CallDynamic(t *testing.T) {
 		_, err := fsCtx.OpenFile(testFS, "/foo", os.O_RDONLY, 0)
 		require.NoError(t, err)
 
-		m, err := s.Instantiate(testCtx, &Module{}, t.Name(), sysCtx)
+		m, err := s.Instantiate(testCtx, &Module{}, t.Name(), sysCtx, nil)
 		require.NoError(t, err)
 
 		// We use side effects to determine if Close in fact called Context.Close (without repeating sys_test.go).
@@ -240,7 +240,7 @@ func TestCallContext_CallDynamic(t *testing.T) {
 		_, err := fsCtx.OpenFile(testFS, path, os.O_RDONLY, 0)
 		require.NoError(t, err)
 
-		m, err := s.Instantiate(testCtx, &Module{}, t.Name(), sysCtx)
+		m, err := s.Instantiate(testCtx, &Module{}, t.Name(), sysCtx, nil)
 		require.NoError(t, err)
 
 		require.EqualError(t, m.Close(testCtx), "error closing")
diff --git a/internal/wasm/global_test.go b/internal/wasm/global_test.go
index a0e5327a..7cb95cb6 100644
--- a/internal/wasm/global_test.go
+++ b/internal/wasm/global_test.go
@@ -296,7 +296,7 @@ func TestPublicModule_Global(t *testing.T) {
 		s := newStore()
 		t.Run(tc.name, func(t *testing.T) {
 			// Instantiate the module and get the export of the above global
-			module, err := s.Instantiate(context.Background(), tc.module, t.Name(), nil)
+			module, err := s.Instantiate(context.Background(), tc.module, t.Name(), nil, nil)
 			require.NoError(t, err)
 
 			if global := module.ExportedGlobal("global"); tc.expected != nil {
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 5643bc1f..ba5fc476 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -80,7 +80,7 @@ func TestModuleInstance_Memory(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			s := newStore()
 
-			instance, err := s.Instantiate(testCtx, tc.input, "test", nil)
+			instance, err := s.Instantiate(testCtx, tc.input, "test", nil, nil)
 			require.NoError(t, err)
 
 			mem := instance.ExportedMemory("memory")
@@ -109,7 +109,7 @@ func TestStore_Instantiate(t *testing.T) {
 	require.NoError(t, err)
 
 	sysCtx := sys.DefaultContext(nil)
-	mod, err := s.Instantiate(testCtx, m, "", sysCtx)
+	mod, err := s.Instantiate(testCtx, m, "", sysCtx, []FunctionTypeID{0})
 	require.NoError(t, err)
 	defer mod.Close(testCtx)
 
@@ -150,7 +150,7 @@ func TestStore_CloseWithExitCode(t *testing.T) {
 				CodeSection:               []*Code{{Body: []byte{OpcodeEnd}}},
 				ExportSection:             []*Export{{Type: ExternTypeFunc, Index: 0, Name: "fn"}},
 				FunctionDefinitionSection: []*FunctionDefinition{{funcType: v_v}},
-			}, importedModuleName, nil)
+			}, importedModuleName, nil, []FunctionTypeID{0})
 			require.NoError(t, err)
 
 			m2, err := s.Instantiate(testCtx, &Module{
@@ -160,7 +160,7 @@ func TestStore_CloseWithExitCode(t *testing.T) {
 				MemoryDefinitionSection: []*MemoryDefinition{{}},
 				GlobalSection:           []*Global{{Type: &GlobalType{}, Init: &ConstantExpression{Opcode: OpcodeI32Const, Data: const1}}},
 				TableSection:            []*Table{{Min: 10}},
-			}, importingModuleName, nil)
+			}, importingModuleName, nil, []FunctionTypeID{0})
 			require.NoError(t, err)
 
 			if tc.testClosed {
@@ -187,7 +187,7 @@ func TestStore_hammer(t *testing.T) {
 	require.NoError(t, err)
 
 	s := newStore()
-	imported, err := s.Instantiate(testCtx, m, importedModuleName, nil)
+	imported, err := s.Instantiate(testCtx, m, importedModuleName, nil, []FunctionTypeID{0})
 	require.NoError(t, err)
 
 	_, ok := s.nameToNode[imported.Name()]
@@ -219,7 +219,7 @@ func TestStore_hammer(t *testing.T) {
 		N = 100
 	}
 	hammer.NewHammer(t, P, N).Run(func(name string) {
-		mod, instantiateErr := s.Instantiate(testCtx, importingModule, name, sys.DefaultContext(nil))
+		mod, instantiateErr := s.Instantiate(testCtx, importingModule, name, sys.DefaultContext(nil), []FunctionTypeID{0})
 		require.NoError(t, instantiateErr)
 		require.NoError(t, mod.Close(testCtx))
 	}, nil)
@@ -241,7 +241,7 @@ func TestStore_hammer_close(t *testing.T) {
 	require.NoError(t, err)
 
 	s := newStore()
-	imported, err := s.Instantiate(testCtx, m, importedModuleName, nil)
+	imported, err := s.Instantiate(testCtx, m, importedModuleName, nil, []FunctionTypeID{0})
 	require.NoError(t, err)
 
 	_, ok := s.nameToNode[imported.Name()]
@@ -267,7 +267,7 @@ func TestStore_hammer_close(t *testing.T) {
 	const instCount = 10000
 	instances := make([]api.Module, instCount)
 	for i := 0; i < instCount; i++ {
-		mod, instantiateErr := s.Instantiate(testCtx, importingModule, strconv.Itoa(i), sys.DefaultContext(nil))
+		mod, instantiateErr := s.Instantiate(testCtx, importingModule, strconv.Itoa(i), sys.DefaultContext(nil), []FunctionTypeID{0})
 		require.NoError(t, instantiateErr)
 		instances[i] = mod
 	}
@@ -301,17 +301,17 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 
 	t.Run("Fails if module name already in use", func(t *testing.T) {
 		s := newStore()
-		_, err = s.Instantiate(testCtx, m, importedModuleName, nil)
+		_, err = s.Instantiate(testCtx, m, importedModuleName, nil, []FunctionTypeID{0})
 		require.NoError(t, err)
 
 		// Trying to register it again should fail
-		_, err = s.Instantiate(testCtx, m, importedModuleName, nil)
+		_, err = s.Instantiate(testCtx, m, importedModuleName, nil, []FunctionTypeID{0})
 		require.EqualError(t, err, "module[imported] has already been instantiated")
 	})
 
 	t.Run("fail resolve import", func(t *testing.T) {
 		s := newStore()
-		_, err = s.Instantiate(testCtx, m, importedModuleName, nil)
+		_, err = s.Instantiate(testCtx, m, importedModuleName, nil, []FunctionTypeID{0})
 		require.NoError(t, err)
 
 		hm := s.nameToNode[importedModuleName]
@@ -325,14 +325,14 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 				// But the second one tries to import uninitialized-module ->
 				{Type: ExternTypeFunc, Module: "non-exist", Name: "fn", DescFunc: 0},
 			},
-		}, importingModuleName, nil)
+		}, importingModuleName, nil, nil)
 		require.EqualError(t, err, "module[non-exist] not instantiated")
 	})
 
 	t.Run("creating engine failed", func(t *testing.T) {
 		s := newStore()
 
-		_, err = s.Instantiate(testCtx, m, importedModuleName, nil)
+		_, err = s.Instantiate(testCtx, m, importedModuleName, nil, []FunctionTypeID{0})
 		require.NoError(t, err)
 
 		hm := s.nameToNode[importedModuleName]
@@ -354,7 +354,7 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 		}
 		importingModule.BuildFunctionDefinitions()
 
-		_, err = s.Instantiate(testCtx, importingModule, importingModuleName, nil)
+		_, err = s.Instantiate(testCtx, importingModule, importingModuleName, nil, []FunctionTypeID{0})
 		require.EqualError(t, err, "some engine creation error")
 	})
 
@@ -363,7 +363,7 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 		engine := s.Engine.(*mockEngine)
 		engine.callFailIndex = 1
 
-		_, err = s.Instantiate(testCtx, m, importedModuleName, nil)
+		_, err = s.Instantiate(testCtx, m, importedModuleName, nil, []FunctionTypeID{0})
 		require.NoError(t, err)
 
 		hm := s.nameToNode[importedModuleName]
@@ -381,7 +381,7 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 		}
 		importingModule.BuildFunctionDefinitions()
 
-		_, err = s.Instantiate(testCtx, importingModule, importingModuleName, nil)
+		_, err = s.Instantiate(testCtx, importingModule, importingModuleName, nil, []FunctionTypeID{0})
 		require.EqualError(t, err, "start function[1] failed: call failed")
 	})
 }
diff --git a/runtime_test.go b/runtime_test.go
index 27517323..18e08e11 100644
--- a/runtime_test.go
+++ b/runtime_test.go
@@ -44,32 +44,32 @@ func TestRuntime_CompileModule(t *testing.T) {
 	tests := []struct {
 		name          string
 		runtime       Runtime
-		wasm          []byte
+		wasm          *wasm.Module
 		moduleBuilder HostModuleBuilder
 		expected      func(CompiledModule)
 	}{
 		{
 			name: "no name section",
-			wasm: binaryencoding.EncodeModule(&wasm.Module{}),
+			wasm: &wasm.Module{},
 		},
 		{
 			name: "empty NameSection.ModuleName",
-			wasm: binaryencoding.EncodeModule(&wasm.Module{NameSection: &wasm.NameSection{}}),
+			wasm: &wasm.Module{NameSection: &wasm.NameSection{}},
 		},
 		{
 			name: "NameSection.ModuleName",
-			wasm: binaryencoding.EncodeModule(&wasm.Module{NameSection: &wasm.NameSection{ModuleName: "test"}}),
+			wasm: &wasm.Module{NameSection: &wasm.NameSection{ModuleName: "test"}},
 			expected: func(compiled CompiledModule) {
 				require.Equal(t, "test", compiled.Name())
 			},
 		},
 		{
 			name: "FunctionSection, but not exported",
-			wasm: binaryencoding.EncodeModule(&wasm.Module{
+			wasm: &wasm.Module{
 				TypeSection:     []*wasm.FunctionType{{Params: []api.ValueType{api.ValueTypeI32}}},
 				FunctionSection: []wasm.Index{0},
 				CodeSection:     []*wasm.Code{{Body: []byte{wasm.OpcodeEnd}}},
-			}),
+			},
 			expected: func(compiled CompiledModule) {
 				require.Nil(t, compiled.ImportedFunctions())
 				require.Zero(t, len(compiled.ExportedFunctions()))
@@ -77,7 +77,7 @@ func TestRuntime_CompileModule(t *testing.T) {
 		},
 		{
 			name: "FunctionSection exported",
-			wasm: binaryencoding.EncodeModule(&wasm.Module{
+			wasm: &wasm.Module{
 				TypeSection:     []*wasm.FunctionType{{Params: []api.ValueType{api.ValueTypeI32}}},
 				FunctionSection: []wasm.Index{0},
 				CodeSection:     []*wasm.Code{{Body: []byte{wasm.OpcodeEnd}}},
@@ -86,7 +86,7 @@ func TestRuntime_CompileModule(t *testing.T) {
 					Name:  "function",
 					Index: 0,
 				}},
-			}),
+			},
 			expected: func(compiled CompiledModule) {
 				require.Nil(t, compiled.ImportedFunctions())
 				f := compiled.ExportedFunctions()["function"]
@@ -95,9 +95,9 @@ func TestRuntime_CompileModule(t *testing.T) {
 		},
 		{
 			name: "MemorySection, but not exported",
-			wasm: binaryencoding.EncodeModule(&wasm.Module{
+			wasm: &wasm.Module{
 				MemorySection: &wasm.Memory{Min: 2, Max: 3, IsMaxEncoded: true},
-			}),
+			},
 			expected: func(compiled CompiledModule) {
 				require.Nil(t, compiled.ImportedMemories())
 				require.Zero(t, len(compiled.ExportedMemories()))
@@ -105,14 +105,14 @@ func TestRuntime_CompileModule(t *testing.T) {
 		},
 		{
 			name: "MemorySection exported",
-			wasm: binaryencoding.EncodeModule(&wasm.Module{
+			wasm: &wasm.Module{
 				MemorySection: &wasm.Memory{Min: 2, Max: 3, IsMaxEncoded: true},
 				ExportSection: []*wasm.Export{{
 					Type:  wasm.ExternTypeMemory,
 					Name:  "memory",
 					Index: 0,
 				}},
-			}),
+			},
 			expected: func(compiled CompiledModule) {
 				require.Nil(t, compiled.ImportedMemories())
 				mem := compiled.ExportedMemories()["memory"]
@@ -124,20 +124,29 @@ func TestRuntime_CompileModule(t *testing.T) {
 		},
 	}
 
-	r := NewRuntime(testCtx)
-	defer r.Close(testCtx)
+	_r := NewRuntime(testCtx)
+	defer _r.Close(testCtx)
+
+	r := _r.(*runtime)
 
 	for _, tt := range tests {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			m, err := r.CompileModule(testCtx, tc.wasm)
+			bin := binaryencoding.EncodeModule(tc.wasm)
+
+			m, err := r.CompileModule(testCtx, bin)
 			require.NoError(t, err)
 			if tc.expected == nil {
 				tc.expected = func(CompiledModule) {}
 			}
 			tc.expected(m)
-			require.Equal(t, r.(*runtime).store.Engine, m.(*compiledModule).compiledEngine)
+			require.Equal(t, r.store.Engine, m.(*compiledModule).compiledEngine)
+
+			// TypeIDs must be assigned to compiledModule.
+			expTypeIDs, err := r.store.GetFunctionTypeIDs(tc.wasm.TypeSection)
+			require.NoError(t, err)
+			require.Equal(t, expTypeIDs, m.(*compiledModule).typeIDs)
 		})
 	}
 }