Rollup merge of #135680 - Zalathar:counters-cleanup, r=compiler-errors

coverage: Clean up a few things after the counters overhaul

Follow-up to #135481. No functional change; this is mostly just deleting or moving code.

1 files changed, 130 insertions, 180 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs
index 1a9323329f6..8d397f63cc7 100644
--- a/compiler/rustc_mir_transform/src/coverage/counters.rs
+++ b/compiler/rustc_mir_transform/src/coverage/counters.rs
@@ -1,10 +1,9 @@
 use std::cmp::Ordering;
-use std::fmt::{self, Debug};
 
 use either::Either;
 use itertools::Itertools;
 use rustc_data_structures::captures::Captures;
-use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
 use rustc_data_structures::graph::DirectedGraph;
 use rustc_index::IndexVec;
 use rustc_index::bit_set::DenseBitSet;
@@ -20,134 +19,163 @@ mod iter_nodes;
 mod node_flow;
 mod union_find;
 
-/// The coverage counter or counter expression associated with a particular
-/// BCB node or BCB edge.
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-enum BcbCounter {
-    Counter { id: CounterId },
-    Expression { id: ExpressionId },
+/// Ensures that each BCB node needing a counter has one, by creating physical
+/// counters or counter expressions for nodes as required.
+pub(super) fn make_bcb_counters(
+    graph: &CoverageGraph,
+    bcb_needs_counter: &DenseBitSet<BasicCoverageBlock>,
+) -> CoverageCounters {
+    // Create the derived graphs that are necessary for subsequent steps.
+    let balanced_graph = BalancedFlowGraph::for_graph(graph, |n| !graph[n].is_out_summable);
+    let merged_graph = MergedNodeFlowGraph::for_balanced_graph(&balanced_graph);
+
+    // Use those graphs to determine which nodes get physical counters, and how
+    // to compute the execution counts of other nodes from those counters.
+    let nodes = make_node_counter_priority_list(graph, balanced_graph);
+    let node_counters = merged_graph.make_node_counters(&nodes);
+
+    // Convert the counters into a form suitable for embedding into MIR.
+    transcribe_counters(&node_counters, bcb_needs_counter)
 }
 
-impl BcbCounter {
-    fn as_term(&self) -> CovTerm {
-        match *self {
-            BcbCounter::Counter { id, .. } => CovTerm::Counter(id),
-            BcbCounter::Expression { id, .. } => CovTerm::Expression(id),
-        }
-    }
+/// Arranges the nodes in `balanced_graph` into a list, such that earlier nodes
+/// take priority in being given a counter expression instead of a physical counter.
+fn make_node_counter_priority_list(
+    graph: &CoverageGraph,
+    balanced_graph: BalancedFlowGraph<&CoverageGraph>,
+) -> Vec<BasicCoverageBlock> {
+    // A "reloop" node has exactly one out-edge, which jumps back to the top
+    // of an enclosing loop. Reloop nodes are typically visited more times
+    // than loop-exit nodes, so try to avoid giving them physical counters.
+    let is_reloop_node = IndexVec::from_fn_n(
+        |node| match graph.successors[node].as_slice() {
+            &[succ] => graph.dominates(succ, node),
+            _ => false,
+        },
+        graph.num_nodes(),
+    );
+
+    let mut nodes = balanced_graph.iter_nodes().rev().collect::<Vec<_>>();
+    // The first node is the sink, which must not get a physical counter.
+    assert_eq!(nodes[0], balanced_graph.sink);
+    // Sort the real nodes, such that earlier (lesser) nodes take priority
+    // in being given a counter expression instead of a physical counter.
+    nodes[1..].sort_by(|&a, &b| {
+        // Start with a dummy `Equal` to make the actual tests line up nicely.
+        Ordering::Equal
+            // Prefer a physical counter for return/yield nodes.
+            .then_with(|| Ord::cmp(&graph[a].is_out_summable, &graph[b].is_out_summable))
+            // Prefer an expression for reloop nodes (see definition above).
+            .then_with(|| Ord::cmp(&is_reloop_node[a], &is_reloop_node[b]).reverse())
+            // Otherwise, prefer a physical counter for dominating nodes.
+            .then_with(|| graph.cmp_in_dominator_order(a, b).reverse())
+    });
+    nodes
 }
 
-impl Debug for BcbCounter {
-    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match self {
-            Self::Counter { id, .. } => write!(fmt, "Counter({:?})", id.index()),
-            Self::Expression { id } => write!(fmt, "Expression({:?})", id.index()),
+// Converts node counters into a form suitable for embedding into MIR.
+fn transcribe_counters(
+    old: &NodeCounters<BasicCoverageBlock>,
+    bcb_needs_counter: &DenseBitSet<BasicCoverageBlock>,
+) -> CoverageCounters {
+    let mut new = CoverageCounters::with_num_bcbs(bcb_needs_counter.domain_size());
+
+    for bcb in bcb_needs_counter.iter() {
+        // Our counter-creation algorithm doesn't guarantee that a counter
+        // expression starts or ends with a positive term, so partition the
+        // counters into "positive" and "negative" lists for easier handling.
+        let (mut pos, mut neg): (Vec<_>, Vec<_>) =
+            old.counter_expr(bcb).iter().partition_map(|&CounterTerm { node, op }| match op {
+                Op::Add => Either::Left(node),
+                Op::Subtract => Either::Right(node),
+            });
+
+        if pos.is_empty() {
+            // If we somehow end up with no positive terms, fall back to
+            // creating a physical counter. There's no known way for this
+            // to happen, but we can avoid an ICE if it does.
+            debug_assert!(false, "{bcb:?} has no positive counter terms");
+            pos = vec![bcb];
+            neg = vec![];
         }
-    }
-}
 
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
-struct BcbExpression {
-    lhs: BcbCounter,
-    op: Op,
-    rhs: BcbCounter,
-}
+        // These intermediate sorts are not strictly necessary, but were helpful
+        // in reducing churn when switching to the current counter-creation scheme.
+        // They also help to slightly decrease the overall size of the expression
+        // table, due to more subexpressions being shared.
+        pos.sort();
+        neg.sort();
+
+        let mut new_counters_for_sites = |sites: Vec<BasicCoverageBlock>| {
+            sites.into_iter().map(|node| new.ensure_phys_counter(node)).collect::<Vec<_>>()
+        };
+        let mut pos = new_counters_for_sites(pos);
+        let mut neg = new_counters_for_sites(neg);
+
+        // These sorts are also not strictly necessary; see above.
+        pos.sort();
+        neg.sort();
+
+        let pos_counter = new.make_sum(&pos).expect("`pos` should not be empty");
+        let new_counter = new.make_subtracted_sum(pos_counter, &neg);
+        new.set_node_counter(bcb, new_counter);
+    }
 
-/// Enum representing either a node or an edge in the coverage graph.
-///
-/// FIXME(#135481): This enum is no longer needed now that we only instrument
-/// nodes and not edges. It can be removed in a subsequent PR.
-#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
-pub(super) enum Site {
-    Node { bcb: BasicCoverageBlock },
+    new
 }
 
 /// Generates and stores coverage counter and coverage expression information
-/// associated with nodes/edges in the BCB graph.
+/// associated with nodes in the coverage graph.
 pub(super) struct CoverageCounters {
     /// List of places where a counter-increment statement should be injected
     /// into MIR, each with its corresponding counter ID.
-    counter_increment_sites: IndexVec<CounterId, Site>,
+    phys_counter_for_node: FxIndexMap<BasicCoverageBlock, CounterId>,
+    next_counter_id: CounterId,
 
     /// Coverage counters/expressions that are associated with individual BCBs.
-    node_counters: IndexVec<BasicCoverageBlock, Option<BcbCounter>>,
+    node_counters: IndexVec<BasicCoverageBlock, Option<CovTerm>>,
 
     /// Table of expression data, associating each expression ID with its
     /// corresponding operator (+ or -) and its LHS/RHS operands.
-    expressions: IndexVec<ExpressionId, BcbExpression>,
+    expressions: IndexVec<ExpressionId, Expression>,
     /// Remember expressions that have already been created (or simplified),
     /// so that we don't create unnecessary duplicates.
-    expressions_memo: FxHashMap<BcbExpression, BcbCounter>,
+    expressions_memo: FxHashMap<Expression, CovTerm>,
 }
 
 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(
-        graph: &CoverageGraph,
-        bcb_needs_counter: &DenseBitSet<BasicCoverageBlock>,
-    ) -> Self {
-        let balanced_graph = BalancedFlowGraph::for_graph(graph, |n| !graph[n].is_out_summable);
-        let merged_graph = MergedNodeFlowGraph::for_balanced_graph(&balanced_graph);
-
-        // A "reloop" node has exactly one out-edge, which jumps back to the top
-        // of an enclosing loop. Reloop nodes are typically visited more times
-        // than loop-exit nodes, so try to avoid giving them physical counters.
-        let is_reloop_node = IndexVec::from_fn_n(
-            |node| match graph.successors[node].as_slice() {
-                &[succ] => graph.dominates(succ, node),
-                _ => false,
-            },
-            graph.num_nodes(),
-        );
-
-        let mut nodes = balanced_graph.iter_nodes().rev().collect::<Vec<_>>();
-        // The first node is the sink, which must not get a physical counter.
-        assert_eq!(nodes[0], balanced_graph.sink);
-        // Sort the real nodes, such that earlier (lesser) nodes take priority
-        // in being given a counter expression instead of a physical counter.
-        nodes[1..].sort_by(|&a, &b| {
-            // Start with a dummy `Equal` to make the actual tests line up nicely.
-            Ordering::Equal
-                // Prefer a physical counter for return/yield nodes.
-                .then_with(|| Ord::cmp(&graph[a].is_out_summable, &graph[b].is_out_summable))
-                // Prefer an expression for reloop nodes (see definition above).
-                .then_with(|| Ord::cmp(&is_reloop_node[a], &is_reloop_node[b]).reverse())
-                // Otherwise, prefer a physical counter for dominating nodes.
-                .then_with(|| graph.cmp_in_dominator_order(a, b).reverse())
-        });
-        let node_counters = merged_graph.make_node_counters(&nodes);
-
-        Transcriber::new(graph.num_nodes(), node_counters).transcribe_counters(bcb_needs_counter)
-    }
-
     fn with_num_bcbs(num_bcbs: usize) -> Self {
         Self {
-            counter_increment_sites: IndexVec::new(),
+            phys_counter_for_node: FxIndexMap::default(),
+            next_counter_id: CounterId::ZERO,
             node_counters: IndexVec::from_elem_n(None, num_bcbs),
             expressions: IndexVec::new(),
             expressions_memo: FxHashMap::default(),
         }
     }
 
-    /// Creates a new physical counter for a BCB node or edge.
-    fn make_phys_counter(&mut self, site: Site) -> BcbCounter {
-        let id = self.counter_increment_sites.push(site);
-        BcbCounter::Counter { id }
+    /// Returns the physical counter for the given node, creating it if necessary.
+    fn ensure_phys_counter(&mut self, bcb: BasicCoverageBlock) -> CovTerm {
+        let id = *self.phys_counter_for_node.entry(bcb).or_insert_with(|| {
+            let id = self.next_counter_id;
+            self.next_counter_id = id + 1;
+            id
+        });
+        CovTerm::Counter(id)
     }
 
-    fn make_expression(&mut self, lhs: BcbCounter, op: Op, rhs: BcbCounter) -> BcbCounter {
-        let new_expr = BcbExpression { lhs, op, rhs };
-        *self.expressions_memo.entry(new_expr).or_insert_with(|| {
+    fn make_expression(&mut self, lhs: CovTerm, op: Op, rhs: CovTerm) -> CovTerm {
+        let new_expr = Expression { lhs, op, rhs };
+        *self.expressions_memo.entry(new_expr.clone()).or_insert_with(|| {
             let id = self.expressions.push(new_expr);
-            BcbCounter::Expression { id }
+            CovTerm::Expression(id)
         })
     }
 
     /// Creates a counter that is the sum of the given counters.
     ///
     /// Returns `None` if the given list of counters was empty.
-    fn make_sum(&mut self, counters: &[BcbCounter]) -> Option<BcbCounter> {
+    fn make_sum(&mut self, counters: &[CovTerm]) -> Option<CovTerm> {
         counters
             .iter()
             .copied()
@@ -155,16 +183,18 @@ impl CoverageCounters {
     }
 
     /// Creates a counter whose value is `lhs - SUM(rhs)`.
-    fn make_subtracted_sum(&mut self, lhs: BcbCounter, rhs: &[BcbCounter]) -> BcbCounter {
+    fn make_subtracted_sum(&mut self, lhs: CovTerm, rhs: &[CovTerm]) -> CovTerm {
         let Some(rhs_sum) = self.make_sum(rhs) else { return lhs };
         self.make_expression(lhs, Op::Subtract, rhs_sum)
     }
 
     pub(super) fn num_counters(&self) -> usize {
-        self.counter_increment_sites.len()
+        let num_counters = self.phys_counter_for_node.len();
+        assert_eq!(num_counters, self.next_counter_id.as_usize());
+        num_counters
     }
 
-    fn set_node_counter(&mut self, bcb: BasicCoverageBlock, counter: BcbCounter) -> BcbCounter {
+    fn set_node_counter(&mut self, bcb: BasicCoverageBlock, counter: CovTerm) -> CovTerm {
         let existing = self.node_counters[bcb].replace(counter);
         assert!(
             existing.is_none(),
@@ -174,16 +204,16 @@ impl CoverageCounters {
     }
 
     pub(super) fn term_for_bcb(&self, bcb: BasicCoverageBlock) -> Option<CovTerm> {
-        self.node_counters[bcb].map(|counter| counter.as_term())
+        self.node_counters[bcb]
     }
 
-    /// Returns an iterator over all the nodes/edges in the coverage graph that
+    /// Returns an iterator over all the nodes in the coverage graph that
     /// should have a counter-increment statement injected into MIR, along with
     /// each site's corresponding counter ID.
     pub(super) fn counter_increment_sites(
         &self,
-    ) -> impl Iterator<Item = (CounterId, Site)> + Captures<'_> {
-        self.counter_increment_sites.iter_enumerated().map(|(id, &site)| (id, site))
+    ) -> impl Iterator<Item = (CounterId, BasicCoverageBlock)> + Captures<'_> {
+        self.phys_counter_for_node.iter().map(|(&site, &id)| (id, site))
     }
 
     /// Returns an iterator over the subset of BCB nodes that have been associated
@@ -193,93 +223,13 @@ impl CoverageCounters {
     ) -> impl Iterator<Item = (BasicCoverageBlock, ExpressionId)> + Captures<'_> {
         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)),
+            Some(CovTerm::Expression(id)) => Some((bcb, id)),
             // This BCB is associated with a counter or nothing, so skip it.
-            Some(BcbCounter::Counter { .. }) | None => None,
+            Some(CovTerm::Counter { .. } | CovTerm::Zero) | None => None,
         })
     }
 
     pub(super) fn into_expressions(self) -> IndexVec<ExpressionId, Expression> {
-        let old_len = self.expressions.len();
-        let expressions = self
-            .expressions
-            .into_iter()
-            .map(|BcbExpression { lhs, op, rhs }| Expression {
-                lhs: lhs.as_term(),
-                op,
-                rhs: rhs.as_term(),
-            })
-            .collect::<IndexVec<ExpressionId, _>>();
-
-        // Expression IDs are indexes into this vector, so make sure we didn't
-        // accidentally invalidate them by changing its length.
-        assert_eq!(old_len, expressions.len());
-        expressions
-    }
-}
-
-struct Transcriber {
-    old: NodeCounters<BasicCoverageBlock>,
-    new: CoverageCounters,
-    phys_counter_for_site: FxHashMap<Site, BcbCounter>,
-}
-
-impl Transcriber {
-    fn new(num_nodes: usize, old: NodeCounters<BasicCoverageBlock>) -> Self {
-        Self {
-            old,
-            new: CoverageCounters::with_num_bcbs(num_nodes),
-            phys_counter_for_site: FxHashMap::default(),
-        }
-    }
-
-    fn transcribe_counters(
-        mut self,
-        bcb_needs_counter: &DenseBitSet<BasicCoverageBlock>,
-    ) -> CoverageCounters {
-        for bcb in bcb_needs_counter.iter() {
-            let site = Site::Node { bcb };
-            let (mut pos, mut neg): (Vec<_>, Vec<_>) =
-                self.old.counter_expr(bcb).iter().partition_map(
-                    |&CounterTerm { node, op }| match op {
-                        Op::Add => Either::Left(node),
-                        Op::Subtract => Either::Right(node),
-                    },
-                );
-
-            if pos.is_empty() {
-                // If we somehow end up with no positive terms, fall back to
-                // creating a physical counter. There's no known way for this
-                // to happen, but we can avoid an ICE if it does.
-                debug_assert!(false, "{site:?} has no positive counter terms");
-                pos = vec![bcb];
-                neg = vec![];
-            }
-
-            pos.sort();
-            neg.sort();
-
-            let mut new_counters_for_sites = |sites: Vec<BasicCoverageBlock>| {
-                sites
-                    .into_iter()
-                    .map(|node| self.ensure_phys_counter(Site::Node { bcb: node }))
-                    .collect::<Vec<_>>()
-            };
-            let mut pos = new_counters_for_sites(pos);
-            let mut neg = new_counters_for_sites(neg);
-
-            pos.sort();
-            neg.sort();
-
-            let pos_counter = self.new.make_sum(&pos).expect("`pos` should not be empty");
-            let new_counter = self.new.make_subtracted_sum(pos_counter, &neg);
-            self.new.set_node_counter(bcb, new_counter);
-        }
-
-        self.new
-    }
-
-    fn ensure_phys_counter(&mut self, site: Site) -> BcbCounter {
-        *self.phys_counter_for_site.entry(site).or_insert_with(|| self.new.make_phys_counter(site))
+        self.expressions
     }
 }