asm: reuse StaticConstPool (#1304)

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

internal/asm/amd64/impl.go

@@ -228,7 +228,7 @@ type AssemblerImpl struct {
 
 	readInstructionAddressNodes []*nodeImpl
 
-	pool *asm.StaticConstPool
+	pool asm.StaticConstPool
 }
 
 func NewAssembler() *AssemblerImpl {
@@ -271,10 +271,12 @@ func (n *nodePool) allocNode() (ret *nodeImpl) {
 
 // Reset implements asm.AssemblerBase.
 func (a *AssemblerImpl) Reset() {
+	pool := a.pool
+	pool.Reset()
 	*a = AssemblerImpl{
 		buf:                         a.buf,
 		nodePool:                    a.nodePool,
-		pool:                        asm.NewStaticConstPool(),
+		pool:                        pool,
 		enablePadding:               a.enablePadding,
 		readInstructionAddressNodes: a.readInstructionAddressNodes[:0],
 		BaseAssemblerImpl: asm.BaseAssemblerImpl{

internal/asm/amd64/impl_staticconst.go

@@ -44,7 +44,7 @@ func (a *AssemblerImpl) maybeFlushConstants(isEndOfFunction bool) {
 			a.buf.Write(c.Raw)
 		}
 
-		a.pool = asm.NewStaticConstPool() // reset
+		a.pool.Reset()
 	}
 }
 

internal/asm/arm64/impl.go

@@ -219,7 +219,7 @@ type AssemblerImpl struct {
 	buf               *bytes.Buffer
 	temporaryRegister asm.Register
 	nodeCount         int
-	pool              *asm.StaticConstPool
+	pool              asm.StaticConstPool
 	// MaxDisplacementForConstantPool is fixed to defaultMaxDisplacementForConstPool
 	// but have it as a field here for testability.
 	MaxDisplacementForConstantPool         int
@@ -267,8 +267,10 @@ func NewAssembler(temporaryRegister asm.Register) *AssemblerImpl {
 // Reset implements asm.AssemblerBase.
 func (a *AssemblerImpl) Reset() {
 	buf, np, tmp := a.buf, a.nodePool, a.temporaryRegister
+	pool := a.pool
+	pool.Reset()
 	*a = AssemblerImpl{
-		buf: buf, nodePool: np, pool: asm.NewStaticConstPool(),
+		buf: buf, nodePool: np, pool: pool,
 		temporaryRegister:   tmp,
 		adrInstructionNodes: a.adrInstructionNodes[:0],
 		relativeJumpNodes:   a.relativeJumpNodes[:0],
@@ -394,7 +396,7 @@ func (a *AssemblerImpl) maybeFlushConstPool(endOfBinary bool) {
 		}
 
 		// After the flush, reset the constant pool.
-		a.pool = asm.NewStaticConstPool()
+		a.pool.Reset()
 	}
 }
 

internal/asm/assembler.go

@@ -90,9 +90,19 @@ type StaticConstPool struct {
 	PoolSizeInBytes int
 }
 
-// NewStaticConstPool returns the pointer to a new StaticConstPool.
-func NewStaticConstPool() *StaticConstPool {
-	return &StaticConstPool{addedConsts: map[*StaticConst]struct{}{}}
+func NewStaticConstPool() StaticConstPool {
+	return StaticConstPool{addedConsts: map[*StaticConst]struct{}{}}
+}
+
+// Reset resets the *StaticConstPool for reuse.
+func (p *StaticConstPool) Reset() {
+	for _, c := range p.Consts {
+		delete(p.addedConsts, c)
+	}
+	// Reuse the slice to avoid re-allocations.
+	p.Consts = p.Consts[:0]
+	p.PoolSizeInBytes = 0
+	p.FirstUseOffsetInBinary = nil
 }
 
 // AddConst adds a *StaticConst into the pool if it's not already added.