wazevo(arm64): removes brtable pointer from instr struct (#2038)

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

internal/engine/wazevo/backend/backend_test.go

@@ -2047,7 +2047,7 @@ L1 (SSA Block: blk0):
 	orr w137?, wzr, #0x6
 	subs wzr, w130?, w137?
 	csel w138?, w137?, w130?, hs
-	br_table_sequence x138?, [L2, L3, L4, L5, L6, L7, L8]
+	br_table_sequence x138?, table_index=0
 L2 (SSA Block: blk7):
 	b L9
 L3 (SSA Block: blk8):
@@ -2088,7 +2088,7 @@ L1 (SSA Block: blk0):
 	orr w8, wzr, #0x6
 	subs wzr, w2, w8
 	csel w8, w8, w2, hs
-	adr x27, #16; ldrsw x8, [x27, x8, UXTW 2]; add x27, x27, x8; br x27; [0x1c 0x20 0x34 0x48 0x5c 0x70 0x84]
+	br_table_sequence x8, table_index=0
 L2 (SSA Block: blk7):
 	b #0x68 (L9)
 L3 (SSA Block: blk8):

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

@@ -3,7 +3,6 @@ package arm64
 import (
 	"fmt"
 	"math"
-	"strings"
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
@@ -26,7 +25,6 @@ type (
 		u1, u2, u3          uint64
 		rd, rm, rn, ra      operand
 		amode               addressMode
-		targets             []uint32
 		kind                instructionKind
 		addedBeforeRegAlloc bool
 	}
@@ -671,10 +669,11 @@ func (i *instruction) asBr(target label) {
 	i.u1 = uint64(target)
 }
 
-func (i *instruction) asBrTableSequence(indexReg regalloc.VReg, targets []uint32) {
+func (i *instruction) asBrTableSequence(indexReg regalloc.VReg, targetIndex, targetCounts int) {
 	i.kind = brTableSequence
 	i.rn = operandNR(indexReg)
-	i.targets = targets
+	i.u1 = uint64(targetIndex)
+	i.u2 = uint64(targetCounts)
 }
 
 func (i *instruction) brTableSequenceOffsetsResolved() {
@@ -1443,28 +1442,8 @@ func (i *instruction) String() (str string) {
 	case adr:
 		str = fmt.Sprintf("adr %s, #%#x", formatVRegSized(i.rd.nr(), 64), int64(i.u1))
 	case brTableSequence:
-		if i.u3 == 0 { // The offsets haven't been resolved yet.
-			labels := make([]string, len(i.targets))
-			for index, l := range i.targets {
-				labels[index] = label(l).String()
-			}
-			str = fmt.Sprintf("br_table_sequence %s, [%s]",
-				formatVRegSized(i.rn.nr(), 64),
-				strings.Join(labels, ", "),
-			)
-		} else {
-			// See encodeBrTableSequence for the encoding.
-			offsets := make([]string, len(i.targets))
-			for index, offset := range i.targets {
-				offsets[index] = fmt.Sprintf("%#x", int32(offset))
-			}
-			str = fmt.Sprintf(
-				`adr %[2]s, #16; ldrsw %[1]s, [%[2]s, %[1]s, UXTW 2]; add %[2]s, %[2]s, %[1]s; br %[2]s; %s`,
-				formatVRegSized(i.rn.nr(), 64),
-				formatVRegSized(tmpRegVReg, 64),
-				offsets,
-			)
-		}
+		targetIndex := i.u1
+		str = fmt.Sprintf("br_table_sequence %s, table_index=%d", formatVRegSized(i.rn.nr(), 64), targetIndex)
 	case exitSequence:
 		str = fmt.Sprintf("exit_sequence %s", formatVRegSized(i.rn.nr(), 64))
 	case atomicRmw:
@@ -2313,7 +2292,7 @@ func (i *instruction) size() int64 {
 		}
 		return 4 + 4 + 16
 	case brTableSequence:
-		return 4*4 + int64(len(i.targets))*4
+		return 4*4 + int64(i.u2)*4
 	default:
 		return 4
 	}

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

@@ -17,11 +17,12 @@ func (m *machine) Encode(ctx context.Context) {
 
 func (m *machine) encode(root *instruction) {
 	for cur := root; cur != nil; cur = cur.next {
-		cur.encode(m.compiler)
+		cur.encode(m)
 	}
 }
 
-func (i *instruction) encode(c backend.Compiler) {
+func (i *instruction) encode(m *machine) {
+	c := m.compiler
 	switch kind := i.kind; kind {
 	case nop0, emitSourceOffsetInfo:
 	case exitSequence:
@@ -341,7 +342,8 @@ func (i *instruction) encode(c backend.Compiler) {
 			vecArrangement(i.u2)),
 		)
 	case brTableSequence:
-		encodeBrTableSequence(c, i.rn.reg(), i.targets)
+		targets := m.jmpTableTargets[i.u1]
+		encodeBrTableSequence(c, i.rn.reg(), targets)
 	case fpuToInt, intToFpu:
 		c.Emit4Bytes(encodeCnvBetweenFloatInt(i))
 	case fpuRR:

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

@@ -1774,15 +1774,17 @@ func TestInstruction_encode(t *testing.T) {
 }
 
 func TestInstruction_encode_call(t *testing.T) {
-	m := &mockCompiler{buf: make([]byte, 128)}
+	_, _, m := newSetupWithMockContext()
+	mock := m.compiler.(*mockCompiler)
+	mock.buf = make([]byte, 128)
 	i := &instruction{}
 	i.asCall(ssa.FuncRef(555), nil)
 	i.encode(m)
-	buf := m.buf[128:]
+	buf := mock.buf[128:]
 	require.Equal(t, "00000094", hex.EncodeToString(buf))
-	require.Equal(t, 1, len(m.relocs))
-	require.Equal(t, ssa.FuncRef(555), m.relocs[0].FuncRef)
-	require.Equal(t, int64(128), m.relocs[0].Offset)
+	require.Equal(t, 1, len(mock.relocs))
+	require.Equal(t, ssa.FuncRef(555), mock.relocs[0].FuncRef)
+	require.Equal(t, int64(128), mock.relocs[0].Offset)
 }
 
 func TestInstruction_encode_br_condflag(t *testing.T) {
@@ -1810,14 +1812,14 @@ func TestInstruction_encode_br_condflag(t *testing.T) {
 		i := &instruction{}
 		i.asCondBr(tc.c.asCond(), label(1), false)
 		i.condBrOffsetResolve(0xf0)
-		m := &mockCompiler{}
+		_, _, m := newSetupWithMockContext()
 		i.encode(m)
 		// Note: for quick iteration we can use golang.org/x/arch package to verify the encoding.
 		// 	but wazero doesn't add even a test dependency to it, so commented out.
 		// inst, err := arm64asm.Decode(m.buf)
 		// require.NoError(t, err)
 		// fmt.Println(inst.String())
-		require.Equal(t, tc.want, hex.EncodeToString(m.buf))
+		require.Equal(t, tc.want, hex.EncodeToString(m.compiler.Buf()))
 	}
 }
 
@@ -2140,14 +2142,14 @@ func TestInstruction_encoding_store_encoding(t *testing.T) {
 			default:
 				t.Fatalf("unknown kind: %v", tc.k)
 			}
-			m := &mockCompiler{}
+			_, _, m := newSetupWithMockContext()
 			i.encode(m)
 			// Note: for quick iteration we can use golang.org/x/arch package to verify the encoding.
 			// 	but wazero doesn't add even a test dependency to it, so commented out.
 			// inst, err := arm64asm.Decode(m.buf)
 			// require.NoError(t, err)
 			// fmt.Println(inst.String())
-			require.Equal(t, tc.want, hex.EncodeToString(m.buf))
+			require.Equal(t, tc.want, hex.EncodeToString(m.compiler.Buf()))
 		})
 	}
 }
@@ -2192,13 +2194,16 @@ func Test_encodeExitSequence(t *testing.T) {
 }
 
 func Test_encodeBrTableSequence(t *testing.T) {
-	m := &mockCompiler{}
-	i := &instruction{kind: brTableSequence, targets: []uint32{1, 2, 3, 4, 5}}
-	encodeBrTableSequence(m, x22VReg, i.targets)
-	encoded := m.Buf()
+	_, _, m := newSetupWithMockContext()
+	i := &instruction{}
+	const tableIndex, tableSize = 5, 10
+	i.asBrTableSequence(x22VReg, tableIndex, tableSize)
+	m.jmpTableTargets = [][]uint32{{}, {}, {}, {}, {}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}}
+	i.encode(m)
+	encoded := m.compiler.Buf()
 	require.Equal(t, i.size(), int64(len(encoded)))
 	require.Equal(t, "9b000010765bb6b87b03168b60031fd6", hex.EncodeToString(encoded[:brTableSequenceOffsetTableBegin]))
-	require.Equal(t, "0100000002000000030000000400000005000000", hex.EncodeToString(encoded[brTableSequenceOffsetTableBegin:]))
+	require.Equal(t, "0100000002000000030000000400000005000000060000000700000008000000090000000a000000", hex.EncodeToString(encoded[brTableSequenceOffsetTableBegin:]))
 }
 
 func Test_encodeUnconditionalBranch(t *testing.T) {

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

@@ -61,13 +61,8 @@ func (m *machine) lowerBrTable(i *ssa.Instruction) {
 
 	brSequence := m.allocateInstr()
 
-	// TODO: reuse the slice!
-	labels := make([]uint32, len(targets))
-	for j, target := range targets {
-		labels[j] = uint32(m.executableContext.GetOrAllocateSSABlockLabel(target))
-	}
-
-	brSequence.asBrTableSequence(adjustedIndex, labels)
+	tableIndex := m.addJmpTableTarget(targets)
+	brSequence.asBrTableSequence(adjustedIndex, tableIndex, len(targets))
 	m.insert(brSequence)
 }
 

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

@@ -31,6 +31,9 @@ type (
 		// condBrRelocs holds the conditional branches which need offset relocation.
 		condBrRelocs []condBrReloc
 
+		// jmpTableTargets holds the labels of the jump table targets.
+		jmpTableTargets [][]uint32
+
 		// spillSlotSize is the size of the stack slot in bytes used for spilling registers.
 		// During the execution of the function, the stack looks like:
 		//
@@ -145,6 +148,7 @@ func (m *machine) Reset() {
 	m.unresolvedAddressModes = m.unresolvedAddressModes[:0]
 	m.maxRequiredStackSizeForCalls = 0
 	m.executableContext.Reset()
+	m.jmpTableTargets = m.jmpTableTargets[:0]
 }
 
 // SetCurrentABI implements backend.Machine SetCurrentABI.
@@ -365,11 +369,13 @@ func (m *machine) resolveRelativeAddresses(ctx context.Context) {
 				cur.condBrOffsetResolve(diff)
 			}
 		case brTableSequence:
-			for i := range cur.targets {
-				l := label(cur.targets[i])
+			tableIndex := cur.u1
+			targets := m.jmpTableTargets[tableIndex]
+			for i := range targets {
+				l := label(targets[i])
 				offsetOfTarget := ectx.LabelPositions[l].BinaryOffset
 				diff := offsetOfTarget - (currentOffset + brTableSequenceOffsetTableBegin)
-				cur.targets[i] = uint32(diff)
+				targets[i] = uint32(diff)
 			}
 			cur.brTableSequenceOffsetsResolved()
 		case emitSourceOffsetInfo:
@@ -494,3 +500,14 @@ func (m *machine) frameSize() int64 {
 	}
 	return s
 }
+
+func (m *machine) addJmpTableTarget(targets []ssa.BasicBlock) (index int) {
+	// TODO: reuse the slice!
+	labels := make([]uint32, len(targets))
+	for j, target := range targets {
+		labels[j] = uint32(m.executableContext.GetOrAllocateSSABlockLabel(target))
+	}
+	index = len(m.jmpTableTargets)
+	m.jmpTableTargets = append(m.jmpTableTargets, labels)
+	return
+}