Repository: tetratelabs/wazero
Pull request #1678 (2023-09-01 06:15:37)

wazevo(regalloc): support for reg spills
This make it pass "skip-stack-guard-page" spectest!!

#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi.go b/internal/engine/wazevo/backend/isa/arm64/abi.go
index 0929e3db..3ab4b05f 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi.go
@@ -339,7 +339,7 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 
 	stackSlotSize := calleeABI.alignedStackSlotSize()
 	if stackSlotSize > 0 {
-		m.insertAddOrSubStackPointer(spVReg, stackSlotSize, false /* == sub */)
+		m.insertAddOrSubStackPointer(spVReg, stackSlotSize, false /* == sub */, false)
 	}
 
 	if m.maxRequiredStackSizeForCalls < stackSlotSize {
@@ -376,17 +376,13 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 	}
 
 	if stackSlotSize > 0 {
-		m.insertAddOrSubStackPointer(spVReg, stackSlotSize, true /* = add */)
+		m.insertAddOrSubStackPointer(spVReg, stackSlotSize, true /* add */, false)
 	}
 }
 
-func (m *machine) insertAddOrSubStackPointer(rd regalloc.VReg, diff int64, add bool) {
-	i := m.allocateAddOrSubStackPointer(rd, diff, add, false)
-	m.insert(i)
-}
-
-func (m *machine) allocateAddOrSubStackPointer(rd regalloc.VReg, diff int64, add bool, afterLowering bool) (alu *instruction) {
+func (m *machine) insertAddOrSubStackPointer(rd regalloc.VReg, diff int64, add bool, afterLowering bool) {
 	if imm12Operand, ok := asImm12Operand(uint64(diff)); ok {
+		var alu *instruction
 		if afterLowering {
 			alu = m.allocateInstrAfterLowering()
 		} else {
@@ -399,8 +395,22 @@ func (m *machine) allocateAddOrSubStackPointer(rd regalloc.VReg, diff int64, add
 			ao = aluOpSub
 		}
 		alu.asALU(ao, operandNR(rd), operandNR(spVReg), imm12Operand, true)
-		return
+		m.insert(alu)
 	} else {
-		panic("TODO: large stack pointer add/sub") // tmpRegVReg can be used for this.
+		m.lowerConstantI64(tmpRegVReg, diff)
+		var alu *instruction
+		if afterLowering {
+			alu = m.allocateInstrAfterLowering()
+		} else {
+			alu = m.allocateInstr()
+		}
+		var ao aluOp
+		if add {
+			ao = aluOpAdd
+		} else {
+			ao = aluOpSub
+		}
+		alu.asALU(ao, operandNR(rd), operandNR(spVReg), operandNR(tmpRegVReg), true)
+		m.insert(alu)
 	}
 }
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
index 16dc99ef..a369b256 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
@@ -51,10 +51,7 @@ func (m *machine) SetupPrologue() {
 			panic(fmt.Errorf("BUG: spill slot size %d is not 16-byte aligned", size))
 		}
 
-		decSP := m.allocateAddOrSubStackPointer(spVReg, size, false, true)
-		cur.next = decSP
-		decSP.prev = cur
-		cur = decSP
+		cur = m.addsAddOrSubStackPointer(cur, spVReg, size, false, true)
 
 		// At this point, the stack looks like:
 		//
@@ -219,10 +216,7 @@ func (m *machine) setupEpilogueAfter(cur *instruction) {
 		//   SP---> +-----------------+
 		//             (low address)
 		//
-		addSP := m.allocateAddOrSubStackPointer(spVReg, s, true, true)
-		cur.next = addSP
-		addSP.prev = cur
-		cur = addSP
+		cur = m.addsAddOrSubStackPointer(cur, spVReg, s, true, true)
 	}
 
 	// Reload the return address (lr).
@@ -409,3 +403,15 @@ func (m *machine) CompileStackGrowCallSequence() []byte {
 	m.encode(m.rootInstr)
 	return m.compiler.Buf()
 }
+
+func (m *machine) addsAddOrSubStackPointer(cur *instruction, rd regalloc.VReg, diff int64, add bool, afterLowering bool) *instruction {
+	m.pendingInstructions = m.pendingInstructions[:0]
+	m.insertAddOrSubStackPointer(rd, diff, add, afterLowering)
+	for _, inserted := range m.pendingInstructions {
+		cur.next = inserted
+		inserted.prev = cur
+		cur = inserted
+	}
+	m.pendingInstructions = m.pendingInstructions[:0]
+	return cur
+}
diff --git a/internal/engine/wazevo/backend/regalloc/assign.go b/internal/engine/wazevo/backend/regalloc/assign.go
index 27b90a3b..ed94bc6c 100644
--- a/internal/engine/wazevo/backend/regalloc/assign.go
+++ b/internal/engine/wazevo/backend/regalloc/assign.go
@@ -113,18 +113,93 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
 		panic("BUG: multiple def instructions must be special cased")
 	}
 
-	a.handleSpills(f, pc, instr, vRegIDToNode, liveNodes, usesSpills, defSpill)
+	a.handleSpills(f, pc, instr, liveNodes, usesSpills, defSpill)
 	a.vs = usesSpills[:0] // for reuse.
 }
 
 func (a *Allocator) handleSpills(
-	f Function, pc programCounter, instr Instr, vRegIDToNode []*node, liveNodes []liveNodeInBlock,
+	f Function, pc programCounter, instr Instr, liveNodes []liveNodeInBlock,
 	usesSpills []VReg, defSpill VReg,
 ) {
-	if len(usesSpills) == 0 && !defSpill.Valid() {
-		return
+	_usesSpills, _defSpill := len(usesSpills) > 0, defSpill.Valid()
+	switch {
+	case !_usesSpills && !_defSpill: // Nothing to do.
+	case !_usesSpills && _defSpill: // Only definition is spilled.
+		a.collectActiveNodesAt(pc+pcDefOffset, liveNodes)
+		a.spillHandler.init(a.nodes1)
+
+		r, evictedNode := a.spillHandler.getUnusedOrEvictReg(defSpill.RegType(), a.regInfo)
+		if evictedNode != nil {
+			evictedNodeV := evictedNode.v.SetRealReg(evictedNode.assignedRealReg())
+			f.StoreRegisterBefore(evictedNodeV, instr)
+			f.ReloadRegisterAfter(evictedNodeV, instr)
+		}
+
+		defSpill = defSpill.SetRealReg(r)
+		instr.AssignDef(defSpill)
+
+		f.StoreRegisterAfter(defSpill, instr)
+
+	case _usesSpills:
+		a.collectActiveNodesAt(pc, liveNodes)
+		a.spillHandler.init(a.nodes1)
+
+		var evicted [3]*node
+		var evictedCount int
+		for i, u := range usesSpills {
+			r, evictedNode := a.spillHandler.getUnusedOrEvictReg(u.RegType(), a.regInfo)
+			if evictedNode != nil {
+				evicted[evictedCount] = evictedNode
+				evictedCount++
+			}
+			usesSpills[i] = u.SetRealReg(r)
+		}
+
+		for i := 0; i < evictedCount; i++ {
+			evictedNode := evicted[i]
+			evictedNodeV := evictedNode.v.SetRealReg(evictedNode.assignedRealReg())
+			f.StoreRegisterBefore(evictedNodeV, instr)
+			f.ReloadRegisterAfter(evictedNodeV, instr)
+		}
+
+		for _, u := range usesSpills {
+			f.ReloadRegisterBefore(u, instr)
+		}
+
+		for useIndex, v := range instr.Uses() {
+			for _, u := range usesSpills {
+				if v.ID() == u.ID() {
+					instr.AssignUse(useIndex, u)
+				}
+			}
+		}
+
+		if _defSpill {
+			// We can reuse the register in usesSpills for the definition.
+			for _, u := range usesSpills {
+				if defSpill.RegType() == u.RegType() {
+					defSpill = defSpill.SetRealReg(u.RealReg())
+					break
+				}
+			}
+
+			if !defSpill.IsRealReg() {
+				// This case, the destination register type is different from the source registers.
+				a.collectActiveNodesAt(pc+pcDefOffset, liveNodes)
+				a.spillHandler.init(a.nodes1)
+				r, evictedNode := a.spillHandler.getUnusedOrEvictReg(defSpill.RegType(), a.regInfo)
+				if evictedNode != nil {
+					evictedNodeV := evictedNode.v.SetRealReg(evictedNode.assignedRealReg())
+					f.StoreRegisterBefore(evictedNodeV, instr)
+					f.ReloadRegisterAfter(evictedNodeV, instr)
+				}
+				defSpill = defSpill.SetRealReg(r)
+			}
+
+			instr.AssignDef(defSpill)
+			f.StoreRegisterAfter(defSpill, instr)
+		}
 	}
-	panic("TODO")
 }
 
 func (a *Allocator) assignIndirectCall(f Function, instr Instr, vRegIDToNode []*node) {
@@ -192,3 +267,21 @@ func (a *Allocator) collectActiveNonRealVRegsAt(pc programCounter, liveNodes []l
 	}
 	a.nodes1 = nodes
 }
+
+func (a *Allocator) collectActiveNodesAt(pc programCounter, liveNodes []liveNodeInBlock) {
+	nodes := a.nodes1[:0]
+	for _, live := range liveNodes {
+		n := live.n
+		if n.assignedRealReg() != RealRegInvalid {
+			r := &n.ranges[live.rangeIndex]
+			if r.begin > pc {
+				// liveNodes are sorted by the start program counter, so we can break here.
+				break
+			}
+			if pc <= r.end { // pc is in the range.
+				nodes = append(nodes, n)
+			}
+		}
+	}
+	a.nodes1 = nodes
+}
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index f003082d..64f14824 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -63,6 +63,7 @@ type (
 		vRegIDToNode [] /* VRegID to */ *node
 		blockInfos   [] /* blockID to */ blockInfo
 		vs           []VReg
+		spillHandler spillHandler
 
 		// Followings are re-used during various places e.g. coloring.
 		realRegSet map[RealReg]struct{}
@@ -666,3 +667,12 @@ func (a *Allocator) recordCopyRelation(dst, src VReg) {
 		srcN.copyToReal = dst.RealReg()
 	}
 }
+
+// assignedRealReg returns either the assigned RealReg to this node or precolored RealReg.
+func (n *node) assignedRealReg() RealReg {
+	r := n.r
+	if r != RealRegInvalid {
+		return r
+	}
+	return n.v.RealReg()
+}
diff --git a/internal/engine/wazevo/backend/regalloc/spill_handler.go b/internal/engine/wazevo/backend/regalloc/spill_handler.go
new file mode 100644
index 00000000..bcd74df8
--- /dev/null
+++ b/internal/engine/wazevo/backend/regalloc/spill_handler.go
@@ -0,0 +1,66 @@
+package regalloc
+
+// spillHandler is a helper to handle the spill and reload of registers at some point in the program.
+type spillHandler struct {
+	activeRegs map[RealReg]spillHandlerRegState
+	deleteTemp []RealReg
+}
+
+type spillHandlerRegState struct {
+	state int
+	node  *node
+}
+
+const (
+	spillHandlerRegStateUsed = iota
+	spillHandlerRegStateEvictable
+	spillHandlerRegStateEvicted
+)
+
+// init initializes the spill handler with the active nodes which are the node alive at a some point in the program.
+func (s *spillHandler) init(activeNodes []*node) {
+	if s.activeRegs == nil {
+		s.activeRegs = make(map[RealReg]spillHandlerRegState)
+	} else {
+		s.deleteTemp = s.deleteTemp[:0]
+		for r := range s.activeRegs {
+			s.deleteTemp = append(s.deleteTemp, r)
+		}
+		for _, r := range s.deleteTemp {
+			delete(s.activeRegs, r)
+		}
+	}
+	for _, n := range activeNodes {
+		s.activeRegs[n.assignedRealReg()] = spillHandlerRegState{node: n, state: spillHandlerRegStateEvictable}
+	}
+}
+
+// getUnusedOrEvictReg returns an unused register of the given type or evicts a register from the active registers.
+func (s *spillHandler) getUnusedOrEvictReg(regType RegType, regInfo *RegisterInfo) (r RealReg, evicted *node) {
+	allocatables := regInfo.AllocatableRegisters[regType]
+	for _, candidate := range allocatables {
+		_, ok := s.activeRegs[candidate]
+		if !ok {
+			r = candidate
+			s.activeRegs[candidate] = spillHandlerRegState{state: spillHandlerRegStateUsed}
+			break
+		} // ok=true meaning that it is either used or evicted.
+	}
+
+	if r == RealRegInvalid {
+		// We need to evict a register from the active registers.
+		for _, candidate := range allocatables {
+			state, ok := s.activeRegs[candidate]
+			if !ok {
+				panic("BUG")
+			}
+			if state.state == spillHandlerRegStateEvictable {
+				evicted = state.node
+				r = candidate
+				s.activeRegs[candidate] = spillHandlerRegState{node: state.node, state: spillHandlerRegStateEvicted}
+				break
+			}
+		}
+	}
+	return
+}


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi_test.go b/internal/engine/wazevo/backend/isa/arm64/abi_test.go
index 7202e0cf..4f621317 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi_test.go
@@ -234,3 +234,28 @@ func TestAbiImpl_init(t *testing.T) {
 		})
 	}
 }
