wazevo: add VRegTable data structure (#1802)

Signed-off-by: Achille Roussel <achille.roussel@gmail.com>

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

@@ -13,7 +13,7 @@ import (
 const xArgRetRegMax, vArgRetRegMax = x7, v7 // x0-x7 & v0-v7.
 
 var regInfo = &regalloc.RegisterInfo{
-	AllocatableRegisters: [regalloc.RegTypeNum][]regalloc.RealReg{
+	AllocatableRegisters: [regalloc.NumRegType][]regalloc.RealReg{
 		// We don't allocate:
 		// - x18: Reserved by the macOS: https://developer.apple.com/documentation/xcode/writing-arm64-code-for-apple-platforms#Respect-the-purpose-of-specific-CPU-registers
 		// - x28: Reserved by Go runtime.

internal/engine/wazevo/backend/regalloc/reg.go

@@ -2,6 +2,7 @@ package regalloc
 
 import (
 	"fmt"
+	"math/bits"
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 )
@@ -13,6 +14,10 @@ type VReg uint64
 // VRegID is the lower 32bit of VReg, which is the pure identifier of VReg without RealReg info.
 type VRegID uint32
 
+const (
+	MaxVRegID = ^VRegID(0)
+)
+
 // RealReg returns the RealReg of this VReg.
 func (v VReg) RealReg() RealReg {
 	return RealReg(v >> 32)
@@ -58,6 +63,130 @@ func (v VReg) Valid() bool {
 	return v.ID() != vRegIDInvalid && v.RegType() != RegTypeInvalid
 }
 
+// VRegIDMinSet is used to collect the minimum ID by type in a collection of
+// virtual registers.
+//
+// We store the min values + 1 so the zero-value of the VRegIDMinSet is valid.
+type VRegIDMinSet [NumRegType]VRegID
+
+func (mins *VRegIDMinSet) Min(t RegType) VRegID {
+	return mins[t] - 1
+}
+
+func (mins *VRegIDMinSet) Observe(v VReg) {
+	if rt, id := v.RegType(), v.ID(); id < (mins[rt] - 1) {
+		mins[rt] = id + 1
+	}
+}
+
+// VRegTable is a data structure designed for fast association of program
+// counters to virtual registers.
+type VRegTable [NumRegType]VRegTypeTable
+
+func (t *VRegTable) Contains(v VReg) bool {
+	return t[v.RegType()].Contains(v.ID())
+}
+
+func (t *VRegTable) Lookup(v VReg) programCounter {
+	return t[v.RegType()].Lookup(v.ID())
+}
+
+func (t *VRegTable) Insert(v VReg, p programCounter) {
+	if v.IsRealReg() {
+		panic("BUG: cannot insert real registers in virtual register table")
+	}
+	t[v.RegType()].Insert(v.ID(), p)
+}
+
+func (t *VRegTable) Range(f func(VReg, programCounter)) {
+	for i := range t {
+		t[i].Range(func(id VRegID, p programCounter) {
+			f(VReg(id).SetRegType(RegType(i)), p)
+		})
+	}
+}
+
+func (t *VRegTable) Reset(minVRegIDs VRegIDMinSet) {
+	for i := range t {
+		t[i].Reset(minVRegIDs.Min(RegType(i)))
+	}
+}
+
+// VRegTypeTable implements a table for virtual registers of a specific type.
+//
+// The parent VRegTable uses 4 instances of this type to maintain a table for
+// each virtual register type.
+//
+// The virtual register type table uses a bitset to accelerate checking whether
+// a virtual register exists in the table. The bitset also helps identify which
+// slots in the program counter array are used, since the table is sparse and
+// some locations may not be occupied and will have the value zero even though
+// zero is a valid program counter value.
+type VRegTypeTable struct {
+	min VRegID
+	set bitset
+	pcs []programCounter
+}
+
+func (t *VRegTypeTable) Contains(id VRegID) bool {
+	return t.set.has(uint(id - t.min))
+}
+
+func (t *VRegTypeTable) Lookup(id VRegID) programCounter {
+	if id := int(id - t.min); t.set.has(uint(id)) {
+		return t.pcs[id]
+	}
+	return -1
+}
+
+func (t *VRegTypeTable) Insert(id VRegID, p programCounter) {
+	if p < 0 {
+		panic("BUG: cannot insert negative program counter in virtual register table")
+	}
+	i := int(id - t.min)
+	if len(t.pcs) <= i {
+		t.pcs = append(t.pcs, make([]programCounter, (i+1)-len(t.pcs))...)
+	}
+	t.set.set(uint(i))
+	t.pcs[i] = p
+}
+
+func (t *VRegTypeTable) Range(f func(VRegID, programCounter)) {
+	t.set.scan(func(i uint) { f(VRegID(i)+t.min, t.pcs[i]) })
+}
+
+func (t *VRegTypeTable) Reset(minVRegID VRegID) {
+	t.min = minVRegID
+	t.set = nil
+	t.pcs = nil
+}
+
+type bitset []uint64
+
+func (b bitset) scan(f func(uint)) {
+	for i, v := range b {
+		for j := uint(i * 64); v != 0; j++ {
+			n := uint(bits.TrailingZeros64(v))
+			j += n
+			v >>= (n + 1)
+			f(j)
+		}
+	}
+}
+
+func (b bitset) has(i uint) bool {
+	index, shift := i/64, i%64
+	return index < uint(len(b)) && ((b[index] & (1 << shift)) != 0)
+}
+
+func (b *bitset) set(i uint) {
+	index, shift := i/64, i%64
+	if index >= uint(len(*b)) {
+		*b = append(*b, make([]uint64, (index+1)-uint(len(*b)))...)
+	}
+	(*b)[index] |= 1 << shift
+}
+
 // RealReg represents a physical register.
 type RealReg byte
 
@@ -95,7 +224,7 @@ const (
 	RegTypeInvalid RegType = iota
 	RegTypeInt
 	RegTypeFloat
-	RegTypeNum
+	NumRegType
 )
 
 // String implements fmt.Stringer.

internal/engine/wazevo/backend/regalloc/reg_test.go

@@ -25,3 +25,33 @@ func Test_FromRealReg(t *testing.T) {
 	require.Equal(t, RealReg(5), r.RealReg())
 	require.Equal(t, VRegID(5), r.ID())
 }
+
+func TestVRegTable(t *testing.T) {
+	min := VRegIDMinSet{}
+	min.Observe(VReg(vRegIDReservedForRealNum + 2))
+	min.Observe(VReg(vRegIDReservedForRealNum + 1))
+	min.Observe(VReg(vRegIDReservedForRealNum + 0))
+
+	table := VRegTable{}
+	table.Reset(min)
+	table.Insert(VReg(vRegIDReservedForRealNum+0), 1)
+	table.Insert(VReg(vRegIDReservedForRealNum+1), 10)
+	table.Insert(VReg(vRegIDReservedForRealNum+2), 100)
+
+	vregs := map[VReg]programCounter{}
+	table.Range(func(v VReg, p programCounter) {
+		vregs[v] = p
+	})
+	require.Equal(t, 3, len(vregs))
+
+	for v, p := range vregs {
+		require.True(t, table.Contains(v))
+		require.Equal(t, p, table.Lookup(v))
+	}
+
+	table.Range(func(v VReg, p programCounter) {
+		require.Equal(t, vregs[v], p)
+		delete(vregs, v)
+	})
+	require.Equal(t, 0, len(vregs))
+}

internal/engine/wazevo/backend/regalloc/regalloc.go

@@ -36,7 +36,7 @@ type (
 	RegisterInfo struct {
 		// AllocatableRegisters is a 2D array of allocatable RealReg, indexed by regTypeNum and regNum.
 		// The order matters: the first element is the most preferred one when allocating.
-		AllocatableRegisters [RegTypeNum][]RealReg
+		AllocatableRegisters [NumRegType][]RealReg
 		CalleeSavedRegisters [RealRegsNumMax]bool
 		CallerSavedRegisters [RealRegsNumMax]bool
 		RealRegToVReg        []VReg
@@ -78,7 +78,7 @@ type (
 		liveOuts map[VReg]struct{}
 		liveIns  map[VReg]struct{}
 		defs     map[VReg]programCounter
-		lastUses map[VReg]programCounter
+		lastUses VRegTable
 		kills    map[VReg]programCounter
 		// Pre-colored real registers can have multiple live ranges in one block.
 		realRegUses [vRegIDReservedForRealNum][]programCounter
@@ -172,6 +172,19 @@ func (a *Allocator) livenessAnalysis(f Function) {
 	for blk := f.PostOrderBlockIteratorBegin(); blk != nil; blk = f.PostOrderBlockIteratorNext() {
 		info := a.blockInfoAt(blk.ID())
 
+		// We have to do a first pass to find the lowest VRegID in the block;
+		// this is used to reduce memory utilization in the VRegTable, which
+		// can avoid allocating memory for registers zero to minVRegID-1.
+		minVRegID := VRegIDMinSet{}
+		for instr := blk.InstrIteratorBegin(); instr != nil; instr = blk.InstrIteratorNext() {
+			for _, use := range instr.Uses() {
+				if !use.IsRealReg() {
+					minVRegID.Observe(use)
+				}
+			}
+		}
+		info.lastUses.Reset(minVRegID)
+
 		var pc programCounter
 		for instr := blk.InstrIteratorBegin(); instr != nil; instr = blk.InstrIteratorNext() {
 			var srcVR, dstVR VReg
@@ -181,7 +194,7 @@ func (a *Allocator) livenessAnalysis(f Function) {
 				if use.IsRealReg() {
 					info.addRealRegUsage(use, pos)
 				} else {
-					info.lastUses[use] = pos
+					info.lastUses.Insert(use, pos)
 				}
 			}
 			for _, def := range instr.Defs() {
@@ -209,7 +222,6 @@ func (a *Allocator) livenessAnalysis(f Function) {
 			}
 			pc += pcStride
 		}
-
 		if wazevoapi.RegAllocLoggingEnabled {
 			fmt.Printf("prepared block info for block[%d]:\n%s\n\n", blk.ID(), info.Format(a.regInfo))
 		}
@@ -231,13 +243,13 @@ func (a *Allocator) livenessAnalysis(f Function) {
 	// Now that we finished gathering liveIns, liveOuts, defs, and lastUses, the only thing left is to construct kills.
 	for blk := f.PostOrderBlockIteratorBegin(); blk != nil; blk = f.PostOrderBlockIteratorNext() { // Order doesn't matter.
 		info := a.blockInfoAt(blk.ID())
-		lastUses, outs := info.lastUses, info.liveOuts
+		outs := info.liveOuts
-		for use, pc := range lastUses {
+		info.lastUses.Range(func(use VReg, pc programCounter) {
 			// Usage without live-outs is a kill.
 			if _, ok := outs[use]; !ok {
 				info.kills[use] = pc
 			}
-		}
+		})
 
 		if wazevoapi.RegAllocLoggingEnabled {
 			fmt.Printf("\nfinalized info for block[%d]:\n%s\n", blk.ID(), info.Format(a.regInfo))
@@ -251,7 +263,7 @@ func (a *Allocator) beginUpAndMarkStack(f Function, v VReg, isPhi bool, phiDefin
 			panic(fmt.Sprintf("block without predecessor must be optimized out by the compiler: %d", blk.ID()))
 		}
 		info := a.blockInfoAt(blk.ID())
-		if _, ok := info.lastUses[v]; !ok {
+		if !info.lastUses.Contains(v) {
 			continue
 		}
 		// TODO: we might want to avoid recursion here.
@@ -463,7 +475,7 @@ func (a *Allocator) Reset() {
 
 func (a *Allocator) allocateBlockInfo(blockID int) *blockInfo {
 	if blockID >= len(a.blockInfos) {
-		a.blockInfos = append(a.blockInfos, make([]blockInfo, blockID+1)...)
+		a.blockInfos = append(a.blockInfos, make([]blockInfo, (blockID+1)-len(a.blockInfos))...)
 	}
 	info := &a.blockInfos[blockID]
 	a.initBlockInfo(info)
@@ -543,11 +555,6 @@ func (a *Allocator) initBlockInfo(i *blockInfo) {
 	} else {
 		resetMap(a, i.defs)
 	}
-	if i.lastUses == nil {
-		i.lastUses = make(map[VReg]programCounter)
-	} else {
-		resetMap(a, i.lastUses)
-	}
 	if i.kills == nil {
 		i.kills = make(map[VReg]programCounter)
 	} else {
@@ -587,9 +594,9 @@ func (i *blockInfo) Format(ri *RegisterInfo) string {
 		buf.WriteString(fmt.Sprintf("%v@%v ", v, pos))
 	}
 	buf.WriteString("\n\tlastUses: ")
-	for v, pos := range i.lastUses {
+	i.lastUses.Range(func(v VReg, pos programCounter) {
 		buf.WriteString(fmt.Sprintf("%v@%v ", v, pos))
-	}
+	})
 	buf.WriteString("\n\tkills: ")
 	for v, pos := range i.kills {
 		buf.WriteString(fmt.Sprintf("%v@%v ", v, pos))

internal/engine/wazevo/backend/regalloc/regalloc_test.go

@@ -6,6 +6,18 @@ import (
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
+func makeVRegTable(vregs map[VReg]programCounter) (table VRegTable) {
+	min := VRegIDMinSet{}
+	for v := range vregs {
+		min.Observe(v)
+	}
+	table.Reset(min)
+	for v, p := range vregs {
+		table.Insert(v, p)
+	}
+	return table
+}
+
 func TestAllocator_livenessAnalysis(t *testing.T) {
 	const realRegID, realRegID2 = 50, 100
 	realReg, realReg2 := FromRealReg(realRegID, RegTypeInt), FromRealReg(realRegID2, RegTypeInt)
@@ -28,11 +40,12 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 			exp: map[int]*blockInfo{
 				0: {
 					defs:     map[VReg]programCounter{2: pcDefOffset + pcStride, 1: pcDefOffset},
-					lastUses: map[VReg]programCounter{1: pcStride + pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{1: pcStride + pcUseOffset}),
 					kills:    map[VReg]programCounter{1: pcStride + pcUseOffset},
 				},
 			},
 		},
+
 		{
 			name: "straight",
 			// b0 -> b1 -> b2
@@ -62,11 +75,11 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 						2:    pcDefOffset,
 						3:    pcStride*2 + pcDefOffset,
 					},
-					lastUses: map[VReg]programCounter{
+					lastUses: makeVRegTable(map[VReg]programCounter{
 						1000: pcStride + pcUseOffset,
 						1:    pcStride*2 + pcUseOffset,
 						2:    pcStride*2 + pcUseOffset,
-					},
+					}),
 					liveOuts: map[VReg]struct{}{3: {}},
 					kills: map[VReg]programCounter{
 						1000: pcStride + pcUseOffset,
@@ -77,9 +90,9 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				1: {
 					liveIns:  map[VReg]struct{}{3: {}},
 					liveOuts: map[VReg]struct{}{3: {}, 4: {}, 5: {}},
-					lastUses: map[VReg]programCounter{
+					lastUses: makeVRegTable(map[VReg]programCounter{
 						3: pcStride + pcUseOffset,
-					},
+					}),
 					defs: map[VReg]programCounter{
 						4: pcStride + pcDefOffset,
 						5: pcStride + pcDefOffset,
@@ -93,7 +106,7 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				},
 				2: {
 					liveIns:  map[VReg]struct{}{3: {}, 4: {}, 5: {}},
-					lastUses: map[VReg]programCounter{3: pcUseOffset, 4: pcUseOffset, 5: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{3: pcUseOffset, 4: pcUseOffset, 5: pcUseOffset}),
 					kills:    map[VReg]programCounter{3: pcUseOffset, 4: pcUseOffset, 5: pcUseOffset},
 				},
 			},
@@ -139,11 +152,12 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 						2:    pcStride*2 + pcDefOffset,
 					},
 					liveOuts: map[VReg]struct{}{1000: {}, 1: {}, 2: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				1: {
 					liveIns:  map[VReg]struct{}{1000: {}, 1: {}},
 					liveOuts: map[VReg]struct{}{1000: {}},
-					lastUses: map[VReg]programCounter{1: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{1: pcUseOffset}),
 					kills:    map[VReg]programCounter{1: pcUseOffset},
 					realRegDefs: [vRegIDReservedForRealNum][]programCounter{
 						realRegID:  {pcDefOffset, pcStride*4 + pcDefOffset},
@@ -157,18 +171,19 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				2: {
 					liveIns:     map[VReg]struct{}{1000: {}, 2: {}},
 					liveOuts:    map[VReg]struct{}{1000: {}},
-					lastUses:    map[VReg]programCounter{2: pcUseOffset},
+					lastUses:    makeVRegTable(map[VReg]programCounter{2: pcUseOffset}),
 					kills:       map[VReg]programCounter{2: pcUseOffset},
 					realRegUses: [vRegIDReservedForRealNum][]programCounter{realRegID2: {pcUseOffset}},
 					realRegDefs: [vRegIDReservedForRealNum][]programCounter{realRegID2: {0}},
 				},
 				3: {
 					liveIns:  map[VReg]struct{}{1000: {}},
-					lastUses: map[VReg]programCounter{1000: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{1000: pcUseOffset}),
 					kills:    map[VReg]programCounter{1000: pcUseOffset},
 				},
 			},
 		},
+
 		{
 			name: "phis",
 			//   0
@@ -203,32 +218,35 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				0: {
 					defs:     map[VReg]programCounter{1000: pcDefOffset, 2000: pcDefOffset, 3000: pcDefOffset},
 					liveOuts: map[VReg]struct{}{1000: {}, 2000: {}, 3000: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				1: {
 					liveIns:  map[VReg]struct{}{2000: {}, 3000: {}},
 					liveOuts: map[VReg]struct{}{phiVReg: {}, 3000: {}},
 					defs:     map[VReg]programCounter{phiVReg: pcDefOffset},
-					lastUses: map[VReg]programCounter{2000: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{2000: pcUseOffset}),
 					kills:    map[VReg]programCounter{2000: pcUseOffset},
 				},
 				2: {
 					liveIns:  map[VReg]struct{}{phiVReg: {}, 3000: {}},
 					liveOuts: map[VReg]struct{}{phiVReg: {}, 3000: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				3: {
 					liveIns:  map[VReg]struct{}{1000: {}, 3000: {}},
 					liveOuts: map[VReg]struct{}{phiVReg: {}, 3000: {}},
 					defs:     map[VReg]programCounter{phiVReg: pcDefOffset},
-					lastUses: map[VReg]programCounter{1000: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{1000: pcUseOffset}),
 					kills:    map[VReg]programCounter{1000: pcUseOffset},
 				},
 				4: {
 					liveIns:  map[VReg]struct{}{phiVReg: {}, 3000: {}},
-					lastUses: map[VReg]programCounter{phiVReg: pcUseOffset, 3000: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{phiVReg: pcUseOffset, 3000: pcUseOffset}),
 					kills:    map[VReg]programCounter{phiVReg: pcUseOffset, 3000: pcUseOffset},
 				},
 			},
 		},
+
 		{
 			name: "loop",
 			// 0 -> 1 -> 2
@@ -274,9 +292,9 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 						1:       pcDefOffset,
 						phiVReg: pcStride + pcDefOffset,
 					},
-					lastUses: map[VReg]programCounter{
+					lastUses: makeVRegTable(map[VReg]programCounter{
 						1: pcStride + pcUseOffset,
-					},
+					}),
 					kills: map[VReg]programCounter{
 						1: pcStride + pcUseOffset,
 					},
@@ -285,32 +303,33 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 					liveIns:  map[VReg]struct{}{phiVReg: {}},
 					liveOuts: map[VReg]struct{}{phiVReg: {}, 9999: {}},
 					defs:     map[VReg]programCounter{9999: pcDefOffset},
-					lastUses: map[VReg]programCounter{},
+					lastUses: makeVRegTable(map[VReg]programCounter{}),
 					kills:    map[VReg]programCounter{},
 				},
 				2: {
 					liveIns:  map[VReg]struct{}{phiVReg: {}, 9999: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
 					defs:     map[VReg]programCounter{100: pcDefOffset},
-					lastUses: map[VReg]programCounter{phiVReg: pcUseOffset, 9999: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{phiVReg: pcUseOffset, 9999: pcUseOffset}),
 					kills:    map[VReg]programCounter{phiVReg: pcUseOffset, 9999: pcUseOffset},
 				},
 				3: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{54321: {}},
 					defs:     map[VReg]programCounter{54321: pcDefOffset},
-					lastUses: map[VReg]programCounter{100: pcStride + pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{100: pcStride + pcUseOffset}),
 					kills:    map[VReg]programCounter{100: pcStride + pcUseOffset},
 				},
 				4: {
 					liveIns:  map[VReg]struct{}{54321: {}},
 					liveOuts: map[VReg]struct{}{phiVReg: {}},
 					defs:     map[VReg]programCounter{phiVReg: pcDefOffset},
-					lastUses: map[VReg]programCounter{54321: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{54321: pcUseOffset}),
 					kills:    map[VReg]programCounter{54321: pcUseOffset},
 				},
 			},
 		},
+
 		{
 			//           -----+
 			//           v    |
@@ -336,23 +355,29 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				0: {
 					defs:     map[VReg]programCounter{99999: pcDefOffset},
 					liveOuts: map[VReg]struct{}{99999: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				1: {
 					liveIns:  map[VReg]struct{}{99999: {}},
 					liveOuts: map[VReg]struct{}{99999: {}},
-					lastUses: map[VReg]programCounter{99999: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{99999: pcUseOffset}),
 				},
 				2: {
 					liveIns:  map[VReg]struct{}{99999: {}},
 					liveOuts: map[VReg]struct{}{99999: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				3: {
 					liveIns:  map[VReg]struct{}{99999: {}},
 					liveOuts: map[VReg]struct{}{99999: {}},
+					lastUses: makeVRegTable(nil),
+				},
+				4: {
+					lastUses: makeVRegTable(nil),
 				},
-				4: {},
 			},
 		},
+
 		//      2
 		//      ^              +----+
 		//      |              v    |
@@ -405,35 +430,48 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				return newMockFunction(b0, b1, b2, b3, b4, b7, b8, b5, b6, b9)
 			},
 			exp: map[int]*blockInfo{
-				0: {},
+				0: {
-				1: {},
+					lastUses: makeVRegTable(nil),
-				2: {},
+				},
+				1: {
+					lastUses: makeVRegTable(nil),
+				},
+				2: {
+					lastUses: makeVRegTable(nil),
+				},
 				3: {
-					liveOuts: map[VReg]struct{}{100: {}},
 					defs:     map[VReg]programCounter{100: pcDefOffset},
+					liveOuts: map[VReg]struct{}{100: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				4: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				5: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
-					lastUses: map[VReg]programCounter{100: pcUseOffset},
+					lastUses: makeVRegTable(map[VReg]programCounter{100: pcUseOffset}),
 				},
 				6: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				7: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
+					lastUses: makeVRegTable(nil),
 				},
 				8: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
+					lastUses: makeVRegTable(nil),
+				},
+				9: {
+					lastUses: makeVRegTable(nil),
 				},
-				9: {},
 			},
 		},
 	} {
@@ -448,7 +486,7 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				initMapInInfo(exp)
 				saved := actual.intervalMng
 				actual.intervalMng = nil // Don't compare intervalManager.
-				require.Equal(t, exp, actual, "\n[exp for block[%d]]\n%s\n[actual for block[%d]]\n%s", blockID, exp, blockID, actual)
+				require.Equal(t, exp, actual, "\n[exp for block[%d]]\n%v\n[actual for block[%d]]\n%v", blockID, exp, blockID, actual)
 				actual.intervalMng = saved
 			}
 
@@ -533,9 +571,6 @@ func initMapInInfo(info *blockInfo) {
 	if info.kills == nil {
 		info.kills = make(map[VReg]programCounter)
 	}
-	if info.lastUses == nil {
-		info.lastUses = make(map[VReg]programCounter)
-	}
 }
 
 func TestNode_assignedRealReg(t *testing.T) {

internal/engine/wazevo/backend/regalloc/spill_handler_test.go

@@ -29,7 +29,7 @@ func TestSpillHandler_getUnusedOrEvictReg(t *testing.T) {
 	require.Equal(t, 4, len(s.activeRegs))
 
 	regInfo := RegisterInfo{
-		AllocatableRegisters: [RegTypeNum][]RealReg{
+		AllocatableRegisters: [NumRegType][]RealReg{
 			RegTypeInt: {
 				RealReg(0xff), // unused.
 				RealReg(0), RealReg(1), RealReg(2),