Repository: tetratelabs/wazero
Pull request #1297 (2023-03-28 05:55:04)

Implements api.Function directly over engine specific callEngine
Consequently, we can not only avoid the unnecessary allocation in `api.Module.ExportFunction`,
but also nuke the `wasm.CallEngine` interface.


Patch:
diff --git a/imports/emscripten/emscripten.go b/imports/emscripten/emscripten.go
index 6bd480da..b013572c 100644
--- a/imports/emscripten/emscripten.go
+++ b/imports/emscripten/emscripten.go
@@ -314,5 +314,5 @@ func callDynamic(ctx context.Context, m *wasm.ModuleInstance, typeID wasm.Functi
 	if err != nil {
 		return nil, err
 	}
-	return m.Function(idx).Call(ctx, params...)
+	return m.Engine.NewFunction(idx).Call(ctx, params...)
 }
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index a67f7e1a..2b2da700 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -50,6 +50,8 @@ type (
 
 	// callEngine holds context per moduleEngine.Call, and shared across all the
 	// function calls originating from the same moduleEngine.Call execution.
+	//
+	// This implements api.Function.
 	callEngine struct {
 		// See note at top of file before modifying this struct.
 
@@ -594,7 +596,7 @@ func (e *moduleEngine) FunctionInstanceReference(funcIndex wasm.Index) wasm.Refe
 	return uintptr(unsafe.Pointer(&e.functions[funcIndex]))
 }
 
-func (e *moduleEngine) NewCallEngine(index wasm.Index) (ce wasm.CallEngine, err error) {
+func (e *moduleEngine) NewFunction(index wasm.Index) api.Function {
 	// Note: The input parameters are pre-validated, so a compiled function is only absent on close. Updates to
 	// code on close aren't locked, neither is this read.
 	compiled := &e.functions[index]
@@ -603,7 +605,7 @@ func (e *moduleEngine) NewCallEngine(index wasm.Index) (ce wasm.CallEngine, err
 	if initialStackSize < compiled.parent.stackPointerCeil {
 		initStackSize = compiled.parent.stackPointerCeil * 2
 	}
-	return e.newCallEngine(initStackSize, compiled), nil
+	return e.newCallEngine(initStackSize, compiled)
 }
 
 // LookupFunction implements the same method as documented on wasm.ModuleEngine.
@@ -640,8 +642,14 @@ func functionFromUintptr(ptr uintptr) *function {
 	return *(**function)(unsafe.Pointer(wrapped))
 }
 
+// Definition implements the same method as documented on wasm.ModuleEngine.
+func (ce *callEngine) Definition() api.FunctionDefinition {
+	return ce.initialFn.def
+}
+
 // Call implements the same method as documented on wasm.ModuleEngine.
-func (ce *callEngine) Call(ctx context.Context, m *wasm.ModuleInstance, params []uint64) (results []uint64, err error) {
+func (ce *callEngine) Call(ctx context.Context, params ...uint64) (results []uint64, err error) {
+	m := ce.initialFn.moduleInstance
 	if ce.fn.parent.withEnsureTermination {
 		select {
 		case <-ctx.Done():
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index aac4b1e2..ba693bd4 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -86,6 +86,8 @@ type moduleEngine struct {
 
 // callEngine holds context per moduleEngine.Call, and shared across all the
 // function calls originating from the same moduleEngine.Call execution.
+//
+// This implements api.Function.
 type callEngine struct {
 	// stack contains the operands.
 	// Note that all the values are represented as uint64.
@@ -741,12 +743,12 @@ func (e *moduleEngine) FunctionInstanceReference(funcIndex wasm.Index) wasm.Refe
 	return uintptr(unsafe.Pointer(&e.functions[funcIndex]))
 }
 
-// NewCallEngine implements the same method as documented on wasm.ModuleEngine.
-func (e *moduleEngine) NewCallEngine(index wasm.Index) (ce wasm.CallEngine, err error) {
+// NewFunction implements the same method as documented on wasm.ModuleEngine.
+func (e *moduleEngine) NewFunction(index wasm.Index) (ce api.Function) {
 	// Note: The input parameters are pre-validated, so a compiled function is only absent on close. Updates to
 	// code on close aren't locked, neither is this read.
 	compiled := &e.functions[index]
-	return e.newCallEngine(compiled), nil
+	return e.newCallEngine(compiled)
 }
 
 // LookupFunction implements the same method as documented on wasm.ModuleEngine.
@@ -770,12 +772,18 @@ func (e *moduleEngine) LookupFunction(t *wasm.TableInstance, typeId wasm.Functio
 	return
 }
 
-// Call implements the same method as documented on wasm.CallEngine.
-func (ce *callEngine) Call(ctx context.Context, m *wasm.ModuleInstance, params []uint64) (results []uint64, err error) {
-	return ce.call(ctx, m, ce.compiled, params)
+// Definition implements the same method as documented on api.Function.
+func (ce *callEngine) Definition() api.FunctionDefinition {
+	return ce.compiled.def
+}
+
+// Call implements the same method as documented on api.Function.
+func (ce *callEngine) Call(ctx context.Context, params ...uint64) (results []uint64, err error) {
+	return ce.call(ctx, ce.compiled, params)
 }
 
-func (ce *callEngine) call(ctx context.Context, m *wasm.ModuleInstance, tf *function, params []uint64) (results []uint64, err error) {
+func (ce *callEngine) call(ctx context.Context, tf *function, params []uint64) (results []uint64, err error) {
+	m := ce.compiled.moduleInstance
 	if ce.compiled.parent.ensureTermination {
 		select {
 		case <-ctx.Done():
diff --git a/internal/wasm/engine.go b/internal/wasm/engine.go
index 6bee3a44..c11e4f37 100644
--- a/internal/wasm/engine.go
+++ b/internal/wasm/engine.go
@@ -3,6 +3,7 @@ package wasm
 import (
 	"context"
 
+	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/experimental"
 )
 
@@ -35,8 +36,8 @@ type Engine interface {
 
 // ModuleEngine implements function calls for a given module.
 type ModuleEngine interface {
-	// NewCallEngine returns a CallEngine for the given FunctionInstance.
-	NewCallEngine(index Index) (CallEngine, error)
+	// NewFunction returns an api.Function for the given function pointed by the given Index.
+	NewFunction(index Index) api.Function
 
 	// ResolveImportedFunction is used to add imported functions needed to make this ModuleEngine fully functional.
 	// 	- `index` is the function Index of this imported function.
@@ -51,10 +52,3 @@ type ModuleEngine interface {
 	// the initialization via ElementSegment.
 	FunctionInstanceReference(funcIndex Index) Reference
 }
-
-// CallEngine implements function calls for a FunctionInstance. It manages its own call frame stack and value stack,
-// internally, and shouldn't be used concurrently.
-type CallEngine interface {
-	// Call invokes a function instance f with given parameters.
-	Call(ctx context.Context, m *ModuleInstance, params []uint64) (results []uint64, err error)
-}
diff --git a/internal/wasm/module_instance.go b/internal/wasm/module_instance.go
index bdf5156b..d7d7fb68 100644
--- a/internal/wasm/module_instance.go
+++ b/internal/wasm/module_instance.go
@@ -51,7 +51,7 @@ func (m *ModuleInstance) closeModuleOnCanceledOrTimeout(ctx context.Context, can
 	}
 }
 
-// CloseWithExitCode closes the module with an exit code based on the type of
+// CloseWithCtxErr closes the module with an exit code based on the type of
 // error reported by the context.
 //
 // If the context's error is unknown or nil, the module does not close.
@@ -159,8 +159,7 @@ func (m *ModuleInstance) ExportedFunction(name string) api.Function {
 	if err != nil {
 		return nil
 	}
-
-	return m.function(exp.Index)
+	return m.Engine.NewFunction(exp.Index)
 }
 
 // ExportedFunctionDefinitions implements the same method as documented on
@@ -175,39 +174,6 @@ func (m *ModuleInstance) ExportedFunctionDefinitions() map[string]api.FunctionDe
 	return result
 }
 
-func (m *ModuleInstance) Function(funcIdx Index) api.Function {
-	if uint32(len(m.Definitions)) < funcIdx {
-		return nil
-	}
-	return m.function(funcIdx)
-}
-
-func (m *ModuleInstance) function(index Index) api.Function {
-	ce, err := m.Engine.NewCallEngine(index)
-	if err != nil {
-		return nil
-	}
-	return &function{index: index, ce: ce, m: m}
-}
-
-// function implements api.Function. This couples FunctionInstance with CallEngine so that
-// it can be used to make function calls originating from the FunctionInstance.
-type function struct {
-	index Index
-	m     *ModuleInstance
-	ce    CallEngine
-}
-
-// Definition implements the same method as documented on api.FunctionDefinition.
-func (f *function) Definition() api.FunctionDefinition {
-	return &f.m.Definitions[f.index]
-}
-
-// Call implements the same method as documented on api.Function.
-func (f *function) Call(ctx context.Context, params ...uint64) (ret []uint64, err error) {
-	return f.ce.Call(ctx, f.m, params)
-}
-
 // GlobalVal is an internal hack to get the lower 64 bits of a global.
 func (m *ModuleInstance) GlobalVal(idx Index) uint64 {
 	return m.Globals[idx].Val
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index 16488d07..d768c7e3 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -376,13 +376,8 @@ func (s *Store) instantiate(
 	// Execute the start function.
 	if module.StartSection != nil {
 		funcIdx := *module.StartSection
-		ce, err := m.Engine.NewCallEngine(funcIdx)
-		if err != nil {
-			return nil, fmt.Errorf("create call engine for start function[%s]: %v",
-				module.funcDesc(SectionIDFunction, funcIdx), err)
-		}
-
-		_, err = ce.Call(ctx, m, nil)
+		ce := m.Engine.NewFunction(funcIdx)
+		_, err = ce.Call(ctx)
 		if exitErr, ok := err.(*sys.ExitError); ok { // Don't wrap an exit error!
 			return nil, exitErr
 		} else if err != nil {
diff --git a/site/content/docs/how_do_compiler_functions_work.md b/site/content/docs/how_do_compiler_functions_work.md
index a9aa3817..db2c44f4 100644
--- a/site/content/docs/how_do_compiler_functions_work.md
+++ b/site/content/docs/how_do_compiler_functions_work.md
@@ -15,7 +15,7 @@ this page. For more general information on architecture, etc., please refer to
 ## Engines
 
 Our [Docs](..) introduce the "engine" concept of wazero. More precisely, there
-are three types of engines, `Engine`, `ModuleEngine` and `CallEngine`. Each has
+are three types of engines, `Engine`, `ModuleEngine` and `callEngine`. Each has
 a different scope and role:
 
 - `Engine` has the same lifetime as `Runtime`. This compiles a `CompiledModule`
@@ -23,7 +23,7 @@ a different scope and role:
 - `ModuleEngine` is a virtual machine with the same lifetime as its [Module][api-module].
   Notably, this binds each [function instance][spec-function-instance] to
   corresponding machine code owned by its `Engine`.
-- `CallEngine` corresponds to an exported [api.Function][api-function] in a
+- `callEngine` is the implementation of [api.Function][api-function] in a
   [Module][api-module]. This implements `Function.Call(...)` by invoking
   machine code corresponding to a function instance in `ModuleEngine` and
   managing the [call stack][call-stack] representing the invocation.
@@ -35,7 +35,7 @@ Here is a diagram showing the relationships of these engines:
      /1:N                   |                                                  |
     /                       |                                                  v
    |     +----------+       v        +----------------+                  +------------+
-   |     |  Engine  |--------------->|  ModuleEngine  |----------------->| CallEngine |
+   |     |  Engine  |--------------->|  ModuleEngine  |----------------->| callEngine |
    |     +----------+                +----------------+                  +------------+
    |          |                               |                            |      |
    .          |                               |                            |      |
@@ -193,7 +193,7 @@ strategy is only used between wasm and the host.
 ## Summary
 
 When an exported wasm function is called, using a wazero API, such as
-`Function.Call()`, wazero allocates a `CallEngine` and starts invocation. This
+`Function.Call()`, wazero allocates a `callEngine` and starts invocation. This
 begins with jumping to machine code compiled from the Wasm binary. When that
 code makes a callback to the host, it exits execution, passing control back to
 `exec_native` which then calls a Go function and resumes the machine code


Test patch:
diff --git a/internal/integration_test/bench/hostfunc_bench_test.go b/internal/integration_test/bench/hostfunc_bench_test.go
index 4de2f45b..fcd5f795 100644
--- a/internal/integration_test/bench/hostfunc_bench_test.go
+++ b/internal/integration_test/bench/hostfunc_bench_test.go
@@ -45,14 +45,11 @@ func BenchmarkHostFunctionCall(b *testing.B) {
 		fn := fn
 
 		b.Run(fn, func(b *testing.B) {
-			ce, err := getCallEngine(m, fn)
-			if err != nil {
-				b.Fatal(err)
-			}
+			ce := getCallEngine(m, fn)
 
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				res, err := ce.Call(testCtx, m, []uint64{offset})
+				res, err := ce.Call(testCtx, offset)
 				if err != nil {
 					b.Fatal(err)
 				}
@@ -73,11 +70,8 @@ func TestBenchmarkFunctionCall(t *testing.T) {
 		require.NoError(t, err)
 	})
 
-	callGoHost, err := getCallEngine(m, callGoHostName)
-	require.NoError(t, err)
-
-	callGoReflectHost, err := getCallEngine(m, callGoReflectHostName)
-	require.NoError(t, err)
+	callGoHost := getCallEngine(m, callGoHostName)
+	callGoReflectHost := getCallEngine(m, callGoReflectHostName)
 
 	require.NotNil(t, callGoHost)
 	require.NotNil(t, callGoReflectHost)
@@ -95,7 +89,7 @@ func TestBenchmarkFunctionCall(t *testing.T) {
 
 	for _, f := range []struct {
 		name string
-		ce   wasm.CallEngine
+		ce   api.Function
 	}{
 		{name: "go", ce: callGoHost},
 		{name: "go-reflect", ce: callGoReflectHost},
@@ -104,7 +98,7 @@ func TestBenchmarkFunctionCall(t *testing.T) {
 		t.Run(f.name, func(t *testing.T) {
 			for _, tc := range tests {
 				binary.LittleEndian.PutUint32(mem[tc.offset:], math.Float32bits(tc.val))
-				res, err := f.ce.Call(context.Background(), m, []uint64{uint64(tc.offset)})
+				res, err := f.ce.Call(context.Background(), uint64(tc.offset))
 				require.NoError(t, err)
 				require.Equal(t, math.Float32bits(tc.val), uint32(res[0]))
 			}
@@ -112,9 +106,9 @@ func TestBenchmarkFunctionCall(t *testing.T) {
 	}
 }
 
-func getCallEngine(m *wasm.ModuleInstance, name string) (ce wasm.CallEngine, err error) {
+func getCallEngine(m *wasm.ModuleInstance, name string) (ce api.Function) {
 	exp := m.Exports[name]
-	ce, err = m.Engine.NewCallEngine(exp.Index)
+	ce = m.Engine.NewFunction(exp.Index)
 	return
 }
 
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index c9a55f12..37a8ee29 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -110,8 +110,8 @@ func RunTestEngine_MemoryGrowInRecursiveCall(t *testing.T, et EngineTester) {
 	inst, err := s.Instantiate(testCtx, m, t.Name(), nil, typeIDs)
 	require.NoError(t, err)
 
-	growFn = inst.Function(2)
-	_, err = inst.Function(1).Call(context.Background())
+	growFn = inst.Engine.NewFunction(2)
+	_, err = inst.Engine.NewFunction(1).Call(context.Background())
 	require.NoError(t, err)
 }
 
@@ -162,26 +162,21 @@ func RunTestModuleEngine_Call(t *testing.T, et EngineTester) {
 
 	// Ensure the base case doesn't fail: A single parameter should work as that matches the function signature.
 	const funcIndex = 0
-	ce, err := me.NewCallEngine(funcIndex)
-	require.NoError(t, err)
+	ce := me.NewFunction(funcIndex)
 
-	results, err := ce.Call(testCtx, module, []uint64{1, 2})
+	results, err := ce.Call(testCtx, 1, 2)
 	require.NoError(t, err)
 	require.Equal(t, []uint64{1, 2}, results)
 
 	t.Run("errs when not enough parameters", func(t *testing.T) {
-		ce, err := me.NewCallEngine(funcIndex)
-		require.NoError(t, err)
-
-		_, err = ce.Call(testCtx, module, nil)
+		ce := me.NewFunction(funcIndex)
+		_, err = ce.Call(testCtx)
 		require.EqualError(t, err, "expected 2 params, but passed 0")
 	})
 
 	t.Run("errs when too many parameters", func(t *testing.T) {
-		ce, err := me.NewCallEngine(funcIndex)
-		require.NoError(t, err)
-
-		_, err = ce.Call(testCtx, module, []uint64{1, 2, 3})
+		ce := me.NewFunction(funcIndex)
+		_, err = ce.Call(testCtx, 1, 2, 3)
 		require.EqualError(t, err, "expected 2 params, but passed 3")
 	})
 }
@@ -286,10 +281,9 @@ func runTestModuleEngine_Call_HostFn_Mem(t *testing.T, et EngineTester, readMem
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			ce, err := importing.Engine.NewCallEngine(tc.fn)
-			require.NoError(t, err)
+			ce := importing.Engine.NewFunction(tc.fn)
 
-			results, err := ce.Call(testCtx, importing, nil)
+			results, err := ce.Call(testCtx)
 			require.NoError(t, err)
 			require.Equal(t, tc.expected, results[0])
 		})
@@ -338,17 +332,15 @@ func runTestModuleEngine_Call_HostFn(t *testing.T, et EngineTester, hostDivBy *w
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			m := tc.module
 			f := tc.fn
 
-			ce, err := tc.module.Engine.NewCallEngine(f)
-			require.NoError(t, err)
+			ce := tc.module.Engine.NewFunction(f)
 
-			results, err := ce.Call(testCtx, m, []uint64{1})
+			results, err := ce.Call(testCtx, 1)
 			require.NoError(t, err)
 			require.Equal(t, uint64(1), results[0])
 
-			results2, err := ce.Call(testCtx, m, []uint64{1})
+			results2, err := ce.Call(testCtx, 1)
 			require.NoError(t, err)
 			require.Equal(t, results, results2)
 
@@ -442,16 +434,13 @@ wasm stack trace:
 	for _, tt := range tests {
 		tc := tt
 		t.Run(tc.name, func(t *testing.T) {
-			m := tc.module
-
-			ce, err := tc.module.Engine.NewCallEngine(tc.fn)
-			require.NoError(t, err)
+			ce := tc.module.Engine.NewFunction(tc.fn)
 
-			_, err = ce.Call(testCtx, m, tc.input)
+			_, err := ce.Call(testCtx, tc.input...)
 			require.EqualError(t, err, tc.expectedErr)
 
 			// Ensure the module still works
-			results, err := ce.Call(testCtx, m, []uint64{1})
+			results, err := ce.Call(testCtx, 1)
 			require.NoError(t, err)
 			require.Equal(t, uint64(1), results[0])
 		})
@@ -534,9 +523,8 @@ func RunTestModuleEngine_Memory(t *testing.T, et EngineTester) {
 	require.Equal(t, make([]byte, wasmPhraseSize), buf)
 
 	// Initialize the memory using Wasm. This copies the test phrase.
-	initCallEngine, err := me.NewCallEngine(init)
-	require.NoError(t, err)
-	_, err = initCallEngine.Call(testCtx, module, nil)
+	initCallEngine := me.NewFunction(init)
+	_, err = initCallEngine.Call(testCtx)
 	require.NoError(t, err)
 
 	// We expect the same []byte read earlier to now include the phrase in wasm.
@@ -566,18 +554,16 @@ func RunTestModuleEngine_Memory(t *testing.T, et EngineTester) {
 	require.Equal(t, hostPhraseTruncated, string(buf2))
 
 	// Now, we need to prove the other direction, that when Wasm changes the capacity, the host's buffer is unaffected.
-	growCallEngine, err := me.NewCallEngine(grow)
-	require.NoError(t, err)
-	_, err = growCallEngine.Call(testCtx, module, []uint64{1})
+	growCallEngine := me.NewFunction(grow)
+	_, err = growCallEngine.Call(testCtx, 1)
 	require.NoError(t, err)
 
 	// The host buffer should still contain the same bytes as before grow
 	require.Equal(t, hostPhraseTruncated, string(buf2))
 
 	// Re-initialize the memory in wasm, which overwrites the region.
-	initCallEngine2, err := me.NewCallEngine(init)
-	require.NoError(t, err)
-	_, err = initCallEngine2.Call(testCtx, module, nil)
+	initCallEngine2 := me.NewFunction(init)
+	_, err = initCallEngine2.Call(testCtx)
 	require.NoError(t, err)
 
 	// The host was not affected because it is a different slice due to "memory.grow" affecting the underlying memory.
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index dcec7e10..3bbeb94e 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -451,9 +451,9 @@ func (e *mockModuleEngine) ResolveImportedFunction(index, importedIndex Index, _
 	e.resolveImportsCalled[index] = importedIndex
 }
 
-// NewCallEngine implements the same method as documented on wasm.ModuleEngine.
-func (e *mockModuleEngine) NewCallEngine(index Index) (CallEngine, error) {
-	return &mockCallEngine{index: index, callFailIndex: e.callFailIndex}, nil
+// NewFunction implements the same method as documented on wasm.ModuleEngine.
+func (e *mockModuleEngine) NewFunction(index Index) api.Function {
+	return &mockCallEngine{index: index, callFailIndex: e.callFailIndex}
 }
 
 // InitializeFuncrefGlobals implements the same method as documented on wasm.ModuleEngine.
@@ -468,8 +468,11 @@ func (e *mockModuleEngine) Name() string {
 func (e *mockModuleEngine) Close(context.Context) {
 }
 
-// Call implements the same method as documented on wasm.ModuleEngine.
-func (ce *mockCallEngine) Call(_ context.Context, _ *ModuleInstance, _ []uint64) (results []uint64, err error) {
+// Call implements the same method as documented on api.Function.
+func (ce *mockCallEngine) Definition() api.FunctionDefinition { return nil }
+
+// Call implements the same method as documented on api.Function.
+func (ce *mockCallEngine) Call(_ context.Context, _ ...uint64) (results []uint64, err error) {
 	if ce.callFailIndex >= 0 && ce.index == Index(ce.callFailIndex) {
 		err = errors.New("call failed")
 		return