wazevo: pass reset func to NewPool (#1796)

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

internal/engine/wazevo/backend/compiler.go

@@ -135,6 +135,7 @@ type compiler struct {
 		dst   regalloc.VReg
 	}
 	vRegSet         []bool
+	vRegIDs         []regalloc.VRegID
 	tempRegs        []regalloc.VReg
 	tmpVals         []ssa.Value
 	ssaTypeOfVRegID [] /* VRegID to */ ssa.Type

internal/engine/wazevo/backend/compiler_lower.go

@@ -160,12 +160,14 @@ func (c *compiler) lowerBlockArguments(args []ssa.Value, succ ssa.BasicBlock) {
 	}
 
 	// Check if there's an overlap among the dsts and srcs in varEdges.
+	c.vRegIDs = c.vRegIDs[:0]
 	for _, edge := range c.varEdges {
 		src := edge[0].ID()
 		if int(src) >= len(c.vRegSet) {
 			c.vRegSet = append(c.vRegSet, make([]bool, src+1)...)
 		}
 		c.vRegSet[src] = true
+		c.vRegIDs = append(c.vRegIDs, src)
 	}
 	separated := true
 	for _, edge := range c.varEdges {
@@ -179,6 +181,9 @@ func (c *compiler) lowerBlockArguments(args []ssa.Value, succ ssa.BasicBlock) {
 			}
 		}
 	}
