diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs index 9a533ea024d..cf9f6981c5a 100644 --- a/compiler/rustc_mir_transform/src/coverage/counters.rs +++ b/compiler/rustc_mir_transform/src/coverage/counters.rs @@ -280,13 +280,12 @@ fn make_bcb_counters(&mut self) { // // 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 @@ fn make_bcb_counters(&mut self) { /// 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 @@ fn make_node_counter_and_out_edge_counters( // 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 @@ fn make_edge_counter_inner( /// 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 { // 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 @@ fn choose_out_edge_for_expression( /// 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 { // If there are no candidates, avoid iterating over the loop stack. @@ -495,14 +489,15 @@ fn find_good_reloop_edge( } // 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 index 930fa129ef2..168262bf01c 100644 --- 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>, pub(crate) successors: IndexVec>, pub(crate) predecessors: IndexVec>, + dominators: Option>, /// 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, + /// A loop header is a node that dominates one or more of its predecessors. + is_loop_header: BitSet, + /// 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>, } impl CoverageGraph { @@ -66,17 +73,38 @@ pub(crate) fn from_mir(mir_body: &mir::Body<'_>) -> Self { 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 @@ pub(crate) fn bcb_from_bb(&self, bb: BasicBlock) -> Option { if bb.index() < self.bb_to_bcb.len() { self.bb_to_bcb[bb] } else { None } } + #[inline(always)] + fn dominators(&self) -> &Dominators { + 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 @@ pub(crate) fn simple_successor( 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 + 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 + Captures<'_> { + self.predecessors[to_bcb].iter().copied().filter(move |&pred| self.dominates(to_bcb, pred)) + } } impl Index for CoverageGraph { @@ -439,15 +502,12 @@ struct TraversalContext { pub(crate) struct TraverseCoverageGraphWithLoops<'a> { basic_coverage_blocks: &'a CoverageGraph, - backedges: IndexVec>, context_stack: Vec, visited: BitSet, } 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 @@ pub(crate) fn new(basic_coverage_blocks: &'a CoverageGraph) -> Self { // 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 { - 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 { @@ -488,7 +538,7 @@ pub(crate) fn next(&mut self) -> Option { } 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 @@ pub(crate) fn unvisited(&self) -> Vec { } } -fn find_loop_backedges( - basic_coverage_blocks: &CoverageGraph, -) -> IndexVec> { - let num_bcbs = basic_coverage_blocks.num_nodes(); - let mut backedges = IndexVec::from_elem_n(Vec::::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,