Rollup merge of #136053 - Zalathar:defer-counters, r=saethlin

coverage: Defer part of counter-creation until codegen

Follow-up to #135481 and #135873.

One of the pleasant properties of the new counter-assignment algorithm is that we can stop partway through the process, store the intermediate state in MIR, and then resume the rest of the algorithm during codegen. This lets it take into account which parts of the control-flow graph were eliminated by MIR opts, resulting in fewer physical counters and simpler counter expressions.

Those improvements end up completely obsoleting much larger chunks of code that were previously responsible for cleaning up the coverage metadata after MIR opts, while also doing a more thorough cleanup job.

(That change also unlocks some further simplifications that I've kept out of this PR to limit its scope.)

1 files changed, 36 insertions, 67 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs
index adb99a75a9e..6f9984d5d0a 100644
--- a/compiler/rustc_mir_transform/src/coverage/counters.rs
+++ b/compiler/rustc_mir_transform/src/coverage/counters.rs
@@ -2,7 +2,6 @@ use std::cmp::Ordering;
 
 use either::Either;
 use itertools::Itertools;
-use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
 use rustc_data_structures::graph::DirectedGraph;
 use rustc_index::IndexVec;
@@ -11,31 +10,35 @@ use rustc_middle::mir::coverage::{CounterId, CovTerm, Expression, ExpressionId,
 
 use crate::coverage::counters::balanced_flow::BalancedFlowGraph;
 use crate::coverage::counters::node_flow::{
-    CounterTerm, NodeCounters, make_node_counters, node_flow_data_for_balanced_graph,
+    CounterTerm, NodeCounters, NodeFlowData, node_flow_data_for_balanced_graph,
 };
 use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
 
 mod balanced_flow;
-mod node_flow;
+pub(crate) mod node_flow;
 mod union_find;
 
-/// 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 {
+/// Struct containing the results of [`prepare_bcb_counters_data`].
+pub(crate) struct BcbCountersData {
+    pub(crate) node_flow_data: NodeFlowData<BasicCoverageBlock>,
+    pub(crate) priority_list: Vec<BasicCoverageBlock>,
+}
+
+/// Analyzes the coverage graph to create intermediate data structures that
+/// will later be used (during codegen) to create physical counters or counter
+/// expressions for each BCB node that needs one.
+pub(crate) fn prepare_bcb_counters_data(graph: &CoverageGraph) -> BcbCountersData {
     // Create the derived graphs that are necessary for subsequent steps.
     let balanced_graph = BalancedFlowGraph::for_graph(graph, |n| !graph[n].is_out_summable);
     let node_flow_data = node_flow_data_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.
+    // Also create a "priority list" of coverage graph nodes, to help determine
+    // which ones get physical counters or counter expressions. This needs to
+    // be done now, because the later parts of the counter-creation process
+    // won't have access to the original coverage graph.
     let priority_list = make_node_flow_priority_list(graph, balanced_graph);
-    let node_counters = make_node_counters(&node_flow_data, &priority_list);
 
-    // Convert the counters into a form suitable for embedding into MIR.
-    transcribe_counters(&node_counters, bcb_needs_counter)
+    BcbCountersData { node_flow_data, priority_list }
 }
 
 /// Arranges the nodes in `balanced_graph` into a list, such that earlier nodes
@@ -74,31 +77,33 @@ fn make_node_flow_priority_list(
 }
 
 // Converts node counters into a form suitable for embedding into MIR.
-fn transcribe_counters(
+pub(crate) fn transcribe_counters(
     old: &NodeCounters<BasicCoverageBlock>,
     bcb_needs_counter: &DenseBitSet<BasicCoverageBlock>,
+    bcbs_seen: &DenseBitSet<BasicCoverageBlock>,
 ) -> CoverageCounters {
     let mut new = CoverageCounters::with_num_bcbs(bcb_needs_counter.domain_size());
 
     for bcb in bcb_needs_counter.iter() {
+        if !bcbs_seen.contains(bcb) {
+            // This BCB's code was removed by MIR opts, so its counter is always zero.
+            new.set_node_counter(bcb, CovTerm::Zero);
+            continue;
+        }
+
         // Our counter-creation algorithm doesn't guarantee that a node's list
         // of terms 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_terms[bcb].iter().partition_map(|&CounterTerm { node, op }| match op {
+        let (mut pos, mut neg): (Vec<_>, Vec<_>) = old.counter_terms[bcb]
+            .iter()
+            // Filter out any BCBs that were removed by MIR opts;
+            // this treats them as having an execution count of 0.
+            .filter(|term| bcbs_seen.contains(term.node))
+            .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![];
-        }
-
         // 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
@@ -116,7 +121,7 @@ fn transcribe_counters(
         pos.sort();
         neg.sort();
 
-        let pos_counter = new.make_sum(&pos).expect("`pos` should not be empty");
+        let pos_counter = new.make_sum(&pos).unwrap_or(CovTerm::Zero);
         let new_counter = new.make_subtracted_sum(pos_counter, &neg);
         new.set_node_counter(bcb, new_counter);
     }
@@ -129,15 +134,15 @@ fn transcribe_counters(
 pub(super) struct CoverageCounters {
     /// List of places where a counter-increment statement should be injected
     /// into MIR, each with its corresponding counter ID.
-    phys_counter_for_node: FxIndexMap<BasicCoverageBlock, CounterId>,
-    next_counter_id: CounterId,
+    pub(crate) phys_counter_for_node: FxIndexMap<BasicCoverageBlock, CounterId>,
+    pub(crate) next_counter_id: CounterId,
 
     /// Coverage counters/expressions that are associated with individual BCBs.
-    node_counters: IndexVec<BasicCoverageBlock, Option<CovTerm>>,
+    pub(crate) 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, Expression>,
+    pub(crate) expressions: IndexVec<ExpressionId, Expression>,
     /// Remember expressions that have already been created (or simplified),
     /// so that we don't create unnecessary duplicates.
     expressions_memo: FxHashMap<Expression, CovTerm>,
@@ -188,12 +193,6 @@ impl CoverageCounters {
         self.make_expression(lhs, Op::Subtract, rhs_sum)
     }
 
-    pub(super) fn num_counters(&self) -> usize {
-        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: CovTerm) -> CovTerm {
         let existing = self.node_counters[bcb].replace(counter);
         assert!(
@@ -202,34 +201,4 @@ impl CoverageCounters {
         );
         counter
     }
-
-    pub(super) fn term_for_bcb(&self, bcb: BasicCoverageBlock) -> Option<CovTerm> {
-        self.node_counters[bcb]
-    }
-
-    /// 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, 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
-    /// with a counter *expression*, along with the ID of that expression.
-    pub(super) fn bcb_nodes_with_coverage_expressions(
-        &self,
-    ) -> 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(CovTerm::Expression(id)) => Some((bcb, id)),
-            // This BCB is associated with a counter or nothing, so skip it.
-            Some(CovTerm::Counter { .. } | CovTerm::Zero) | None => None,
-        })
-    }
-
-    pub(super) fn into_expressions(self) -> IndexVec<ExpressionId, Expression> {
-        self.expressions
-    }
 }