compiler(arm64): avoid conditional jumps to the trap handler (#1524)

Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>

internal/engine/compiler/impl_arm64.go

@@ -380,42 +380,16 @@ func (c *arm64Compiler) compileReturnFunction() error {
 }
 
 func (c *arm64Compiler) compileMaybeExitFromNativeCode(skipCondition asm.Instruction, status nativeCallStatusCode) {
-	if target := c.compiledTrapTargets[status]; target != nil {
+	skip := c.assembler.CompileJump(skipCondition)
-		// We've already compiled this.
+	c.compileExitFromNativeCode(status)
-		// Invert the condition to jump into the appropriate target.
+	c.assembler.SetJumpTargetOnNext(skip)
-		var trapCondition asm.Instruction
-		switch skipCondition {
-		case arm64.BCONDEQ:
-			trapCondition = arm64.BCONDNE
-		case arm64.BCONDNE:
-			trapCondition = arm64.BCONDEQ
-		case arm64.BCONDLO:
-			trapCondition = arm64.BCONDHS
-		case arm64.BCONDHS:
-			trapCondition = arm64.BCONDLO
-		case arm64.BCONDLS:
-			trapCondition = arm64.BCONDHI
-		case arm64.BCONDHI:
-			trapCondition = arm64.BCONDLS
-		case arm64.BCONDVS:
-			trapCondition = arm64.BCONDVC
-		case arm64.BCONDVC:
-			trapCondition = arm64.BCONDVS
-		default:
-			panic("BUG: couldn't invert condition")
-		}
-		c.assembler.CompileJump(trapCondition).AssignJumpTarget(target)
-	} else {
-		skip := c.assembler.CompileJump(skipCondition)
-		c.compileExitFromNativeCode(status)
-		c.assembler.SetJumpTargetOnNext(skip)
-	}
 }
 
 // compileExitFromNativeCode adds instructions to give the control back to ce.exec with the given status code.
 func (c *arm64Compiler) compileExitFromNativeCode(status nativeCallStatusCode) {
 	if target := c.compiledTrapTargets[status]; target != nil {
 		c.assembler.CompileJump(arm64.B).AssignJumpTarget(target)
+		return
 	}
 
 	switch status {

internal/engine/compiler/impl_arm64_test.go

@@ -236,6 +236,52 @@ func TestArm64Compiler_getSavedTemporaryLocationStack(t *testing.T) {
 	})
 }
 
+// https://github.com/tetratelabs/wazero/issues/1522
+func TestArm64Compiler_LargeTrapOffsets(t *testing.T) {
+	env := newCompilerEnvironment()
+	compiler := env.requireNewCompiler(t, &wasm.FunctionType{}, newCompiler, &wazeroir.CompilationResult{
+		Types: []wasm.FunctionType{{}},
+	})
+	err := compiler.compilePreamble()
+	require.NoError(t, err)
+
+	one := operationPtr(wazeroir.NewOperationConstI32(uint32(1)))
+	five := operationPtr(wazeroir.NewOperationConstI32(uint32(5)))
+	div := operationPtr(wazeroir.NewOperationDiv(wazeroir.SignedTypeInt32))
+
+	// Place the offset value.
+	err = compiler.compileConstI32(one)
+	require.NoError(t, err)
+
+	// Repeat enough times that jump labels are not within (-524288, 524287).
+	// Relative offset -2097164/4(=-524291).
+	// At the time of writing, 52429 is empirically the value that starts
+	// triggering the bug on arm64. We impose an arbitrarily higher value
+	// to account for possible future improvement to the number of instructions
+	// we emit.
+	for i := 0; i < 80_000; i++ {
+		err = compiler.compileConstI32(five)
+		require.NoError(t, err)
+
+		err = compiler.compileDiv(div)
+		require.NoError(t, err)
+	}
+
+	err = compiler.compileReturnFunction()
+	require.NoError(t, err)
+
+	code := asm.CodeSegment{}
+	defer func() { require.NoError(t, code.Unmap()) }()
+
+	// Generate the code under test and run.
+	_, err = compiler.compile(code.NextCodeSection())
+	require.NoError(t, err)
+
+	env.exec(code.Bytes())
+
+	require.Equal(t, nativeCallStatusCodeReturned.String(), env.compilerStatus().String())
+}
+
 // compile implements compilerImpl.setStackPointerCeil for the amd64 architecture.
 func (c *arm64Compiler) setStackPointerCeil(v uint64) {
 	c.stackPointerCeil = v