Repository: tetratelabs/wazero
Pull request #667 (2022-06-29 03:24:07)

compiler: fix stackPointerCeil calculation
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/internal/engine/compiler/compiler_value_location.go b/internal/engine/compiler/compiler_value_location.go
index 325698af..0c582e6f 100644
--- a/internal/engine/compiler/compiler_value_location.go
+++ b/internal/engine/compiler/compiler_value_location.go
@@ -201,10 +201,12 @@ func (v *runtimeValueLocationStack) push(loc *runtimeValueLocation) {
 	} else {
 		v.stack[v.sp] = loc
 	}
+	v.sp++
+	// stackPointerCeil must be set after sp is incremented since
+	// we skip the stack grow if len(valueStack) >= basePointer+stackPointerCeil.
 	if v.sp > v.stackPointerCeil {
 		v.stackPointerCeil = v.sp
 	}
-	v.sp++
 }
 
 func (v *runtimeValueLocationStack) pop() (loc *runtimeValueLocation) {


Test patch:
diff --git a/internal/engine/compiler/compiler_initialization_test.go b/internal/engine/compiler/compiler_initialization_test.go
index 3cb1b08e..aac9bd8c 100644
--- a/internal/engine/compiler/compiler_initialization_test.go
+++ b/internal/engine/compiler/compiler_initialization_test.go
@@ -201,40 +201,62 @@ func TestCompiler_compileMaybeGrowValueStack(t *testing.T) {
 			})
 		}
 	})
+
+	var defaultValueStackLen = uint64(initialValueStackSize)
 	t.Run("grow", func(t *testing.T) {
-		env := newCompilerEnvironment()
-		compiler := env.requireNewCompiler(t, newCompiler, nil)
-
-		err := compiler.compileMaybeGrowValueStack()
-		require.NoError(t, err)
-
-		// On the return from grow value stack, we simply return.
-		err = compiler.compileReturnFunction()
-		require.NoError(t, err)
-
-		stackPointerCeil := uint64(6)
-		compiler.setStackPointerCeil(stackPointerCeil)
-		valueStackLen := uint64(len(env.stack()))
-		stackBasePointer := valueStackLen - 5 // Ceil > valueStackLen - stackBasePointer = need to grow!
-		env.setValueStackBasePointer(stackBasePointer)
-
-		// Generate and run the code under test.
-		code, _, err := compiler.compile()
-		require.NoError(t, err)
-		env.exec(code)
-
-		// Check if the call exits with builtin function call status.
-		require.Equal(t, nativeCallStatusCodeCallBuiltInFunction, env.compilerStatus())
-
-		// Reenter from the return address.
-		returnAddress := env.callFrameStackPeek().returnAddress
-		require.True(t, returnAddress != 0, "returnAddress was non-zero %d", returnAddress)
-		nativecall(
-			returnAddress, uintptr(unsafe.Pointer(env.callEngine())),
-			uintptr(unsafe.Pointer(env.module())),
-		)
-
-		// Check the result. This should be "Returned".
-		require.Equal(t, nativeCallStatusCodeReturned, env.compilerStatus())
+		tests := []struct {
+			name             string
+			stackPointerCeil uint64
+			stackBasePointer uint64
+		}{
+			{
+				name:             "ceil=6/sbp=len-5",
+				stackPointerCeil: 6,
+				stackBasePointer: defaultValueStackLen - 5,
+			},
+			{
+				name:             "ceil=10000/sbp=0",
+				stackPointerCeil: 10000,
+				stackBasePointer: 0,
+			},
+		}
+
+		for _, tc := range tests {
+			tc := tc
+			t.Run(tc.name, func(t *testing.T) {
+				env := newCompilerEnvironment()
+				compiler := env.requireNewCompiler(t, newCompiler, nil)
+
+				err := compiler.compileMaybeGrowValueStack()
+				require.NoError(t, err)
+
+				// On the return from grow value stack, we simply return.
+				err = compiler.compileReturnFunction()
+				require.NoError(t, err)
+
+				// Generate code under test with the given stackPointerCeil.
+				compiler.setStackPointerCeil(tc.stackPointerCeil)
+				code, _, err := compiler.compile()
+				require.NoError(t, err)
+
+				// And run the code with the specified stackBasePointer.
+				env.setValueStackBasePointer(tc.stackBasePointer)
+				env.exec(code)
+
+				// Check if the call exits with builtin function call status.
+				require.Equal(t, nativeCallStatusCodeCallBuiltInFunction, env.compilerStatus())
+
+				// Reenter from the return address.
+				returnAddress := env.callFrameStackPeek().returnAddress
+				require.True(t, returnAddress != 0, "returnAddress was zero %d", returnAddress)
+				nativecall(
+					returnAddress, uintptr(unsafe.Pointer(env.callEngine())),
+					uintptr(unsafe.Pointer(env.module())),
+				)
+
+				// Check the result. This should be "Returned".
+				require.Equal(t, nativeCallStatusCodeReturned, env.compilerStatus())
+			})
+		}
 	})
 }
diff --git a/internal/engine/compiler/compiler_value_location_test.go b/internal/engine/compiler/compiler_value_location_test.go
index f71fe70d..48a227fe 100644
--- a/internal/engine/compiler/compiler_value_location_test.go
+++ b/internal/engine/compiler/compiler_value_location_test.go
@@ -57,7 +57,7 @@ func TestRuntimeValueLocationStack_basic(t *testing.T) {
 	for i := 0; i < 1000; i++ {
 		s.pop()
 	}
-	require.Equal(t, uint64(1001), s.stackPointerCeil)
+	require.Equal(t, uint64(1002), s.stackPointerCeil)
 }
 
 func TestRuntimeValueLocationStack_takeFreeRegister(t *testing.T) {