From 506c922f0b4244fa84483488030fbc8b5dc086e3 Mon Sep 17 00:00:00 2001
From: Takeshi Yoneda <t.y.mathetake@gmail.com>
Date: Wed, 5 Jun 2024 14:33:58 -0700
Subject: [PATCH] ssa: reuses slice on basicBlock.loopNestingForestChildren
 (#2232)

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>
---
 internal/engine/wazevo/ssa/basic_block.go   | 9 ++++++---
 internal/engine/wazevo/ssa/builder.go       | 4 ++++
 internal/engine/wazevo/ssa/pass_cfg.go      | 4 ++--
 internal/engine/wazevo/ssa/pass_cfg_test.go | 2 +-
 4 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/internal/engine/wazevo/ssa/basic_block.go b/internal/engine/wazevo/ssa/basic_block.go
index f9dab4eab7..2e7069dca4 100644
--- a/internal/engine/wazevo/ssa/basic_block.go
+++ b/internal/engine/wazevo/ssa/basic_block.go
@@ -108,7 +108,7 @@ type (
 
 		// loopNestingForestChildren holds the children of this block in the loop nesting forest.
 		// Non-empty if and only if this block is a loop header (i.e. loopHeader=true)
-		loopNestingForestChildren []BasicBlock
+		loopNestingForestChildren wazevoapi.VarLength[BasicBlock]
 
 		// reversePostOrder is used to sort all the blocks in the function in reverse post order.
 		// This is used in builder.LayoutBlocks.
@@ -128,6 +128,9 @@ type (
 	}
 )
 
+// basicBlockVarLengthNil is the default nil value for basicBlock.loopNestingForestChildren.
+var basicBlockVarLengthNil = wazevoapi.NewNilVarLength[BasicBlock]()
+
 const basicBlockIDReturnBlock = 0xffffffff
 
 // Name implements BasicBlock.Name.
@@ -271,7 +274,7 @@ func resetBasicBlock(bb *basicBlock) {
 	bb.unknownValues = bb.unknownValues[:0]
 	bb.lastDefinitions = wazevoapi.ResetMap(bb.lastDefinitions)
 	bb.reversePostOrder = -1
-	bb.loopNestingForestChildren = bb.loopNestingForestChildren[:0]
+	bb.loopNestingForestChildren = basicBlockVarLengthNil
 	bb.loopHeader = false
 	bb.sibling = nil
 	bb.child = nil
@@ -364,7 +367,7 @@ func (bb *basicBlock) String() string {
 
 // LoopNestingForestChildren implements BasicBlock.LoopNestingForestChildren.
 func (bb *basicBlock) LoopNestingForestChildren() []BasicBlock {
-	return bb.loopNestingForestChildren
+	return bb.loopNestingForestChildren.View()
 }
 
 // LoopHeader implements BasicBlock.LoopHeader.
diff --git a/internal/engine/wazevo/ssa/builder.go b/internal/engine/wazevo/ssa/builder.go
index 09d9967c87..0671b2a7c8 100644
--- a/internal/engine/wazevo/ssa/builder.go
+++ b/internal/engine/wazevo/ssa/builder.go
@@ -135,6 +135,7 @@ func NewBuilder() Builder {
 	return &builder{
 		instructionsPool:               wazevoapi.NewPool[Instruction](resetInstruction),
 		basicBlocksPool:                wazevoapi.NewPool[basicBlock](resetBasicBlock),
+		varLengthBasicBlockPool:        wazevoapi.NewVarLengthPool[BasicBlock](),
 		varLengthPool:                  wazevoapi.NewVarLengthPool[Value](),
 		valueAnnotations:               make(map[ValueID]string),
 		signatures:                     make(map[SignatureID]*Signature),
@@ -177,6 +178,8 @@ type builder struct {
 	dominators []*basicBlock
 	sparseTree dominatorSparseTree
 
+	varLengthBasicBlockPool wazevoapi.VarLengthPool[BasicBlock]
+
 	// loopNestingForestRoots are the roots of the loop nesting forest.
 	loopNestingForestRoots []BasicBlock
 
@@ -219,6 +222,7 @@ func (b *builder) Init(s *Signature) {
 	b.instructionsPool.Reset()
 	b.basicBlocksPool.Reset()
 	b.varLengthPool.Reset()
+	b.varLengthBasicBlockPool.Reset()
 	b.donePreBlockLayoutPasses = false
 	b.doneBlockLayout = false
 	b.donePostBlockLayoutPasses = false
diff --git a/internal/engine/wazevo/ssa/pass_cfg.go b/internal/engine/wazevo/ssa/pass_cfg.go
index 50cb9c4750..1bec185ee5 100644
--- a/internal/engine/wazevo/ssa/pass_cfg.go
+++ b/internal/engine/wazevo/ssa/pass_cfg.go
@@ -180,7 +180,7 @@ func passBuildLoopNestingForest(b *builder) {
 			b.loopNestingForestRoots = append(b.loopNestingForestRoots, blk)
 		} else if n == ent {
 		} else if n.loopHeader {
-			n.loopNestingForestChildren = append(n.loopNestingForestChildren, blk)
+			n.loopNestingForestChildren = n.loopNestingForestChildren.Append(&b.varLengthBasicBlockPool, blk)
 		}
 	}
 
@@ -193,7 +193,7 @@ func passBuildLoopNestingForest(b *builder) {
 
 func printLoopNestingForest(root *basicBlock, depth int) {
 	fmt.Println(strings.Repeat("\t", depth), "loop nesting forest root:", root.ID())
-	for _, child := range root.loopNestingForestChildren {
+	for _, child := range root.loopNestingForestChildren.View() {
 		fmt.Println(strings.Repeat("\t", depth+1), "child:", child.ID())
 		if child.LoopHeader() {
 			printLoopNestingForest(child.(*basicBlock), depth+2)
diff --git a/internal/engine/wazevo/ssa/pass_cfg_test.go b/internal/engine/wazevo/ssa/pass_cfg_test.go
index f6c48e6efa..ade6eeb960 100644
--- a/internal/engine/wazevo/ssa/pass_cfg_test.go
+++ b/internal/engine/wazevo/ssa/pass_cfg_test.go
@@ -669,7 +669,7 @@ func TestBuildLoopNestingForest(t *testing.T) {
 			for expBlkID, blk := range blocks {
 				expChildren := tc.expLNF.children[expBlkID]
 				var actualChildren []BasicBlockID
-				for _, child := range blk.loopNestingForestChildren {
+				for _, child := range blk.loopNestingForestChildren.View() {
 					actualChildren = append(actualChildren, child.(*basicBlock).id)
 				}
 				sort.Slice(actualChildren, func(i, j int) bool {