Auto merge of #132091 - Zalathar:graph-loops, r=saethlin

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

This opens up the possibility of modifying or removing the custom graph traversal used in coverage counter creation, without losing access to the heuristics that care about a node's enclosing loops.

Actually changing the traversal is left for future work, because this PR on its own doesn't change the emitted coverage mappings at all.

Author: bors

Diff for: 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 for: 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,