arm64: account for imported functions when encoding relocation islands (#2346)

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

Author: evacchi

internal/engine/wazevo/backend/isa/amd64/machine.go

@@ -2196,7 +2196,7 @@ func (m *machine) Encode(ctx context.Context) (err error) {
 }
 
 // ResolveRelocations implements backend.Machine.
-func (m *machine) ResolveRelocations(refToBinaryOffset []int, binary []byte, relocations []backend.RelocationInfo, _ []int) {
+func (m *machine) ResolveRelocations(refToBinaryOffset []int, _ int, binary []byte, relocations []backend.RelocationInfo, _ []int) {
 	for _, r := range relocations {
 		offset := r.Offset
 		calleeFnOffset := refToBinaryOffset[r.FuncRef]

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

@@ -21,7 +21,7 @@ const (
 
 	// trampolineIslandInterval is the range of the trampoline island.
 	// Half of the range is used for the trampoline island, and the other half is used for the function.
-	trampolineIslandInterval = maxUnconditionalBranchOffset / 2
+	trampolineIslandInterval = (maxUnconditionalBranchOffset - 1) / 2
 
 	// maxNumFunctions explicitly specifies the maximum number of functions that can be allowed in a single executable.
 	maxNumFunctions = trampolineIslandInterval >> 6
@@ -42,12 +42,13 @@ func (m *machine) CallTrampolineIslandInfo(numFunctions int) (interval, size int
 // ResolveRelocations implements backend.Machine ResolveRelocations.
 func (m *machine) ResolveRelocations(
 	refToBinaryOffset []int,
+	importedFns int,
 	executable []byte,
 	relocations []backend.RelocationInfo,
 	callTrampolineIslandOffsets []int,
 ) {
 	for _, islandOffset := range callTrampolineIslandOffsets {
-		encodeCallTrampolineIsland(refToBinaryOffset, islandOffset, executable)
+		encodeCallTrampolineIsland(refToBinaryOffset, importedFns, islandOffset, executable)
 	}
 
 	for _, r := range relocations {
@@ -71,11 +72,15 @@ func (m *machine) ResolveRelocations(
 // encodeCallTrampolineIsland encodes a trampoline island for the given functions.
 // Each island consists of a trampoline instruction sequence for each function.
 // Each trampoline instruction sequence consists of 4 instructions + 32-bit immediate.
-func encodeCallTrampolineIsland(refToBinaryOffset []int, islandOffset int, executable []byte) {
-	for i := 0; i < len(refToBinaryOffset); i++ {
+func encodeCallTrampolineIsland(refToBinaryOffset []int, importedFns int, islandOffset int, executable []byte) {
+	// We skip the imported functions: they don't need trampolines
+	// and are not accounted for.
+	binaryOffsets := refToBinaryOffset[importedFns:]
+
+	for i := 0; i < len(binaryOffsets); i++ {
 		trampolineOffset := islandOffset + trampolineCallSize*i
 
-		fnOffset := refToBinaryOffset[i]
+		fnOffset := binaryOffsets[i]
 		diff := fnOffset - (trampolineOffset + 16)
 		if diff > math.MaxInt32 || diff < math.MinInt32 {
 			// This case even amd64 can't handle. 4GB is too big.

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

@@ -20,7 +20,8 @@ func Test_encodeCallTrampolineIsland(t *testing.T) {
 	executable := make([]byte, 16*1000)
 	islandOffset := 160
 	refToBinaryOffset := []int{0: 0, 1: 16, 2: 1600, 3: 16000}
-	encodeCallTrampolineIsland(refToBinaryOffset, islandOffset, executable)
+
+	encodeCallTrampolineIsland(refToBinaryOffset, 0, islandOffset, executable)
 	for i := 0; i < len(refToBinaryOffset); i++ {
 		offset := islandOffset + trampolineCallSize*i
 		instrs := executable[offset : offset+trampolineCallSize-4]
@@ -29,6 +30,23 @@ func Test_encodeCallTrampolineIsland(t *testing.T) {
 		imm := binary.LittleEndian.Uint32(executable[offset+trampolineCallSize-4:])
 		require.Equal(t, uint32(refToBinaryOffset[ssa.FuncRef(i)]-(offset+16)), imm)
 	}
+
+	executable = executable[:]
+	// We test again, assuming that the first offset is an imported function.
+	const importedFns = 1
+	encodeCallTrampolineIsland(refToBinaryOffset, importedFns, islandOffset, executable)
+	for i := 0; i < len(refToBinaryOffset[importedFns:]); i++ {
+		offset := islandOffset + trampolineCallSize*i
+		instrs := executable[offset : offset+trampolineCallSize-4]
+		// Instructions are always the same except for the last immediate.
+		require.Equal(t, "9b0000106b0380b97b030b8b60031fd6", hex.EncodeToString(instrs))
+		imm := binary.LittleEndian.Uint32(executable[offset+trampolineCallSize-4:])
+		require.Equal(t, uint32(refToBinaryOffset[ssa.FuncRef(i+importedFns)]-(offset+16)), imm)
+	}
+	// If the encoding is incorrect, then this offset should contain zeroes.
+	finalOffset := islandOffset + trampolineCallSize*len(refToBinaryOffset)
+	instrs := executable[finalOffset : finalOffset+trampolineCallSize]
+	require.Equal(t, "0000000000000000000000000000000000000000", hex.EncodeToString(instrs))
 }
 
 func Test_searchTrampolineIsland(t *testing.T) {

internal/engine/wazevo/backend/machine.go

@@ -77,11 +77,13 @@ type (
 
 		// ResolveRelocations resolves the relocations after emitting machine code.
 		//  * refToBinaryOffset: the map from the function reference (ssa.FuncRef) to the executable offset.
+		//  * importedFns: the max index of the imported functions at the beginning of refToBinaryOffset
 		//  * executable: the binary to resolve the relocations.
 		//  * relocations: the relocations to resolve.
 		//  * callTrampolineIslandOffsets: the offsets of the trampoline islands in the executable.
 		ResolveRelocations(
 			refToBinaryOffset []int,
+			importedFns int,
 			executable []byte,
 			relocations []RelocationInfo,
 			callTrampolineIslandOffsets []int,

internal/engine/wazevo/backend/machine_test.go

@@ -55,7 +55,7 @@ func (m mockMachine) CompileGoFunctionTrampoline(wazevoapi.ExitCode, *ssa.Signat
 func (m mockMachine) Encode(context.Context) (err error) { return }
 
 // ResolveRelocations implements Machine.ResolveRelocations.
-func (m mockMachine) ResolveRelocations([]int, []byte, []RelocationInfo, []int) {}
+func (m mockMachine) ResolveRelocations([]int, int, []byte, []RelocationInfo, []int) {}
 
 // PostRegAlloc implements Machine.SetupPrologue.
 func (m mockMachine) PostRegAlloc() {}

internal/engine/wazevo/engine.go

@@ -314,7 +314,7 @@ func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listene
 
 	// Resolve relocations for local function calls.
 	if len(rels) > 0 {
-		machine.ResolveRelocations(refToBinaryOffset, executable, rels, callTrampolineIslandOffsets)
+		machine.ResolveRelocations(refToBinaryOffset, importedFns, executable, rels, callTrampolineIslandOffsets)
 	}
 
 	if runtime.GOARCH == "arm64" {