asm: moved Instruction.Symbol and Instruction.Reference into metadata

This commit adds per-instruction metadata, which is sparse in nature,
meaning that most instructions don't have metadata but some do. Metadata
is linked to an instruction via a pointer which is nil by default, thus
if an instruction has no metadata, at most one pointer worth of space is
unused which is desirable since we expect the size of the metadata
struct to grow in future commits.

The metadata struct is copy-on-write, so multiple instructions can share
the same metadata object after being copied, but as soon as a
modification on one of the instruction's metadata is performed, the
metadata object is copied so writes to a copied instruction doesn't
effect the original instruction.

This commit also deprecates `Instruction.Sym` in favor of
`Instruction.WithSymbol`. This follows the convention of prefixing
func names with `With...` if they have value receivers to make it clear
to callers that this doesn't modify the instruction it is being called
on.

Author: Dylan Reimerink

asm/dsl_test.go

@@ -25,21 +25,21 @@ func TestDSL(t *testing.T) {
 			OpCode: 0x04, Dst: R1, Constant: 22,
 		}},
 		{"JSGT.Imm", JSGT.Imm(R1, 4, "foo"), Instruction{
-			OpCode: 0x65, Dst: R1, Constant: 4, Offset: -1, Reference: "foo",
-		}},
+			OpCode: 0x65, Dst: R1, Constant: 4, Offset: -1,
+		}.WithReference("foo")},
 		{"JSGT.Imm32", JSGT.Imm32(R1, -2, "foo"), Instruction{
-			OpCode: 0x66, Dst: R1, Constant: -2, Offset: -1, Reference: "foo",
-		}},
+			OpCode: 0x66, Dst: R1, Constant: -2, Offset: -1,
+		}.WithReference("foo")},
 		{"JSLT.Reg", JSLT.Reg(R1, R2, "foo"), Instruction{
-			OpCode: 0xcd, Dst: R1, Src: R2, Offset: -1, Reference: "foo",
-		}},
+			OpCode: 0xcd, Dst: R1, Src: R2, Offset: -1,
+		}.WithReference("foo")},
 		{"JSLT.Reg32", JSLT.Reg32(R1, R3, "foo"), Instruction{
-			OpCode: 0xce, Dst: R1, Src: R3, Offset: -1, Reference: "foo",
-		}},
+			OpCode: 0xce, Dst: R1, Src: R3, Offset: -1,
+		}.WithReference("foo")},
 	}
 
 	for _, tc := range testcases {
-		if tc.have != tc.want {
+		if !tc.have.equal(tc.want) {
 			t.Errorf("%s: have %v, want %v", tc.name, tc.have, tc.want)
 		}
 	}

asm/instruction.go

@@ -35,17 +35,39 @@ type Instruction struct {
 	Offset   int16
 	Constant int64
 
-	// Reference denotes a reference (e.g. a jump) to another symbol.
-	Reference string
+	// Metadata contains optional metadata about this instruction
+	metadata *metadata
+}
+
+// WithSymbol marks the Instruction as a Symbol, which other Instructions
+// can point to using corresponding calls to WithReference.
+func (ins Instruction) WithSymbol(name string) Instruction {
+	if (ins.metadata != nil && ins.metadata.symbol == name) ||
+		(ins.metadata == nil && name == "") {
+		return ins
+	}
 
-	// Symbol denotes an instruction at the start of a function body.
-	Symbol string
+	ins.metadata = ins.metadata.copy()
+	ins.metadata.symbol = name
+	return ins
 }
 
 // Sym creates a symbol.
+//
+// Deprecated: use WithSymbol instead.
 func (ins Instruction) Sym(name string) Instruction {
-	ins.Symbol = name
-	return ins
+	return ins.WithSymbol(name)
+}
+
+// Symbol returns the value ins has been marked with using WithSymbol,
+// otherwise returns an empty string. A symbol is often an Instruction
+// at the start of a function body.
+func (ins Instruction) Symbol() string {
+	if ins.metadata == nil {
+		return ""
+	}
+
+	return ins.metadata.symbol
 }
 
 // Unmarshal decodes a BPF instruction.
@@ -299,16 +321,63 @@ func (ins Instruction) Format(f fmt.State, c rune) {
 	}
 
 ref:
-	if ins.Reference != "" {
-		fmt.Fprintf(f, " <%s>", ins.Reference)
+	if ins.Reference() != "" {
+		fmt.Fprintf(f, " <%s>", ins.Reference())
 	}
 }
 
+func (ins Instruction) equal(other Instruction) bool {
+	return ins.OpCode == other.OpCode &&
+		ins.Dst == other.Dst &&
+		ins.Src == other.Src &&
+		ins.Offset == other.Offset &&
+		ins.Constant == other.Constant
+}
+
 // Size returns the amount of bytes ins would occupy in binary form.
 func (ins Instruction) Size() uint64 {
 	return uint64(InstructionSize * ins.OpCode.rawInstructions())
 }
 
+// WithReference makes ins reference another Symbol or map by name.
+func (ins Instruction) WithReference(ref string) Instruction {
+	if (ins.metadata != nil && ins.metadata.reference == ref) ||
+		(ins.metadata == nil && ref == "") {
+		return ins
+	}
+
+	ins.metadata = ins.metadata.copy()
+	ins.metadata.reference = ref
+	return ins
+}
+
+// Reference returns the Symbol or map name referenced by ins, if any.
+func (ins Instruction) Reference() string {
+	if ins.metadata == nil {
+		return ""
+	}
+
+	return ins.metadata.reference
+}
+
+// metadata holds metadata about an Instruction.
+type metadata struct {
+	// reference denotes a reference (e.g. a jump) to another symbol.
+	reference string
+	// symbol denotes an instruction at the start of a function body.
+	symbol string
+}
+
+// copy returns a copy of metadata.
+// Always returns a valid pointer, even when called on a nil metadata.
+func (m *metadata) copy() *metadata {
+	var copy metadata
+	if m != nil {
+		copy = *m
+	}
+	return &copy
+}
+
 // Instructions is an eBPF program.
 type Instructions []Instruction
 
@@ -342,7 +411,7 @@ func (insns Instructions) Name() string {
 	if len(insns) == 0 {
 		return ""
 	}
-	return insns[0].Symbol
+	return insns[0].Symbol()
 }
 
 func (insns Instructions) String() string {
@@ -369,7 +438,7 @@ func (insns Instructions) RewriteMapPtr(symbol string, fd int) error {
 	found := false
 	for i := range insns {
 		ins := &insns[i]
-		if ins.Reference != symbol {
+		if ins.Reference() != symbol {
 			continue
 		}
 
@@ -393,15 +462,15 @@ func (insns Instructions) SymbolOffsets() (map[string]int, error) {
 	offsets := make(map[string]int)
 
 	for i, ins := range insns {
-		if ins.Symbol == "" {
+		if ins.Symbol() == "" {
 			continue
 		}
 
-		if _, ok := offsets[ins.Symbol]; ok {
-			return nil, fmt.Errorf("duplicate symbol %s", ins.Symbol)
+		if _, ok := offsets[ins.Symbol()]; ok {
+			return nil, fmt.Errorf("duplicate symbol %s", ins.Symbol())
 		}
 
-		offsets[ins.Symbol] = i
+		offsets[ins.Symbol()] = i
 	}
 
 	return offsets, nil
@@ -418,15 +487,15 @@ func (insns Instructions) FunctionReferences() map[string]bool {
 			continue
 		}
 
-		if ins.Reference == "" {
+		if ins.Reference() == "" {
 			continue
 		}
 
 		if !ins.IsFunctionReference() {
 			continue
 		}
 
-		calls[ins.Reference] = true
+		calls[ins.Reference()] = true
 	}
 
 	return calls
@@ -438,11 +507,11 @@ func (insns Instructions) ReferenceOffsets() map[string][]int {
 	offsets := make(map[string][]int)
 
 	for i, ins := range insns {
-		if ins.Reference == "" {
+		if ins.Reference() == "" {
 			continue
 		}
 
-		offsets[ins.Reference] = append(offsets[ins.Reference], i)
+		offsets[ins.Reference()] = append(offsets[ins.Reference()], i)
 	}
 
 	return offsets
@@ -493,8 +562,8 @@ func (insns Instructions) Format(f fmt.State, c rune) {
 
 	iter := insns.Iterate()
 	for iter.Next() {
-		if iter.Ins.Symbol != "" {
-			fmt.Fprintf(f, "%s%s:\n", symIndent, iter.Ins.Symbol)
+		if iter.Ins.Symbol() != "" {
+			fmt.Fprintf(f, "%s%s:\n", symIndent, iter.Ins.Symbol())
 		}
 		fmt.Fprintf(f, "%s%*d: %v\n", indent, offsetWidth, iter.Offset, iter.Ins)
 	}
@@ -552,15 +621,15 @@ func (insns Instructions) resolveFunctionReferences() error {
 	for iter.Next() {
 		ins := iter.Ins
 
-		if ins.Symbol == "" {
+		if ins.Symbol() == "" {
 			continue
 		}
 
-		if _, ok := symbolOffsets[ins.Symbol]; ok {
-			return fmt.Errorf("duplicate symbol %s", ins.Symbol)
+		if _, ok := symbolOffsets[ins.Symbol()]; ok {
+			return fmt.Errorf("duplicate symbol %s", ins.Symbol())
 		}
 
-		symbolOffsets[ins.Symbol] = iter.Offset
+		symbolOffsets[ins.Symbol()] = iter.Offset
 	}
 
 	// Find all instructions tagged as references to other symbols.
@@ -572,23 +641,23 @@ func (insns Instructions) resolveFunctionReferences() error {
 		offset := iter.Offset
 		ins := iter.Ins
 
-		if ins.Reference == "" {
+		if ins.Reference() == "" {
 			continue
 		}
 
 		switch {
 		case ins.IsFunctionReference() && ins.Constant == -1:
-			symOffset, ok := symbolOffsets[ins.Reference]
+			symOffset, ok := symbolOffsets[ins.Reference()]
 			if !ok {
-				return fmt.Errorf("%s at insn %d: symbol %q: %w", ins.OpCode, i, ins.Reference, ErrUnsatisfiedProgramReference)
+				return fmt.Errorf("%s at insn %d: symbol %q: %w", ins.OpCode, i, ins.Reference(), ErrUnsatisfiedProgramReference)
 			}
 
 			ins.Constant = int64(symOffset - offset - 1)
 
 		case ins.OpCode.Class().IsJump() && ins.Offset == -1:
-			symOffset, ok := symbolOffsets[ins.Reference]
+			symOffset, ok := symbolOffsets[ins.Reference()]
 			if !ok {
-				return fmt.Errorf("%s at insn %d: symbol %q: %w", ins.OpCode, i, ins.Reference, ErrUnsatisfiedProgramReference)
+				return fmt.Errorf("%s at insn %d: symbol %q: %w", ins.OpCode, i, ins.Reference(), ErrUnsatisfiedProgramReference)
 			}
 
 			ins.Offset = int16(symOffset - offset - 1)

asm/instruction_test.go

@@ -151,10 +151,9 @@ func TestInstructionLoadMapValue(t *testing.T) {
 
 func TestInstructionsRewriteMapPtr(t *testing.T) {
 	insns := Instructions{
-		LoadMapPtr(R1, 0),
+		LoadMapPtr(R1, 0).WithReference("good"),
 		Return(),
 	}
-	insns[0].Reference = "good"
 
 	if err := insns.RewriteMapPtr("good", 1); err != nil {
 		t.Fatal(err)
@@ -181,7 +180,7 @@ func TestInstructionsRewriteMapPtr(t *testing.T) {
 // program is stringified.
 func ExampleInstructions_Format() {
 	insns := Instructions{
-		FnMapLookupElem.Call().Sym("my_func"),
+		FnMapLookupElem.Call().WithSymbol("my_func"),
 		LoadImm(R0, 42, DWord),
 		Return(),
 	}

asm/jump.go

@@ -63,47 +63,43 @@ func (op JumpOp) Op(source Source) OpCode {
 // Imm compares 64 bit dst to 64 bit value (sign extended), and adjusts PC by offset if the condition is fulfilled.
 func (op JumpOp) Imm(dst Register, value int32, label string) Instruction {
 	return Instruction{
-		OpCode:    op.opCode(JumpClass, ImmSource),
-		Dst:       dst,
-		Offset:    -1,
-		Constant:  int64(value),
-		Reference: label,
-	}
+		OpCode:   op.opCode(JumpClass, ImmSource),
+		Dst:      dst,
+		Offset:   -1,
+		Constant: int64(value),
+	}.WithReference(label)
 }
 
 // Imm32 compares 32 bit dst to 32 bit value, and adjusts PC by offset if the condition is fulfilled.
 // Requires kernel 5.1.
 func (op JumpOp) Imm32(dst Register, value int32, label string) Instruction {
 	return Instruction{
-		OpCode:    op.opCode(Jump32Class, ImmSource),
-		Dst:       dst,
-		Offset:    -1,
-		Constant:  int64(value),
-		Reference: label,
-	}
+		OpCode:   op.opCode(Jump32Class, ImmSource),
+		Dst:      dst,
+		Offset:   -1,
+		Constant: int64(value),
+	}.WithReference(label)
 }
 
 // Reg compares 64 bit dst to 64 bit src, and adjusts PC by offset if the condition is fulfilled.
 func (op JumpOp) Reg(dst, src Register, label string) Instruction {
 	return Instruction{
-		OpCode:    op.opCode(JumpClass, RegSource),
-		Dst:       dst,
-		Src:       src,
-		Offset:    -1,
-		Reference: label,
-	}
+		OpCode: op.opCode(JumpClass, RegSource),
+		Dst:    dst,
+		Src:    src,
+		Offset: -1,
+	}.WithReference(label)
 }
 
 // Reg32 compares 32 bit dst to 32 bit src, and adjusts PC by offset if the condition is fulfilled.
 // Requires kernel 5.1.
 func (op JumpOp) Reg32(dst, src Register, label string) Instruction {
 	return Instruction{
-		OpCode:    op.opCode(Jump32Class, RegSource),
-		Dst:       dst,
-		Src:       src,
-		Offset:    -1,
-		Reference: label,
-	}
+		OpCode: op.opCode(Jump32Class, RegSource),
+		Dst:    dst,
+		Src:    src,
+		Offset: -1,
+	}.WithReference(label)
 }
 
 func (op JumpOp) opCode(class Class, source Source) OpCode {
@@ -118,16 +114,14 @@ func (op JumpOp) opCode(class Class, source Source) OpCode {
 func (op JumpOp) Label(label string) Instruction {
 	if op == Call {
 		return Instruction{
-			OpCode:    OpCode(JumpClass).SetJumpOp(Call),
-			Src:       PseudoCall,
-			Constant:  -1,
-			Reference: label,
-		}
+			OpCode:   OpCode(JumpClass).SetJumpOp(Call),
+			Src:      PseudoCall,
+			Constant: -1,
+		}.WithReference(label)
 	}
 
 	return Instruction{
-		OpCode:    OpCode(JumpClass).SetJumpOp(op),
-		Offset:    -1,
-		Reference: label,
-	}
+		OpCode: OpCode(JumpClass).SetJumpOp(op),
+		Offset: -1,
+	}.WithReference(label)
 }