Moved Instruction.Symbol and Instruction.Reference into per-instruction metadata

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)
 }