coverage: Store a graph reference in the graph traversal struct

Having to keep passing in a graph reference was a holdover from when the graph
was partly mutated during traversal. As of rust-lang#114354 that is no longer necessary,
so we can simplify the traversal code by storing a graph reference as a field
in `TraverseCoverageGraphWithLoops`.

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/counters.rs

@@ -245,13 +245,13 @@ impl<'a> MakeBcbCounters<'a> {
         // the loop. The `traversal` state includes a `context_stack`, providing a way to know if
         // the current BCB is in one or more nested loops or not.
         let mut traversal = TraverseCoverageGraphWithLoops::new(&self.basic_coverage_blocks);
-        while let Some(bcb) = traversal.next(self.basic_coverage_blocks) {
+        while let Some(bcb) = traversal.next() {
             if bcb_has_coverage_spans(bcb) {
                 debug!("{:?} has at least one coverage span. Get or make its counter", bcb);
                 let branching_counter_operand = self.get_or_make_counter_operand(bcb)?;
 
                 if self.bcb_needs_branch_counters(bcb) {
-                    self.make_branch_counters(&mut traversal, bcb, branching_counter_operand)?;
+                    self.make_branch_counters(&traversal, bcb, branching_counter_operand)?;
                 }
             } else {
                 debug!(
@@ -274,7 +274,7 @@ impl<'a> MakeBcbCounters<'a> {
 
     fn make_branch_counters(
         &mut self,
-        traversal: &mut TraverseCoverageGraphWithLoops,
+        traversal: &TraverseCoverageGraphWithLoops<'_>,
         branching_bcb: BasicCoverageBlock,
         branching_counter_operand: Operand,
     ) -> Result<(), Error> {
@@ -507,7 +507,7 @@ impl<'a> MakeBcbCounters<'a> {
     /// found, select any branch.
     fn choose_preferred_expression_branch(
         &self,
-        traversal: &TraverseCoverageGraphWithLoops,
+        traversal: &TraverseCoverageGraphWithLoops<'_>,
         branches: &[BcbBranch],
     ) -> BcbBranch {
         let good_reloop_branch = self.find_good_reloop_branch(traversal, &branches);
@@ -537,7 +537,7 @@ impl<'a> MakeBcbCounters<'a> {
     /// will tend to be executed more times than a loop-exit branch.
     fn find_good_reloop_branch(
         &self,
-        traversal: &TraverseCoverageGraphWithLoops,
+        traversal: &TraverseCoverageGraphWithLoops<'_>,
         branches: &[BcbBranch],
     ) -> Option<BcbBranch> {
         // Consider each loop on the current traversal context stack, top-down.

compiler/rustc_mir_transform/src/coverage/graph.rs

@@ -393,14 +393,16 @@ pub(super) struct TraversalContext {
     worklist: VecDeque<BasicCoverageBlock>,
 }
 
-pub(super) struct TraverseCoverageGraphWithLoops {
+pub(super) struct TraverseCoverageGraphWithLoops<'a> {
+    basic_coverage_blocks: &'a CoverageGraph,
+
     backedges: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>,
     context_stack: Vec<TraversalContext>,
     visited: BitSet<BasicCoverageBlock>,
 }
 
-impl TraverseCoverageGraphWithLoops {
-    pub fn new(basic_coverage_blocks: &CoverageGraph) -> Self {
+impl<'a> TraverseCoverageGraphWithLoops<'a> {
+    pub(super) fn new(basic_coverage_blocks: &'a CoverageGraph) -> Self {
         let backedges = find_loop_backedges(basic_coverage_blocks);
 
         let worklist = VecDeque::from([basic_coverage_blocks.start_node()]);
@@ -411,7 +413,7 @@ impl TraverseCoverageGraphWithLoops {
         // of the stack as loops are entered, and popped off of the stack when a loop's worklist is
         // exhausted.
         let visited = BitSet::new_empty(basic_coverage_blocks.num_nodes());
-        Self { backedges, context_stack, visited }
+        Self { basic_coverage_blocks, backedges, context_stack, visited }
     }
 
     /// For each loop on the loop context stack (top-down), yields a list of BCBs
@@ -424,7 +426,7 @@ impl TraverseCoverageGraphWithLoops {
             .map(|(from_bcbs, _to_bcb)| from_bcbs.as_slice())
     }
 
-    pub fn next(&mut self, basic_coverage_blocks: &CoverageGraph) -> Option<BasicCoverageBlock> {
+    pub(super) fn next(&mut self) -> Option<BasicCoverageBlock> {
         debug!(
             "TraverseCoverageGraphWithLoops::next - context_stack: {:?}",
             self.context_stack.iter().rev().collect::<Vec<_>>()
@@ -445,7 +447,7 @@ impl TraverseCoverageGraphWithLoops {
                         worklist: VecDeque::new(),
                     });
                 }
-                self.extend_worklist(basic_coverage_blocks, bcb);
+                self.extend_worklist(bcb);
                 return Some(bcb);
             } else {
                 // Strip contexts with empty worklists from the top of the stack
@@ -456,11 +458,8 @@ impl TraverseCoverageGraphWithLoops {
         None
     }
 
-    pub fn extend_worklist(
-        &mut self,
-        basic_coverage_blocks: &CoverageGraph,
-        bcb: BasicCoverageBlock,
-    ) {
+    pub fn extend_worklist(&mut self, bcb: BasicCoverageBlock) {
+        let Self { basic_coverage_blocks, .. } = *self;
         let successors = &basic_coverage_blocks.successors[bcb];
         debug!("{:?} has {} successors:", bcb, successors.len());
         for &successor in successors {

compiler/rustc_mir_transform/src/coverage/tests.rs

@@ -628,7 +628,7 @@ fn test_traverse_coverage_with_loops() {
     let basic_coverage_blocks = graph::CoverageGraph::from_mir(&mir_body);
     let mut traversed_in_order = Vec::new();
     let mut traversal = graph::TraverseCoverageGraphWithLoops::new(&basic_coverage_blocks);
-    while let Some(bcb) = traversal.next(&basic_coverage_blocks) {
+    while let Some(bcb) = traversal.next() {
         traversed_in_order.push(bcb);
     }