Repository: tetratelabs/wazero (go)

Issue:
Too large offset for ADR
**Describe the bug**
I'm compiling a production wasm module that has many, very large regular expressions (around 9kb) and I'm getting the following error when compiling in debug mode only
```
error compiling wasm func[.<regex::exec::ExecNoSync as regex::re_trait::RegularExpression>::captures_read_at::hdbe37ea0d10b827b]: failed to compile: BUG: too large offset for ADR
```
In release mode, I'm not getting the same error
```
[profile.release]
lto = "thin"
opt-level = 's'
strip = "symbols"
```

**To Reproduce**
Sadly, the wasm module is a closed-source project. If it is helpful, I can try to come up with a reproduction, but I wanted to check if there was a quick answer first.

**Environment (please complete the relevant information):**
- Go version: go version go1.19 darwin/arm64
- wazero Version: v1.0.0-pre.3
- Host architecture: arm64
- Runtime mode: compiler


Pull request description:
arm64: adds support for large in-block jumps with ADR
Fixes #847 

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

Patch:
diff --git a/internal/asm/arm64/impl.go b/internal/asm/arm64/impl.go
index bf5c102b..97c775c1 100644
--- a/internal/asm/arm64/impl.go
+++ b/internal/asm/arm64/impl.go
@@ -5,7 +5,6 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
-	"math"
 
 	"github.com/tetratelabs/wazero/internal/asm"
 )
@@ -1835,25 +1834,23 @@ func (a *AssemblerImpl) encodeADR(n *nodeImpl) (err error) {
 		}
 
 		offset := targetNode.OffsetInBinary() - n.OffsetInBinary()
-		if offset > math.MaxUint8 {
-			// We could support up to 20-bit integer, but byte should be enough for our impl.
-			// If the necessity comes up, we could fix the below to support larger offsets.
-			return fmt.Errorf("BUG: too large offset for ADR: %d", offset)
+		if i64 := int64(offset); i64 >= 1<<20 || i64 < -1<<20 {
+			// We could support offset over 20-bit range by special casing them here,
+			// but 20-bit range should be enough for our impl. If the necessity comes up,
+			// we could add the special casing here to support arbitrary large offset.
+			return fmt.Errorf("BUG: too large offset for ADR: %#x", offset)
 		}
 
-		// Now ready to write an offset byte.
-		v := byte(offset)
-
 		adrInstructionBytes := code[n.OffsetInBinary() : n.OffsetInBinary()+4]
-		// According to the binary format of ADR instruction in arm64:
+		// According to the binary format of ADR instruction:
 		// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/ADR--Form-PC-relative-address-?lang=en
-		//
-		// The 0 to 1 bits live on 29 to 30 bits of the instruction.
-		adrInstructionBytes[3] |= (v & 0b00000011) << 5
-		// The 2 to 4 bits live on 5 to 7 bits of the instruction.
-		adrInstructionBytes[0] |= (v & 0b00011100) << 3
-		// The 5 to 7 bits live on 8 to 10 bits of the instruction.
-		adrInstructionBytes[1] |= (v & 0b11100000) >> 5
+		adrInstructionBytes[3] |= byte(offset & 0b00000011 << 5)
+		offset >>= 2
+		adrInstructionBytes[0] |= byte(offset << 5)
+		offset >>= 3
+		adrInstructionBytes[1] |= byte(offset)
+		offset >>= 8
+		adrInstructionBytes[2] |= byte(offset)
 		return nil
 	})
 	return