+
+func TestMachine_insertAddOrSubStackPointer(t *testing.T) {
+	t.Run("add/small", func(t *testing.T) {
+		_, _, m := newSetupWithMockContext()
+		m.insertAddOrSubStackPointer(spVReg, 0x10, true, false)
+		require.Equal(t, `add sp, sp, #0x10`, formatEmittedInstructionsInCurrentBlock(m))
+	})
+	t.Run("add/large stack", func(t *testing.T) {
+		_, _, m := newSetupWithMockContext()
+		m.insertAddOrSubStackPointer(spVReg, 0xffffffff_8, true, false)
+		require.Equal(t, `orr x27, xzr, #0xffffffff8
+add sp, sp, x27`, formatEmittedInstructionsInCurrentBlock(m))
+	})
+	t.Run("sub/small", func(t *testing.T) {
+		_, _, m := newSetupWithMockContext()
+		m.insertAddOrSubStackPointer(spVReg, 0x10, false, false)
+		require.Equal(t, `sub sp, sp, #0x10`, formatEmittedInstructionsInCurrentBlock(m))
+	})
+	t.Run("sub/large stack", func(t *testing.T) {
+		_, _, m := newSetupWithMockContext()
+		m.insertAddOrSubStackPointer(spVReg, 0xffffffff_8, false, false)
+		require.Equal(t, `orr x27, xzr, #0xffffffff8
+sub sp, sp, x27`, formatEmittedInstructionsInCurrentBlock(m))
+	})
+}
diff --git a/internal/engine/wazevo/backend/regalloc/api_test.go b/internal/engine/wazevo/backend/regalloc/api_test.go
index 1e8cb6a6..c9fbd78d 100644
--- a/internal/engine/wazevo/backend/regalloc/api_test.go
+++ b/internal/engine/wazevo/backend/regalloc/api_test.go
@@ -9,15 +9,15 @@ import (
 type (
 	// mockFunction implements Function.
 	mockFunction struct {
-		iter   int
-		blocks []*mockBlock
-		storeRegisterAfter, reloadRegisterAfter,
-		storeRegisterBefore, reloadRegisterBefore []storeOrReloadInfo
+		iter            int
+		blocks          []*mockBlock
+		befores, afters []storeOrReloadInfo
 	}
 
 	storeOrReloadInfo struct {
-		v     VReg
-		instr Instr
+		reload bool
+		v      VReg
+		instr  Instr
 	}
 
 	// mockBlock implements Block.
@@ -109,22 +109,22 @@ func (m *mockInstr) asIndirectCall() *mockInstr {
 
 // StoreRegisterAfter implements Function.StoreRegisterAfter.
 func (m *mockFunction) StoreRegisterAfter(v VReg, instr Instr) {
-	m.storeRegisterAfter = append(m.storeRegisterAfter, storeOrReloadInfo{v, instr})
+	m.afters = append(m.afters, storeOrReloadInfo{false, v, instr})
 }
 
 // ReloadRegisterBefore implements Function.ReloadRegisterBefore.
 func (m *mockFunction) ReloadRegisterBefore(v VReg, instr Instr) {
-	m.reloadRegisterBefore = append(m.reloadRegisterBefore, storeOrReloadInfo{v, instr})
+	m.befores = append(m.befores, storeOrReloadInfo{true, v, instr})
 }
 
 // StoreRegisterBefore implements Function.StoreRegisterBefore.
 func (m *mockFunction) StoreRegisterBefore(v VReg, instr Instr) {
-	m.storeRegisterBefore = append(m.storeRegisterBefore, storeOrReloadInfo{v, instr})
+	m.befores = append(m.befores, storeOrReloadInfo{false, v, instr})
 }
 
 // ReloadRegisterAfter implements Function.ReloadRegisterAfter.
 func (m *mockFunction) ReloadRegisterAfter(v VReg, instr Instr) {
-	m.reloadRegisterAfter = append(m.reloadRegisterAfter, storeOrReloadInfo{v, instr})
+	m.afters = append(m.afters, storeOrReloadInfo{true, v, instr})
 }
 
 // ClobberedRegisters implements Function.ClobberedRegisters.
@@ -227,7 +227,7 @@ func (m *mockBlock) Entry() bool { return m._entry }
 // AssignUses implements Instr.
 func (m *mockInstr) AssignUse(index int, reg VReg) {
 	if index >= len(m.uses) {
-		m.uses = append(m.uses, make([]VReg, 100)...)
+		m.uses = append(m.uses, make([]VReg, 5)...)
 	}
 	m.uses[index] = reg
 }
diff --git a/internal/engine/wazevo/backend/regalloc/assign_test.go b/internal/engine/wazevo/backend/regalloc/assign_test.go
index 18154c1a..a93d0976 100644
--- a/internal/engine/wazevo/backend/regalloc/assign_test.go
+++ b/internal/engine/wazevo/backend/regalloc/assign_test.go
@@ -22,8 +22,8 @@ func TestAllocator_assignRegistersPerInstr(t *testing.T) {
 		f := newMockFunction(blk)
 		a.assignRegistersPerInstr(f, pc, call, nil, liveNodes)
 
-		require.Equal(t, 2, len(f.storeRegisterBefore))
-		require.Equal(t, 2, len(f.reloadRegisterAfter))
+		require.Equal(t, 2, len(f.befores))
+		require.Equal(t, 2, len(f.afters))
 	})
 	t.Run("call_indirect/func_ptr not spilled", func(t *testing.T) {
 		a := NewAllocator(&RegisterInfo{CallerSavedRegisters: map[RealReg]struct{}{1: {}, 3: {}, 0xff: {}}})
@@ -45,8 +45,8 @@ func TestAllocator_assignRegistersPerInstr(t *testing.T) {
 		f := newMockFunction(blk)
 		a.assignRegistersPerInstr(f, pc, callInd, []*node{0: functionPtrLiveNode.n}, liveNodes)
 
-		require.Equal(t, 2, len(f.storeRegisterBefore))
-		require.Equal(t, 2, len(f.reloadRegisterAfter))
+		require.Equal(t, 2, len(f.befores))
+		require.Equal(t, 2, len(f.afters))
 		require.True(t, callInd.uses[0].IsRealReg())
 		require.Equal(t, functionPtrVReg.SetRealReg(0xff), callInd.uses[0])
 	})
@@ -71,11 +71,10 @@ func TestAllocator_assignRegistersPerInstr(t *testing.T) {
 			0: {r: RealRegInvalid},
 		}, liveNodes)
 
-		require.Equal(t, 2, len(f.storeRegisterBefore))
-		require.Equal(t, 2, len(f.reloadRegisterAfter))
-		require.Equal(t, 1, len(f.reloadRegisterBefore))
-		require.Equal(t, callInd, f.reloadRegisterBefore[0].instr)
-		require.Equal(t, functionPtrVReg.SetRealReg(0xbb), f.reloadRegisterBefore[0].v)
+		require.Equal(t, 3, len(f.befores))
+		require.Equal(t, 2, len(f.afters))
+		require.Equal(t, callInd, f.befores[2].instr)
+		require.Equal(t, functionPtrVReg.SetRealReg(0xbb), f.befores[2].v)
 		require.True(t, callInd.uses[0].IsRealReg())
 		require.Equal(t, functionPtrVReg.SetRealReg(0xbb), callInd.uses[0])
 	})