+	for _, id := range c.vRegIDs {
+		c.vRegSet[id] = false // reset for the next use.
+	}
 
 	if separated {
 		// If there's no overlap, we can simply move the source to destination.

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

@@ -112,8 +112,8 @@ const (
 // NewBackend returns a new backend for arm64.
 func NewBackend() backend.Machine {
 	m := &machine{
-		instrPool:         wazevoapi.NewPool[instruction](),
-		labelPositionPool: wazevoapi.NewPool[labelPosition](),
+		instrPool:         wazevoapi.NewPool[instruction](resetInstruction),
+		labelPositionPool: wazevoapi.NewPool[labelPosition](resetLabelPosition),
 		labelPositions:    make(map[label]*labelPosition),
 		spillSlots:        make(map[regalloc.VRegID]int64),
 		nextLabel:         invalidLabel,
@@ -243,10 +243,13 @@ func (m *machine) insertBrTargetLabel() label {
 
 func (m *machine) allocateLabelPosition() *labelPosition {
 	l := m.labelPositionPool.Allocate()
-	*l = labelPosition{}
 	return l
 }
 
+func resetLabelPosition(l *labelPosition) {
+	*l = labelPosition{}
+}
+
 func (m *machine) FlushPendingInstructions() {
 	l := len(m.pendingInstructions)
 	if l == 0 {
@@ -277,15 +280,18 @@ func (l label) String() string {
 // allocateInstr allocates an instruction.
 func (m *machine) allocateInstr() *instruction {
 	instr := m.instrPool.Allocate()
-	*instr = instruction{}
 	if !m.regAllocStarted {
 		instr.addedBeforeRegAlloc = true
 	}
 	return instr
 }
 
+func resetInstruction(i *instruction) {
+	*i = instruction{}
+}
+
 func (m *machine) allocateNop() *instruction {
-	instr := m.instrPool.Allocate()
+	instr := m.allocateInstr()
 	instr.asNop0()
 	return instr
 }

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

@@ -57,8 +57,7 @@ func TestMachine_resolveAddressingMode(t *testing.T) {
 	})
 
 	t.Run("tmp reg", func(t *testing.T) {
-		// 0x89705f4136b4a598
-		m := &machine{instrPool: wazevoapi.NewPool[instruction]()}
+		m := &machine{instrPool: wazevoapi.NewPool[instruction](resetInstruction)}
 		root := &instruction{kind: udf}
 		i := &instruction{prev: root}
 		i.asULoad(operandNR(x17VReg), addressMode{

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

@@ -32,7 +32,7 @@ func (t *intervalTree) reset() {
 
 func newIntervalTree() *intervalTree {
 	return &intervalTree{
-		allocator: wazevoapi.NewPool[intervalTreeNode](),
+		allocator: wazevoapi.NewPool[intervalTreeNode](resetIntervalTreeNode),
 		intervals: make(map[uint64]*intervalTreeNode),
 	}
 }
@@ -48,16 +48,24 @@ type intervalTreeNode struct {
 	// TODO: color for red-black balancing.
 }
 
+func resetIntervalTreeNode(i *intervalTreeNode) {
+	i.begin = 0
+	i.end = 0
+	i.nodes = i.nodes[:0]
+	i.maxEnd = 0
+	i.neighbors = i.neighbors[:0]
+	i.left = nil
+	i.right = nil
+}
+
 func (i *intervalTreeNode) insert(t *intervalTree, n *node, begin, end programCounter) *intervalTreeNode {
 	if i == nil {
 		intervalNode := t.allocator.Allocate()
-		intervalNode.right = nil
-		intervalNode.left = nil
 		intervalNode.nodes = append(intervalNode.nodes, n)
 		intervalNode.maxEnd = end
 		intervalNode.begin = begin
 		intervalNode.end = end
-		key := uint64(begin) | uint64(end)<<32
+		key := intervalTreeNodeKey(begin, end)
 		t.intervals[key] = intervalNode
 		return intervalNode
 	}

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

@@ -21,7 +21,7 @@ import (
 func NewAllocator(allocatableRegs *RegisterInfo) Allocator {
 	a := Allocator{
 		regInfo:  allocatableRegs,
-		nodePool: wazevoapi.NewPool[node](),
+		nodePool: wazevoapi.NewPool[node](resetNode),
 		phis:     make(map[VReg]Block),
 	}
 	for _, regs := range allocatableRegs.AllocatableRegisters {
@@ -487,16 +487,22 @@ func (a *Allocator) getOrAllocateNode(v VReg) (n *node) {
 	return
 }
 
-func (a *Allocator) allocateNode() (n *node) {
-	n = a.nodePool.Allocate()
-	n.id = a.nodePool.Allocated() - 1
+func resetNode(n *node) {
+	n.r = RealRegInvalid
+	n.v = VRegInvalid
 	n.ranges = n.ranges[:0]
+	n.neighbors = n.neighbors[:0]
 	n.copyFromVReg = nil
 	n.copyToVReg = nil
 	n.copyFromReal = RealRegInvalid
 	n.copyToReal = RealRegInvalid
-	n.neighbors = n.neighbors[:0]
+	n.degree = 0
 	n.visited = false
+}
+
+func (a *Allocator) allocateNode() (n *node) {
+	n = a.nodePool.Allocate()
+	n.id = a.nodePool.Allocated() - 1
 	return
 }
 

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

@@ -445,7 +445,10 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				actual := &a.blockInfos[blockID]
 				exp := tc.exp[blockID]
 				initMapInInfo(exp)
+				saved := actual.intervalTree
+				actual.intervalTree = nil // Don't compare intervalTree.
 				require.Equal(t, exp, actual, "\n[exp for block[%d]]\n%s\n[actual for block[%d]]\n%s", blockID, exp, blockID, actual)
+				actual.intervalTree = saved
 			}
 
 			// Sanity check: buildLiveRanges should not panic.
@@ -517,9 +520,6 @@ func TestAllocator_recordCopyRelation(t *testing.T) {
 }
 
 func initMapInInfo(info *blockInfo) {
-	if info.intervalTree == nil {
-		info.intervalTree = newIntervalTree()
-	}
 	if info.liveIns == nil {
 		info.liveIns = make(map[VReg]struct{})
 	}

internal/engine/wazevo/ssa/basic_block.go

@@ -245,7 +245,7 @@ func (bb *basicBlock) Tail() *Instruction {
 }
 
 // reset resets the basicBlock to its initial state so that it can be reused for another function.
-func (bb *basicBlock) reset() {
+func resetBasicBlock(bb *basicBlock) {
 	bb.params = bb.params[:0]
 	bb.rootInstr, bb.currentInstr = nil, nil
 	bb.preds = bb.preds[:0]

internal/engine/wazevo/ssa/builder.go

@@ -126,8 +126,8 @@ type Builder interface {
 // NewBuilder returns a new Builder implementation.
 func NewBuilder() Builder {
 	return &builder{
-		instructionsPool:               wazevoapi.NewPool[Instruction](),
-		basicBlocksPool:                wazevoapi.NewPool[basicBlock](),
+		instructionsPool:               wazevoapi.NewPool[Instruction](resetInstruction),
+		basicBlocksPool:                wazevoapi.NewPool[basicBlock](resetBasicBlock),
 		valueAnnotations:               make(map[ValueID]string),
 		signatures:                     make(map[SignatureID]*Signature),
 		blkVisited:                     make(map[*basicBlock]int),
@@ -196,7 +196,7 @@ func (b *builder) ReturnBlock() BasicBlock {
 // Init implements Builder.Reset.
 func (b *builder) Init(s *Signature) {
 	b.currentSignature = s
-	b.returnBlk.reset()
+	resetBasicBlock(b.returnBlk)
 	b.instructionsPool.Reset()
 	b.basicBlocksPool.Reset()
 	b.donePasses = false
@@ -211,7 +211,6 @@ func (b *builder) Init(s *Signature) {
 
 	for i := 0; i < b.basicBlocksPool.Allocated(); i++ {
 		blk := b.basicBlocksPool.View(i)
-		blk.reset()
 		delete(b.blkVisited, blk)
 	}
 	b.basicBlocksPool.Reset()
@@ -250,7 +249,6 @@ func (b *builder) AnnotateValue(value Value, a string) {
 // AllocateInstruction implements Builder.AllocateInstruction.
 func (b *builder) AllocateInstruction() *Instruction {
 	instr := b.instructionsPool.Allocate()
-	instr.reset()
 	return instr
 }
 

internal/engine/wazevo/ssa/instructions.go

@@ -74,8 +74,8 @@ func (i *Instruction) Lowered() bool {
 	return i.alreadyLowered
 }
 
-// reset resets this instruction to the initial state.
-func (i *Instruction) reset() {
+// resetInstruction resets this instruction to the initial state.
+func resetInstruction(i *Instruction) {
 	*i = Instruction{}
 	i.v = ValueInvalid
 	i.v2 = ValueInvalid

internal/engine/wazevo/wazevoapi/pool.go

@@ -6,12 +6,15 @@ const poolPageSize = 128
 // This is useful to avoid unnecessary allocations.
 type Pool[T any] struct {
 	pages            []*[poolPageSize]T
+	resetFn          func(*T)
 	allocated, index int
 }
 
 // NewPool returns a new Pool.
-func NewPool[T any]() Pool[T] {
+// resetFn is called when a new T is allocated in Pool.Allocate.
+func NewPool[T any](resetFn func(*T)) Pool[T] {
 	var ret Pool[T]
+	ret.resetFn = resetFn
 	ret.Reset()
 	return ret
 }
@@ -36,6 +39,7 @@ func (p *Pool[T]) Allocate() *T {
 		p.index = 0
 	}
 	ret := &p.pages[len(p.pages)-1][p.index]
+	p.resetFn(ret)
 	p.index++
 	p.allocated++
 	return ret
@@ -49,13 +53,6 @@ func (p *Pool[T]) View(i int) *T {
 
 // Reset resets the pool.
 func (p *Pool[T]) Reset() {
-	for _, ns := range p.pages {
-		pages := ns[:]
-		for i := range pages {
-			var v T
-			pages[i] = v
-		}
-	}
 	p.pages = p.pages[:0]
 	p.index = poolPageSize
 	p.allocated = 0