Repository: tetratelabs/wazero
Pull request #1420 (2023-05-01 02:42:17)

interpreter: refactors around FunctionDefinition
I forgot committing the change of interpreter in #1419 


Patch:
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index 0f6153f1..2411a313 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -29,17 +29,17 @@ var callStackCeiling = 2000
 
 // engine is an interpreter implementation of wasm.Engine
 type engine struct {
-	enabledFeatures api.CoreFeatures
-	codes           map[wasm.ModuleID][]code // guarded by mutex.
-	mux             sync.RWMutex
+	enabledFeatures   api.CoreFeatures
+	compiledFunctions map[wasm.ModuleID][]compiledFunction // guarded by mutex.
+	mux               sync.RWMutex
 	// labelAddressResolutionCache is the temporary cache used to map LabelKind -> FrameID -> the index to the body.
 	labelAddressResolutionCache [wazeroir.LabelKindNum][]uint64
 }
 
 func NewEngine(_ context.Context, enabledFeatures api.CoreFeatures, _ filecache.Cache) wasm.Engine {
 	return &engine{
-		enabledFeatures: enabledFeatures,
-		codes:           map[wasm.ModuleID][]code{},
+		enabledFeatures:   enabledFeatures,
+		compiledFunctions: map[wasm.ModuleID][]compiledFunction{},
 	}
 }
 