@@ -164,3 +163,223 @@ func TestAllocator_activeNonRealVRegsAt(t *testing.T) {
 		})
 	}
 }
+
+func TestAllocator_handleSpills(t *testing.T) {
+	requireInsertedInst := func(t *testing.T, f *mockFunction, before bool, index int, instr Instr, reload bool, v VReg) {
+		lists := f.afters
+		if before {
+			lists = f.befores
+		}
+		actual := lists[index]
+		require.Equal(t, instr, actual.instr)
+		require.Equal(t, reload, actual.reload)
+		require.Equal(t, v, actual.v)
+	}
+
+	t.Run("no spills", func(t *testing.T) {
+		a := NewAllocator(&RegisterInfo{})
+		a.handleSpills(nil, 0, nil, nil, nil, VRegInvalid)
+	})
+	t.Run("only def / evicted / Real reg backed VReg", func(t *testing.T) {
+		const pc = 5
+		liveNodes := []liveNodeInBlock{
+			// Real reg backed VReg.
+			{n: &node{r: RealRegInvalid, v: VReg(1).SetRealReg(0xa), ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0xb), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0xc), v: 0xc, ranges: []liveRange{{begin: pc, end: 20}}}},
+		}
+
+		a := NewAllocator(&RegisterInfo{
+			AllocatableRegisters: [3][]RealReg{RegTypeInt: {0xa, 0xb, 0xc}}, // Only live nodes are allocatable.
+		})
+
+		f := newMockFunction(newMockBlock(0).entry())
+
+		vr := VReg(100).SetRegType(RegTypeInt)
+		instr := newMockInstr().def(vr)
+		a.handleSpills(f, pc, instr, liveNodes, nil, vr)
+		require.Equal(t, 1, len(instr.defs))
+		require.Equal(t, RealReg(0xa), instr.defs[0].RealReg())
+
+		require.Equal(t, 1, len(f.befores))
+		requireInsertedInst(t, f, true, 0, instr, false, liveNodes[0].n.v)
+
+		require.Equal(t, 2, len(f.afters))
+		requireInsertedInst(t, f, false, 0, instr, true, liveNodes[0].n.v)
+		requireInsertedInst(t, f, false, 1, instr, false, vr.SetRealReg(0xa))
+	})
+
+	t.Run("only def / evicted", func(t *testing.T) {
+		const pc = 5
+		liveNodes := []liveNodeInBlock{
+			{n: &node{r: RealReg(0xb), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0xc), v: 0xc, ranges: []liveRange{{begin: pc, end: 20}}}},
+		}
+
+		a := NewAllocator(&RegisterInfo{
+			AllocatableRegisters: [3][]RealReg{RegTypeInt: {0xb, 0xc}}, // Only live nodes are allocatable.
+		})
+
+		f := newMockFunction(newMockBlock(0).entry())
+
+		vr := VReg(100).SetRegType(RegTypeInt)
+		instr := newMockInstr().def(vr)
+		a.handleSpills(f, pc, instr, liveNodes, nil, vr)
+		require.Equal(t, 1, len(instr.defs))
+		require.Equal(t, RealReg(0xb), instr.defs[0].RealReg())
+
+		require.Equal(t, 1, len(f.befores))
+		requireInsertedInst(t, f, true, 0, instr, false, liveNodes[0].n.v.SetRealReg(0xb))
+
+		require.Equal(t, 2, len(f.afters))
+		requireInsertedInst(t, f, false, 0, instr, true, liveNodes[0].n.v.SetRealReg(0xb))
+		requireInsertedInst(t, f, false, 1, instr, false, vr.SetRealReg(0xb))
+	})
+
+	t.Run("only def / not evicted", func(t *testing.T) {
+		const pc = 5
+		liveNodes := []liveNodeInBlock{
+			{n: &node{r: RealReg(0xb), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0xc), v: 0xc, ranges: []liveRange{{begin: pc, end: 20}}}},
+		}
+
+		a := NewAllocator(&RegisterInfo{
+			AllocatableRegisters: [3][]RealReg{RegTypeInt: {0xb, 0xc, 0xff /* free */}},
+		})
+
+		f := newMockFunction(newMockBlock(0).entry())
+
+		vr := VReg(100).SetRegType(RegTypeInt)
+		instr := newMockInstr().def(vr)
+		a.handleSpills(f, pc, instr, liveNodes, nil, vr)
+		require.Equal(t, 1, len(instr.defs))
+		require.Equal(t, RealReg(0xff), instr.defs[0].RealReg())
+
+		require.Equal(t, 0, len(f.befores))
+		require.Equal(t, 1, len(f.afters))
+		requireInsertedInst(t, f, false, 0, instr, false, vr.SetRealReg(0xff))
+	})
+
+	t.Run("uses and def / not evicted / def same type", func(t *testing.T) {
+		const pc = 5
+		liveNodes := []liveNodeInBlock{
+			{n: &node{r: RealReg(0xb), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0xc), v: 0xc, ranges: []liveRange{{begin: pc, end: 20}}}},
+		}
+
+		a := NewAllocator(&RegisterInfo{
+			AllocatableRegisters: [3][]RealReg{
+				RegTypeInt:   {0xb, 0xc, 0xaa /* free */},
+				RegTypeFloat: {0xbb /* free */},
+			},
+		})
+
+		f := newMockFunction(newMockBlock(0).entry())
+		u1, u2, u3 := VReg(100).SetRegType(RegTypeInt),
+			VReg(101).SetRegType(RegTypeInt).SetRealReg(0x88), // This one isn't spilled.
+			VReg(102).SetRegType(RegTypeFloat)
+		d1 := VReg(104).SetRegType(RegTypeFloat)
+		instr := newMockInstr().use(u1, u2, u3).def(d1)
+		a.handleSpills(f, pc, instr, liveNodes, []VReg{u1, u3}, d1)
+		require.Equal(t, []VReg{u1.SetRealReg(0xaa), u2, u3.SetRealReg(0xbb)}, instr.uses)
+		require.Equal(t, []VReg{d1.SetRealReg(0xbb)}, instr.defs)
+
+		require.Equal(t, 2, len(f.befores))
+		requireInsertedInst(t, f, true, 0, instr, true, u1.SetRealReg(0xaa))
+		requireInsertedInst(t, f, true, 1, instr, true, u3.SetRealReg(0xbb))
+
+		require.Equal(t, 1, len(f.afters))
+		requireInsertedInst(t, f, false, 0, instr, false, d1.SetRealReg(0xbb))
+	})
+
+	t.Run("uses and def / not evicted / def different type", func(t *testing.T) {
+		const pc = 5
+		liveNodes := []liveNodeInBlock{
+			{n: &node{r: RealReg(0xb), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0xff), v: 0xb, ranges: []liveRange{{begin: pc, end: 20}}}},
+		}
+
+		a := NewAllocator(&RegisterInfo{
+			AllocatableRegisters: [3][]RealReg{
+				RegTypeInt:   {0xb, 0xaa /* free */},
+				RegTypeFloat: {0xff},
+			},
+		})
+
+		f := newMockFunction(newMockBlock(0).entry())
+		u1 := VReg(100).SetRegType(RegTypeInt)
+		d1 := VReg(104).SetRegType(RegTypeFloat)
+		instr := newMockInstr().use(u1).def(d1)
+		a.handleSpills(f, pc, instr, liveNodes, []VReg{u1}, d1)
+		require.Equal(t, []VReg{u1.SetRealReg(0xaa)}, instr.uses)
+		require.Equal(t, []VReg{d1.SetRealReg(0xff)}, instr.defs)
+
+		require.Equal(t, 2, len(f.befores))
+		requireInsertedInst(t, f, true, 0, instr, true, u1.SetRealReg(0xaa))
+		requireInsertedInst(t, f, true, 1, instr, false, liveNodes[1].n.v.SetRealReg(0xff))
+		require.Equal(t, 2, len(f.afters))
+		requireInsertedInst(t, f, false, 0, instr, true, liveNodes[1].n.v.SetRealReg(0xff))
+		requireInsertedInst(t, f, false, 1, instr, false, d1.SetRealReg(0xff))
+	})
+
+	t.Run("uses and def / evicted / def different type", func(t *testing.T) {
+		const pc = 5
+		liveNodes := []liveNodeInBlock{
+			{n: &node{r: RealReg(0xb), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0xc), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0xff), v: 0xb, ranges: []liveRange{{begin: pc, end: 20}}}},
+		}
+
+		a := NewAllocator(&RegisterInfo{
+			AllocatableRegisters: [3][]RealReg{
+				RegTypeInt:   {0xb, 0xc},
+				RegTypeFloat: {0xff},
+			},
+		})
+
+		f := newMockFunction(newMockBlock(0).entry())
+		u1 := VReg(100).SetRegType(RegTypeInt)
+		d1 := VReg(104).SetRegType(RegTypeFloat)
+		instr := newMockInstr().use(u1).def(d1)
+		a.handleSpills(f, pc, instr, liveNodes, []VReg{u1}, d1)
+		require.Equal(t, []VReg{u1.SetRealReg(0xb)}, instr.uses)
+		require.Equal(t, []VReg{d1.SetRealReg(0xff)}, instr.defs)
+
+		require.Equal(t, 3, len(f.befores))
+		requireInsertedInst(t, f, true, 0, instr, false, liveNodes[0].n.v.SetRealReg(0xb))
+		requireInsertedInst(t, f, true, 1, instr, true, u1.SetRealReg(0xb))
+		requireInsertedInst(t, f, true, 2, instr, false, liveNodes[2].n.v.SetRealReg(0xff))
+		require.Equal(t, 3, len(f.afters))
+		requireInsertedInst(t, f, false, 0, instr, true, liveNodes[0].n.v.SetRealReg(0xb))
+		requireInsertedInst(t, f, false, 1, instr, true, liveNodes[2].n.v.SetRealReg(0xff))
+		requireInsertedInst(t, f, false, 2, instr, false, d1.SetRealReg(0xff))
+	})
+}
+
+func TestAllocator_collectActiveNodesAt(t *testing.T) {
+	t.Run("no live nodes", func(t *testing.T) {
+		a := NewAllocator(&RegisterInfo{})
+		a.nodes1 = []*node{{r: 1}, {r: 2}} // Must be cleared.
+		a.collectActiveNodesAt(0, nil)
+		require.Equal(t, 0, len(a.nodes1))
+	})
+
+	t.Run("lives", func(t *testing.T) {
+		const pc = 5
+		liveNodes := []liveNodeInBlock{
+			{n: &node{r: RealReg(0xff), v: 0xa, ranges: []liveRange{{begin: 0, end: pc - 1}}}},
+			{n: &node{r: RealReg(0x1), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0x2), v: 0xa, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0x4), v: 0xb, ranges: []liveRange{{begin: pc, end: 20}}}},
+			{n: &node{r: RealReg(0xff), v: 0xa, ranges: []liveRange{{begin: 1000, end: 2000000}}}},
+		}
+
+		a := NewAllocator(&RegisterInfo{})
+		a.nodes1 = []*node{{r: 1}, {r: 2}} // Must be cleared.
+		a.collectActiveNodesAt(pc, liveNodes)
+		require.Equal(t, 3, len(a.nodes1))
+		require.Equal(t, liveNodes[1].n, a.nodes1[0])
+		require.Equal(t, liveNodes[2].n, a.nodes1[1])
+		require.Equal(t, liveNodes[3].n, a.nodes1[2])
+	})
+}
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc_test.go b/internal/engine/wazevo/backend/regalloc/regalloc_test.go
index 38b9b9e5..77c4a8e4 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc_test.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc_test.go
@@ -706,3 +706,9 @@ func initMapInInfo(info *blockInfo) {
 		info.realRegDefs = make(map[VReg][]programCounter)
 	}
 }
