compiler: uses slice instead of map for refToBinaryOffset (#2195)

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

Author: mathetake

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

@@ -2059,7 +2059,7 @@ func (m *machine) Encode(ctx context.Context) (err error) {
 }
 
 // ResolveRelocations implements backend.Machine.
-func (m *machine) ResolveRelocations(refToBinaryOffset map[ssa.FuncRef]int, binary []byte, relocations []backend.RelocationInfo, _ []int) {
+func (m *machine) ResolveRelocations(refToBinaryOffset []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

@@ -7,7 +7,6 @@ import (
 	"sort"
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
-	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 )
 
 const (
@@ -42,7 +41,7 @@ func (m *machine) CallTrampolineIslandInfo(numFunctions int) (interval, size int
 
 // ResolveRelocations implements backend.Machine ResolveRelocations.
 func (m *machine) ResolveRelocations(
-	refToBinaryOffset map[ssa.FuncRef]int,
+	refToBinaryOffset []int,
 	executable []byte,
 	relocations []backend.RelocationInfo,
 	callTrampolineIslandOffsets []int,
@@ -72,11 +71,11 @@ 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 map[ssa.FuncRef]int, islandOffset int, executable []byte) {
+func encodeCallTrampolineIsland(refToBinaryOffset []int, islandOffset int, executable []byte) {
 	for i := 0; i < len(refToBinaryOffset); i++ {
 		trampolineOffset := islandOffset + trampolineCallSize*i
 
-		fnOffset := refToBinaryOffset[ssa.FuncRef(i)]
+		fnOffset := refToBinaryOffset[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

@@ -19,7 +19,7 @@ func Test_maxNumFunctions(t *testing.T) {
 func Test_encodeCallTrampolineIsland(t *testing.T) {
 	executable := make([]byte, 16*1000)
 	islandOffset := 160
-	refToBinaryOffset := map[ssa.FuncRef]int{0: 0, 1: 16, 2: 1600, 3: 16000}
+	refToBinaryOffset := []int{0: 0, 1: 16, 2: 1600, 3: 16000}
 	encodeCallTrampolineIsland(refToBinaryOffset, islandOffset, executable)
 	for i := 0; i < len(refToBinaryOffset); i++ {
 		offset := islandOffset + trampolineCallSize*i

internal/engine/wazevo/backend/machine.go

@@ -59,12 +59,12 @@ type (
 		PostRegAlloc()
 
 		// ResolveRelocations resolves the relocations after emitting machine code.
-		//  * refToBinaryOffset: the map from the function reference to the executable offset.
+		//  * refToBinaryOffset: the map from the function reference (ssa.FuncRef) to the executable offset.
 		//  * executable: the binary to resolve the relocations.
 		//  * relocations: the relocations to resolve.
 		//  * callTrampolineIslandOffsets: the offsets of the trampoline islands in the executable.
 		ResolveRelocations(
-			refToBinaryOffset map[ssa.FuncRef]int,
+			refToBinaryOffset []int,
 			executable []byte,
 			relocations []RelocationInfo,
 			callTrampolineIslandOffsets []int,

internal/engine/wazevo/backend/machine_test.go

@@ -56,7 +56,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(map[ssa.FuncRef]int, []byte, []RelocationInfo, []int) {}
+func (m mockMachine) ResolveRelocations([]int, []byte, []RelocationInfo, []int) {}
 
 // PostRegAlloc implements Machine.SetupPrologue.
 func (m mockMachine) PostRegAlloc() {}

internal/engine/wazevo/engine.go

@@ -188,10 +188,6 @@ func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listene
 		executables:       &executables{},
 	}
 
-	// rels is a list of relocations to be resolved. This is reused for each compilation to avoid allocation.
-	rels := make([]backend.RelocationInfo, 0)
-	refToBinaryOffset := map[ssa.FuncRef]int{}
-
 	if module.IsHostModule {
 		return e.compileHostModule(ctx, module, listeners)
 	}
@@ -201,6 +197,9 @@ func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listene
 		return cm, nil
 	}
 
+	rels := make([]backend.RelocationInfo, 0)
+	refToBinaryOffset := make([]int, importedFns+localFns)
+
 	if wazevoapi.DeterministicCompilationVerifierEnabled {
 		// The compilation must be deterministic regardless of the order of functions being compiled.
 		wazevoapi.DeterministicCompilationVerifierRandomizeIndexes(ctx)