Merge pull request #18431 from lnicola/sync-from-rust

minor: Sync from downstream

1 files changed, 116 insertions, 131 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs
index 94088156756..cf9f6981c5a 100644
--- a/compiler/rustc_mir_transform/src/coverage/counters.rs
+++ b/compiler/rustc_mir_transform/src/coverage/counters.rs
@@ -5,7 +5,6 @@ use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::graph::DirectedGraph;
 use rustc_index::IndexVec;
 use rustc_index::bit_set::BitSet;
-use rustc_middle::bug;
 use rustc_middle::mir::coverage::{CounterId, CovTerm, Expression, ExpressionId, Op};
 use tracing::{debug, debug_span, instrument};
 
@@ -58,13 +57,13 @@ pub(super) struct CoverageCounters {
     counter_increment_sites: IndexVec<CounterId, CounterIncrementSite>,
 
     /// Coverage counters/expressions that are associated with individual BCBs.
-    bcb_counters: IndexVec<BasicCoverageBlock, Option<BcbCounter>>,
+    node_counters: IndexVec<BasicCoverageBlock, Option<BcbCounter>>,
     /// Coverage counters/expressions that are associated with the control-flow
     /// edge between two BCBs.
     ///
     /// We currently don't iterate over this map, but if we do in the future,
     /// switch it back to `FxIndexMap` to avoid query stability hazards.
-    bcb_edge_counters: FxHashMap<(BasicCoverageBlock, BasicCoverageBlock), BcbCounter>,
+    edge_counters: FxHashMap<(BasicCoverageBlock, BasicCoverageBlock), BcbCounter>,
 
     /// Table of expression data, associating each expression ID with its
     /// corresponding operator (+ or -) and its LHS/RHS operands.
@@ -78,20 +77,20 @@ impl CoverageCounters {
     /// Ensures that each BCB node needing a counter has one, by creating physical
     /// counters or counter expressions for nodes and edges as required.
     pub(super) fn make_bcb_counters(
-        basic_coverage_blocks: &CoverageGraph,
+        graph: &CoverageGraph,
         bcb_needs_counter: &BitSet<BasicCoverageBlock>,
     ) -> Self {
-        let mut counters = MakeBcbCounters::new(basic_coverage_blocks, bcb_needs_counter);
-        counters.make_bcb_counters();
+        let mut builder = CountersBuilder::new(graph, bcb_needs_counter);
+        builder.make_bcb_counters();
 
-        counters.coverage_counters
+        builder.counters
     }
 
     fn with_num_bcbs(num_bcbs: usize) -> Self {
         Self {
             counter_increment_sites: IndexVec::new(),
-            bcb_counters: IndexVec::from_elem_n(None, num_bcbs),
-            bcb_edge_counters: FxHashMap::default(),
+            node_counters: IndexVec::from_elem_n(None, num_bcbs),
+            edge_counters: FxHashMap::default(),
             expressions: IndexVec::new(),
             expressions_memo: FxHashMap::default(),
         }
@@ -104,24 +103,18 @@ impl CoverageCounters {
         BcbCounter::Counter { id }
     }
 
-    /// Creates a new physical counter attached a BCB node.
-    /// The node must not already have a counter.
+    /// Creates a new physical counter for a BCB node.
     fn make_phys_node_counter(&mut self, bcb: BasicCoverageBlock) -> BcbCounter {
-        let counter = self.make_counter_inner(CounterIncrementSite::Node { bcb });
-        debug!(?bcb, ?counter, "node gets a physical counter");
-        self.set_bcb_counter(bcb, counter)
+        self.make_counter_inner(CounterIncrementSite::Node { bcb })
     }
 
-    /// Creates a new physical counter attached to a BCB edge.
-    /// The edge must not already have a counter.
+    /// Creates a new physical counter for a BCB edge.
     fn make_phys_edge_counter(
         &mut self,
         from_bcb: BasicCoverageBlock,
         to_bcb: BasicCoverageBlock,
     ) -> BcbCounter {
-        let counter = self.make_counter_inner(CounterIncrementSite::Edge { from_bcb, to_bcb });
-        debug!(?from_bcb, ?to_bcb, ?counter, "edge gets a physical counter");
-        self.set_bcb_edge_counter(from_bcb, to_bcb, counter)
+        self.make_counter_inner(CounterIncrementSite::Edge { from_bcb, to_bcb })
     }
 
     fn make_expression(&mut self, lhs: BcbCounter, op: Op, rhs: BcbCounter) -> BcbCounter {
@@ -193,35 +186,31 @@ impl CoverageCounters {
         self.counter_increment_sites.len()
     }
 
-    fn set_bcb_counter(&mut self, bcb: BasicCoverageBlock, counter_kind: BcbCounter) -> BcbCounter {
-        if let Some(replaced) = self.bcb_counters[bcb].replace(counter_kind) {
-            bug!(
-                "attempt to set a BasicCoverageBlock coverage counter more than once; \
-                {bcb:?} already had counter {replaced:?}",
-            );
-        } else {
-            counter_kind
-        }
+    fn set_node_counter(&mut self, bcb: BasicCoverageBlock, counter: BcbCounter) -> BcbCounter {
+        let existing = self.node_counters[bcb].replace(counter);
+        assert!(
+            existing.is_none(),
+            "node {bcb:?} already has a counter: {existing:?} => {counter:?}"
+        );
+        counter
     }
 
-    fn set_bcb_edge_counter(
+    fn set_edge_counter(
         &mut self,
         from_bcb: BasicCoverageBlock,
         to_bcb: BasicCoverageBlock,
-        counter_kind: BcbCounter,
+        counter: BcbCounter,
     ) -> BcbCounter {
-        if let Some(replaced) = self.bcb_edge_counters.insert((from_bcb, to_bcb), counter_kind) {
-            bug!(
-                "attempt to set an edge counter more than once; from_bcb: \
-                {from_bcb:?} already had counter {replaced:?}",
-            );
-        } else {
-            counter_kind
-        }
+        let existing = self.edge_counters.insert((from_bcb, to_bcb), counter);
+        assert!(
+            existing.is_none(),
+            "edge ({from_bcb:?} -> {to_bcb:?}) already has a counter: {existing:?} => {counter:?}"
+        );
+        counter
     }
 
     pub(super) fn term_for_bcb(&self, bcb: BasicCoverageBlock) -> Option<CovTerm> {
-        self.bcb_counters[bcb].map(|counter| counter.as_term())
+        self.node_counters[bcb].map(|counter| counter.as_term())
     }
 
     /// Returns an iterator over all the nodes/edges in the coverage graph that
@@ -238,7 +227,7 @@ impl CoverageCounters {
     pub(super) fn bcb_nodes_with_coverage_expressions(
         &self,
     ) -> impl Iterator<Item = (BasicCoverageBlock, ExpressionId)> + Captures<'_> {
-        self.bcb_counters.iter_enumerated().filter_map(|(bcb, &counter_kind)| match counter_kind {
+        self.node_counters.iter_enumerated().filter_map(|(bcb, &counter)| match counter {
             // Yield the BCB along with its associated expression ID.
             Some(BcbCounter::Expression { id }) => Some((bcb, id)),
             // This BCB is associated with a counter or nothing, so skip it.
@@ -265,22 +254,20 @@ impl CoverageCounters {
     }
 }
 
-/// Helper struct that allows counter creation to inspect the BCB graph.
-struct MakeBcbCounters<'a> {
-    coverage_counters: CoverageCounters,
-    basic_coverage_blocks: &'a CoverageGraph,
+/// Helper struct that allows counter creation to inspect the BCB graph, and
+/// the set of nodes that need counters.
+struct CountersBuilder<'a> {
+    counters: CoverageCounters,
+    graph: &'a CoverageGraph,
     bcb_needs_counter: &'a BitSet<BasicCoverageBlock>,
 }
 
-impl<'a> MakeBcbCounters<'a> {
-    fn new(
-        basic_coverage_blocks: &'a CoverageGraph,
-        bcb_needs_counter: &'a BitSet<BasicCoverageBlock>,
-    ) -> Self {
-        assert_eq!(basic_coverage_blocks.num_nodes(), bcb_needs_counter.domain_size());
+impl<'a> CountersBuilder<'a> {
+    fn new(graph: &'a CoverageGraph, bcb_needs_counter: &'a BitSet<BasicCoverageBlock>) -> Self {
+        assert_eq!(graph.num_nodes(), bcb_needs_counter.domain_size());
         Self {
-            coverage_counters: CoverageCounters::with_num_bcbs(basic_coverage_blocks.num_nodes()),
-            basic_coverage_blocks,
+            counters: CoverageCounters::with_num_bcbs(graph.num_nodes()),
+            graph,
             bcb_needs_counter,
         }
     }
@@ -293,13 +280,12 @@ impl<'a> MakeBcbCounters<'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.
-        let mut traversal = TraverseCoverageGraphWithLoops::new(self.basic_coverage_blocks);
+        // 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);
             }
         }
 
@@ -312,36 +298,41 @@ impl<'a> MakeBcbCounters<'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);
 
-        let successors = self.basic_coverage_blocks.successors[from_bcb].as_slice();
+        let successors = self.graph.successors[from_bcb].as_slice();
 
         // If this node's out-edges won't sum to the node's counter,
         // then there's no reason to create edge counters here.
-        if !self.basic_coverage_blocks[from_bcb].is_out_summable {
+        if !self.graph[from_bcb].is_out_summable {
             return;
         }
 
-        // Determine the set of out-edges that don't yet have edge counters.
-        let candidate_successors = self.basic_coverage_blocks.successors[from_bcb]
+        // When choosing which out-edge should be given a counter expression, ignore edges that
+        // already have counters, or could use the existing counter of their target node.
+        let out_edge_has_counter = |to_bcb| {
+            if self.counters.edge_counters.contains_key(&(from_bcb, to_bcb)) {
+                return true;
+            }
+            self.graph.sole_predecessor(to_bcb) == Some(from_bcb)
+                && self.counters.node_counters[to_bcb].is_some()
+        };
+
+        // Determine the set of out-edges that could benefit from being given an expression.
+        let candidate_successors = self.graph.successors[from_bcb]
             .iter()
             .copied()
-            .filter(|&to_bcb| self.edge_has_no_counter(from_bcb, to_bcb))
+            .filter(|&to_bcb| !out_edge_has_counter(to_bcb))
             .collect::<Vec<_>>();
         debug!(?candidate_successors);
 
         // 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(expression_to_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;
         };
@@ -353,43 +344,44 @@ impl<'a> MakeBcbCounters<'a> {
             .iter()
             .copied()
             // Skip the chosen edge, since we'll calculate its count from this sum.
-            .filter(|&to_bcb| to_bcb != expression_to_bcb)
+            .filter(|&edge_target_bcb| edge_target_bcb != target_bcb)
             .map(|to_bcb| self.get_or_make_edge_counter(from_bcb, to_bcb))
             .collect::<Vec<_>>();
-        let Some(sum_of_all_other_out_edges) =
-            self.coverage_counters.make_sum(&other_out_edge_counters)
+        let Some(sum_of_all_other_out_edges) = self.counters.make_sum(&other_out_edge_counters)
         else {
             return;
         };
 
         // Now create an expression for the chosen edge, by taking the counter
         // for its source node and subtracting the sum of its sibling out-edges.
-        let expression = self.coverage_counters.make_expression(
-            node_counter,
-            Op::Subtract,
-            sum_of_all_other_out_edges,
-        );
+        let expression =
+            self.counters.make_expression(node_counter, Op::Subtract, sum_of_all_other_out_edges);
 
-        debug!("{expression_to_bcb:?} gets an expression: {expression:?}");
-        if let Some(sole_pred) = self.basic_coverage_blocks.sole_predecessor(expression_to_bcb) {
-            // This edge normally wouldn't get its own counter, so attach the expression
-            // to its target node instead, so that `edge_has_no_counter` can see it.
-            assert_eq!(sole_pred, from_bcb);
-            self.coverage_counters.set_bcb_counter(expression_to_bcb, expression);
-        } else {
-            self.coverage_counters.set_bcb_edge_counter(from_bcb, expression_to_bcb, expression);
-        }
+        debug!("{target_bcb:?} gets an expression: {expression:?}");
+        self.counters.set_edge_counter(from_bcb, target_bcb, expression);
     }
 
     #[instrument(level = "debug", skip(self))]
     fn get_or_make_node_counter(&mut self, bcb: BasicCoverageBlock) -> BcbCounter {
         // If the BCB already has a counter, return it.
-        if let Some(counter_kind) = self.coverage_counters.bcb_counters[bcb] {
-            debug!("{bcb:?} already has a counter: {counter_kind:?}");
-            return counter_kind;
+        if let Some(counter) = self.counters.node_counters[bcb] {
+            debug!("{bcb:?} already has a counter: {counter:?}");
+            return counter;
+        }
+
+        let counter = self.make_node_counter_inner(bcb);
+        self.counters.set_node_counter(bcb, counter)
+    }
+
+    fn make_node_counter_inner(&mut self, bcb: BasicCoverageBlock) -> BcbCounter {
+        // If the node's sole in-edge already has a counter, use that.
+        if let Some(sole_pred) = self.graph.sole_predecessor(bcb)
+            && let Some(&edge_counter) = self.counters.edge_counters.get(&(sole_pred, bcb))
+        {
+            return edge_counter;
         }
 
-        let predecessors = self.basic_coverage_blocks.predecessors[bcb].as_slice();
+        let predecessors = self.graph.predecessors[bcb].as_slice();
 
         // Handle cases where we can't compute a node's count from its in-edges:
         // - START_BCB has no in-edges, so taking the sum would panic (or be wrong).
@@ -398,7 +390,9 @@ impl<'a> MakeBcbCounters<'a> {
         //   leading to infinite recursion.
         if predecessors.len() <= 1 || predecessors.contains(&bcb) {
             debug!(?bcb, ?predecessors, "node has <=1 predecessors or is its own predecessor");
-            return self.coverage_counters.make_phys_node_counter(bcb);
+            let counter = self.counters.make_phys_node_counter(bcb);
+            debug!(?bcb, ?counter, "node gets a physical counter");
+            return counter;
         }
 
         // A BCB with multiple incoming edges can compute its count by ensuring that counters
@@ -408,13 +402,11 @@ impl<'a> MakeBcbCounters<'a> {
             .copied()
             .map(|from_bcb| self.get_or_make_edge_counter(from_bcb, bcb))
             .collect::<Vec<_>>();
-        let sum_of_in_edges: BcbCounter = self
-            .coverage_counters
-            .make_sum(&in_edge_counters)
-            .expect("there must be at least one in-edge");
+        let sum_of_in_edges: BcbCounter =
+            self.counters.make_sum(&in_edge_counters).expect("there must be at least one in-edge");
 
         debug!("{bcb:?} gets a new counter (sum of predecessor counters): {sum_of_in_edges:?}");
-        self.coverage_counters.set_bcb_counter(bcb, sum_of_in_edges)
+        sum_of_in_edges
     }
 
     #[instrument(level = "debug", skip(self))]
@@ -423,9 +415,24 @@ impl<'a> MakeBcbCounters<'a> {
         from_bcb: BasicCoverageBlock,
         to_bcb: BasicCoverageBlock,
     ) -> BcbCounter {
+        // If the edge already has a counter, return it.
+        if let Some(&counter) = self.counters.edge_counters.get(&(from_bcb, to_bcb)) {
+            debug!("Edge {from_bcb:?}->{to_bcb:?} already has a counter: {counter:?}");
+            return counter;
+        }
+
+        let counter = self.make_edge_counter_inner(from_bcb, to_bcb);
+        self.counters.set_edge_counter(from_bcb, to_bcb, counter)
+    }
+
+    fn make_edge_counter_inner(
+        &mut self,
+        from_bcb: BasicCoverageBlock,
+        to_bcb: BasicCoverageBlock,
+    ) -> BcbCounter {
         // If the target node has exactly one in-edge (i.e. this one), then just
         // use the node's counter, since it will have the same value.
-        if let Some(sole_pred) = self.basic_coverage_blocks.sole_predecessor(to_bcb) {
+        if let Some(sole_pred) = self.graph.sole_predecessor(to_bcb) {
             assert_eq!(sole_pred, from_bcb);
             // This call must take care not to invoke `get_or_make_edge` for
             // this edge, since that would result in infinite recursion!
@@ -434,33 +441,27 @@ impl<'a> MakeBcbCounters<'a> {
 
         // If the source node has exactly one out-edge (i.e. this one) and would have
         // the same execution count as that edge, then just use the node's counter.
-        if let Some(simple_succ) = self.basic_coverage_blocks.simple_successor(from_bcb) {
+        if let Some(simple_succ) = self.graph.simple_successor(from_bcb) {
             assert_eq!(simple_succ, to_bcb);
             return self.get_or_make_node_counter(from_bcb);
         }
 
-        // If the edge already has a counter, return it.
-        if let Some(&counter_kind) =
-            self.coverage_counters.bcb_edge_counters.get(&(from_bcb, to_bcb))
-        {
-            debug!("Edge {from_bcb:?}->{to_bcb:?} already has a counter: {counter_kind:?}");
-            return counter_kind;
-        }
-
         // Make a new counter to count this edge.
-        self.coverage_counters.make_phys_edge_counter(from_bcb, to_bcb)
+        let counter = self.counters.make_phys_edge_counter(from_bcb, to_bcb);
+        debug!(?from_bcb, ?to_bcb, ?counter, "edge gets a physical counter");
+        counter
     }
 
     /// Given a set of candidate out-edges (represented by their successor node),
     /// 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);
         }
@@ -479,7 +480,7 @@ impl<'a> MakeBcbCounters<'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.
@@ -488,14 +489,15 @@ impl<'a> MakeBcbCounters<'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.basic_coverage_blocks.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);
@@ -508,21 +510,4 @@ impl<'a> MakeBcbCounters<'a> {
 
         None
     }
-
-    #[inline]
-    fn edge_has_no_counter(
-        &self,
-        from_bcb: BasicCoverageBlock,
-        to_bcb: BasicCoverageBlock,
-    ) -> bool {
-        let edge_counter =
-            if let Some(sole_pred) = self.basic_coverage_blocks.sole_predecessor(to_bcb) {
-                assert_eq!(sole_pred, from_bcb);
-                self.coverage_counters.bcb_counters[to_bcb]
-            } else {
-                self.coverage_counters.bcb_edge_counters.get(&(from_bcb, to_bcb)).copied()
-            };
-
-        edge_counter.is_none()
-    }
 }