asm/amd64: avoids allocations in static const offset resolution (#1384)

internal/asm/amd64/impl.go

@@ -39,8 +39,11 @@ type nodeImpl struct {
 	// jumpOriginsHead true if this is the target of all the nodes in the singly linked list jumpOrigins,
 	// and can be traversed from this nodeImpl's forwardJumpOrigins.
 	forwardJumpTarget bool
-
-	staticConst *asm.StaticConst
+	// staticConstReferrersAdded true if this node is already added into AssemblerImpl.staticConstReferrers.
+	// Only used when staticConst is not nil. Through re-assembly, we might end up adding multiple times which causes unnecessary
+	// allocations, so we use this flag to do it once.
+	staticConstReferrersAdded bool
+	staticConst               *asm.StaticConst
 }
 
 type nodeFlag byte
@@ -217,22 +220,33 @@ func (o operandTypes) String() string {
 	return fmt.Sprintf("from:%s,to:%s", o.src, o.dst)
 }
 
-// AssemblerImpl implements Assembler.
-type AssemblerImpl struct {
-	nodePool *nodePool
-	asm.BaseAssemblerImpl
-	enablePadding   bool
-	root, current   *nodeImpl
-	buf             *bytes.Buffer
-	forceReAssemble bool
-	// MaxDisplacementForConstantPool is fixed to defaultMaxDisplacementForConstantPool
-	// but have it as an exported field here for testability.
-	MaxDisplacementForConstantPool int
+type (
+	// AssemblerImpl implements Assembler.
+	AssemblerImpl struct {
+		nodePool *nodePool
+		asm.BaseAssemblerImpl
+		enablePadding   bool
+		root, current   *nodeImpl
+		buf             *bytes.Buffer
+		forceReAssemble bool
+		// MaxDisplacementForConstantPool is fixed to defaultMaxDisplacementForConstantPool
+		// but have it as an exported field here for testability.
+		MaxDisplacementForConstantPool int
 
-	readInstructionAddressNodes []*nodeImpl
+		readInstructionAddressNodes []*nodeImpl
+		// staticConstReferrers maintains the list of static const referrers which requires the
+		// offset resolution after finalizing the binary layout.
+		staticConstReferrers []staticConstReferrer
+		pool                 asm.StaticConstPool
+	}
 
-	pool asm.StaticConstPool
-}
+	// staticConstReferrer represents a referrer of a asm.StaticConst.
+	staticConstReferrer struct {
+		n *nodeImpl
+		// instLen is the encoded length of the instruction for `n`.
+		instLen int
+	}
+)
 
 func NewAssembler() *AssemblerImpl {
 	return &AssemblerImpl{
@@ -295,6 +309,7 @@ func (a *AssemblerImpl) Reset() {
 		pool:                        pool,
 		enablePadding:               a.enablePadding,
 		readInstructionAddressNodes: a.readInstructionAddressNodes[:0],
+		staticConstReferrers:        a.staticConstReferrers[:0],
 		BaseAssemblerImpl: asm.BaseAssemblerImpl{
 			SetBranchTargetOnNextNodes: a.SetBranchTargetOnNextNodes[:0],
 			JumpTableEntries:           a.JumpTableEntries[:0],
@@ -404,6 +419,17 @@ func (a *AssemblerImpl) Assemble() ([]byte, error) {
 		}
 	}
 
+	// Now that we've finished the layout, fill out static consts offsets.
+	for i := range a.staticConstReferrers {
+		ref := &a.staticConstReferrers[i]
+		n, instLen := ref.n, ref.instLen
+		// Calculate the displacement between the RIP (the offset _after_ n) and the static constant.
+		displacement := int(n.staticConst.OffsetInBinary) - int(n.OffsetInBinary()) - instLen
+		// The offset must be stored at the 4 bytes from the tail of this n. See AssemblerImpl.encodeStaticConstImpl for detail.
+		displacementOffsetInInstruction := n.OffsetInBinary() + uint64(instLen-4)
+		binary.LittleEndian.PutUint32(code[displacementOffsetInInstruction:], uint32(int32(displacement)))
+	}
+
 	if err := a.FinalizeJumpTableEntry(code); err != nil {
 		return nil, err
 	}

internal/asm/amd64/impl_staticconst.go

@@ -1,7 +1,6 @@
 package amd64
 
 import (
-	"encoding/binary"
 	"fmt"
 	"math"
 
@@ -137,32 +136,36 @@ func (a *AssemblerImpl) encodeStaticConstImpl(n *nodeImpl, opcode []byte, rex Re
 
 	rexPrefix |= rex
 
-	var inst []byte
-	n.staticConst.AddOffsetFinalizedCallback(func(offsetOfConstInBinary uint64) {
-		bin := a.buf.Bytes()
-		displacement := int(offsetOfConstInBinary) - int(n.OffsetInBinary()) - len(inst)
-		displacementOffsetInInstruction := n.OffsetInBinary() + uint64(len(inst)-4)
-		binary.LittleEndian.PutUint32(bin[displacementOffsetInInstruction:], uint32(int32(displacement)))
-	})
+	var instLen int
+	if mandatoryPrefix != 0 {
+		a.buf.WriteByte(mandatoryPrefix)
+		instLen++
+	}
+
+	if rexPrefix != RexPrefixNone {
+		a.buf.WriteByte(rexPrefix)
+		instLen++
+	}
+
+	a.buf.Write(opcode)
+	instLen += len(opcode)
 
 	// https://wiki.osdev.org/X86-64_Instruction_Encoding#32.2F64-bit_addressing
 	modRM := 0b00_000_101 | // Indicate "[RIP + 32bit displacement]" encoding.
 		(reg3Bits << 3) // Place the reg on ModRM:reg.
+	a.buf.WriteByte(modRM)
+	instLen++
 
-	if mandatoryPrefix != 0 {
-		inst = append(inst, mandatoryPrefix)
+	// Preserve 4 bytes for displacement which will be filled after we finalize the location.
+	for i := 0; i < 4; i++ {
+		a.buf.WriteByte(0)
 	}
+	instLen += 4
 
-	if rexPrefix != RexPrefixNone {
-		inst = append(inst, rexPrefix)
+	if !n.staticConstReferrersAdded {
+		a.staticConstReferrers = append(a.staticConstReferrers, staticConstReferrer{n: n, instLen: instLen})
+		n.staticConstReferrersAdded = true
 	}
-
-	inst = append(inst, opcode...)
-	inst = append(inst, modRM,
-		0x0, 0x0, 0x0, 0x0, // Preserve 4 bytes for displacement.
-	)
-
-	a.buf.Write(inst)
 	return
 }