Repository: tetratelabs/wazero
Pull request #1456 (2023-05-11 14:38:18)

compiler: avoid some unsafe.Pointers
These seem unneeded, so best to avoid.

ncruces: Do we know the recent source of test failures?
achille-roussel: Looks like it probably was https://github.com/tetratelabs/wazero/pull/1453, weird that the tests passed in that PR 🤔 

It seems like a linter issue, one of the fields on the compiler engine isn't used anymore, do you want to address as part of this change?
ncruces: Fixed one of the CI failures, the other one is pretty out of scope.
Trying to address it elsewhere.

Patch:
diff --git a/internal/engine/compiler/arch.go b/internal/engine/compiler/arch.go
index 3eb1cd3f..81b9efc3 100644
--- a/internal/engine/compiler/arch.go
+++ b/internal/engine/compiler/arch.go
@@ -11,10 +11,9 @@ var newArchContext func() archContext
 
 // nativecall is used by callEngine.execWasmFunction and the entrypoint to enter the compiled native code.
 // codeSegment is the pointer to the initial instruction of the compiled native code.
-// ce is "*callEngine" as uintptr.
 //
 // Note: this is implemented in per-arch Go assembler file. For example, arch_amd64.s implements this for amd64.
-func nativecall(codeSegment, ce uintptr, moduleInstanceAddress *wasm.ModuleInstance)
+func nativecall(codeSegment uintptr, ce *callEngine, moduleInstanceAddress *wasm.ModuleInstance)
 
 // registerNameFn is used for debugging purpose to have register symbols in the string of runtimeValueLocation.
 var registerNameFn func(register asm.Register) string
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index 9595fec6..3c12755c 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -129,11 +129,6 @@ type (
 		// initialFn is the initial function for this call engine.
 		initialFn *function
 
-		// ctx is the context.Context passed to all the host function calls.
-		// This is modified when there's a function listener call, otherwise it's always the context.Context
-		// passed to the Call API.
-		ctx context.Context
-
 		// stackIterator provides a way to iterate over the stack for Listeners.
 		// It is setup and valid only during a call to a Listener hook.
 		stackIterator stackIterator
@@ -543,7 +538,7 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, listeners
 			cmp.Init(typ, nil, lsn != nil)
 			withGoFunc = true
 			if body, err = compileGoDefinedHostFunction(cmp); err != nil {
-				def := module.FunctionDefinitionSection[funcIndex+importedFuncs]
+				def := module.FunctionDefinition(funcIndex + importedFuncs)
 				return fmt.Errorf("error compiling host go func[%s]: %w", def.DebugName(), err)
 			}
 			compiledFn.goFunc = codeSeg.GoFunc
@@ -556,7 +551,7 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, listeners
 
 			body, compiledFn.stackPointerCeil, compiledFn.sourceOffsetMap, err = compileWasmFunction(cmp, ir)
 			if err != nil {
-				def := module.FunctionDefinitionSection[funcIndex+importedFuncs]
+				def := module.FunctionDefinition(funcIndex + importedFuncs)
 				return fmt.Errorf("error compiling wasm func[%s]: %w", def.DebugName(), err)
 			}
 		}