Test patch:
diff --git a/internal/asm/arm64/impl_4_test.go b/internal/asm/arm64/impl_4_test.go
index e4a60915..ca2c60be 100644
--- a/internal/asm/arm64/impl_4_test.go
+++ b/internal/asm/arm64/impl_4_test.go
@@ -2,6 +2,7 @@ package arm64
 
 import (
 	"encoding/hex"
+	"fmt"
 	"math"
 	"testing"
 
@@ -662,17 +663,65 @@ func TestAssemblerImpl_EncodeMemoryToRegister(t *testing.T) {
 
 func TestAssemblerImpl_encodeReadInstructionAddress(t *testing.T) {
 	t.Run("ok", func(t *testing.T) {
-		const targetBeforeInstruction, dstReg = RET, RegR23
-		a := NewAssembler(asm.NilRegister)
+		tests := []struct {
+			name                   string
+			numDummyInstructions   int
+			expADRInstructionBytes []byte
+		}{
+			{
+				name:                   "< 8-bit offset",
+				numDummyInstructions:   1,
+				expADRInstructionBytes: []byte{0x77, 0x0, 0x0, 0x10},
+			},
+			{
+				name:                   "> 8-bit offset",
+				numDummyInstructions:   5000,
+				expADRInstructionBytes: []byte{0x57, 0x71, 0x2, 0x10},
+			},
+		}
+
+		for _, tc := range tests {
+			tc := tc
+			t.Run(tc.name, func(t *testing.T) {
+				const targetBeforeInstruction, dstReg = RET, RegR23
+				a := NewAssembler(asm.NilRegister)
 
-		a.CompileReadInstructionAddress(dstReg, targetBeforeInstruction)
-		a.CompileConstToRegister(MOVD, 1000, RegR10) // Dummy
-		a.CompileJumpToRegister(targetBeforeInstruction, RegR25)
-		a.CompileConstToRegister(MOVD, 1000, RegR10) // Target.
+				a.CompileReadInstructionAddress(dstReg, targetBeforeInstruction)
+				adrInst := a.Current
+				for i := 0; i < tc.numDummyInstructions; i++ {
+					a.CompileJumpToRegister(B, RegR5)
+				}
+				a.CompileJumpToRegister(targetBeforeInstruction, RegR25)
+				a.CompileConstToRegister(MOVD, 0x3e8, RegR10) // Target.
+				target := a.Current
 
-		actual, err := a.Assemble()
-		require.NoError(t, err)
-		require.Equal(t, []byte{0x77, 0x0, 0x0, 0x10, 0xa, 0x7d, 0x80, 0xd2, 0x20, 0x3, 0x5f, 0xd6, 0xa, 0x7d, 0x80, 0xd2}, actual)
+				actual, err := a.Assemble()
+				require.NoError(t, err)
+				// The binary should start with ADR instruction.
+				require.Equal(t, tc.expADRInstructionBytes, actual[:4], hex.EncodeToString(actual))
+				// Then, follow the dummy B instructions.
+				pos := 4
+				for i := 0; i < tc.numDummyInstructions; i++ {
+					require.Equal(t,
+						// A0 00 1F D6    br   x5
+						[]byte{0xa0, 0x0, 0x1f, 0xd6},
+						actual[pos:pos+4], hex.EncodeToString(actual))
+					pos += 4
+				}
+				// And targetBeforeInstruction follows: "20 03 5F D6    ret  x25"
+				require.Equal(t, []byte{0x20, 0x03, 0x5F, 0xd6},
+					actual[pos:pos+4], hex.EncodeToString(actual))
+
+				// After that, we end with the target instruction "movz x10, #0x3e8"
+				pos += 4
+				require.Equal(t, []byte{0xa, 0x7d, 0x80, 0xd2},
+					actual[pos:pos+4], hex.EncodeToString(actual))
+				fmt.Println(hex.EncodeToString(actual))
+
+				require.Equal(t, uint64(4+tc.numDummyInstructions*4+4),
+					target.offsetInBinaryField-adrInst.offsetInBinaryField)
+			})
+		}
 	})
 
 	t.Run("not found", func(t *testing.T) {
@@ -683,25 +732,34 @@ func TestAssemblerImpl_encodeReadInstructionAddress(t *testing.T) {
 		require.EqualError(t, err, "BUG: target instruction NOP not found for ADR")
 	})
 	t.Run("offset too large", func(t *testing.T) {
-		a := NewAssembler(asm.NilRegister)
-		a.CompileReadInstructionAddress(RegR27, RET)
-		a.CompileJumpToRegister(RET, RegR25)
-		a.CompileConstToRegister(MOVD, 1000, RegR10)
+		for _, offset := range []int64{
+			1 << 20,
+			-(1 << 20) - 1,
+			math.MaxInt64, math.MinInt64,
+		} {
+			u64 := uint64(offset)
+			t.Run(fmt.Sprintf("offset=%#b", u64), func(t *testing.T) {
+				a := NewAssembler(asm.NilRegister)
+				a.CompileReadInstructionAddress(RegR27, RET)
+				a.CompileJumpToRegister(RET, RegR25)
+				a.CompileConstToRegister(MOVD, 1000, RegR10)
 
-		for n := a.Root; n != nil; n = n.next {
-			n.offsetInBinaryField = uint64(a.Buf.Len())
+				for n := a.Root; n != nil; n = n.next {
+					n.offsetInBinaryField = uint64(a.Buf.Len())
 
-			err := a.encodeNode(n)
-			require.NoError(t, err)
-		}
+					err := a.encodeNode(n)
+					require.NoError(t, err)
+				}
 
-		require.Equal(t, 1, len(a.OnGenerateCallbacks))
-		cb := a.OnGenerateCallbacks[0]
+				require.Equal(t, 1, len(a.OnGenerateCallbacks))
+				cb := a.OnGenerateCallbacks[0]
 
-		targetNode := a.Current
-		targetNode.offsetInBinaryField = uint64(math.MaxInt64)
+				targetNode := a.Current
+				targetNode.offsetInBinaryField = u64
 
-		err := cb(nil)
-		require.EqualError(t, err, "BUG: too large offset for ADR: 9223372036854775807")
+				err := cb(nil)
+				require.EqualError(t, err, fmt.Sprintf("BUG: too large offset for ADR: %#x", u64))
+			})
+		}
 	})
 }