From 509827a23f68dd9d99c554738e8c095236974907 Mon Sep 17 00:00:00 2001
From: Achille <achille.roussel@gmail.com>
Date: Tue, 9 May 2023 08:10:12 -0700
Subject: [PATCH] experimental: optimize function listener context stack
 (#1445)

---
 internal/engine/compiler/engine.go            | 20 +++------
 internal/engine/compiler/engine_bench_test.go | 44 +++++++++++++++++++
 internal/engine/compiler/engine_test.go       | 10 ++---
 3 files changed, 56 insertions(+), 18 deletions(-)
 create mode 100644 internal/engine/compiler/engine_bench_test.go

diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index fcc9c6ea..55c13e57 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -135,7 +135,7 @@ type (
 		ctx context.Context
 		// contextStack is a stack of contexts which is pushed and popped by function listeners.
 		// This is used and modified when there are function listeners.
-		contextStack *contextStack
+		contextStack []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.
@@ -144,14 +144,6 @@ type (
 		ensureTermination bool
 	}
 
-	// contextStack is a stack of context.Context.
-	contextStack struct {
-		// See note at top of file before modifying this struct.
-
-		self context.Context
-		prev *contextStack
-	}
-
 	// moduleContext holds the per-function call specific module information.
 	// This is subject to be manipulated from compiled native code whenever we make function calls.
 	moduleContext struct {
@@ -1175,9 +1167,8 @@ func (ce *callEngine) builtinFunctionFunctionListenerBefore(ctx context.Context,
 
 	params := ce.stack[base : base+fn.funcType.ParamNumInUint64]
 	listerCtx := fn.parent.listener.Before(ctx, mod, fn.definition(), params, &ce.stackIterator)
-	prevStackTop := ce.contextStack
-	ce.contextStack = &contextStack{self: ctx, prev: prevStackTop}
 
+	ce.contextStack = append(ce.contextStack, ctx)
 	ce.ctx = listerCtx
 	ce.stackIterator.clear()
 }
@@ -1185,8 +1176,11 @@ func (ce *callEngine) builtinFunctionFunctionListenerBefore(ctx context.Context,
 func (ce *callEngine) builtinFunctionFunctionListenerAfter(ctx context.Context, mod api.Module, fn *function) {
 	base := int(ce.stackBasePointerInBytes >> 3)
 	fn.parent.listener.After(ctx, mod, fn.definition(), nil, ce.stack[base:base+fn.funcType.ResultNumInUint64])
-	ce.ctx = ce.contextStack.self
-	ce.contextStack = ce.contextStack.prev
+
+	i := len(ce.contextStack) - 1
+	ce.ctx = ce.contextStack[i]
+	ce.contextStack[i] = nil
+	ce.contextStack = ce.contextStack[:i]
 }
 
 func compileGoDefinedHostFunction(cmp compiler) (body []byte, err error) {
diff --git a/internal/engine/compiler/engine_bench_test.go b/internal/engine/compiler/engine_bench_test.go
new file mode 100644
index 00000000..23ca605a
--- /dev/null
+++ b/internal/engine/compiler/engine_bench_test.go
@@ -0,0 +1,44 @@
+package compiler
+
+import (
+	"context"
+	"testing"
+
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/experimental"
+	"github.com/tetratelabs/wazero/internal/wasm"
+)
+
+func BenchmarkCallEngine_builtinFunctionFunctionListener(b *testing.B) {
+	f := &function{
+		funcType: &wasm.FunctionType{ParamNumInUint64: 3},
+		parent: &compiledFunction{
+			listener: mockListener{
+				before: func(context.Context, api.Module, api.FunctionDefinition, []uint64, experimental.StackIterator) context.Context {
+					return context.Background()
+				},
+				after: func(context.Context, api.Module, api.FunctionDefinition, error, []uint64) {
+				},
+			},
+			index: 0,
+			parent: &compiledModule{
+				source: &wasm.Module{
+					FunctionDefinitionSection: []wasm.FunctionDefinition{{}},
+				},
+			},
+		},
+	}
+
+	ce := &callEngine{
+		ctx:          context.Background(),
+		stack:        []uint64{0, 1, 2, 3, 4, 0, 0, 0},
+		stackContext: stackContext{stackBasePointerInBytes: 16},
+	}
+
+	module := new(wasm.ModuleInstance)
+
+	for i := 0; i < b.N; i++ {
+		ce.builtinFunctionFunctionListenerBefore(ce.ctx, module, f)
+		ce.builtinFunctionFunctionListenerAfter(ce.ctx, module, f)
+	}
+}
diff --git a/internal/engine/compiler/engine_test.go b/internal/engine/compiler/engine_test.go
index 46daa269..396fa777 100644
--- a/internal/engine/compiler/engine_test.go
+++ b/internal/engine/compiler/engine_test.go
@@ -605,13 +605,13 @@ func TestCallEngine_builtinFunctionFunctionListenerBefore(t *testing.T) {
 	ce := &callEngine{
 		ctx: currentContext, stack: []uint64{0, 1, 2, 3, 4, 0, 0, 0},
 		stackContext: stackContext{stackBasePointerInBytes: 16},
-		contextStack: &contextStack{self: prevContext},
+		contextStack: []context.Context{prevContext},
 	}
 	ce.builtinFunctionFunctionListenerBefore(ce.ctx, &wasm.ModuleInstance{}, f)
 
 	// Contexts must be stacked.
-	require.Equal(t, currentContext, ce.contextStack.self)
-	require.Equal(t, prevContext, ce.contextStack.prev.self)
+	require.Equal(t, currentContext, ce.contextStack[1])
+	require.Equal(t, prevContext, ce.contextStack[0])
 }
 
 func TestCallEngine_builtinFunctionFunctionListenerAfter(t *testing.T) {
@@ -635,12 +635,12 @@ func TestCallEngine_builtinFunctionFunctionListenerAfter(t *testing.T) {
 	ce := &callEngine{
 		ctx: currentContext, stack: []uint64{0, 1, 2, 3, 4, 5},
 		stackContext: stackContext{stackBasePointerInBytes: 40},
-		contextStack: &contextStack{self: prevContext},
+		contextStack: []context.Context{prevContext},
 	}
 	ce.builtinFunctionFunctionListenerAfter(ce.ctx, &wasm.ModuleInstance{}, f)
 
 	// Contexts must be popped.
-	require.Nil(t, ce.contextStack)
+	require.Equal(t, 0, len(ce.contextStack))
 	require.Equal(t, prevContext, ce.ctx)
 }