@@ -50,30 +50,30 @@ func (e *engine) Close() (err error) {
 
 // CompiledModuleCount implements the same method as documented on wasm.Engine.
 func (e *engine) CompiledModuleCount() uint32 {
-	return uint32(len(e.codes))
+	return uint32(len(e.compiledFunctions))
 }
 
 // DeleteCompiledModule implements the same method as documented on wasm.Engine.
 func (e *engine) DeleteCompiledModule(m *wasm.Module) {
-	e.deleteCodes(m)
+	e.deleteCompiledFunctions(m)
 }
 
-func (e *engine) deleteCodes(module *wasm.Module) {
+func (e *engine) deleteCompiledFunctions(module *wasm.Module) {
 	e.mux.Lock()
 	defer e.mux.Unlock()
-	delete(e.codes, module.ID)
+	delete(e.compiledFunctions, module.ID)
 }
 
-func (e *engine) addCodes(module *wasm.Module, fs []code) {
+func (e *engine) addCompiledFunctions(module *wasm.Module, fs []compiledFunction) {
 	e.mux.Lock()
 	defer e.mux.Unlock()
-	e.codes[module.ID] = fs
+	e.compiledFunctions[module.ID] = fs
 }
 
-func (e *engine) getCodes(module *wasm.Module) (fs []code, ok bool) {
+func (e *engine) getCompiledFunctions(module *wasm.Module) (fs []compiledFunction, ok bool) {
 	e.mux.RLock()
 	defer e.mux.RUnlock()
-	fs, ok = e.codes[module.ID]
+	fs, ok = e.compiledFunctions[module.ID]
 	return
 }
 
@@ -101,15 +101,15 @@ type callEngine struct {
 	// frames are the function call stack.
 	frames []*callFrame
 
-	// compiled is the initial function for this call engine.
-	compiled *function
+	// f is the initial function for this call engine.
+	f *function
 
 	// stackiterator for Listeners to walk frames and stack.
 	stackIterator stackIterator
 }
 
 func (e *moduleEngine) newCallEngine(compiled *function) *callEngine {
-	return &callEngine{compiled: compiled}
+	return &callEngine{f: compiled}
 }
 
 func (ce *callEngine) pushValue(v uint64) {
@@ -187,22 +187,21 @@ type callFrame struct {
 	base int
 }
 
-type code struct {
+type compiledFunction struct {
 	source              *wasm.Module
 	body                []wazeroir.UnionOperation
 	listener            experimental.FunctionListener
 	offsetsInWasmBinary []uint64
 	hostFn              interface{}
 	ensureTermination   bool
+	index               wasm.Index
 }
 
 type function struct {
-	index          wasm.Index
 	funcType       *wasm.FunctionType
-	def            api.FunctionDefinition
 	moduleInstance *wasm.ModuleInstance
 	typeID         wasm.FunctionTypeID
-	parent         *code
+	parent         *compiledFunction
 }
 
 // functionFromUintptr resurrects the original *function from the given uintptr
@@ -259,10 +258,10 @@ func (si *stackIterator) Next() bool {
 
 // FunctionDefinition implements experimental.StackIterator.
 func (si *stackIterator) FunctionDefinition() api.FunctionDefinition {
-	return si.fn.def
+	return si.fn.definition()
 }
 
-// Args implements experimental.StackIterator.
+// Parameters implements experimental.StackIterator.
 func (si *stackIterator) Parameters() []uint64 {
 	paramsCount := si.fn.funcType.ParamNumInUint64
 	top := len(si.stack)
@@ -274,15 +273,16 @@ const callFrameStackSize = 0
 
 // CompileModule implements the same method as documented on wasm.Engine.
 func (e *engine) CompileModule(_ context.Context, module *wasm.Module, listeners []experimental.FunctionListener, ensureTermination bool) error {
-	if _, ok := e.getCodes(module); ok { // cache hit!
+	if _, ok := e.getCompiledFunctions(module); ok { // cache hit!
 		return nil
 	}
 
-	funcs := make([]code, len(module.FunctionSection))
+	funcs := make([]compiledFunction, len(module.FunctionSection))
 	irCompiler, err := wazeroir.NewCompiler(e.enabledFeatures, callFrameStackSize, module, ensureTermination)
 	if err != nil {
 		return err
 	}
+	imported := module.ImportFunctionCount
 	for i := range module.CodeSection {
 		var lsn experimental.FunctionListener
 		if i < len(listeners) {
@@ -309,8 +309,9 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, listeners
 		compiled.source = module
 		compiled.ensureTermination = ensureTermination
 		compiled.listener = lsn
+		compiled.index = imported + uint32(i)
 	}
-	e.addCodes(module, funcs)
+	e.addCompiledFunctions(module, funcs)
 	return nil
 }
 
@@ -321,7 +322,7 @@ func (e *engine) NewModuleEngine(module *wasm.Module, instance *wasm.ModuleInsta
 		functions:    make([]function, len(module.FunctionSection)+int(module.ImportFunctionCount)),
 	}
 
-	codes, ok := e.getCodes(module)
+	codes, ok := e.getCompiledFunctions(module)
 	if !ok {
 		return nil, errors.New("source module must be compiled before instantiation")
 	}
@@ -331,11 +332,9 @@ func (e *engine) NewModuleEngine(module *wasm.Module, instance *wasm.ModuleInsta
 		offset := i + int(module.ImportFunctionCount)
 		typeIndex := module.FunctionSection[i]
 		me.functions[offset] = function{
-			index:          wasm.Index(offset),
 			moduleInstance: instance,
 			typeID:         instance.TypeIDs[typeIndex],
 			funcType:       &module.TypeSection[typeIndex],
-			def:            &module.FunctionDefinitionSection[offset],
 			parent:         c,
 		}
 	}
@@ -343,7 +342,7 @@ func (e *engine) NewModuleEngine(module *wasm.Module, instance *wasm.ModuleInsta
 }
 
 // lowerIR lowers the wazeroir operations to engine friendly struct.
-func (e *engine) lowerIR(ir *wazeroir.CompilationResult, ret *code) error {
+func (e *engine) lowerIR(ir *wazeroir.CompilationResult, ret *compiledFunction) error {
 	// Copy the body from the result.
 	ret.body = make([]wazeroir.UnionOperation, len(ir.Operations))
 	copy(ret.body, ir.Operations)
@@ -411,7 +410,6 @@ func (e *engine) setLabelAddress(op *uint64, label wazeroir.Label) {
 func (e *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
 	imported := importedModuleEngine.(*moduleEngine)
 	e.functions[index] = imported.functions[indexInImportedModule]
-	e.functions[index].index = index
 }
 
 // FunctionInstanceReference implements the same method as documented on wasm.ModuleEngine.
@@ -451,12 +449,17 @@ func (e *moduleEngine) LookupFunction(t *wasm.TableInstance, typeId wasm.Functio
 
 // Definition implements the same method as documented on api.Function.
 func (ce *callEngine) Definition() api.FunctionDefinition {
-	return ce.compiled.def
+	return ce.f.definition()
+}
+
+func (f *function) definition() api.FunctionDefinition {
+	compiled := f.parent
+	return &compiled.source.FunctionDefinitionSection[compiled.index]
 }
 
 // Call implements the same method as documented on api.Function.
 func (ce *callEngine) Call(ctx context.Context, params ...uint64) (results []uint64, err error) {
-	ft := ce.compiled.funcType
+	ft := ce.f.funcType
 	if n := ft.ParamNumInUint64; n != len(params) {
 		return nil, fmt.Errorf("expected %d params, but passed %d", n, len(params))
 	}
@@ -465,7 +468,7 @@ func (ce *callEngine) Call(ctx context.Context, params ...uint64) (results []uin
 
 // CallWithStack implements the same method as documented on api.Function.
 func (ce *callEngine) CallWithStack(ctx context.Context, stack []uint64) error {
-	params, results, err := wasm.SplitCallStack(ce.compiled.funcType, stack)
+	params, results, err := wasm.SplitCallStack(ce.f.funcType, stack)
 	if err != nil {
 		return err
 	}
@@ -474,8 +477,8 @@ func (ce *callEngine) CallWithStack(ctx context.Context, stack []uint64) error {
 }
 
 func (ce *callEngine) call(ctx context.Context, params, results []uint64) (_ []uint64, err error) {
-	m := ce.compiled.moduleInstance
-	if ce.compiled.parent.ensureTermination {
+	m := ce.f.moduleInstance
+	if ce.f.parent.ensureTermination {
 		select {
 		case <-ctx.Done():
 			// If the provided context is already done, close the call context
@@ -500,17 +503,17 @@ func (ce *callEngine) call(ctx context.Context, params, results []uint64) (_ []u
 
 	ce.pushValues(params)
 
-	if ce.compiled.parent.ensureTermination {
+	if ce.f.parent.ensureTermination {
 		done := m.CloseModuleOnCanceledOrTimeout(ctx)
 		defer done()
 	}
 
-	ce.callFunction(ctx, m, ce.compiled)
+	ce.callFunction(ctx, m, ce.f)
 
 	// This returns a safe copy of the results, instead of a slice view. If we
 	// returned a re-slice, the caller could accidentally or purposefully
 	// corrupt the stack of subsequent calls.
-	ft := ce.compiled.funcType
+	ft := ce.f.funcType
 	if results == nil && ft.ResultNumInUint64 > 0 {
 		results = make([]uint64, ft.ResultNumInUint64)
 	}
@@ -527,7 +530,7 @@ func (ce *callEngine) recoverOnCall(v interface{}) (err error) {
 	for i := 0; i < frameCount; i++ {
 		frame := ce.popFrame()
 		f := frame.f
-		def := f.def
+		def := f.definition()
 		var sources []string
 		if parent := frame.f.parent; parent.body != nil && len(parent.offsetsInWasmBinary) > 0 {
 			sources = parent.source.DWARFLines.Line(parent.offsetsInWasmBinary[frame.pc])
@@ -552,12 +555,12 @@ func (ce *callEngine) callFunction(ctx context.Context, m *wasm.ModuleInstance,
 }
 
 func (ce *callEngine) callGoFunc(ctx context.Context, m *wasm.ModuleInstance, f *function, stack []uint64) {
-	def, typ := f.def, f.funcType
+	typ := f.funcType
 	lsn := f.parent.listener
 	if lsn != nil {
 		params := stack[:typ.ParamNumInUint64]
 		ce.stackIterator.reset(ce.stack, ce.frames, f)
-		ctx = lsn.Before(ctx, m, def, params, &ce.stackIterator)
+		ctx = lsn.Before(ctx, m, f.definition(), params, &ce.stackIterator)
 		ce.stackIterator.clear()
 	}
 	frame := &callFrame{f: f, base: len(ce.stack)}
@@ -575,7 +578,7 @@ func (ce *callEngine) callGoFunc(ctx context.Context, m *wasm.ModuleInstance, f
 	if lsn != nil {
 		// TODO: This doesn't get the error due to use of panic to propagate them.
 		results := stack[:typ.ResultNumInUint64]
-		lsn.After(ctx, m, def, nil, results)
+		lsn.After(ctx, m, f.definition(), nil, results)
 	}
 }
 
@@ -4048,7 +4051,7 @@ func i32Abs(v uint32) uint32 {
 }
 
 func (ce *callEngine) callNativeFuncWithListener(ctx context.Context, m *wasm.ModuleInstance, f *function, fnl experimental.FunctionListener) context.Context {
-	def, typ := &f.moduleInstance.Definitions[f.index], f.funcType
+	def, typ := f.definition(), f.funcType
 
 	ce.stackIterator.reset(ce.stack, ce.frames, f)
 	ctx = fnl.Before(ctx, m, def, ce.peekValues(len(typ.Params)), &ce.stackIterator)


Test patch:
diff --git a/internal/engine/interpreter/interpreter_test.go b/internal/engine/interpreter/interpreter_test.go
index 7f03b9e0..ef84b825 100644
--- a/internal/engine/interpreter/interpreter_test.go
+++ b/internal/engine/interpreter/interpreter_test.go
@@ -362,7 +362,7 @@ func TestInterpreter_NonTrappingFloatToIntConversion(t *testing.T) {
 					ce := &callEngine{}
 					f := &function{
 						moduleInstance: &wasm.ModuleInstance{Engine: &moduleEngine{}},
-						parent:         &code{body: body},
+						parent:         &compiledFunction{body: body},
 					}
 					ce.callNativeFunc(testCtx, &wasm.ModuleInstance{}, f)
 
@@ -425,7 +425,7 @@ func TestInterpreter_CallEngine_callNativeFunc_signExtend(t *testing.T) {
 				ce := &callEngine{}
 				f := &function{
 					moduleInstance: &wasm.ModuleInstance{Engine: &moduleEngine{}},
-					parent: &code{body: []wazeroir.UnionOperation{
+					parent: &compiledFunction{body: []wazeroir.UnionOperation{
 						{Kind: wazeroir.OperationKindConstI32, U1: uint64(uint32(tc.in))},
 						{Kind: translateToIROperationKind(tc.opcode)},
 						{Kind: wazeroir.OperationKindBr, U1: uint64(math.MaxUint64)},
@@ -479,7 +479,7 @@ func TestInterpreter_CallEngine_callNativeFunc_signExtend(t *testing.T) {
 				ce := &callEngine{}
 				f := &function{
 					moduleInstance: &wasm.ModuleInstance{Engine: &moduleEngine{}},
-					parent: &code{body: []wazeroir.UnionOperation{
+					parent: &compiledFunction{body: []wazeroir.UnionOperation{
 						{Kind: wazeroir.OperationKindConstI64, U1: uint64(tc.in)},
 						{Kind: translateToIROperationKind(tc.opcode)},
 						{Kind: wazeroir.OperationKindBr, U1: uint64(math.MaxUint64)},
@@ -520,7 +520,7 @@ func TestInterpreter_Compile(t *testing.T) {
 		require.EqualError(t, err, "handling instruction: apply stack failed for call: reading immediates: EOF")
 
 		// On the compilation failure, all the compiled functions including succeeded ones must be released.
-		_, ok := e.codes[errModule.ID]
+		_, ok := e.compiledFunctions[errModule.ID]
 		require.False(t, ok)
 	})
 	t.Run("ok", func(t *testing.T) {
@@ -540,34 +540,34 @@ func TestInterpreter_Compile(t *testing.T) {
 		err := e.CompileModule(testCtx, okModule, nil, false)
 		require.NoError(t, err)
 
-		compiled, ok := e.codes[okModule.ID]
+		compiled, ok := e.compiledFunctions[okModule.ID]
 		require.True(t, ok)
 		require.Equal(t, len(okModule.FunctionSection), len(compiled))
 
-		_, ok = e.codes[okModule.ID]
+		_, ok = e.compiledFunctions[okModule.ID]
 		require.True(t, ok)
 	})
 }
 
-func TestEngine_CachedcodesPerModule(t *testing.T) {
+func TestEngine_CachedCompiledFunctionPerModule(t *testing.T) {
 	e := et.NewEngine(api.CoreFeaturesV1).(*engine)
-	exp := []code{
+	exp := []compiledFunction{
 		{body: []wazeroir.UnionOperation{}},
 		{body: []wazeroir.UnionOperation{}},
 	}
 	m := &wasm.Module{}
 
-	e.addCodes(m, exp)
+	e.addCompiledFunctions(m, exp)
 
-	actual, ok := e.getCodes(m)
+	actual, ok := e.getCompiledFunctions(m)
 	require.True(t, ok)
 	require.Equal(t, len(exp), len(actual))
 	for i := range actual {
 		require.Equal(t, exp[i], actual[i])
 	}
 
-	e.deleteCodes(m)
-	_, ok = e.getCodes(m)
+	e.deleteCompiledFunctions(m)
+	_, ok = e.getCompiledFunctions(m)
 	require.False(t, ok)
 }