@@ -695,7 +690,7 @@ func (ce *callEngine) Definition() api.FunctionDefinition {
 
 func (f *function) definition() api.FunctionDefinition {
 	compiled := f.parent
-	return &compiled.parent.source.FunctionDefinitionSection[compiled.index]
+	return compiled.parent.source.FunctionDefinition(compiled.index)
 }
 
 // Call implements the same method as documented on wasm.ModuleEngine.
@@ -998,7 +993,7 @@ func (ce *callEngine) execWasmFunction(ctx context.Context, m *wasm.ModuleInstan
 entry:
 	{
 		// Call into the native code.
-		nativecall(codeAddr, uintptr(unsafe.Pointer(ce)), modAddr)
+		nativecall(codeAddr, ce, modAddr)
 
 		// Check the status code from Compiler code.
 		switch status := ce.exitContext.statusCode; status {
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index a820005c..1cf33b0b 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -332,7 +332,7 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, listeners
 			}
 			err = e.lowerIR(ir, compiled)
 			if err != nil {
-				def := module.FunctionDefinitionSection[uint32(i)+module.ImportFunctionCount]
+				def := module.FunctionDefinition(uint32(i) + module.ImportFunctionCount)
 				return fmt.Errorf("failed to lower func[%s] to wazeroir: %w", def.DebugName(), err)
 			}
 		}
@@ -484,7 +484,7 @@ func (ce *callEngine) Definition() api.FunctionDefinition {
 
 func (f *function) definition() api.FunctionDefinition {
 	compiled := f.parent
-	return &compiled.source.FunctionDefinitionSection[compiled.index]
+	return compiled.source.FunctionDefinition(compiled.index)
 }
 
 // Call implements the same method as documented on api.Function.
diff --git a/internal/wasm/function_definition.go b/internal/wasm/function_definition.go
index 51d530a7..c5f6e912 100644
--- a/internal/wasm/function_definition.go
+++ b/internal/wasm/function_definition.go
@@ -10,11 +10,8 @@ import (
 //
 // Note: Unlike ExportedFunctions, there is no unique constraint on imports.
 func (m *Module) ImportedFunctions() (ret []api.FunctionDefinition) {
-	for i := range m.FunctionDefinitionSection {
-		d := &m.FunctionDefinitionSection[i]
-		if d.importDesc != nil {
-			ret = append(ret, d)
-		}
+	for i := uint32(0); i < m.ImportFunctionCount; i++ {
+		ret = append(ret, m.FunctionDefinition(i))
 	}
 	return
 }
@@ -22,21 +19,30 @@ func (m *Module) ImportedFunctions() (ret []api.FunctionDefinition) {
 // ExportedFunctions returns the definitions of each exported function.
 func (m *Module) ExportedFunctions() map[string]api.FunctionDefinition {
 	ret := map[string]api.FunctionDefinition{}
-	for i := range m.FunctionDefinitionSection {
-		d := &m.FunctionDefinitionSection[i]
-		for _, e := range d.exportNames {
-			ret[e] = d
+	for i := range m.ExportSection {
+		exp := &m.ExportSection[i]
+		if exp.Type == ExternTypeFunc {
+			d := m.FunctionDefinition(exp.Index)
+			ret[exp.Name] = d
 		}
 	}
 	return ret
 }
 
-// BuildFunctionDefinitions generates function metadata that can be parsed from
+// FunctionDefinition returns the FunctionDefinition for the given `index`.
+func (m *Module) FunctionDefinition(index Index) *FunctionDefinition {
+	// TODO: function initialization is lazy, but bulk. Make it per function.
+	m.buildFunctionDefinitions()
+	return &m.FunctionDefinitionSection[index]
+}
+
+// buildFunctionDefinitions generates function metadata that can be parsed from
 // the module. This must be called after all validation.
-//
-// Note: This is exported for tests who don't use wazero.Runtime or
-// NewHostModule to compile the module.
-func (m *Module) BuildFunctionDefinitions() {
+func (m *Module) buildFunctionDefinitions() {
+	m.functionDefinitionSectionInitOnce.Do(m.buildFunctionDefinitionsOnce)
+}
+
+func (m *Module) buildFunctionDefinitionsOnce() {
 	var moduleName string
 	var functionNames NameMap
 	var localNames, resultNames IndirectNameMap
diff --git a/internal/wasm/host.go b/internal/wasm/host.go
index 400bcea3..bb9c40f5 100644
--- a/internal/wasm/host.go
+++ b/internal/wasm/host.go
@@ -73,7 +73,6 @@ func NewHostModule(
 	// TODO: refactor engines so that we can properly cache compiled machine codes for host modules.
 	m.AssignModuleID([]byte(fmt.Sprintf("@@@@@@@@%p", m)), // @@@@@@@@ = any 8 bytes different from Wasm header.
 		false, false)
-	m.BuildFunctionDefinitions()
 	return
 }
 
diff --git a/internal/wasm/module.go b/internal/wasm/module.go
index 57063bec..ba2c032a 100644
--- a/internal/wasm/module.go
+++ b/internal/wasm/module.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"sort"
 	"strings"
+	"sync"
 
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/internal/ieee754"
@@ -169,10 +170,13 @@ type Module struct {
 	// IsHostModule true if this is the host module, false otherwise.
 	IsHostModule bool
 
-	// FunctionDefinitionSection is a wazero-specific section built on Validate.
+	// functionDefinitionSectionInitOnce guards FunctionDefinitionSection so that it is initialized exactly once.
+	functionDefinitionSectionInitOnce sync.Once
+
+	// FunctionDefinitionSection is a wazero-specific section.
 	FunctionDefinitionSection []FunctionDefinition
 
-	// MemoryDefinitionSection is a wazero-specific section built on Validate.
+	// MemoryDefinitionSection is a wazero-specific section.
 	MemoryDefinitionSection []MemoryDefinition
 
 	// DWARFLines is used to emit DWARF based stack trace. This is created from the multiple custom sections
@@ -212,28 +216,28 @@ func boolToByte(b bool) (ret byte) {
 	return
 }
 
-// TypeOfFunction returns the wasm.SectionIDType index for the given function space index or nil.
-// Note: The function index is preceded by imported functions.
-// TODO: Returning nil should be impossible when decode results are validated. Validate decode before back-filling tests.
-func (m *Module) TypeOfFunction(funcIdx Index) *FunctionType {
-	typeSectionLength := uint32(len(m.TypeSection))
-	if typeSectionLength == 0 {
-		return nil
-	}
-	funcImportCount := Index(0)
-	for i := range m.ImportSection {
-		imp := &m.ImportSection[i]
-		if imp.Type == ExternTypeFunc {
-			if funcIdx == funcImportCount {
+// typeOfFunction returns the wasm.FunctionType for the given function space index or nil.
+func (m *Module) typeOfFunction(funcIdx Index) *FunctionType {
+	typeSectionLength, importedFunctionCount := uint32(len(m.TypeSection)), m.ImportFunctionCount
+	if funcIdx < importedFunctionCount {
+		// Imports are not exclusively functions. This is the current function index in the loop.
+		cur := Index(0)
+		for i := range m.ImportSection {
+			imp := &m.ImportSection[i]
+			if imp.Type != ExternTypeFunc {
+				continue
+			}
+			if funcIdx == cur {
 				if imp.DescFunc >= typeSectionLength {
 					return nil
 				}
 				return &m.TypeSection[imp.DescFunc]
 			}
-			funcImportCount++
+			cur++
 		}
 	}
-	funcSectionIdx := funcIdx - funcImportCount
+
+	funcSectionIdx := funcIdx - m.ImportFunctionCount
 	if funcSectionIdx >= uint32(len(m.FunctionSection)) {
 		return nil
 	}
@@ -296,7 +300,7 @@ func (m *Module) validateStartSection() error {
 	// TODO: this should be verified during decode so that errors have the correct source positions
 	if m.StartSection != nil {
 		startIndex := *m.StartSection
-		ft := m.TypeOfFunction(startIndex)
+		ft := m.typeOfFunction(startIndex)
 		if ft == nil { // TODO: move this check to decoder so that a module can never be decoded invalidly
 			return fmt.Errorf("invalid start function: func[%d] has an invalid type", startIndex)
 		}
diff --git a/internal/wasm/module_instance.go b/internal/wasm/module_instance.go
index a11e9a83..20cdcd96 100644
--- a/internal/wasm/module_instance.go
+++ b/internal/wasm/module_instance.go
@@ -203,7 +203,7 @@ func (m *ModuleInstance) ExportedFunctionDefinitions() map[string]api.FunctionDe
 	result := map[string]api.FunctionDefinition{}
 	for name, exp := range m.Exports {
 		if exp.Type == ExternTypeFunc {
-			result[name] = &m.Definitions[exp.Index]
+			result[name] = m.Source.FunctionDefinition(exp.Index)
 		}
 	}
 	return result
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index 8134c50a..264242cc 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -127,8 +127,8 @@ type (
 		//
 		// Note: This is currently only used for spectests and will be nil in most cases.
 		aliases []string
-		// Definitions is derived from *Module, and is constructed during compilation phrase.
-		Definitions []FunctionDefinition
+		// Source is a pointer to the Module from which this ModuleInstance derives.
+		Source *Module
 	}
 
 	// DataInstance holds bytes corresponding to the data segment in a module.
@@ -321,7 +321,7 @@ func (s *Store) instantiate(
 	sysCtx *internalsys.Context,
 	typeIDs []FunctionTypeID,
 ) (m *ModuleInstance, err error) {
-	m = &ModuleInstance{ModuleName: name, TypeIDs: typeIDs, Sys: sysCtx, s: s, Definitions: module.FunctionDefinitionSection}
+	m = &ModuleInstance{ModuleName: name, TypeIDs: typeIDs, Sys: sysCtx, s: s, Source: module}
 
 	m.Tables = make([]*TableInstance, int(module.ImportTableCount)+len(module.TableSection))
 	m.Globals = make([]*GlobalInstance, int(module.ImportGlobalCount)+len(module.GlobalSection))
@@ -396,9 +396,10 @@ func (m *ModuleInstance) resolveImports(module *Module) (err error) {
 			switch i.Type {
 			case ExternTypeFunc:
 				expectedType := &module.TypeSection[i.DescFunc]
-				actual := &importedModule.Definitions[imported.Index]
-				if !actual.Functype.EqualsSignature(expectedType.Params, expectedType.Results) {
-					err = errorInvalidImport(i, fmt.Errorf("signature mismatch: %s != %s", expectedType, actual.Functype))
+				src := importedModule.Source
+				actual := src.typeOfFunction(imported.Index)
+				if !actual.EqualsSignature(expectedType.Params, expectedType.Results) {
+					err = errorInvalidImport(i, fmt.Errorf("signature mismatch: %s != %s", expectedType, actual))
 					return
 				}
 
diff --git a/runtime.go b/runtime.go
index b651861c..789069d8 100644
--- a/runtime.go
+++ b/runtime.go
@@ -215,8 +215,8 @@ func (r *runtime) CompileModule(ctx context.Context, binary []byte) (CompiledMod
 		return nil, err
 	}
 
-	// Now that the module is validated, cache the function and memory definitions.
-	internal.BuildFunctionDefinitions()
+	// Now that the module is validated, cache the memory definitions.
+	// TODO: lazy initialization of memory definition.
 	internal.BuildMemoryDefinitions()
 
 	c := &compiledModule{module: internal, compiledEngine: r.store.Engine}
@@ -249,7 +249,7 @@ func buildFunctionListeners(ctx context.Context, internal *wasm.Module) ([]exper
 	importCount := internal.ImportFunctionCount
 	listeners := make([]experimentalapi.FunctionListener, len(internal.FunctionSection))
 	for i := 0; i < len(listeners); i++ {
-		listeners[i] = factory.NewFunctionListener(&internal.FunctionDefinitionSection[uint32(i)+importCount])
+		listeners[i] = factory.NewFunctionListener(internal.FunctionDefinition(uint32(i) + importCount))
 	}
 	return listeners, nil
 }


Test patch:
diff --git a/experimental/logging/log_listener_test.go b/experimental/logging/log_listener_test.go
index 8234f199..284cbb94 100644
--- a/experimental/logging/log_listener_test.go
+++ b/experimental/logging/log_listener_test.go
@@ -274,8 +274,7 @@ func Test_loggingListener(t *testing.T) {
 			} else {
 				m.CodeSection = []wasm.Code{{Body: []byte{wasm.OpcodeEnd}}}
 			}
-			m.BuildFunctionDefinitions()
-			def := &m.FunctionDefinitionSection[0]
+			def := m.FunctionDefinition(0)
 			l := lf.NewFunctionListener(def)
 
 			out.Reset()
@@ -309,10 +308,9 @@ func Test_loggingListener_indentation(t *testing.T) {
 			FunctionNames: wasm.NameMap{{Index: 0, Name: "fn1"}, {Index: 1, Name: "fn2"}},
 		},
 	}
-	m.BuildFunctionDefinitions()
-	def1 := &m.FunctionDefinitionSection[0]
+	def1 := m.FunctionDefinition(0)
 	l1 := lf.NewFunctionListener(def1)
-	def2 := &m.FunctionDefinitionSection[1]
+	def2 := m.FunctionDefinition(1)
 	l2 := lf.NewFunctionListener(def2)
 
 	l1.Before(testCtx, nil, def1, []uint64{}, nil)
diff --git a/internal/engine/compiler/compiler_bench_test.go b/internal/engine/compiler/compiler_bench_test.go
index 71f8cfb2..88b76462 100644
--- a/internal/engine/compiler/compiler_bench_test.go
+++ b/internal/engine/compiler/compiler_bench_test.go
@@ -114,8 +114,7 @@ func (j *compilerEnv) execBench(b *testing.B, codeSegment []byte) {
 	for i := 0; i < b.N; i++ {
 		nativecall(
 			uintptr(unsafe.Pointer(&codeSegment[0])),
-			uintptr(unsafe.Pointer(j.ce)),
-			j.moduleInstance,
+			j.ce, j.moduleInstance,
 		)
 	}
 	b.StopTimer()
diff --git a/internal/engine/compiler/compiler_controlflow_test.go b/internal/engine/compiler/compiler_controlflow_test.go
index 069a728c..4240e34a 100644
--- a/internal/engine/compiler/compiler_controlflow_test.go
+++ b/internal/engine/compiler/compiler_controlflow_test.go
@@ -37,10 +37,7 @@ func TestCompiler_compileHostFunction(t *testing.T) {
 
 	// Re-enter the return address.
 	require.NotEqual(t, uintptr(0), uintptr(env.ce.returnAddress))
-	nativecall(env.ce.returnAddress,
-		uintptr(unsafe.Pointer(env.callEngine())),
-		env.module(),
-	)
+	nativecall(env.ce.returnAddress, env.callEngine(), env.module())
 
 	// After that, the code must exit with returned status.
 	require.Equal(t, nativeCallStatusCodeReturned, env.compilerStatus())
diff --git a/internal/engine/compiler/compiler_initialization_test.go b/internal/engine/compiler/compiler_initialization_test.go
index c2c153f7..8f92ca87 100644
--- a/internal/engine/compiler/compiler_initialization_test.go
+++ b/internal/engine/compiler/compiler_initialization_test.go
@@ -249,10 +249,7 @@ func TestCompiler_compileMaybeGrowStack(t *testing.T) {
 				// Reenter from the return address.
 				returnAddress := env.ce.returnAddress
 				require.True(t, returnAddress != 0, "returnAddress was zero %d", returnAddress)
-				nativecall(
-					returnAddress, uintptr(unsafe.Pointer(env.callEngine())),
-					env.module(),
-				)
+				nativecall(returnAddress, env.callEngine(), env.module())
 
 				// Check the result. This should be "Returned".
 				require.Equal(t, nativeCallStatusCodeReturned, env.compilerStatus())
diff --git a/internal/engine/compiler/compiler_memory_test.go b/internal/engine/compiler/compiler_memory_test.go
index 7dea2cbb..b3aadab3 100644
--- a/internal/engine/compiler/compiler_memory_test.go
+++ b/internal/engine/compiler/compiler_memory_test.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"math"
 	"testing"
-	"unsafe"
 
 	"github.com/tetratelabs/wazero/internal/testing/require"
 	"github.com/tetratelabs/wazero/internal/wasm"
@@ -41,7 +40,7 @@ func TestCompiler_compileMemoryGrow(t *testing.T) {
 	// Reenter from the return address.
 	nativecall(
 		env.ce.returnAddress,
-		uintptr(unsafe.Pointer(env.callEngine())),
+		env.callEngine(),
 		env.module(),
 	)
 
diff --git a/internal/engine/compiler/compiler_test.go b/internal/engine/compiler/compiler_test.go
index a2ff5d7b..505c034b 100644
--- a/internal/engine/compiler/compiler_test.go
+++ b/internal/engine/compiler/compiler_test.go
@@ -213,8 +213,7 @@ func (j *compilerEnv) exec(machineCode []byte) {
 
 	nativecall(
 		uintptr(unsafe.Pointer(&executable[0])),
-		uintptr(unsafe.Pointer(j.ce)),
-		j.moduleInstance,
+		j.ce, j.moduleInstance,
 	)
 }
 
diff --git a/internal/engine/compiler/engine_test.go b/internal/engine/compiler/engine_test.go
index 7938a52e..302bb1a0 100644
--- a/internal/engine/compiler/engine_test.go
+++ b/internal/engine/compiler/engine_test.go
@@ -219,7 +219,6 @@ func TestCompiler_CompileModule(t *testing.T) {
 			},
 			ID: wasm.ModuleID{},
 		}
-		errModule.BuildFunctionDefinitions()
 
 		e := et.NewEngine(api.CoreFeaturesV1).(*engine)
 		err := e.CompileModule(testCtx, errModule, nil, false)
@@ -336,7 +335,6 @@ func TestCompiler_SliceAllocatedOnHeap(t *testing.T) {
 		},
 		ID: wasm.ModuleID{1},
 	}
-	m.BuildFunctionDefinitions()
 
 	err = s.Engine.CompileModule(testCtx, m, nil, false)
 	require.NoError(t, err)
@@ -382,11 +380,10 @@ func ptrAsUint64(f *function) uint64 {
 }
 
 func TestCallEngine_deferredOnCall(t *testing.T) {
-	vv := &wasm.FunctionType{}
 	s := &wasm.Module{
-		FunctionDefinitionSection: []wasm.FunctionDefinition{
-			{Debugname: "1", Functype: vv}, {Debugname: "2", Functype: vv}, {Debugname: "3", Functype: vv},
-		},
+		FunctionSection: []wasm.Index{0, 1, 2},
+		CodeSection:     []wasm.Code{{}, {}, {}},
+		TypeSection:     []wasm.FunctionType{{}, {}, {}},
 	}
 	f1 := &function{
 		funcType: &wasm.FunctionType{ParamNumInUint64: 2},
@@ -429,9 +426,9 @@ func TestCallEngine_deferredOnCall(t *testing.T) {
 	err := ce.deferredOnCall(context.Background(), &wasm.ModuleInstance{}, errors.New("some error"))
 	require.EqualError(t, err, `some error (recovered by wazero)
 wasm stack trace:
-	3()
-	2()
-	1()`)
+	.$2()
+	.$1()
+	.$0()`)
 
 	// After recover, the state of callEngine must be reset except that the underlying slices must be intact
 	// for the subsequent calls to avoid additional allocations on each call.
@@ -597,7 +594,9 @@ func TestCallEngine_builtinFunctionFunctionListenerBefore(t *testing.T) {
 			},
 			index: 0,
 			parent: &compiledModule{source: &wasm.Module{
-				FunctionDefinitionSection: []wasm.FunctionDefinition{{}},
+				FunctionSection: []wasm.Index{0},
+				CodeSection:     []wasm.Code{{}},
+				TypeSection:     []wasm.FunctionType{{}},
 			}},
 		},
 	}
@@ -621,7 +620,9 @@ func TestCallEngine_builtinFunctionFunctionListenerAfter(t *testing.T) {
 			},
 			index: 0,
 			parent: &compiledModule{source: &wasm.Module{
-				FunctionDefinitionSection: []wasm.FunctionDefinition{{}},
+				FunctionSection: []wasm.Index{0},
+				CodeSection:     []wasm.Code{{}},
+				TypeSection:     []wasm.FunctionType{{}},
 			}},
 		},
 	}
diff --git a/internal/engine/interpreter/interpreter_test.go b/internal/engine/interpreter/interpreter_test.go
index 81a41f30..26efed33 100644
--- a/internal/engine/interpreter/interpreter_test.go
+++ b/internal/engine/interpreter/interpreter_test.go
@@ -514,7 +514,6 @@ func TestInterpreter_Compile(t *testing.T) {
 			},
 			ID: wasm.ModuleID{},
 		}
-		errModule.BuildFunctionDefinitions()
 
 		err := e.CompileModule(testCtx, errModule, nil, false)
 		require.EqualError(t, err, "handling instruction: apply stack failed for call: reading immediates: EOF")
diff --git a/internal/integration_test/bench/hostfunc_bench_test.go b/internal/integration_test/bench/hostfunc_bench_test.go
index e69c4bad..cd0a597c 100644
--- a/internal/integration_test/bench/hostfunc_bench_test.go
+++ b/internal/integration_test/bench/hostfunc_bench_test.go
@@ -172,12 +172,8 @@ func setupHostCallBench(requireNoError func(error)) *wasm.ModuleInstance {
 		},
 		ID: wasm.ModuleID{1, 2, 3, 4, 5},
 	}
-	hostModule.BuildFunctionDefinitions()
 
-	host := &wasm.ModuleInstance{
-		ModuleName: "host", TypeIDs: []wasm.FunctionTypeID{0},
-		Definitions: hostModule.FunctionDefinitionSection,
-	}
+	host := &wasm.ModuleInstance{ModuleName: "host", TypeIDs: []wasm.FunctionTypeID{0}}
 	host.Exports = hostModule.Exports
 
 	err := eng.CompileModule(testCtx, hostModule, nil, false)
@@ -214,14 +210,10 @@ func setupHostCallBench(requireNoError func(error)) *wasm.ModuleInstance {
 		ID:            wasm.ModuleID{1},
 	}
 
-	importingModule.BuildFunctionDefinitions()
 	err = eng.CompileModule(testCtx, importingModule, nil, false)
 	requireNoError(err)
 
-	importing := &wasm.ModuleInstance{
-		TypeIDs:     []wasm.FunctionTypeID{0},
-		Definitions: importingModule.FunctionDefinitionSection,
-	}
+	importing := &wasm.ModuleInstance{TypeIDs: []wasm.FunctionTypeID{0}}
 	importing.Exports = importingModule.Exports
 
 	importingMe, err := eng.NewModuleEngine(importingModule, importing)
diff --git a/internal/integration_test/spectest/spectest.go b/internal/integration_test/spectest/spectest.go
index 144a44a1..253e6fd0 100644
--- a/internal/integration_test/spectest/spectest.go
+++ b/internal/integration_test/spectest/spectest.go
@@ -329,7 +329,6 @@ func addSpectestModule(t *testing.T, ctx context.Context, s *wasm.Store, enabled
 
 	maybeSetMemoryCap(mod)
 	mod.BuildMemoryDefinitions()
-	mod.BuildFunctionDefinitions()
 
 	err = mod.Validate(enabledFeatures)
 	require.NoError(t, err)
@@ -403,7 +402,6 @@ func Run(t *testing.T, testDataFS embed.FS, ctx context.Context, fc filecache.Ca
 
 						maybeSetMemoryCap(mod)
 						mod.BuildMemoryDefinitions()
-						mod.BuildFunctionDefinitions()
 						err = s.Engine.CompileModule(ctx, mod, nil, false)
 						require.NoError(t, err, msg)
 
@@ -544,7 +542,6 @@ func Run(t *testing.T, testDataFS embed.FS, ctx context.Context, fc filecache.Ca
 							mod.AssignModuleID(buf, false, false)
 
 							maybeSetMemoryCap(mod)
-							mod.BuildFunctionDefinitions()
 							err = s.Engine.CompileModule(ctx, mod, nil, false)
 							require.NoError(t, err, msg)
 
@@ -581,7 +578,6 @@ func requireInstantiationError(t *testing.T, ctx context.Context, s *wasm.Store,
 
 	maybeSetMemoryCap(mod)
 	mod.BuildMemoryDefinitions()
-	mod.BuildFunctionDefinitions()
 	err = s.Engine.CompileModule(ctx, mod, nil, false)
 	if err != nil {
 		return
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index 2f715c4c..b9dbbb77 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -109,7 +109,6 @@ func RunTestEngineMemoryGrowInRecursiveCall(t *testing.T, et EngineTester) {
 		ImportSection:   []wasm.Import{{Module: hostModuleName, Name: hostFnName, DescFunc: 0}},
 		ImportPerModule: map[string][]*wasm.Import{hostModuleName: {{Module: hostModuleName, Name: hostFnName, DescFunc: 0}}},
 	}
-	m.BuildFunctionDefinitions()
 	m.BuildMemoryDefinitions()
 
 	err = s.Engine.CompileModule(testCtx, m, nil, false)
@@ -155,16 +154,12 @@ func RunTestModuleEngineCall(t *testing.T, et EngineTester) {
 		},
 	}
 
-	m.BuildFunctionDefinitions()
 	listeners := buildFunctionListeners(et.ListenerFactory(), m)
 	err := e.CompileModule(testCtx, m, listeners, false)
 	require.NoError(t, err)
 
 	// To use the function, we first need to add it to a module.
-	module := &wasm.ModuleInstance{
-		ModuleName: t.Name(), TypeIDs: []wasm.FunctionTypeID{0},
-		Definitions: m.FunctionDefinitionSection,
-	}
+	module := &wasm.ModuleInstance{ModuleName: t.Name(), TypeIDs: []wasm.FunctionTypeID{0}}
 
 	// Compile the module
 	me, err := e.NewModuleEngine(m, module)
@@ -212,16 +207,12 @@ func RunTestModuleEngineCallWithStack(t *testing.T, et EngineTester) {
 		},
 	}
 
-	m.BuildFunctionDefinitions()
 	listeners := buildFunctionListeners(et.ListenerFactory(), m)
 	err := e.CompileModule(testCtx, m, listeners, false)
 	require.NoError(t, err)
 
 	// To use the function, we first need to add it to a module.
-	module := &wasm.ModuleInstance{
-		ModuleName: t.Name(), TypeIDs: []wasm.FunctionTypeID{0},
-		Definitions: m.FunctionDefinitionSection,
-	}
+	module := &wasm.ModuleInstance{ModuleName: t.Name(), TypeIDs: []wasm.FunctionTypeID{0}}
 
 	// Compile the module
 	me, err := e.NewModuleEngine(m, module)
@@ -257,7 +248,6 @@ func RunTestModuleEngineLookupFunction(t *testing.T, et EngineTester) {
 		},
 	}
 
-	mod.BuildFunctionDefinitions()
 	err := e.CompileModule(testCtx, mod, nil, false)
 	require.NoError(t, err)
 	m := &wasm.ModuleInstance{
@@ -653,17 +643,15 @@ func RunTestModuleEngineBeforeListenerStackIterator(t *testing.T, et EngineTeste
 		},
 		ID: wasm.ModuleID{0},
 	}
-	m.BuildFunctionDefinitions()
 
 	listeners := buildFunctionListeners(fnListener, m)
 	err := e.CompileModule(testCtx, m, listeners, false)
 	require.NoError(t, err)
 
 	module := &wasm.ModuleInstance{
-		ModuleName:  t.Name(),
-		TypeIDs:     []wasm.FunctionTypeID{0, 1, 2, 3},
-		Definitions: m.FunctionDefinitionSection,
-		Exports:     exportMap(m),
+		ModuleName: t.Name(),
+		TypeIDs:    []wasm.FunctionTypeID{0, 1, 2, 3},
+		Exports:    exportMap(m),
 	}
 
 	me, err := e.NewModuleEngine(m, module)
@@ -773,17 +761,15 @@ func RunTestModuleEngineBeforeListenerGlobals(t *testing.T, et EngineTester) {
 		},
 		ID: wasm.ModuleID{0},
 	}
-	m.BuildFunctionDefinitions()
 
 	listeners := buildFunctionListeners(fnListener, m)
 	err := e.CompileModule(testCtx, m, listeners, false)
 	require.NoError(t, err)
 
 	module := &wasm.ModuleInstance{
-		ModuleName:  t.Name(),
-		TypeIDs:     []wasm.FunctionTypeID{0, 1, 2, 3},
-		Definitions: m.FunctionDefinitionSection,
-		Exports:     exportMap(m),
+		ModuleName: t.Name(),
+		TypeIDs:    []wasm.FunctionTypeID{0, 1, 2, 3},
+		Exports:    exportMap(m),
 		Globals: []*wasm.GlobalInstance{
 			{Val: 100, Type: wasm.GlobalType{ValType: wasm.ValueTypeI32, Mutable: true}},
 			{Val: 200, Type: wasm.GlobalType{ValType: wasm.ValueTypeI32, Mutable: true}},
@@ -953,21 +939,19 @@ func RunTestModuleEngineStackIteratorOffset(t *testing.T, et EngineTester) {
 	m.CodeSection[0].BodyOffsetInCodeSection = f1offset
 	m.CodeSection[1].BodyOffsetInCodeSection = f2offset
 
-	m.BuildFunctionDefinitions()
-
 	listeners := buildFunctionListeners(fnListener, m)
 	err = e.CompileModule(testCtx, m, listeners, false)
 	require.NoError(t, err)
 
 	module := &wasm.ModuleInstance{
-		ModuleName:  t.Name(),
-		TypeIDs:     []wasm.FunctionTypeID{0, 1, 2},
-		Definitions: m.FunctionDefinitionSection,
-		Exports:     exportMap(m),
+		ModuleName: t.Name(),
+		TypeIDs:    []wasm.FunctionTypeID{0, 1, 2},
+		Exports:    exportMap(m),
 		Globals: []*wasm.GlobalInstance{
 			{Val: 100, Type: wasm.GlobalType{ValType: wasm.ValueTypeI32, Mutable: true}},
 			{Val: 200, Type: wasm.GlobalType{ValType: wasm.ValueTypeI32, Mutable: true}},
 		},
+		Source: m,
 	}
 
 	me, err := e.NewModuleEngine(m, module)
@@ -1064,7 +1048,6 @@ func RunTestModuleEngineMemory(t *testing.T, et EngineTester) {
 			{Name: "init", Type: wasm.ExternTypeFunc, Index: 1},
 		},
 	}
-	m.BuildFunctionDefinitions()
 	listeners := buildFunctionListeners(et.ListenerFactory(), m)
 
 	err := e.CompileModule(testCtx, m, listeners, false)
@@ -1076,7 +1059,6 @@ func RunTestModuleEngineMemory(t *testing.T, et EngineTester) {
 		MemoryInstance: wasm.NewMemoryInstance(m.MemorySection),
 		DataInstances:  []wasm.DataInstance{m.DataSection[0].Init},
 		TypeIDs:        []wasm.FunctionTypeID{0, 1},
-		Definitions:    m.FunctionDefinitionSection,
 	}
 	memory := module.MemoryInstance
 
@@ -1196,14 +1178,10 @@ func setupCallTests(t *testing.T, e wasm.Engine, divBy *wasm.Code, fnlf experime
 		},
 		ID: wasm.ModuleID{0},
 	}
-	hostModule.BuildFunctionDefinitions()
 	lns := buildFunctionListeners(fnlf, hostModule)
 	err := e.CompileModule(testCtx, hostModule, lns, false)
 	require.NoError(t, err)
-	host := &wasm.ModuleInstance{
-		ModuleName: hostModule.NameSection.ModuleName, TypeIDs: []wasm.FunctionTypeID{0},
-		Definitions: hostModule.FunctionDefinitionSection,
-	}
+	host := &wasm.ModuleInstance{ModuleName: hostModule.NameSection.ModuleName, TypeIDs: []wasm.FunctionTypeID{0}}
 	host.Exports = exportMap(hostModule)
 
 	hostME, err := e.NewModuleEngine(hostModule, host)
@@ -1233,14 +1211,12 @@ func setupCallTests(t *testing.T, e wasm.Engine, divBy *wasm.Code, fnlf experime
 		},
 		ID: wasm.ModuleID{1},
 	}
-	importedModule.BuildFunctionDefinitions()
 	lns = buildFunctionListeners(fnlf, importedModule)
 	err = e.CompileModule(testCtx, importedModule, lns, false)
 	require.NoError(t, err)
 
 	imported := &wasm.ModuleInstance{
 		ModuleName: importedModule.NameSection.ModuleName, TypeIDs: []wasm.FunctionTypeID{0},
-		Definitions: importedModule.FunctionDefinitionSection,
 	}
 	imported.Exports = exportMap(importedModule)
 
@@ -1268,16 +1244,12 @@ func setupCallTests(t *testing.T, e wasm.Engine, divBy *wasm.Code, fnlf experime
 		},
 		ID: wasm.ModuleID{2},
 	}
-	importingModule.BuildFunctionDefinitions()
 	lns = buildFunctionListeners(fnlf, importingModule)
 	err = e.CompileModule(testCtx, importingModule, lns, false)
 	require.NoError(t, err)
 
 	// Add the exported function.
-	importing := &wasm.ModuleInstance{
-		ModuleName: importingModule.NameSection.ModuleName, TypeIDs: []wasm.FunctionTypeID{0},
-		Definitions: importingModule.FunctionDefinitionSection,
-	}
+	importing := &wasm.ModuleInstance{ModuleName: importingModule.NameSection.ModuleName, TypeIDs: []wasm.FunctionTypeID{0}}
 	importing.Exports = exportMap(importingModule)
 
 	// Compile the importing module
@@ -1304,13 +1276,9 @@ func setupCallMemTests(t *testing.T, e wasm.Engine, readMem *wasm.Code) *wasm.Mo
 		},
 		ID: wasm.ModuleID{0},
 	}
-	hostModule.BuildFunctionDefinitions()
 	err := e.CompileModule(testCtx, hostModule, nil, false)
 	require.NoError(t, err)
-	host := &wasm.ModuleInstance{
-		ModuleName: hostModule.NameSection.ModuleName, TypeIDs: []wasm.FunctionTypeID{0},
-		Definitions: hostModule.FunctionDefinitionSection,
-	}
+	host := &wasm.ModuleInstance{ModuleName: hostModule.NameSection.ModuleName, TypeIDs: []wasm.FunctionTypeID{0}}
 	host.Exports = exportMap(hostModule)
 
 	hostMe, err := e.NewModuleEngine(hostModule, host)
@@ -1341,15 +1309,11 @@ func setupCallMemTests(t *testing.T, e wasm.Engine, readMem *wasm.Code) *wasm.Mo
 		MemorySection: &wasm.Memory{Min: 1},
 		ID:            wasm.ModuleID{1},
 	}
-	importingModule.BuildFunctionDefinitions()
 	err = e.CompileModule(testCtx, importingModule, nil, false)
 	require.NoError(t, err)
 
 	// Add the exported function.
-	importing := &wasm.ModuleInstance{
-		ModuleName: importingModule.NameSection.ModuleName, TypeIDs: []wasm.FunctionTypeID{0},
-		Definitions: importingModule.FunctionDefinitionSection,
-	}
+	importing := &wasm.ModuleInstance{ModuleName: importingModule.NameSection.ModuleName, TypeIDs: []wasm.FunctionTypeID{0}}
 	// Note: adds imported functions readMemFn and callReadMemFn at index 0 and 1.
 	importing.Exports = exportMap(importingModule)
 
@@ -1379,7 +1343,7 @@ func buildFunctionListeners(factory experimental.FunctionListenerFactory, m *was
 	listeners := make([]experimental.FunctionListener, len(m.FunctionSection))
 	importCount := m.ImportFunctionCount
 	for i := 0; i < len(listeners); i++ {
-		listeners[i] = factory.NewFunctionListener(&m.FunctionDefinitionSection[uint32(i)+importCount])
+		listeners[i] = factory.NewFunctionListener(m.FunctionDefinition(uint32(i) + importCount))
 	}
 	return listeners
 }
diff --git a/internal/wasm/function_definition_test.go b/internal/wasm/function_definition_test.go
index 99c71f55..6e6e92f2 100644
--- a/internal/wasm/function_definition_test.go
+++ b/internal/wasm/function_definition_test.go
@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/internal/testing/hammer"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
@@ -253,7 +254,25 @@ func TestModule_BuildFunctionDefinitions(t *testing.T) {
 	for _, tc := range tests {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
-			tc.m.BuildFunctionDefinitions()
+			tc.m.buildFunctionDefinitions()
+			require.Equal(t, tc.expected, tc.m.FunctionDefinitionSection)
+			require.Equal(t, tc.expectedImports, tc.m.ImportedFunctions())
+			require.Equal(t, tc.expectedExports, tc.m.ExportedFunctions())
+		})
+	}
+
+	// Execute the same tests with n=`concurrentCount` goroutines invoking `buildFunctionDefinitions()` at once.
+	const nGoroutines = 100
+	const nIterations = 10
+	for _, tc := range tests {
+		tc := tc
+		testName := tc.name + " (concurrent)"
+		t.Run(testName, func(t *testing.T) {
+			hammer.NewHammer(t, nGoroutines, nIterations).
+				Run(func(name string) {
+					tc.m.buildFunctionDefinitions()
+				}, nil)
+
 			require.Equal(t, tc.expected, tc.m.FunctionDefinitionSection)
 			require.Equal(t, tc.expectedImports, tc.m.ImportedFunctions())
 			require.Equal(t, tc.expectedExports, tc.m.ExportedFunctions())
diff --git a/internal/wasm/module_test.go b/internal/wasm/module_test.go
index dadb2bda..77eb830d 100644
--- a/internal/wasm/module_test.go
+++ b/internal/wasm/module_test.go
@@ -376,8 +376,9 @@ func TestModule_validateStartSection(t *testing.T) {
 	t.Run("imported valid func", func(t *testing.T) {
 		index := Index(1)
 		m := Module{
-			StartSection: &index,
-			TypeSection:  []FunctionType{{}, {Results: []ValueType{ValueTypeI32}}},
+			StartSection:        &index,
+			TypeSection:         []FunctionType{{}, {Results: []ValueType{ValueTypeI32}}},
+			ImportFunctionCount: 2,
 			ImportSection: []Import{
 				{Type: ExternTypeFunc, DescFunc: 1},
 				// import with index 1 is global but this should be skipped when searching imported functions.
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 9046ab02..383757af 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -211,7 +211,6 @@ func TestStore_hammer(t *testing.T) {
 			{Type: ExternTypeFunc, Module: importedModuleName, Name: "fn", DescFunc: 0},
 		},
 	}
-	importingModule.BuildFunctionDefinitions()
 
 	// Concurrent instantiate, close should test if locks work on the ns. If they don't, we should see leaked modules
 	// after all of these complete, or an error raised.
@@ -271,7 +270,6 @@ func TestStore_hammer_close(t *testing.T) {
 			{Type: ExternTypeFunc, Module: importedModuleName, Name: "fn", DescFunc: 0},
 		},
 	}
-	importingModule.BuildFunctionDefinitions()
 
 	const instCount = 10000
 	instances := make([]api.Module, instCount)
@@ -371,7 +369,6 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 				{Type: ExternTypeFunc, Module: importedModuleName, Name: "fn", DescFunc: 0},
 			},
 		}
-		importingModule.BuildFunctionDefinitions()
 
 		_, err = s.Instantiate(testCtx, importingModule, importingModuleName, nil, []FunctionTypeID{0})
 		require.EqualError(t, err, "some engine creation error")
@@ -399,7 +396,6 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 				{Type: ExternTypeFunc, Module: importedModuleName, Name: "fn", DescFunc: 0},
 			},
 		}
-		importingModule.BuildFunctionDefinitions()
 
 		_, err = s.Instantiate(testCtx, importingModule, importingModuleName, nil, []FunctionTypeID{0})
 		require.EqualError(t, err, "start function[1] failed: call failed")
@@ -692,19 +688,21 @@ func Test_resolveImports(t *testing.T) {
 					"":   {Type: ExternTypeFunc, Index: 4},
 				},
 				ModuleName: moduleName,
-				Definitions: []FunctionDefinition{
-					{},
-					{},
-					{Functype: &FunctionType{Params: []ValueType{i32}, Results: []ValueType{ValueTypeV128}}},
-					{},
-					{Functype: &FunctionType{Params: []ValueType{ExternTypeFunc}, Results: []ValueType{}}},
+				Source: &Module{
+					FunctionSection: []Index{0, 0, 1, 0, 0},
+					TypeSection: []FunctionType{
+						{Params: []ValueType{ExternTypeFunc}},
+						{Params: []ValueType{i32}, Results: []ValueType{ValueTypeV128}},
+					},
 				},
 			}
+
 			module := &Module{
 				TypeSection: []FunctionType{
 					{Params: []ValueType{i32}, Results: []ValueType{ValueTypeV128}},
 					{Params: []ValueType{ExternTypeFunc}},
 				},
+				ImportFunctionCount: 2,
 				ImportPerModule: map[string][]*Import{
 					moduleName: {
 						{Module: moduleName, Name: name, Type: ExternTypeFunc, DescFunc: 0, IndexPerType: 0},
@@ -713,7 +711,7 @@ func Test_resolveImports(t *testing.T) {
 				},
 			}
 
-			m := &ModuleInstance{Engine: &mockModuleEngine{resolveImportsCalled: map[Index]Index{}}, s: s}
+			m := &ModuleInstance{Engine: &mockModuleEngine{resolveImportsCalled: map[Index]Index{}}, s: s, Source: module}
 			err := m.resolveImports(module)
 			require.NoError(t, err)
 
@@ -727,9 +725,14 @@ func Test_resolveImports(t *testing.T) {
 				Exports: map[string]*Export{
 					name: {Type: ExternTypeFunc, Index: 0},
 				},
-				ModuleName:  moduleName,
-				TypeIDs:     []FunctionTypeID{123435},
-				Definitions: []FunctionDefinition{{Functype: &FunctionType{}}},
+				ModuleName: moduleName,
+				TypeIDs:    []FunctionTypeID{123435},
+				Source: &Module{
+					FunctionSection: []Index{0},
+					TypeSection: []FunctionType{
+						{Params: []ValueType{}},
+					},
+				},
 			}
 			module := &Module{
 				TypeSection: []FunctionType{{Results: []ValueType{ValueTypeF32}}},
@@ -738,7 +741,7 @@ func Test_resolveImports(t *testing.T) {
 				},
 			}
 
-			m := &ModuleInstance{Engine: &mockModuleEngine{resolveImportsCalled: map[Index]Index{}}, s: s}
+			m := &ModuleInstance{Engine: &mockModuleEngine{resolveImportsCalled: map[Index]Index{}}, s: s, Source: module}
 			err := m.resolveImports(module)
 			require.EqualError(t, err, "import func[test.target]: signature mismatch: v_f32 != v_v")
 		})