BPF optimizer: Refactor removal operations across optimizers.

This introduces a `codeRemoval` type which acts as a tracker for removal of
instructions during execution of an optimizer.

This simplifies optimizers that use it, notably to avoid having to do
backwards iterations in all of them.

No heap allocations are performed unless code is actually removed, and only
one allocation is performed when code is to be removed.

PiperOrigin-RevId: 595575409

pkg/bpf/optimizer.go

@@ -136,6 +136,35 @@ func optimizeUnconditionalJumps(insns []Instruction) ([]Instruction, bool) {
 	return insns, changed
 }
 
+// codeRemoval efficiently tracks indexes to remove from instructions.
+type codeRemoval struct {
+	insns    []Instruction
+	toRemove []int
+}
+
+// MarkRemoved adds a new instruction index to be removed.
+func (cr *codeRemoval) MarkRemoved(index int) {
+	if cr.toRemove == nil {
+		cr.toRemove = make([]int, 0, len(cr.insns))
+	}
+	cr.toRemove = append(cr.toRemove, index)
+}
+
+// Apply returns the set of instructions after removing marked indexes,
+// along with a boolean representing whether any instruction was removed.
+func (cr *codeRemoval) Apply() ([]Instruction, bool) {
+	if len(cr.toRemove) == 0 {
+		return cr.insns, false
+	}
+	sort.Ints(cr.toRemove)
+	for i := len(cr.toRemove) - 1; i >= 0; i-- {
+		pc := cr.toRemove[i]
+		cr.insns = append(cr.insns[:pc], cr.insns[pc+1:]...)
+		decrementJumps(cr.insns, pc)
+	}
+	return cr.insns, true
+}
+
 // decrementJumps decrements all jumps within `insns` that are jumping to an
 // instruction with index larger than `target`, the index of an
 // instruction that just got removed (i.e. `target` now points to the
@@ -172,21 +201,14 @@ func decrementJumps(insns []Instruction, target int) {
 // instructions forward. This may seem silly but it can happen due to other
 // optimizations in this file which decrement jump target indexes.
 func removeZeroInstructionJumps(insns []Instruction) ([]Instruction, bool) {
-	changed := false
-	for pc := 0; pc < len(insns); pc++ {
-		ins := insns[pc]
+	removal := codeRemoval{insns: insns}
+	for pc, ins := range insns {
 		if !ins.IsUnconditionalJump() || ins.K != 0 {
 			continue
 		}
-		insns = append(insns[:pc], insns[pc+1:]...)
-		decrementJumps(insns, pc)
-		changed = true
-
-		// Rewind back one instruction, in case the instruction now at `pc`
-		// is also a zero-instruction unconditional jump.
-		pc--
+		removal.MarkRemoved(pc)
 	}
-	return insns, changed
+	return removal.Apply()
 }
 
 // removeDeadCode removes instructions which are unreachable.
@@ -229,30 +251,14 @@ func removeDeadCode(insns []Instruction) ([]Instruction, bool) {
 		}
 	}
 
-	// Now scan for unreachable code.
-	var unreachable []int
-	for i := 0; i < len(reachable); i++ {
-		if !reachable[i] {
-			unreachable = append(unreachable, i)
+	// Now remove unreachable code.
+	removal := codeRemoval{insns: insns}
+	for pc := range insns {
+		if !reachable[pc] {
+			removal.MarkRemoved(pc)
 		}
 	}
-
-	// And finally cull unreachable code.
-	for u := 0; u < len(unreachable); u++ {
-		i := unreachable[u]
-		// Remove the instruction at this index:
-		insns = append(insns[:i], insns[i+1:]...)
-
-		// Rewrite all previous jumps which would have straddled over this instruction:
-		decrementJumps(insns, i)
-
-		// And decrement all future unreachable indexes, since we just shortened `insns` by one:
-		for u2 := u + 1; u2 < len(unreachable); u2++ {
-			unreachable[u2] = unreachable[u2] - 1
-		}
-	}
-
-	return insns, len(unreachable) > 0
+	return removal.Apply()
 }
 
 // optimizeJumpsToReturn replaces unconditional jumps that go to return
@@ -298,13 +304,8 @@ func removeRedundantLoads(insns []Instruction) ([]Instruction, bool) {
 	}
 
 	// Now look for redundant load instructions.
-	// We iterate backwards here so that we can remove instructions as we go
-	// without a past iteration interfering with the results of the current
-	// iteration. This is relying on the property that BPF programs may only
-	// flow forwards (no backwards jumps).
-	changed := false
-	for pc := len(insns) - 1; pc >= 0; pc-- {
-		ins := insns[pc]
+	removal := codeRemoval{insns: insns}
+	for pc, ins := range insns {
 		if ins.OpCode&instructionClassMask != Ld {
 			continue
 		}
@@ -336,12 +337,10 @@ func removeRedundantLoads(insns []Instruction) ([]Instruction, bool) {
 			}
 		}
 		if lastModifiedA != -1 && insns[pc].Equal(insns[lastModifiedA]) {
-			insns = append(insns[:pc], insns[pc+1:]...)
-			decrementJumps(insns, pc)
-			changed = true
+			removal.MarkRemoved(pc)
 		}
 	}
-	return insns, changed
+	return removal.Apply()
 }
 
 // jumpRewriteOperation rewrites a jump target.