Repository: tetratelabs/wazero
Pull request #1902 (2024-01-05 22:30:13)

wazevo: shares regalloc impls for amd64


Patch:
diff --git a/internal/engine/wazevo/backend/compiler.go b/internal/engine/wazevo/backend/compiler.go
index 2c0e3cb8..87cc0c31 100644
--- a/internal/engine/wazevo/backend/compiler.go
+++ b/internal/engine/wazevo/backend/compiler.go
@@ -19,7 +19,6 @@ func newCompiler(_ context.Context, mach Machine, builder ssa.Builder) *compiler
 	c := &compiler{
 		mach: mach, ssaBuilder: builder,
 		nextVRegID: regalloc.VRegIDNonReservedBegin,
-		regAlloc:   regalloc.NewAllocator(mach.RegisterInfo()),
 	}
 	mach.SetCompiler(c)
 	return c
@@ -122,7 +121,6 @@ type compiler struct {
 	ssaValueRefCounts []int
 	// returnVRegs is the list of virtual registers that store the return values.
 	returnVRegs  []regalloc.VReg
-	regAlloc     regalloc.Allocator
 	varEdges     [][2]regalloc.VReg
 	varEdgeTypes []ssa.Type
 	constEdges   []struct {
@@ -179,8 +177,7 @@ func (c *compiler) Compile(ctx context.Context) ([]byte, []RelocationInfo, error
 
 // RegAlloc implements Compiler.RegAlloc.
 func (c *compiler) RegAlloc() {
-	regAllocFn := c.mach.Function()
-	c.regAlloc.DoAllocation(regAllocFn)
+	c.mach.RegAlloc()
 }
 
 // Finalize implements Compiler.Finalize.
@@ -290,7 +287,6 @@ func (c *compiler) Init() {
 	c.mach.Reset()
 	c.varEdges = c.varEdges[:0]
 	c.constEdges = c.constEdges[:0]
-	c.regAlloc.Reset()
 	c.buf = c.buf[:0]
 	c.sourceOffsets = c.sourceOffsets[:0]
 	c.relocations = c.relocations[:0]
diff --git a/internal/engine/wazevo/backend/compiler_lower.go b/internal/engine/wazevo/backend/compiler_lower.go
index ddc7e679..676c1a72 100644
--- a/internal/engine/wazevo/backend/compiler_lower.go
+++ b/internal/engine/wazevo/backend/compiler_lower.go
@@ -33,7 +33,8 @@ func (c *compiler) lowerBlocks() {
 
 func (c *compiler) lowerBlock(blk ssa.BasicBlock) {
 	mach := c.mach
-	mach.StartBlock(blk)
+	ectx := mach.ExecutableContext()
+	ectx.StartBlock(blk)
 
 	// We traverse the instructions in reverse order because we might want to lower multiple
 	// instructions together.
@@ -59,7 +60,6 @@ func (c *compiler) lowerBlock(blk ssa.BasicBlock) {
 	}
 
 	// Now start lowering the non-branching instructions.
-	ectx := mach.ExecutableContext()
 	for ; cur != nil; cur = cur.Prev() {
 		c.setCurrentGroupID(cur.GroupID())
 		if cur.Lowered() {
diff --git a/internal/engine/wazevo/backend/executable_context.go b/internal/engine/wazevo/backend/executable_context.go
index a89488ea..a1a9658c 100644
--- a/internal/engine/wazevo/backend/executable_context.go
+++ b/internal/engine/wazevo/backend/executable_context.go
@@ -16,6 +16,11 @@ type ExecutableContext interface {
 	// LinkAdjacentBlocks is called after finished lowering all blocks in order to create one single instruction list.
 	LinkAdjacentBlocks(prev, next ssa.BasicBlock)
 
+	// StartBlock is called when the compilation of the given block is started.
+	// The order of this being called is the reverse post order of the ssa.BasicBlock(s) as we iterate with
+	// ssa.Builder BlockIteratorReversePostOrderBegin and BlockIteratorReversePostOrderEnd.
+	StartBlock(ssa.BasicBlock)
+
 	// EndBlock is called when the compilation of the current block is finished.
 	EndBlock()
 
@@ -80,10 +85,10 @@ func (e *ExecutableContextT[Instr]) StartLoweringFunction(max ssa.BasicBlockID)
 	}
 }
 
-func (e *ExecutableContextT[Instr]) StartBlock(blk ssa.BasicBlock) (l Label, pos *LabelPosition[Instr]) {
+func (e *ExecutableContextT[Instr]) StartBlock(blk ssa.BasicBlock) {
 	e.CurrentSSABlk = blk
 
-	l = e.SsaBlockIDToLabels[e.CurrentSSABlk.ID()]
+	l := e.SsaBlockIDToLabels[e.CurrentSSABlk.ID()]
 	if l == LabelInvalid {
 		l = e.AllocateLabel()
 		e.SsaBlockIDToLabels[blk.ID()] = l
@@ -99,7 +104,7 @@ func (e *ExecutableContextT[Instr]) StartBlock(blk ssa.BasicBlock) (l Label, pos
 	}
 	e.OrderedBlockLabels = append(e.OrderedBlockLabels, labelPos)
 	labelPos.Begin, labelPos.End = end, end
-	return l, labelPos
+	labelPos.SB = blk
 }
 
 // EndBlock implements ExecutableContext.
@@ -191,6 +196,7 @@ func (e *ExecutableContextT[T]) LinkAdjacentBlocks(prev, next ssa.BasicBlock) {
 
 // LabelPosition represents the regions of the generated code which the label represents.
 type LabelPosition[Instr any] struct {
+	SB           ssa.BasicBlock
 	L            Label
 	Begin, End   *Instr
 	BinarySize   int64
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 321907fb..26d20878 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -18,7 +18,8 @@ func NewBackend() backend.Machine {
 		asNop,
 	)
 	return &machine{
-		ectx: ectx,
+		ectx:     ectx,
+		regAlloc: regalloc.NewAllocator(regInfo),
 	}
 }
 
@@ -28,6 +29,8 @@ type machine struct {
 	ectx                     *backend.ExecutableContextT[instruction]
 	stackBoundsCheckDisabled bool
 
+	regAlloc regalloc.Allocator
+
 	// abis maps ssa.SignatureID to the ABI implementation.
 	abis       []abiImpl
 	currentABI *abiImpl
@@ -48,20 +51,11 @@ func (m *machine) DisableStackCheck() { m.stackBoundsCheckDisabled = true }
 // SetCompiler implements backend.Machine.
 func (m *machine) SetCompiler(compiler backend.Compiler) { m.c = compiler }
 
-// RegisterInfo implements backend.Machine.
-func (m *machine) RegisterInfo() *regalloc.RegisterInfo { return regInfo }
-
 // InitializeABI implements backend.Machine.
 func (m *machine) InitializeABI(sig *ssa.Signature) {
 	m.currentABI = m.getOrCreateFunctionABI(sig)
 }
 
-// StartBlock implements backend.Machine.
-func (m *machine) StartBlock(block ssa.BasicBlock) {
-	// TODO implement me
-	panic("implement me")
-}
-
 // LowerSingleBranch implements backend.Machine.
 func (m *machine) LowerSingleBranch(b *ssa.Instruction) {
 	// TODO implement me
@@ -80,12 +74,6 @@ func (m *machine) LowerInstr(instruction *ssa.Instruction) {
 	panic("implement me")
 }
 
-// FlushPendingInstructions implements backend.Machine.
-func (m *machine) FlushPendingInstructions() {
-	// TODO implement me
-	panic("implement me")
-}
-
 // InsertMove implements backend.Machine.
 func (m *machine) InsertMove(dst, src regalloc.VReg, typ ssa.Type) {
 	// TODO implement me
@@ -110,12 +98,6 @@ func (m *machine) Format() string {
 	panic("implement me")
 }
 
-// Function implements backend.Machine.
-func (m *machine) Function() regalloc.Function {
-	// TODO implement me
-	panic("implement me")
-}
-
 // SetupPrologue implements backend.Machine.
 func (m *machine) SetupPrologue() {
 	// TODO implement me
@@ -163,3 +145,6 @@ func (m *machine) CompileEntryPreamble(signature *ssa.Signature) []byte {
 	// TODO implement me
 	panic("implement me")
 }
+
+// RegAlloc implements backend.Machine.
+func (m *machine) RegAlloc() { panic("implement me") }
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr.go b/internal/engine/wazevo/backend/isa/arm64/instr.go
index 80195569..bcd1566f 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr.go
@@ -63,6 +63,21 @@ func (i *instruction) IsReturn() bool {
 	return i.kind == ret
 }
 
+// Next implements regalloc.Instr Next.
+func (i *instruction) Next() regalloc.Instr {
+	return i.next
+}
+
+// Prev implements regalloc.Instr Prev.
+func (i *instruction) Prev() regalloc.Instr {
+	return i.prev
+}
+
+// AddedBeforeRegAlloc implements regalloc.Instr AddedBeforeRegAlloc.
+func (i *instruction) AddedBeforeRegAlloc() bool {
+	return i.addedBeforeRegAlloc
+}
+
 type defKind byte
 
 const (
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine.go b/internal/engine/wazevo/backend/isa/arm64/machine.go
index 9581f3ec..bdf88ae2 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine.go
@@ -20,7 +20,8 @@ type (
 		// abis maps ssa.SignatureID to the ABI implementation.
 		abis []functionABI
 
-		regAllocFn regAllocFunctionImpl
+		regAlloc   regalloc.Allocator
+		regAllocFn *backend.RegAllocFunction[*instruction, *machine]
 
 		// addendsWorkQueue is used during address lowering, defined here for reuse.
 		addendsWorkQueue queue[ssa.Value]
@@ -102,9 +103,8 @@ func NewBackend() backend.Machine {
 	m := &machine{
 		spillSlots:        make(map[regalloc.VRegID]int64),
 		executableContext: newExecutableContext(),
+		regAlloc:          regalloc.NewAllocator(regInfo),
 	}
-	m.regAllocFn.m = m
-	m.regAllocFn.labelToRegAllocBlockIndex = make(map[label]int)
 	return m
 }
 
@@ -117,9 +117,21 @@ func (m *machine) ExecutableContext() backend.ExecutableContext {
 	return m.executableContext
 }
 
+// RegAlloc implements backend.Machine Function.
+func (m *machine) RegAlloc() {
+	rf := m.regAllocFn
+	for _, pos := range m.executableContext.OrderedBlockLabels {
+		rf.AddBlock(pos.SB, pos.L, pos.Begin, pos.End)
+	}
+
+	m.regAllocStarted = true
+	m.regAlloc.DoAllocation(rf)
+	// Now that we know the final spill slot size, we must align spillSlotSize to 16 bytes.
+	m.spillSlotSize = (m.spillSlotSize + 15) &^ 15
+}
+
 // Reset implements backend.Machine.
 func (m *machine) Reset() {
-	m.regAllocStarted = false
 	m.clobberedRegs = m.clobberedRegs[:0]
 	for key := range m.spillSlots {
 		m.clobberedRegs = append(m.clobberedRegs, regalloc.VReg(key))
@@ -128,7 +140,9 @@ func (m *machine) Reset() {
 		delete(m.spillSlots, regalloc.VRegID(key))
 	}
 	m.clobberedRegs = m.clobberedRegs[:0]
-	m.regAllocFn.reset()
+	m.regAllocStarted = false
+	m.regAlloc.Reset()
+	m.regAllocFn.Reset()
 	m.spillSlotSize = 0
 	m.unresolvedAddressModes = m.unresolvedAddressModes[:0]
 	m.maxRequiredStackSizeForCalls = 0
@@ -148,12 +162,7 @@ func (m *machine) DisableStackCheck() {
 // SetCompiler implements backend.Machine.
 func (m *machine) SetCompiler(ctx backend.Compiler) {
 	m.compiler = ctx
-}
-
-// StartBlock implements backend.Machine.
-func (m *machine) StartBlock(blk ssa.BasicBlock) {
-	l, pos := m.executableContext.StartBlock(blk)
-	m.regAllocFn.addBlock(blk, l, pos)
+	m.regAllocFn = backend.NewRegAllocFunction[*instruction, *machine](m, ctx.SSABuilder(), ctx)
 }
 
 func (m *machine) insert(i *instruction) {
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
index 5bff98d1..1c8793b7 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
@@ -1,113 +1,20 @@
 package arm64
 
-// This file implements the interfaces required for register allocations. See regalloc/api.go.
+// This file implements the interfaces required for register allocations. See backend.RegAllocFunctionMachine.
 
 import (
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 )
 
-type (
-	// regAllocFunctionImpl implements regalloc.Function.
-	regAllocFunctionImpl struct {
-		m *machine
-		// iter is the iterator for reversePostOrderBlocks
-		iter                   int
-		reversePostOrderBlocks []regAllocBlockImpl
-		// labelToRegAllocBlockIndex maps label to the index of reversePostOrderBlocks.
-		labelToRegAllocBlockIndex map[label]int
-		loopNestingForestRoots    []ssa.BasicBlock
-	}
-
-	// regAllocBlockImpl implements regalloc.Block.
-	regAllocBlockImpl struct {
-		// f is the function this instruction belongs to. Used to reuse the regAllocFunctionImpl.predsSlice slice for Defs() and Uses().
-		f                         *regAllocFunctionImpl
-		sb                        ssa.BasicBlock
-		l                         label
-		pos                       *labelPosition
-		loopNestingForestChildren []ssa.BasicBlock
-		cur                       *instruction
-		id                        int
-		cachedLastInstr           regalloc.Instr
-	}
-)
-
-func (f *regAllocFunctionImpl) addBlock(sb ssa.BasicBlock, l label, pos *labelPosition) {
-	i := len(f.reversePostOrderBlocks)
-	f.reversePostOrderBlocks = append(f.reversePostOrderBlocks, regAllocBlockImpl{
-		f:   f,
-		sb:  sb,
-		l:   l,
-		pos: pos,
-		id:  int(sb.ID()),
-	})
-	f.labelToRegAllocBlockIndex[l] = i
-}
-
-func (f *regAllocFunctionImpl) reset() {
-	f.reversePostOrderBlocks = f.reversePostOrderBlocks[:0]
-	f.iter = 0
-}
-
-var (
-	_ regalloc.Function = (*regAllocFunctionImpl)(nil)
-	_ regalloc.Block    = (*regAllocBlockImpl)(nil)
-)
-
-// PostOrderBlockIteratorBegin implements regalloc.Function PostOrderBlockIteratorBegin.
-func (f *regAllocFunctionImpl) PostOrderBlockIteratorBegin() regalloc.Block {
-	f.iter = len(f.reversePostOrderBlocks) - 1
-	return f.PostOrderBlockIteratorNext()
-}
-
-// PostOrderBlockIteratorNext implements regalloc.Function PostOrderBlockIteratorNext.
-func (f *regAllocFunctionImpl) PostOrderBlockIteratorNext() regalloc.Block {
-	if f.iter < 0 {
-		return nil
-	}
-	b := &f.reversePostOrderBlocks[f.iter]
-	f.iter--
-	return b
-}
-
-// ReversePostOrderBlockIteratorBegin implements regalloc.Function ReversePostOrderBlockIteratorBegin.
-func (f *regAllocFunctionImpl) ReversePostOrderBlockIteratorBegin() regalloc.Block {
-	f.iter = 0
-	return f.ReversePostOrderBlockIteratorNext()
-}
-
-// ReversePostOrderBlockIteratorNext implements regalloc.Function ReversePostOrderBlockIteratorNext.
-func (f *regAllocFunctionImpl) ReversePostOrderBlockIteratorNext() regalloc.Block {
-	if f.iter >= len(f.reversePostOrderBlocks) {
-		return nil
-	}
-	b := &f.reversePostOrderBlocks[f.iter]
-	f.iter++
-	return b
-}
-
-// ClobberedRegisters implements regalloc.Function ClobberedRegisters.
-func (f *regAllocFunctionImpl) ClobberedRegisters(regs []regalloc.VReg) {
-	m := f.m
+// ClobberedRegisters implements backend.RegAllocFunctionMachine.
+func (m *machine) ClobberedRegisters(regs []regalloc.VReg) {
 	m.clobberedRegs = append(m.clobberedRegs[:0], regs...)
 }
 
-// StoreRegisterBefore implements regalloc.Function StoreRegisterBefore.
-func (f *regAllocFunctionImpl) StoreRegisterBefore(v regalloc.VReg, instr regalloc.Instr) {
-	m := f.m
-	m.insertStoreRegisterAt(v, instr.(*instruction), false)
-}
-
-// SwapAtEndOfBlock implements regalloc.Function SwapAtEndOfBlock.
-func (f *regAllocFunctionImpl) SwapAtEndOfBlock(x1, x2, tmp regalloc.VReg, block regalloc.Block) {
-	blk := block.(*regAllocBlockImpl)
-	cur := blk.LastInstr().(*instruction)
-	cur = cur.prev
-	f.m.swap(cur, x1, x2, tmp)
-}
-
-func (m *machine) swap(cur *instruction, x1, x2, tmp regalloc.VReg) {
+// Swap implements backend.RegAllocFunctionMachine.
+func (m *machine) Swap(cur *instruction, x1, x2, tmp regalloc.VReg) {
 	prevNext := cur.next
 	var mov1, mov2, mov3 *instruction
 	if x1.RegType() == regalloc.RegTypeInt {
@@ -125,12 +32,12 @@ func (m *machine) swap(cur *instruction, x1, x2, tmp regalloc.VReg) {
 		if !tmp.Valid() {
 			r2 := x2.RealReg()
 			// Temporarily spill x1 to stack.
-			cur = m.insertStoreRegisterAt(x1, cur, true).prev
+			cur = m.InsertStoreRegisterAt(x1, cur, true).prev
 			// Then move x2 to x1.
 			cur = linkInstr(cur, m.allocateInstr().asFpuMov128(x1, x2))
 			linkInstr(cur, prevNext)
 			// Then reload the original value on x1 from stack to r2.
-			m.insertReloadRegisterAt(x1.SetRealReg(r2), cur, true)
+			m.InsertReloadRegisterAt(x1.SetRealReg(r2), cur, true)
 		} else {
 			mov1 = m.allocateInstr().asFpuMov128(tmp, x1)
 			mov2 = m.allocateInstr().asFpuMov128(x1, x2)
@@ -143,10 +50,8 @@ func (m *machine) swap(cur *instruction, x1, x2, tmp regalloc.VReg) {
 	}
 }
 
-// InsertMoveBefore implements regalloc.Function InsertMoveBefore.
-func (f *regAllocFunctionImpl) InsertMoveBefore(dst, src regalloc.VReg, instr regalloc.Instr) {
-	m := f.m
-
+// InsertMoveBefore implements backend.RegAllocFunctionMachine.
+func (m *machine) InsertMoveBefore(dst, src regalloc.VReg, instr *instruction) {
 	typ := src.RegType()
 	if typ != dst.RegType() {
 		panic("BUG: src and dst must have the same type")
@@ -159,178 +64,19 @@ func (f *regAllocFunctionImpl) InsertMoveBefore(dst, src regalloc.VReg, instr re
 		mov.asFpuMov128(dst, src)
 	}
 
-	cur := instr.(*instruction).prev
+	cur := instr.prev
 	prevNext := cur.next
 	cur = linkInstr(cur, mov)
 	linkInstr(cur, prevNext)
 }
 
-// StoreRegisterAfter implements regalloc.Function StoreRegisterAfter.
-func (f *regAllocFunctionImpl) StoreRegisterAfter(v regalloc.VReg, instr regalloc.Instr) {
-	m := f.m
-	m.insertStoreRegisterAt(v, instr.(*instruction), true)
-}
-
-// ReloadRegisterBefore implements regalloc.Function ReloadRegisterBefore.
-func (f *regAllocFunctionImpl) ReloadRegisterBefore(v regalloc.VReg, instr regalloc.Instr) {
-	m := f.m
-	m.insertReloadRegisterAt(v, instr.(*instruction), false)
-}
-
-// ReloadRegisterAfter implements regalloc.Function ReloadRegisterAfter.
-func (f *regAllocFunctionImpl) ReloadRegisterAfter(v regalloc.VReg, instr regalloc.Instr) {
-	m := f.m
-	m.insertReloadRegisterAt(v, instr.(*instruction), true)
-}
-
-// Done implements regalloc.Function Done.
-func (f *regAllocFunctionImpl) Done() {
-	m := f.m
-	// Now that we know the final spill slot size, we must align spillSlotSize to 16 bytes.
-	m.spillSlotSize = (m.spillSlotSize + 15) &^ 15
-}
-
-// ID implements regalloc.Block ID.
-func (r *regAllocBlockImpl) ID() int {
-	return r.id
-}
-
-// Preds implements regalloc.Block Preds.
-func (r *regAllocBlockImpl) Preds() int {
-	return r.sb.Preds()
-}
-
-// Pred implements regalloc.Block Pred.
-func (r *regAllocBlockImpl) Pred(i int) regalloc.Block {
-	sb := r.sb
-	pred := sb.Pred(i)
-	l := r.f.m.executableContext.SsaBlockIDToLabels[pred.ID()]
-	index := r.f.labelToRegAllocBlockIndex[l]
-	return &r.f.reversePostOrderBlocks[index]
-}
-
-// Succs implements regalloc.Block Succs.
-func (r *regAllocBlockImpl) Succs() int {
-	return r.sb.Succs()
-}
-
-// Succ implements regalloc.Block Succ.
-func (r *regAllocBlockImpl) Succ(i int) regalloc.Block {
-	sb := r.sb
-	succ := sb.Succ(i)
-	if succ.ReturnBlock() {
-		return nil
-	}
-	l := r.f.m.executableContext.SsaBlockIDToLabels[succ.ID()]
-	index := r.f.labelToRegAllocBlockIndex[l]
-	return &r.f.reversePostOrderBlocks[index]
+// SSABlockLabel implements backend.RegAllocFunctionMachine.
+func (m *machine) SSABlockLabel(id ssa.BasicBlockID) backend.Label {
+	return m.executableContext.SsaBlockIDToLabels[id]
 }
 
-// LoopHeader implements regalloc.Block LoopHeader.
-func (r *regAllocBlockImpl) LoopHeader() bool {
-	return r.sb.LoopHeader()
-}
-
-// LoopNestingForestRoots implements regalloc.Function LoopNestingForestRoots.
-func (f *regAllocFunctionImpl) LoopNestingForestRoots() int {
-	f.loopNestingForestRoots = f.m.compiler.SSABuilder().LoopNestingForestRoots()
-	return len(f.loopNestingForestRoots)
-}
-
-// LoopNestingForestRoot implements regalloc.Function LoopNestingForestRoot.
-func (f *regAllocFunctionImpl) LoopNestingForestRoot(i int) regalloc.Block {
-	blk := f.loopNestingForestRoots[i]
-	l := f.m.executableContext.SsaBlockIDToLabels[blk.ID()]
-	index := f.labelToRegAllocBlockIndex[l]
-	return &f.reversePostOrderBlocks[index]
-}
-
-// LoopNestingForestChildren implements regalloc.Block LoopNestingForestChildren.
-func (r *regAllocBlockImpl) LoopNestingForestChildren() int {
-	r.loopNestingForestChildren = r.sb.LoopNestingForestChildren()
-	return len(r.loopNestingForestChildren)
-}
-
-// LoopNestingForestChild implements regalloc.Block LoopNestingForestChild.
-func (r *regAllocBlockImpl) LoopNestingForestChild(i int) regalloc.Block {
-	blk := r.loopNestingForestChildren[i]
-	l := r.f.m.executableContext.SsaBlockIDToLabels[blk.ID()]
-	index := r.f.labelToRegAllocBlockIndex[l]
-	return &r.f.reversePostOrderBlocks[index]
-}
-
-// InstrIteratorBegin implements regalloc.Block InstrIteratorBegin.
-func (r *regAllocBlockImpl) InstrIteratorBegin() regalloc.Instr {
-	r.cur = r.pos.Begin
-	return r.cur
-}
-
-// InstrIteratorNext implements regalloc.Block InstrIteratorNext.
-func (r *regAllocBlockImpl) InstrIteratorNext() regalloc.Instr {
-	for {
-		if r.cur == r.pos.End {
-			return nil
-		}
-		instr := r.cur.next
-		r.cur = instr
-		if instr == nil {
-			return nil
-		} else if instr.addedBeforeRegAlloc {
-			// Only concerned about the instruction added before regalloc.
-			return instr
-		}
-	}
-}
-
-// InstrRevIteratorBegin implements regalloc.Block InstrRevIteratorBegin.
-func (r *regAllocBlockImpl) InstrRevIteratorBegin() regalloc.Instr {
-	r.cur = r.pos.End
-	return r.cur
-}
-
-// InstrRevIteratorNext implements regalloc.Block InstrRevIteratorNext.
-func (r *regAllocBlockImpl) InstrRevIteratorNext() regalloc.Instr {
-	for {
-		if r.cur == r.pos.Begin {
-			return nil
-		}
-		instr := r.cur.prev
-		r.cur = instr
-		if instr == nil {
-			return nil
-		} else if instr.addedBeforeRegAlloc {
-			// Only concerned about the instruction added before regalloc.
-			return instr
-		}
-	}
-}
-
-// BlockParams implements regalloc.Block BlockParams.
-func (r *regAllocBlockImpl) BlockParams(regs *[]regalloc.VReg) []regalloc.VReg {
-	c := r.f.m.compiler
-	*regs = (*regs)[:0]
-	for i := 0; i < r.sb.Params(); i++ {
-		v := c.VRegOf(r.sb.Param(i))
-		*regs = append(*regs, v)
-	}
-	return *regs
-}
-
-// Entry implements regalloc.Block Entry.
-func (r *regAllocBlockImpl) Entry() bool { return r.sb.EntryBlock() }
-
-// RegisterInfo implements backend.Machine.
-func (m *machine) RegisterInfo() *regalloc.RegisterInfo {
-	return regInfo
-}
-
-// Function implements backend.Machine Function.
-func (m *machine) Function() regalloc.Function {
-	m.regAllocStarted = true
-	return &m.regAllocFn
-}
-
-func (m *machine) insertStoreRegisterAt(v regalloc.VReg, instr *instruction, after bool) *instruction {
+// InsertStoreRegisterAt implements backend.RegAllocFunctionMachine.
+func (m *machine) InsertStoreRegisterAt(v regalloc.VReg, instr *instruction, after bool) *instruction {
 	if !v.IsRealReg() {
 		panic("BUG: VReg must be backed by real reg to be stored")
 	}
@@ -354,7 +100,8 @@ func (m *machine) insertStoreRegisterAt(v regalloc.VReg, instr *instruction, aft
 	return linkInstr(cur, prevNext)
 }
 
-func (m *machine) insertReloadRegisterAt(v regalloc.VReg, instr *instruction, after bool) *instruction {
+// InsertReloadRegisterAt implements backend.RegAllocFunctionMachine.
+func (m *machine) InsertReloadRegisterAt(v regalloc.VReg, instr *instruction, after bool) *instruction {
 	if !v.IsRealReg() {
 		panic("BUG: VReg must be backed by real reg to be stored")
 	}
@@ -387,47 +134,19 @@ func (m *machine) insertReloadRegisterAt(v regalloc.VReg, instr *instruction, af
 	return linkInstr(cur, prevNext)
 }
 
-// LowestCommonAncestor implements regalloc.Function LowestCommonAncestor.
-func (f *regAllocFunctionImpl) LowestCommonAncestor(blk1, blk2 regalloc.Block) regalloc.Block {
-	ret := f.m.compiler.SSABuilder().LowestCommonAncestor(blk1.(*regAllocBlockImpl).sb, blk2.(*regAllocBlockImpl).sb)
-	l := f.m.executableContext.SsaBlockIDToLabels[ret.ID()]
-	index := f.labelToRegAllocBlockIndex[l]
-	return &f.reversePostOrderBlocks[index]
-}
-
-func (f *regAllocFunctionImpl) Idom(blk regalloc.Block) regalloc.Block {
-	builder := f.m.compiler.SSABuilder()
-	idom := builder.Idom(blk.(*regAllocBlockImpl).sb)
-	if idom == nil {
-		panic("BUG: idom must not be nil")
-	}
-	l := f.m.executableContext.SsaBlockIDToLabels[idom.ID()]
-	index := f.labelToRegAllocBlockIndex[l]
-	return &f.reversePostOrderBlocks[index]
-}
-
-// FirstInstr implements regalloc.Block FirstInstr.
-func (r *regAllocBlockImpl) FirstInstr() regalloc.Instr {
-	return r.pos.Begin
-}
-
-// LastInstr implements regalloc.Block LastInstr.
-func (r *regAllocBlockImpl) LastInstr() regalloc.Instr {
-	if r.cachedLastInstr == nil {
-		cur := r.pos.End
-		for cur.kind == nop0 {
-			cur = cur.prev
-			if cur == r.pos.Begin {
-				r.cachedLastInstr = r.pos.End
-				return r.cachedLastInstr
-			}
-		}
-		switch cur.kind {
-		case br:
-			r.cachedLastInstr = cur
-		default:
-			r.cachedLastInstr = r.pos.End
+// LastInstrForInsertion implements backend.RegAllocFunctionMachine.
+func (m *machine) LastInstrForInsertion(begin, end *instruction) *instruction {
+	cur := end
+	for cur.kind == nop0 {
+		cur = cur.prev
+		if cur == begin {
+			return end
 		}
 	}
-	return r.cachedLastInstr
+	switch cur.kind {
+	case br:
+		return cur
+	default:
+		return end
+	}
 }
diff --git a/internal/engine/wazevo/backend/machine.go b/internal/engine/wazevo/backend/machine.go
index 4cdea558..34dad417 100644
--- a/internal/engine/wazevo/backend/machine.go
+++ b/internal/engine/wazevo/backend/machine.go
@@ -16,10 +16,6 @@ type (
 		// DisableStackCheck disables the stack check for the current compilation for debugging/testing.
 		DisableStackCheck()
 
-		// RegisterInfo returns the set of registers that can be used for register allocation.
-		// This is only called once, and the result is shared across all compilations.
-		RegisterInfo() *regalloc.RegisterInfo
-
 		// InitializeABI initializes the FunctionABI for the given signature.
 		InitializeABI(sig *ssa.Signature)
 
@@ -27,11 +23,6 @@ type (
 		// This is only called once per Machine, i.e. before the first compilation.
 		SetCompiler(Compiler)
 
-		// StartBlock is called when the compilation of the given block is started.
-		// The order of this being called is the reverse post order of the ssa.BasicBlock(s) as we iterate with
-		// ssa.Builder BlockIteratorReversePostOrderBegin and BlockIteratorReversePostOrderEnd.
-		StartBlock(ssa.BasicBlock)
-
 		// LowerSingleBranch is called when the compilation of the given single branch is started.
 		LowerSingleBranch(b *ssa.Instruction)
 
@@ -61,8 +52,8 @@ type (
 		// This is only for testing purpose.
 		Format() string
 
-		// Function returns the currently compiled state as regalloc.Function so that we can perform register allocation.
-		Function() regalloc.Function
+		// RegAlloc does the register allocation after lowering.
+		RegAlloc()
 
 		// SetupPrologue inserts the prologue after register allocations.
 		SetupPrologue()
diff --git a/internal/engine/wazevo/backend/regalloc.go b/internal/engine/wazevo/backend/regalloc.go
new file mode 100644
index 00000000..223fa4a6
--- /dev/null
+++ b/internal/engine/wazevo/backend/regalloc.go
@@ -0,0 +1,320 @@
+package backend
+
+import (
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
+)
+
+// RegAllocFunctionMachine is the interface for the machine specific logic that will be used in RegAllocFunction.
+type RegAllocFunctionMachine[I regalloc.InstrConstraint] interface {
+	// InsertMoveBefore inserts the move instruction from src to dst before the given instruction.
+	InsertMoveBefore(dst, src regalloc.VReg, instr I)
+	// InsertStoreRegisterAt inserts the instruction(s) to store the given virtual register at the given instruction.
+	// If after is true, the instruction(s) will be inserted after the given instruction, otherwise before.
+	InsertStoreRegisterAt(v regalloc.VReg, instr I, after bool) I
+	// InsertReloadRegisterAt inserts the instruction(s) to reload the given virtual register at the given instruction.
+	// If after is true, the instruction(s) will be inserted after the given instruction, otherwise before.
+	InsertReloadRegisterAt(v regalloc.VReg, instr I, after bool) I
+	// ClobberedRegisters is called when the register allocation is done and the clobbered registers are known.
+	ClobberedRegisters(regs []regalloc.VReg)
+	// Swap swaps the two virtual registers after the given instruction.
+	Swap(cur I, x1, x2, tmp regalloc.VReg)
+	// LastInstrForInsertion implements LastInstrForInsertion of regalloc.Function. See its comment for details.
+	LastInstrForInsertion(begin, end I) I
+	// SSABlockLabel returns the label of the given ssa.BasicBlockID.
+	SSABlockLabel(id ssa.BasicBlockID) Label
+}
+
+type (
+	// RegAllocFunction implements regalloc.Function.
+	RegAllocFunction[I regalloc.InstrConstraint, m RegAllocFunctionMachine[I]] struct {
+		m   m
+		ssb ssa.Builder
+		c   Compiler
+		// iter is the iterator for reversePostOrderBlocks
+		iter                   int
+		reversePostOrderBlocks []RegAllocBlock[I, m]
+		// labelToRegAllocBlockIndex maps label to the index of reversePostOrderBlocks.
+		labelToRegAllocBlockIndex map[Label]int
+		loopNestingForestRoots    []ssa.BasicBlock
+	}
+
+	// RegAllocBlock implements regalloc.Block.
+	RegAllocBlock[I regalloc.InstrConstraint, m RegAllocFunctionMachine[I]] struct {
+		// f is the function this instruction belongs to. Used to reuse the regAllocFunctionImpl.predsSlice slice for Defs() and Uses().
+		f                           *RegAllocFunction[I, m]
+		sb                          ssa.BasicBlock
+		l                           Label
+		begin, end                  I
+		loopNestingForestChildren   []ssa.BasicBlock
+		cur                         I
+		id                          int
+		cachedLastInstrForInsertion I
+	}
+)
+
+// NewRegAllocFunction returns a new RegAllocFunction.
+func NewRegAllocFunction[I regalloc.InstrConstraint, M RegAllocFunctionMachine[I]](m M, ssb ssa.Builder, c Compiler) *RegAllocFunction[I, M] {
+	return &RegAllocFunction[I, M]{
+		m:                         m,
+		ssb:                       ssb,
+		c:                         c,
+		labelToRegAllocBlockIndex: make(map[Label]int),
+	}
+}
+
+// AddBlock adds a new block to the function.
+func (f *RegAllocFunction[I, M]) AddBlock(sb ssa.BasicBlock, l Label, begin, end I) {
+	i := len(f.reversePostOrderBlocks)
+	f.reversePostOrderBlocks = append(f.reversePostOrderBlocks, RegAllocBlock[I, M]{
+		f:     f,
+		sb:    sb,
+		l:     l,
+		begin: begin,
+		end:   end,
+		id:    int(sb.ID()),
+	})
+	f.labelToRegAllocBlockIndex[l] = i
+}
+
+// Reset resets the function for the next compilation.
+func (f *RegAllocFunction[I, M]) Reset() {
+	f.reversePostOrderBlocks = f.reversePostOrderBlocks[:0]
+	f.iter = 0
+}
+
+// StoreRegisterAfter implements regalloc.Function StoreRegisterAfter.
+func (f *RegAllocFunction[I, M]) StoreRegisterAfter(v regalloc.VReg, instr regalloc.Instr) {
+	m := f.m
+	m.InsertStoreRegisterAt(v, instr.(I), true)
+}
+
+// ReloadRegisterBefore implements regalloc.Function ReloadRegisterBefore.
+func (f *RegAllocFunction[I, M]) ReloadRegisterBefore(v regalloc.VReg, instr regalloc.Instr) {
+	m := f.m
+	m.InsertReloadRegisterAt(v, instr.(I), false)
+}
+
+// ReloadRegisterAfter implements regalloc.Function ReloadRegisterAfter.
+func (f *RegAllocFunction[I, M]) ReloadRegisterAfter(v regalloc.VReg, instr regalloc.Instr) {
+	m := f.m
+	m.InsertReloadRegisterAt(v, instr.(I), true)
+}
+
+// StoreRegisterBefore implements regalloc.Function StoreRegisterBefore.
+func (f *RegAllocFunction[I, M]) StoreRegisterBefore(v regalloc.VReg, instr regalloc.Instr) {
+	m := f.m
+	m.InsertStoreRegisterAt(v, instr.(I), false)
+}
+
+// ClobberedRegisters implements regalloc.Function ClobberedRegisters.
+func (f *RegAllocFunction[I, M]) ClobberedRegisters(regs []regalloc.VReg) {
+	f.m.ClobberedRegisters(regs)
+}
+
+// SwapAtEndOfBlock implements regalloc.Function SwapAtEndOfBlock.
+func (f *RegAllocFunction[I, M]) SwapAtEndOfBlock(x1, x2, tmp regalloc.VReg, blk regalloc.Block) {
+	cur := blk.LastInstrForInsertion()
+	f.m.Swap(cur.Prev().(I), x1, x2, tmp)
+}
+
+// PostOrderBlockIteratorBegin implements regalloc.Function PostOrderBlockIteratorBegin.
+func (f *RegAllocFunction[I, M]) PostOrderBlockIteratorBegin() regalloc.Block {
+	f.iter = len(f.reversePostOrderBlocks) - 1
+	return f.PostOrderBlockIteratorNext()
+}
+
+// PostOrderBlockIteratorNext implements regalloc.Function PostOrderBlockIteratorNext.
+func (f *RegAllocFunction[I, M]) PostOrderBlockIteratorNext() regalloc.Block {
+	if f.iter < 0 {
+		return nil
+	}
+	b := &f.reversePostOrderBlocks[f.iter]
+	f.iter--
+	return b
+}
+
+// ReversePostOrderBlockIteratorBegin implements regalloc.Function ReversePostOrderBlockIteratorBegin.
+func (f *RegAllocFunction[I, M]) ReversePostOrderBlockIteratorBegin() regalloc.Block {
+	f.iter = 0
+	return f.ReversePostOrderBlockIteratorNext()
+}
+
+// ReversePostOrderBlockIteratorNext implements regalloc.Function ReversePostOrderBlockIteratorNext.
+func (f *RegAllocFunction[I, M]) ReversePostOrderBlockIteratorNext() regalloc.Block {
+	if f.iter >= len(f.reversePostOrderBlocks) {
+		return nil
+	}
+	b := &f.reversePostOrderBlocks[f.iter]
+	f.iter++
+	return b
+}
+
+// LoopNestingForestRoots implements regalloc.Function LoopNestingForestRoots.
+func (f *RegAllocFunction[I, M]) LoopNestingForestRoots() int {
+	f.loopNestingForestRoots = f.ssb.LoopNestingForestRoots()
+	return len(f.loopNestingForestRoots)
+}
+
+// LoopNestingForestRoot implements regalloc.Function LoopNestingForestRoot.
+func (f *RegAllocFunction[I, M]) LoopNestingForestRoot(i int) regalloc.Block {
+	blk := f.loopNestingForestRoots[i]
+	l := f.m.SSABlockLabel(blk.ID())
+	index := f.labelToRegAllocBlockIndex[l]
+	return &f.reversePostOrderBlocks[index]
+}
+
+// InsertMoveBefore implements regalloc.Function InsertMoveBefore.
+func (f *RegAllocFunction[I, M]) InsertMoveBefore(dst, src regalloc.VReg, instr regalloc.Instr) {
+	f.m.InsertMoveBefore(dst, src, instr.(I))
+}
+
+// LowestCommonAncestor implements regalloc.Function LowestCommonAncestor.
+func (f *RegAllocFunction[I, M]) LowestCommonAncestor(blk1, blk2 regalloc.Block) regalloc.Block {
+	ret := f.ssb.LowestCommonAncestor(blk1.(*RegAllocBlock[I, M]).sb, blk2.(*RegAllocBlock[I, M]).sb)
+	l := f.m.SSABlockLabel(ret.ID())
+	index := f.labelToRegAllocBlockIndex[l]
+	return &f.reversePostOrderBlocks[index]
+}
+
+// Idom implements regalloc.Function Idom.
+func (f *RegAllocFunction[I, M]) Idom(blk regalloc.Block) regalloc.Block {
+	builder := f.ssb
+	idom := builder.Idom(blk.(*RegAllocBlock[I, M]).sb)
+	if idom == nil {
+		panic("BUG: idom must not be nil")
+	}
+	l := f.m.SSABlockLabel(idom.ID())
+	index := f.labelToRegAllocBlockIndex[l]
+	return &f.reversePostOrderBlocks[index]
+}
+
+// ID implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) ID() int { return r.id }
+
+// BlockParams implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) BlockParams(regs *[]regalloc.VReg) []regalloc.VReg {
+	c := r.f.c
+	*regs = (*regs)[:0]
+	for i := 0; i < r.sb.Params(); i++ {
+		v := c.VRegOf(r.sb.Param(i))
+		*regs = append(*regs, v)
+	}
+	return *regs
+}
+
+// InstrIteratorBegin implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) InstrIteratorBegin() regalloc.Instr {
+	r.cur = r.begin
+	return r.cur
+}
+
+// InstrIteratorNext implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) InstrIteratorNext() regalloc.Instr {
+	for {
+		if r.cur == r.end {
+			return nil
+		}
+		instr := r.cur.Next()
+		r.cur = instr.(I)
+		if instr == nil {
+			return nil
+		} else if instr.AddedBeforeRegAlloc() {
+			// Only concerned about the instruction added before regalloc.
+			return instr
+		}
+	}
+}
+
+// InstrRevIteratorBegin implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) InstrRevIteratorBegin() regalloc.Instr {
+	r.cur = r.end
+	return r.cur
+}
+
+// InstrRevIteratorNext implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) InstrRevIteratorNext() regalloc.Instr {
+	for {
+		if r.cur == r.begin {
+			return nil
+		}
+		instr := r.cur.Prev()
+		r.cur = instr.(I)
+		if instr == nil {
+			return nil
+		} else if instr.AddedBeforeRegAlloc() {
+			// Only concerned about the instruction added before regalloc.
+			return instr
+		}
+	}
+}
+
+// FirstInstr implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) FirstInstr() regalloc.Instr {
+	return r.begin
+}
+
+// EndInstr implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) EndInstr() regalloc.Instr {
+	return r.end
+}
+
+// LastInstrForInsertion implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) LastInstrForInsertion() regalloc.Instr {
+	var nil I
+	if r.cachedLastInstrForInsertion == nil {
+		r.cachedLastInstrForInsertion = r.f.m.LastInstrForInsertion(r.begin, r.end)
+	}
+	return r.cachedLastInstrForInsertion
+}
+
+// Preds implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) Preds() int { return r.sb.Preds() }
+
+// Pred implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) Pred(i int) regalloc.Block {
+	sb := r.sb
+	pred := sb.Pred(i)
+	l := r.f.m.SSABlockLabel(pred.ID())
+	index := r.f.labelToRegAllocBlockIndex[l]
+	return &r.f.reversePostOrderBlocks[index]
+}
+
+// Entry implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) Entry() bool { return r.sb.EntryBlock() }
+
+// Succs implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) Succs() int {
+	return r.sb.Succs()
+}
+
+// Succ implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) Succ(i int) regalloc.Block {
+	sb := r.sb
+	succ := sb.Succ(i)
+	if succ.ReturnBlock() {
+		return nil
+	}
+	l := r.f.m.SSABlockLabel(succ.ID())
+	index := r.f.labelToRegAllocBlockIndex[l]
+	return &r.f.reversePostOrderBlocks[index]
+}
+
+// LoopHeader implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) LoopHeader() bool {
+	return r.sb.LoopHeader()
+}
+
+// LoopNestingForestChildren implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) LoopNestingForestChildren() int {
+	r.loopNestingForestChildren = r.sb.LoopNestingForestChildren()
+	return len(r.loopNestingForestChildren)
+}
+
+// LoopNestingForestChild implements regalloc.Block.
+func (r *RegAllocBlock[I, m]) LoopNestingForestChild(i int) regalloc.Block {
+	blk := r.loopNestingForestChildren[i]
+	l := r.f.m.SSABlockLabel(blk.ID())
+	index := r.f.labelToRegAllocBlockIndex[l]
+	return &r.f.reversePostOrderBlocks[index]
+}
diff --git a/internal/engine/wazevo/backend/regalloc/api.go b/internal/engine/wazevo/backend/regalloc/api.go
index 81240cfd..941808ad 100644
--- a/internal/engine/wazevo/backend/regalloc/api.go
+++ b/internal/engine/wazevo/backend/regalloc/api.go
@@ -25,8 +25,6 @@ type (
 		ReversePostOrderBlockIteratorNext() Block
 		// ClobberedRegisters tell the clobbered registers by this function.
 		ClobberedRegisters([]VReg)
-		// Done tells the implementation that register allocation is done, and it can finalize the stack
-		Done()
 		// LoopNestingForestRoots returns the number of roots of the loop nesting forest in a function.
 		LoopNestingForestRoots() int
 		// LoopNestingForestRoot returns the i-th root of the loop nesting forest in a function.
@@ -70,12 +68,14 @@ type (
 		InstrRevIteratorNext() Instr
 		// FirstInstr returns the fist instruction in this block where instructions will be inserted after it.
 		FirstInstr() Instr
-		// LastInstr returns the last instruction in this block where instructions will be inserted before it.
+		// EndInstr returns the end instruction in this block.
+		EndInstr() Instr
+		// LastInstrForInsertion returns the last instruction in this block where instructions will be inserted before it.
 		// Such insertions only happen when we need to insert spill/reload instructions to adjust the merge edges.
 		// If the very last instruction is the unconditional branching, then the returned instruction is the one before it.
 		// Note that at the time of register allocation, all the critical edges are already split, so there is no need
 		// to worry about the case where branching instruction has multiple successors.
-		LastInstr() Instr
+		LastInstrForInsertion() Instr
 		// Preds returns the number of predecessors of this block in the CFG.
 		Preds() int
 		// Pred returns the i-th predecessor of this block in the CFG.
@@ -97,7 +97,10 @@ type (
 	// Instr is an instruction in a block, abstracting away the underlying ISA.
 	Instr interface {
 		fmt.Stringer
-
+		// Next returns the next instruction in the same block.
+		Next() Instr
+		// Prev returns the previous instruction in the same block.
+		Prev() Instr
 		// Defs returns the virtual registers defined by this instruction.
 		Defs(*[]VReg) []VReg
 		// Uses returns the virtual registers used by this instruction.
@@ -120,5 +123,13 @@ type (
 		IsIndirectCall() bool
 		// IsReturn returns true if this instruction is a return instruction.
 		IsReturn() bool
+		// AddedBeforeRegAlloc returns true if this instruction is added before register allocation.
+		AddedBeforeRegAlloc() bool
+	}
+
+	// InstrConstraint is an interface for arch-specific instruction constraints.
+	InstrConstraint interface {
+		comparable
+		Instr
 	}
 )
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index d2d212fe..b7f2da1b 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -318,7 +318,6 @@ func (a *Allocator) DoAllocation(f Function) {
 	a.livenessAnalysis(f)
 	a.alloc(f)
 	a.determineCalleeSavedRealRegs(f)
-	f.Done()
 }
 
 func (a *Allocator) determineCalleeSavedRealRegs(f Function) {
@@ -809,11 +808,11 @@ func (a *Allocator) reconcileEdge(f Function,
 			// This case is that the desired value is on the stack, but currentVReg is on the target register.
 			// We need to move the current value to the stack, and reload the desired value.
 			// TODO: we can do better here.
-			f.StoreRegisterBefore(currentVReg.SetRealReg(r), pred.LastInstr())
+			f.StoreRegisterBefore(currentVReg.SetRealReg(r), pred.LastInstrForInsertion())
 			delete(currentOccupantsRev, currentVReg)
 
 			s.getVRegState(desiredVReg).recordReload(f, pred)
-			f.ReloadRegisterBefore(desiredVReg.SetRealReg(r), pred.LastInstr())
+			f.ReloadRegisterBefore(desiredVReg.SetRealReg(r), pred.LastInstrForInsertion())
 			currentOccupants.add(r, desiredVReg)
 			currentOccupantsRev[desiredVReg] = r
 			return
@@ -849,12 +848,12 @@ func (a *Allocator) reconcileEdge(f Function,
 			f.InsertMoveBefore(
 				FromRealReg(r, typ),
 				desiredVReg.SetRealReg(currentReg),
-				pred.LastInstr(),
+				pred.LastInstrForInsertion(),
 			)
 			currentOccupants.remove(currentReg)
 		} else {
 			s.getVRegState(desiredVReg).recordReload(f, pred)
-			f.ReloadRegisterBefore(desiredVReg.SetRealReg(r), pred.LastInstr())
+			f.ReloadRegisterBefore(desiredVReg.SetRealReg(r), pred.LastInstrForInsertion())
 		}
 		currentOccupantsRev[desiredVReg] = r
 		currentOccupants.add(r, desiredVReg)


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc_test.go b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc_test.go
index 396aabd1..91c0768d 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc_test.go
@@ -8,56 +8,6 @@ import (
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
-func TestRegAllocFunctionImpl_addBlock(t *testing.T) {
-	ssab := ssa.NewBuilder()
-	sb1, sb2 := ssab.AllocateBasicBlock(), ssab.AllocateBasicBlock()
-	p1, p2 := &labelPosition{}, &labelPosition{}
-
-	f := &regAllocFunctionImpl{labelToRegAllocBlockIndex: map[label]int{}}
-	f.addBlock(sb1, label(10), p1)
-	f.addBlock(sb2, label(20), p2)
-
-	require.Equal(t, 2, len(f.labelToRegAllocBlockIndex))
-	require.Equal(t, 2, len(f.reversePostOrderBlocks))
-
-	rb1, rb2 := &f.reversePostOrderBlocks[0], &f.reversePostOrderBlocks[1]
-	require.Equal(t, f, rb1.f)
-	require.Equal(t, f, rb2.f)
-
-	require.Equal(t, p1, rb1.pos)
-	require.Equal(t, p2, rb2.pos)
-
-	require.Equal(t, label(10), rb1.l)
-	require.Equal(t, label(20), rb2.l)
-
-	require.Equal(t, sb1, rb1.sb)
-	require.Equal(t, sb2, rb2.sb)
-}
-
-func TestRegAllocFunctionImpl_PostOrderBlockIterator(t *testing.T) {
-	f := &regAllocFunctionImpl{reversePostOrderBlocks: []regAllocBlockImpl{{}, {}, {}}}
-	blk := f.PostOrderBlockIteratorBegin()
-	require.Equal(t, blk, &f.reversePostOrderBlocks[2])
-	blk = f.PostOrderBlockIteratorNext()
-	require.Equal(t, blk, &f.reversePostOrderBlocks[1])
-	blk = f.PostOrderBlockIteratorNext()
-	require.Equal(t, blk, &f.reversePostOrderBlocks[0])
-	blk = f.PostOrderBlockIteratorNext()
-	require.Nil(t, blk)
-}
-
-func TestRegAllocFunctionImpl_ReversePostOrderBlockIterator(t *testing.T) {
-	f := &regAllocFunctionImpl{reversePostOrderBlocks: []regAllocBlockImpl{{}, {}, {}}}
-	blk := f.ReversePostOrderBlockIteratorBegin()
-	require.Equal(t, blk, &f.reversePostOrderBlocks[0])
-	blk = f.ReversePostOrderBlockIteratorNext()
-	require.Equal(t, blk, &f.reversePostOrderBlocks[1])
-	blk = f.ReversePostOrderBlockIteratorNext()
-	require.Equal(t, blk, &f.reversePostOrderBlocks[2])
-	blk = f.ReversePostOrderBlockIteratorNext()
-	require.Nil(t, blk)
-}
-
 func TestRegAllocFunctionImpl_ReloadRegisterAfter(t *testing.T) {
 	ctx, _, m := newSetupWithMockContext()
 
@@ -66,9 +16,8 @@ func TestRegAllocFunctionImpl_ReloadRegisterAfter(t *testing.T) {
 	i1.next = i2
 	i2.prev = i1
 
-	f := &regAllocFunctionImpl{m: m}
-	f.ReloadRegisterAfter(x1VReg, i1)
-	f.ReloadRegisterAfter(v1VReg, i1)
+	m.InsertReloadRegisterAt(x1VReg, i1, true)
+	m.InsertReloadRegisterAt(v1VReg, i1, true)
 
 	require.NotEqual(t, i1, i2.prev)
 	require.NotEqual(t, i1.next, i2)
@@ -96,9 +45,8 @@ func TestRegAllocFunctionImpl_StoreRegisterBefore(t *testing.T) {
 	i1.next = i2
 	i2.prev = i1
 
-	f := &regAllocFunctionImpl{m: m}
-	f.StoreRegisterBefore(x1VReg, i2)
-	f.StoreRegisterBefore(v1VReg, i2)
+	m.InsertStoreRegisterAt(x1VReg, i2, false)
+	m.InsertStoreRegisterAt(v1VReg, i2, false)
 
 	require.NotEqual(t, i1, i2.prev)
 	require.NotEqual(t, i1.next, i2)
@@ -177,11 +125,11 @@ func TestMachine_insertStoreRegisterAt(t *testing.T) {
 					i2.prev = i1
 
 					if after {
-						m.insertStoreRegisterAt(v1VReg, i1, after)
-						m.insertStoreRegisterAt(x1VReg, i1, after)
+						m.InsertStoreRegisterAt(v1VReg, i1, after)
+						m.InsertStoreRegisterAt(x1VReg, i1, after)
 					} else {
-						m.insertStoreRegisterAt(x1VReg, i2, after)
-						m.insertStoreRegisterAt(v1VReg, i2, after)
+						m.InsertStoreRegisterAt(x1VReg, i2, after)
+						m.InsertStoreRegisterAt(v1VReg, i2, after)
 					}
 					m.executableContext.RootInstr = i1
 					require.Equal(t, tc.expected, m.Format())
@@ -250,11 +198,11 @@ func TestMachine_insertReloadRegisterAt(t *testing.T) {
 					i2.prev = i1
 
 					if after {
-						m.insertReloadRegisterAt(v1VReg, i1, after)
-						m.insertReloadRegisterAt(x1VReg, i1, after)
+						m.InsertReloadRegisterAt(v1VReg, i1, after)
+						m.InsertReloadRegisterAt(x1VReg, i1, after)
 					} else {
-						m.insertReloadRegisterAt(x1VReg, i2, after)
-						m.insertReloadRegisterAt(v1VReg, i2, after)
+						m.InsertReloadRegisterAt(x1VReg, i2, after)
+						m.InsertReloadRegisterAt(v1VReg, i2, after)
 					}
 					m.executableContext.RootInstr = i1
 
@@ -265,14 +213,13 @@ func TestMachine_insertReloadRegisterAt(t *testing.T) {
 	}
 }
 
-func TestRegAllocFunctionImpl_ClobberedRegisters(t *testing.T) {
+func TestRegMachine_ClobberedRegisters(t *testing.T) {
 	_, _, m := newSetupWithMockContext()
-	f := &regAllocFunctionImpl{m: m}
-	f.ClobberedRegisters([]regalloc.VReg{v19VReg, v19VReg, v19VReg, v19VReg})
+	m.ClobberedRegisters([]regalloc.VReg{v19VReg, v19VReg, v19VReg, v19VReg})
 	require.Equal(t, []regalloc.VReg{v19VReg, v19VReg, v19VReg, v19VReg}, m.clobberedRegs)
 }
 
-func TestMachine_swap(t *testing.T) {
+func TestMachineMachineswap(t *testing.T) {
 	for _, tc := range []struct {
 		x1, x2, tmp regalloc.VReg
 		expected    string
@@ -337,7 +284,7 @@ func TestMachine_swap(t *testing.T) {
 			cur.next = i2
 			i2.prev = cur
 
-			m.swap(cur, tc.x1, tc.x2, tc.tmp)
+			m.Swap(cur, tc.x1, tc.x2, tc.tmp)
 			m.executableContext.RootInstr = cur
 
 			require.Equal(t, tc.expected, m.Format())
diff --git a/internal/engine/wazevo/backend/isa/arm64/util_test.go b/internal/engine/wazevo/backend/isa/arm64/util_test.go
index 1be5a76d..a033db70 100644
--- a/internal/engine/wazevo/backend/isa/arm64/util_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/util_test.go
@@ -61,7 +61,7 @@ type mockCompiler struct {
 	buf         []byte
 }
 
-func (m *mockCompiler) SSABuilder() ssa.Builder { panic("TODO") }
+func (m *mockCompiler) SSABuilder() ssa.Builder { return nil }
 
 func (m *mockCompiler) LoopNestingForestRoots() []ssa.BasicBlock { panic("TODO") }
 
diff --git a/internal/engine/wazevo/backend/machine_test.go b/internal/engine/wazevo/backend/machine_test.go
index 666c1c7b..45ed0f22 100644
--- a/internal/engine/wazevo/backend/machine_test.go
+++ b/internal/engine/wazevo/backend/machine_test.go
@@ -25,6 +25,8 @@ type mockMachine struct {
 	rinfo                  *regalloc.RegisterInfo
 }
 
+func (m mockMachine) RegAlloc() { panic("implement me") }
+
 func (m mockMachine) LowerParams(params []ssa.Value) { panic("implement me") }
 
 func (m mockMachine) LowerReturns(returns []ssa.Value) { panic("implement me") }
diff --git a/internal/engine/wazevo/backend/regalloc/api_test.go b/internal/engine/wazevo/backend/regalloc/api_test.go
index 6fd25e69..0bb4fad7 100644
--- a/internal/engine/wazevo/backend/regalloc/api_test.go
+++ b/internal/engine/wazevo/backend/regalloc/api_test.go
@@ -36,6 +36,7 @@ type (
 
 	// mockInstr implements Instr.
 	mockInstr struct {
+		next, prev                 *mockInstr
 		defs, uses                 []VReg
 		isCopy, isCall, isIndirect bool
 	}
@@ -58,6 +59,14 @@ func (m *mockFunction) loopNestingForestRoots(blocks ...*mockBlock) {
 }
 
 func newMockBlock(id int, instructions ...*mockInstr) *mockBlock {
+	if len(instructions) > 0 {
+		instructions[0].prev = nil
+		for i := 1; i < len(instructions); i++ {
+			instructions[i].prev = instructions[i-1]
+			instructions[i-1].next = instructions[i]
+		}
+		instructions[len(instructions)-1].next = nil
+	}
 	return &mockBlock{id: id, instructions: instructions}
 }
 
@@ -260,6 +269,9 @@ func (m *mockInstr) Defs(ret *[]VReg) []VReg {
 	return *ret
 }
 
+// AddedBeforeRegAlloc implements Instr.
+func (m *mockInstr) AddedBeforeRegAlloc() bool { return true }
+
 // Uses implements Instr.
 func (m *mockInstr) Uses(ret *[]VReg) []VReg {
 	*ret = append((*ret)[:0], m.uses...)
@@ -278,6 +290,12 @@ func (m *mockInstr) IsIndirectCall() bool { return m.isIndirect }
 // IsReturn implements Instr.
 func (m *mockInstr) IsReturn() bool { return false }
 
+// Next implements Instr.
+func (m *mockInstr) Next() Instr { return m.next }
+
+// Prev implements Instr.
+func (m *mockInstr) Prev() Instr { return m.prev }
+
 // Entry implements Entry.
 func (m *mockBlock) Entry() bool { return m._entry }
 
@@ -328,7 +346,24 @@ func (m *mockBlock) LoopNestingForestChild(i int) Block {
 	return m.lnfChildren[i]
 }
 
-func (m *mockBlock) LastInstr() Instr {
+func (m *mockBlock) BeginInstr() Instr {
+	if len(m.instructions) == 0 {
+		return nil
+	}
+	return m.instructions[0]
+}
+
+func (m *mockBlock) EndInstr() Instr {
+	if len(m.instructions) == 0 {
+		return nil
+	}
+	return m.instructions[len(m.instructions)-1]
+}
+
+func (m *mockBlock) LastInstrForInsertion() Instr {
+	if len(m.instructions) == 0 {
+		return nil
+	}
 	return m.instructions[len(m.instructions)-1]
 }