+
+func TestNode_assignedRealReg(t *testing.T) {
+	require.Equal(t, RealRegInvalid, (&node{}).assignedRealReg())
+	require.Equal(t, RealReg(100), (&node{r: 100}).assignedRealReg())
+	require.Equal(t, RealReg(200), (&node{v: VReg(1).SetRealReg(200)}).assignedRealReg())
+}
diff --git a/internal/engine/wazevo/backend/regalloc/spill_handler_test.go b/internal/engine/wazevo/backend/regalloc/spill_handler_test.go
new file mode 100644
index 00000000..c82c1141
--- /dev/null
+++ b/internal/engine/wazevo/backend/regalloc/spill_handler_test.go
@@ -0,0 +1,62 @@
+package regalloc
+
+import (
+	"testing"
+
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+func TestSpillHandler_init(t *testing.T) {
+	var s spillHandler
+	activeNodes := []*node{{r: RealReg(0)}, {r: RealReg(1)}, {r: RealReg(2)}, {v: VReg(10).SetRealReg(100)}}
+	s.init(activeNodes)
+	require.Equal(t, 4, len(s.activeRegs))
+	for _, n := range activeNodes {
+		require.Equal(t, spillHandlerRegState{node: n, state: spillHandlerRegStateEvictable}, s.activeRegs[n.assignedRealReg()])
+	}
+}
+
+func TestSpillHandler_getUnusedOrEvictReg(t *testing.T) {
+	var s spillHandler
+	activeNodes := []*node{{r: RealReg(0)}, {r: RealReg(1)}, {r: RealReg(2)}}
+	s.init(activeNodes)
+	require.Equal(t, 3, len(s.activeRegs))
+
+	regInfo := RegisterInfo{
+		AllocatableRegisters: [RegTypeNum][]RealReg{
+			RegTypeInt: {
+				RealReg(0xff), // unused.
+				RealReg(0), RealReg(1), RealReg(2),
+			},
+			RegTypeFloat: {RealReg(0xaa)},
+		},
+	}
+
+	// Get unused register.
+	r, evicted := s.getUnusedOrEvictReg(RegTypeInt, &regInfo)
+	require.Equal(t, RealReg(0xff), r)
+	require.Nil(t, evicted)
+	require.Equal(t, 4, len(s.activeRegs))
+	require.Equal(t, spillHandlerRegState{state: spillHandlerRegStateUsed}, s.activeRegs[RealReg(0xff)])
+
+	// Get unused register.
+	r, evicted = s.getUnusedOrEvictReg(RegTypeFloat, &regInfo)
+	require.Equal(t, RealReg(0xaa), r)
+	require.Nil(t, evicted)
+	require.Equal(t, 5, len(s.activeRegs))
+	require.Equal(t, spillHandlerRegState{state: spillHandlerRegStateUsed}, s.activeRegs[RealReg(0xaa)])
+
+	// Evict register.
+	r, evicted = s.getUnusedOrEvictReg(RegTypeInt, &regInfo)
+	require.Equal(t, RealReg(0), r)
+	require.Equal(t, activeNodes[0], evicted)
+	require.Equal(t, 5, len(s.activeRegs))
+	require.Equal(t, spillHandlerRegState{node: activeNodes[0], state: spillHandlerRegStateEvicted}, s.activeRegs[RealReg(0)])
+
+	// Evict again.
+	r, evicted = s.getUnusedOrEvictReg(RegTypeInt, &regInfo)
+	require.Equal(t, RealReg(1), r)
+	require.Equal(t, activeNodes[1], evicted)
+	require.Equal(t, 5, len(s.activeRegs))
+	require.Equal(t, spillHandlerRegState{node: activeNodes[1], state: spillHandlerRegStateEvicted}, s.activeRegs[RealReg(1)])
+}
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index d709061d..39336765 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -99,7 +99,7 @@ func TestSpectestV1(t *testing.T) {
 		{name: "stack"},
 		{name: "store"},
 		{name: "switch"},
-		//{name: "skip-stack-guard-page"},
+		{name: "skip-stack-guard-page"},
 		{name: "token"},
 		{name: "type"},
 		{name: "traps"},