coverage: Don't rely on the custom traversal to find enclosing loops

Zalathar · Zalathar · commit 19b2142d1880 · 2024-10-26T09:38:47.000+11:00
diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs
@@ -280,13 +280,12 @@ impl<'a> CountersBuilder<'a> {
         //
         // The traversal tries to ensure that, when a loop is encountered, all
         // nodes within the loop are visited before visiting any nodes outside
-        // the loop. It also keeps track of which loop(s) the traversal is
-        // currently inside.
+        // the loop.
         let mut traversal = TraverseCoverageGraphWithLoops::new(self.graph);
         while let Some(bcb) = traversal.next() {
             let _span = debug_span!("traversal", ?bcb).entered();
             if self.bcb_needs_counter.contains(bcb) {
-                self.make_node_counter_and_out_edge_counters(&traversal, bcb);
+                self.make_node_counter_and_out_edge_counters(bcb);
             }
         }
 
@@ -299,12 +298,8 @@ impl<'a> CountersBuilder<'a> {
 
     /// Make sure the given node has a node counter, and then make sure each of
     /// its out-edges has an edge counter (if appropriate).
-    #[instrument(level = "debug", skip(self, traversal))]
-    fn make_node_counter_and_out_edge_counters(
-        &mut self,
-        traversal: &TraverseCoverageGraphWithLoops<'_>,
-        from_bcb: BasicCoverageBlock,
-    ) {
+    #[instrument(level = "debug", skip(self))]
+    fn make_node_counter_and_out_edge_counters(&mut self, from_bcb: BasicCoverageBlock) {
         // First, ensure that this node has a counter of some kind.
         // We might also use that counter to compute one of the out-edge counters.
         let node_counter = self.get_or_make_node_counter(from_bcb);
@@ -337,8 +332,7 @@ impl<'a> CountersBuilder<'a> {
 
         // If there are out-edges without counters, choose one to be given an expression
         // (computed from this node and the other out-edges) instead of a physical counter.
-        let Some(target_bcb) =
-            self.choose_out_edge_for_expression(traversal, &candidate_successors)
+        let Some(target_bcb) = self.choose_out_edge_for_expression(from_bcb, &candidate_successors)
         else {
             return;
         };
@@ -462,12 +456,12 @@ impl<'a> CountersBuilder<'a> {
     /// choose one to be given a counter expression instead of a physical counter.
     fn choose_out_edge_for_expression(
         &self,
-        traversal: &TraverseCoverageGraphWithLoops<'_>,
+        from_bcb: BasicCoverageBlock,
         candidate_successors: &[BasicCoverageBlock],
     ) -> Option<BasicCoverageBlock> {
         // Try to find a candidate that leads back to the top of a loop,
         // because reloop edges tend to be executed more times than loop-exit edges.
-        if let Some(reloop_target) = self.find_good_reloop_edge(traversal, &candidate_successors) {
+        if let Some(reloop_target) = self.find_good_reloop_edge(from_bcb, &candidate_successors) {
             debug!("Selecting reloop target {reloop_target:?} to get an expression");
             return Some(reloop_target);
         }
@@ -486,7 +480,7 @@ impl<'a> CountersBuilder<'a> {
     /// for them to be able to avoid a physical counter increment.
     fn find_good_reloop_edge(
         &self,
-        traversal: &TraverseCoverageGraphWithLoops<'_>,
+        from_bcb: BasicCoverageBlock,
         candidate_successors: &[BasicCoverageBlock],
     ) -> Option<BasicCoverageBlock> {
         // If there are no candidates, avoid iterating over the loop stack.
@@ -495,14 +489,15 @@ impl<'a> CountersBuilder<'a> {
         }
 
         // Consider each loop on the current traversal context stack, top-down.
-        for reloop_bcbs in traversal.reloop_bcbs_per_loop() {
+        for loop_header_node in self.graph.loop_headers_containing(from_bcb) {
             // Try to find a candidate edge that doesn't exit this loop.
             for &target_bcb in candidate_successors {
                 // An edge is a reloop edge if its target dominates any BCB that has
                 // an edge back to the loop header. (Otherwise it's an exit edge.)
-                let is_reloop_edge = reloop_bcbs
-                    .iter()
-                    .any(|&reloop_bcb| self.graph.dominates(target_bcb, reloop_bcb));
+                let is_reloop_edge = self
+                    .graph
+                    .reloop_predecessors(loop_header_node)
+                    .any(|reloop_bcb| self.graph.dominates(target_bcb, reloop_bcb));
                 if is_reloop_edge {
                     // We found a good out-edge to be given an expression.
                     return Some(target_bcb);
diff --git a/compiler/rustc_mir_transform/src/coverage/graph.rs b/compiler/rustc_mir_transform/src/coverage/graph.rs
@@ -1,10 +1,11 @@
 use std::cmp::Ordering;
 use std::collections::VecDeque;
+use std::iter;
 use std::ops::{Index, IndexMut};
 
 use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::FxHashSet;
-use rustc_data_structures::graph::dominators::{self, Dominators};
+use rustc_data_structures::graph::dominators::Dominators;
 use rustc_data_structures::graph::{self, DirectedGraph, StartNode};
 use rustc_index::IndexVec;
 use rustc_index::bit_set::BitSet;
@@ -20,11 +21,17 @@ pub(crate) struct CoverageGraph {
     bb_to_bcb: IndexVec<BasicBlock, Option<BasicCoverageBlock>>,
     pub(crate) successors: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>,
     pub(crate) predecessors: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>,
+
     dominators: Option<Dominators<BasicCoverageBlock>>,
     /// Allows nodes to be compared in some total order such that _if_
     /// `a` dominates `b`, then `a < b`. If neither node dominates the other,
     /// their relative order is consistent but arbitrary.
     dominator_order_rank: IndexVec<BasicCoverageBlock, u32>,
+    /// A loop header is a node that dominates one or more of its predecessors.
+    is_loop_header: BitSet<BasicCoverageBlock>,
+    /// For each node, the loop header node of its nearest enclosing loop.
+    /// This forms a linked list that can be traversed to find all enclosing loops.
+    enclosing_loop_header: IndexVec<BasicCoverageBlock, Option<BasicCoverageBlock>>,
 }
 
 impl CoverageGraph {
@@ -66,17 +73,38 @@ impl CoverageGraph {
             predecessors,
             dominators: None,
             dominator_order_rank: IndexVec::from_elem_n(0, num_nodes),
+            is_loop_header: BitSet::new_empty(num_nodes),
+            enclosing_loop_header: IndexVec::from_elem_n(None, num_nodes),
         };
         assert_eq!(num_nodes, this.num_nodes());
 
-        this.dominators = Some(dominators::dominators(&this));
+        // Set the dominators first, because later init steps rely on them.
+        this.dominators = Some(graph::dominators::dominators(&this));
 
-        // The dominator rank of each node is just its index in a reverse-postorder traversal.
-        let reverse_post_order = graph::iterate::reverse_post_order(&this, this.start_node());
+        // Iterate over all nodes, such that dominating nodes are visited before
+        // the nodes they dominate. Either preorder or reverse postorder is fine.
+        let dominator_order = graph::iterate::reverse_post_order(&this, this.start_node());
         // The coverage graph is created by traversal, so all nodes are reachable.
-        assert_eq!(reverse_post_order.len(), this.num_nodes());
-        for (rank, bcb) in (0u32..).zip(reverse_post_order) {
+        assert_eq!(dominator_order.len(), this.num_nodes());
+        for (rank, bcb) in (0u32..).zip(dominator_order) {
+            // The dominator rank of each node is its index in a dominator-order traversal.
             this.dominator_order_rank[bcb] = rank;
+
+            // A node is a loop header if it dominates any of its predecessors.
+            if this.reloop_predecessors(bcb).next().is_some() {
+                this.is_loop_header.insert(bcb);
+            }
+
+            // If the immediate dominator is a loop header, that's our enclosing loop.
+            // Otherwise, inherit the immediate dominator's enclosing loop.
+            // (Dominator order ensures that we already processed the dominator.)
+            if let Some(dom) = this.dominators().immediate_dominator(bcb) {
+                this.enclosing_loop_header[bcb] = this
+                    .is_loop_header
+                    .contains(dom)
+                    .then_some(dom)
+                    .or_else(|| this.enclosing_loop_header[dom]);
+            }
         }
 
         // The coverage graph's entry-point node (bcb0) always starts with bb0,
@@ -172,9 +200,14 @@ impl CoverageGraph {
         if bb.index() < self.bb_to_bcb.len() { self.bb_to_bcb[bb] } else { None }
     }
 
+    #[inline(always)]
+    fn dominators(&self) -> &Dominators<BasicCoverageBlock> {
+        self.dominators.as_ref().unwrap()
+    }
+
     #[inline(always)]
     pub(crate) fn dominates(&self, dom: BasicCoverageBlock, node: BasicCoverageBlock) -> bool {
-        self.dominators.as_ref().unwrap().dominates(dom, node)
+        self.dominators().dominates(dom, node)
     }
 
     #[inline(always)]
@@ -214,6 +247,36 @@ impl CoverageGraph {
             None
         }
     }
+
+    /// For each loop that contains the given node, yields the "loop header"
+    /// node representing that loop, from innermost to outermost. If the given
+    /// node is itself a loop header, it is yielded first.
+    pub(crate) fn loop_headers_containing(
+        &self,
+        bcb: BasicCoverageBlock,
+    ) -> impl Iterator<Item = BasicCoverageBlock> + Captures<'_> {
+        let self_if_loop_header = self.is_loop_header.contains(bcb).then_some(bcb).into_iter();
+
+        let mut curr = Some(bcb);
+        let strictly_enclosing = iter::from_fn(move || {
+            let enclosing = self.enclosing_loop_header[curr?];
+            curr = enclosing;
+            enclosing
+        });
+
+        self_if_loop_header.chain(strictly_enclosing)
+    }
+
+    /// For the given node, yields the subset of its predecessor nodes that
+    /// it dominates. If that subset is non-empty, the node is a "loop header",
+    /// and each of those predecessors represents an in-edge that jumps back to
+    /// the top of its loop.
+    pub(crate) fn reloop_predecessors(
+        &self,
+        to_bcb: BasicCoverageBlock,
+    ) -> impl Iterator<Item = BasicCoverageBlock> + Captures<'_> {
+        self.predecessors[to_bcb].iter().copied().filter(move |&pred| self.dominates(to_bcb, pred))
+    }
 }
 
 impl Index<BasicCoverageBlock> for CoverageGraph {
@@ -439,15 +502,12 @@ struct TraversalContext {
 pub(crate) struct TraverseCoverageGraphWithLoops<'a> {
     basic_coverage_blocks: &'a CoverageGraph,
 
-    backedges: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>,
     context_stack: Vec<TraversalContext>,
     visited: BitSet<BasicCoverageBlock>,
 }
 
 impl<'a> TraverseCoverageGraphWithLoops<'a> {
     pub(crate) 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()]);
         let context_stack = vec![TraversalContext { loop_header: None, worklist }];
 
@@ -456,17 +516,7 @@ impl<'a> TraverseCoverageGraphWithLoops<'a> {
         // 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 { basic_coverage_blocks, backedges, context_stack, visited }
-    }
-
-    /// For each loop on the loop context stack (top-down), yields a list of BCBs
-    /// within that loop that have an outgoing edge back to the loop header.
-    pub(crate) fn reloop_bcbs_per_loop(&self) -> impl Iterator<Item = &[BasicCoverageBlock]> {
-        self.context_stack
-            .iter()
-            .rev()
-            .filter_map(|context| context.loop_header)
-            .map(|header_bcb| self.backedges[header_bcb].as_slice())
+        Self { basic_coverage_blocks, context_stack, visited }
     }
 
     pub(crate) fn next(&mut self) -> Option<BasicCoverageBlock> {
@@ -488,7 +538,7 @@ impl<'a> TraverseCoverageGraphWithLoops<'a> {
             }
             debug!("Visiting {bcb:?}");
 
-            if self.backedges[bcb].len() > 0 {
+            if self.basic_coverage_blocks.is_loop_header.contains(bcb) {
                 debug!("{bcb:?} is a loop header! Start a new TraversalContext...");
                 self.context_stack
                     .push(TraversalContext { loop_header: Some(bcb), worklist: VecDeque::new() });
@@ -566,29 +616,6 @@ impl<'a> TraverseCoverageGraphWithLoops<'a> {
     }
 }
 
-fn find_loop_backedges(
-    basic_coverage_blocks: &CoverageGraph,
-) -> IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>> {
-    let num_bcbs = basic_coverage_blocks.num_nodes();
-    let mut backedges = IndexVec::from_elem_n(Vec::<BasicCoverageBlock>::new(), num_bcbs);
-
-    // Identify loops by their backedges.
-    for (bcb, _) in basic_coverage_blocks.iter_enumerated() {
-        for &successor in &basic_coverage_blocks.successors[bcb] {
-            if basic_coverage_blocks.dominates(successor, bcb) {
-                let loop_header = successor;
-                let backedge_from_bcb = bcb;
-                debug!(
-                    "Found BCB backedge: {:?} -> loop_header: {:?}",
-                    backedge_from_bcb, loop_header
-                );
-                backedges[loop_header].push(backedge_from_bcb);
-            }
-        }
-    }
-    backedges
-}
-
 fn short_circuit_preorder<'a, 'tcx, F, Iter>(
     body: &'a mir::Body<'tcx>,
     filtered_successors: F,
