arm64: account for imported funcs in the branch to a relocation island

internal/engine/wazevo/backend/isa/arm64/machine_relocation.go

@@ -59,7 +59,10 @@ func (m *machine) ResolveRelocations(
 		if diff < minUnconditionalBranchOffset || diff > maxUnconditionalBranchOffset {
 			// Find the near trampoline island from callTrampolineIslandOffsets.
 			islandOffset := searchTrampolineIsland(callTrampolineIslandOffsets, int(instrOffset))
-			islandTargetOffset := islandOffset + trampolineCallSize*int(r.FuncRef)
+			// Imported functions don't need trampolines, so we ignore them when we compute the offset
+			// (see also encodeCallTrampolineIsland)
+			funcOffset := int(r.FuncRef) - importedFns
+			islandTargetOffset := islandOffset + trampolineCallSize*funcOffset
 			diff = int64(islandTargetOffset) - (instrOffset)
 			if diff < minUnconditionalBranchOffset || diff > maxUnconditionalBranchOffset {
 				panic("BUG in trampoline placement")

internal/engine/wazevo/backend/isa/arm64/machine_relocation_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/hex"
 	"testing"
 
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
@@ -49,6 +50,121 @@ func Test_encodeCallTrampolineIsland(t *testing.T) {
 	require.Equal(t, "0000000000000000000000000000000000000000", hex.EncodeToString(instrs))
 }
 
+func Test_ResolveRelocations_importedFns(t *testing.T) {
+	tests := []struct {
+		name          string
+		refToOffset   []int
+		importedFns   int
+		relocations   []backend.RelocationInfo
+		islandOffsets []int
+	}{
+		{
+			name:        "Single imported function",
+			refToOffset: []int{10, 20, 200 * 1024 * 1024},
+			importedFns: 1,
+			relocations: []backend.RelocationInfo{
+				{Offset: 100, FuncRef: 2},
+			},
+			islandOffsets: []int{500},
+		},
+		{
+			name:        "Multiple imported functions",
+			refToOffset: []int{10, 20, 30, 40, 200 * 1024 * 1024},
+			importedFns: 3,
+			relocations: []backend.RelocationInfo{
+				{Offset: 100, FuncRef: 4},
+			},
+			islandOffsets: []int{500},
+		},
+		{
+			name:        "Multiple relocations with imported functions",
+			refToOffset: []int{10, 20, 30, 40, 200 * 1024 * 1024, 250 * 1024 * 1024},
+			importedFns: 2,
+			relocations: []backend.RelocationInfo{
+				{Offset: 100, FuncRef: 4},
+				{Offset: 200, FuncRef: 5},
+			},
+			islandOffsets: []int{500},
+		},
+		{
+			name:        "Many imported functions",
+			refToOffset: []int{10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 200 * 1024 * 1024},
+			importedFns: 9,
+			relocations: []backend.RelocationInfo{
+				{Offset: 120, FuncRef: 10},
+			},
+			islandOffsets: []int{600},
+		},
+		{
+			name:        "Multiple islands",
+			refToOffset: []int{10, 20, 30, 40, 200 * 1024 * 1024, 300 * 1024 * 1024},
+			importedFns: 3,
+			relocations: []backend.RelocationInfo{
+				{Offset: 100, FuncRef: 4},
+				{Offset: 200*1024*1024 - 1000, FuncRef: 5},
+			},
+			islandOffsets: []int{500, 200*1024*1024 - 500},
+		},
+		{
+			name:        "Multiple islands, no imported functions",
+			refToOffset: []int{10, 20, 30, 40, 200 * 1024 * 1024, 300 * 1024 * 1024},
+			importedFns: 0,
+			relocations: []backend.RelocationInfo{
+				{Offset: 100, FuncRef: 4},
+				{Offset: 200*1024*1024 - 1000, FuncRef: 5},
+			},
+			islandOffsets: []int{500, 200*1024*1024 - 500},
+		},
+	}
+
+	extractTargetInstruction := func(executable []byte, branchInstrOffset int64) uint32 {
+		bytes := executable[branchInstrOffset : branchInstrOffset+4]
+		return uint32(bytes[0]) | uint32(bytes[1])<<8 | uint32(bytes[2])<<16 | uint32(bytes[3])<<24
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create large executable buffer to handle far offsets
+			executable := make([]byte, 310*1024*1024)
+
+			m := &machine{}
+			m.ResolveRelocations(
+				tt.refToOffset,
+				tt.importedFns,
+				executable,
+				tt.relocations,
+				tt.islandOffsets,
+			)
+
+			for _, reloc := range tt.relocations {
+				relocOffset := reloc.Offset
+
+				// Extract the actual target instruction from the generated branch instruction
+				actualInstruction := extractTargetInstruction(executable, relocOffset)
+
+				// First find which island this relocation would use
+				funcOffset := tt.refToOffset[reloc.FuncRef]
+				diff := int64(funcOffset) - relocOffset
+
+				// Only test relocations that need trampolines
+				if diff < minUnconditionalBranchOffset || diff > maxUnconditionalBranchOffset {
+					// Find the nearest trampoline island
+					islandOffset := searchTrampolineIsland(tt.islandOffsets, int(relocOffset))
+
+					expectedIslandOffset := islandOffset + trampolineCallSize*(int(reloc.FuncRef)-tt.importedFns)
+					expectedBranchOffset := int64(expectedIslandOffset) - relocOffset
+					expectedInstruction := encodeUnconditionalBranch(true, expectedBranchOffset)
+
+					require.Equal(t, expectedInstruction, actualInstruction)
+
+				} else {
+					t.Logf("Skipping relocation that doesn't need a trampoline: %v", reloc)
+				}
+			}
+		})
+	}
+}
+
 func Test_searchTrampolineIsland(t *testing.T) {
 	offsets := []int{16, 32, 48}
 	require.Equal(t, 32, searchTrampolineIsland(offsets, 30))