coverage: Defer part of counter-creation until codegen

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen/covfun.rs

@@ -54,7 +54,7 @@ pub(crate) fn prepare_covfun_record<'tcx>(
     let fn_cov_info = tcx.instance_mir(instance.def).function_coverage_info.as_deref()?;
     let ids_info = tcx.coverage_ids_info(instance.def)?;
 
-    let expressions = prepare_expressions(fn_cov_info, ids_info, is_used);
+    let expressions = prepare_expressions(ids_info, is_used);
 
     let mut covfun = CovfunRecord {
         mangled_function_name: tcx.symbol_name(instance).name,
@@ -76,11 +76,7 @@ pub(crate) fn prepare_covfun_record<'tcx>(
 }
 
 /// Convert the function's coverage-counter expressions into a form suitable for FFI.
-fn prepare_expressions(
-    fn_cov_info: &FunctionCoverageInfo,
-    ids_info: &CoverageIdsInfo,
-    is_used: bool,
-) -> Vec<ffi::CounterExpression> {
+fn prepare_expressions(ids_info: &CoverageIdsInfo, is_used: bool) -> Vec<ffi::CounterExpression> {
     // If any counters or expressions were removed by MIR opts, replace their
     // terms with zero.
     let counter_for_term = |term| {
@@ -95,7 +91,7 @@ fn prepare_expressions(
     // producing the final coverage map, so there's no need to do the same
     // thing on the Rust side unless we're confident we can do much better.
     // (See `CounterExpressionsMinimizer` in `CoverageMappingWriter.cpp`.)
-    fn_cov_info
+    ids_info
         .expressions
         .iter()
         .map(move |&Expression { lhs, op, rhs }| ffi::CounterExpression {
@@ -142,8 +138,7 @@ fn fill_region_tables<'tcx>(
         // If the mapping refers to counters/expressions that were removed by
         // MIR opts, replace those occurrences with zero.
         let counter_for_bcb = |bcb: BasicCoverageBlock| -> ffi::Counter {
-            let term =
-                fn_cov_info.term_for_bcb[bcb].expect("every BCB in a mapping was given a term");
+            let term = ids_info.term_for_bcb[bcb].expect("every BCB in a mapping was given a term");
             let term = if is_zero_term(term) { CovTerm::Zero } else { term };
             ffi::Counter::from_term(term)
         };

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -160,17 +160,10 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
             CoverageKind::SpanMarker | CoverageKind::BlockMarker { .. } => unreachable!(
                 "marker statement {kind:?} should have been removed by CleanupPostBorrowck"
             ),
-            CoverageKind::CounterIncrement { id } => {
-                // The number of counters passed to `llvm.instrprof.increment` might
-                // be smaller than the number originally inserted by the instrumentor,
-                // if some high-numbered counters were removed by MIR optimizations.
-                // If so, LLVM's profiler runtime will use fewer physical counters.
+            CoverageKind::VirtualCounter { bcb }
+                if let Some(&id) = ids_info.phys_counter_for_node.get(&bcb) =>
+            {
                 let num_counters = ids_info.num_counters_after_mir_opts();
-                assert!(
-                    num_counters as usize <= function_coverage_info.num_counters,
-                    "num_counters disagreement: query says {num_counters} but function info only has {}",
-                    function_coverage_info.num_counters
-                );
 
                 let fn_name = bx.get_pgo_func_name_var(instance);
                 let hash = bx.const_u64(function_coverage_info.function_source_hash);
@@ -182,10 +175,8 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
                 );
                 bx.instrprof_increment(fn_name, hash, num_counters, index);
             }
-            CoverageKind::ExpressionUsed { id: _ } => {
-                // Expression-used statements are markers that are handled by
-                // `coverage_ids_info`, so there's nothing to codegen here.
-            }
+            // If a BCB doesn't have an associated physical counter, there's nothing to codegen.
+            CoverageKind::VirtualCounter { .. } => {}
             CoverageKind::CondBitmapUpdate { index, decision_depth } => {
                 let cond_bitmap = coverage_cx
                     .try_get_mcdc_condition_bitmap(&instance, decision_depth)

compiler/rustc_codegen_llvm/src/lib.rs

@@ -13,6 +13,7 @@
 #![feature(extern_types)]
 #![feature(file_buffered)]
 #![feature(hash_raw_entry)]
+#![feature(if_let_guard)]
 #![feature(impl_trait_in_assoc_type)]
 #![feature(iter_intersperse)]
 #![feature(let_chains)]

compiler/rustc_middle/src/mir/coverage.rs

@@ -2,8 +2,9 @@
 
 use std::fmt::{self, Debug, Formatter};
 
-use rustc_index::IndexVec;
+use rustc_data_structures::fx::FxIndexMap;
 use rustc_index::bit_set::DenseBitSet;
+use rustc_index::{Idx, IndexVec};
 use rustc_macros::{HashStable, TyDecodable, TyEncodable};
 use rustc_span::Span;
 
@@ -103,23 +104,12 @@ pub enum CoverageKind {
     /// Should be erased before codegen (at some point after `InstrumentCoverage`).
     BlockMarker { id: BlockMarkerId },
 
-    /// Marks the point in MIR control flow represented by a coverage counter.
+    /// Marks its enclosing basic block with the ID of the coverage graph node
+    /// that it was part of during the `InstrumentCoverage` MIR pass.
     ///
-    /// This is eventually lowered to `llvm.instrprof.increment` in LLVM IR.
-    ///
-    /// If this statement does not survive MIR optimizations, any mappings that
-    /// refer to this counter can have those references simplified to zero.
-    CounterIncrement { id: CounterId },
-
-    /// Marks the point in MIR control-flow represented by a coverage expression.
-    ///
-    /// If this statement does not survive MIR optimizations, any mappings that
-    /// refer to this expression can have those references simplified to zero.
-    ///
-    /// (This is only inserted for expression IDs that are directly used by
-    /// mappings. Intermediate expressions with no direct mappings are
-    /// retained/zeroed based on whether they are transitively used.)
-    ExpressionUsed { id: ExpressionId },
+    /// During codegen, this might be lowered to `llvm.instrprof.increment` or
+    /// to a no-op, depending on the outcome of counter-creation.
+    VirtualCounter { bcb: BasicCoverageBlock },
 
     /// Marks the point in MIR control flow represented by a evaluated condition.
     ///
@@ -138,8 +128,7 @@ impl Debug for CoverageKind {
         match self {
             SpanMarker => write!(fmt, "SpanMarker"),
             BlockMarker { id } => write!(fmt, "BlockMarker({:?})", id.index()),
-            CounterIncrement { id } => write!(fmt, "CounterIncrement({:?})", id.index()),
-            ExpressionUsed { id } => write!(fmt, "ExpressionUsed({:?})", id.index()),
+            VirtualCounter { bcb } => write!(fmt, "VirtualCounter({bcb:?})"),
             CondBitmapUpdate { index, decision_depth } => {
                 write!(fmt, "CondBitmapUpdate(index={:?}, depth={:?})", index, decision_depth)
             }
@@ -208,11 +197,12 @@ pub struct FunctionCoverageInfo {
     pub function_source_hash: u64,
     pub body_span: Span,
 
-    pub num_counters: usize,
-    pub expressions: IndexVec<ExpressionId, Expression>,
+    /// Used in conjunction with `priority_list` to create physical counters
+    /// and counter expressions, after MIR optimizations.
+    pub node_flow_data: NodeFlowData<BasicCoverageBlock>,
+    pub priority_list: Vec<BasicCoverageBlock>,
 
     pub mappings: Vec<Mapping>,
-    pub term_for_bcb: IndexVec<BasicCoverageBlock, Option<CovTerm>>,
 
     pub mcdc_bitmap_bits: usize,
     /// The depth of the deepest decision is used to know how many
@@ -281,15 +271,18 @@ pub struct MCDCDecisionSpan {
     pub num_conditions: usize,
 }
 
-/// Summarizes coverage IDs inserted by the `InstrumentCoverage` MIR pass
-/// (for compiler option `-Cinstrument-coverage`), after MIR optimizations
-/// have had a chance to potentially remove some of them.
+/// Contains information needed during codegen, obtained by inspecting the
+/// function's MIR after MIR optimizations.
 ///
-/// Used by the `coverage_ids_info` query.
+/// Returned by the `coverage_ids_info` query.
 #[derive(Clone, TyEncodable, TyDecodable, Debug, HashStable)]
 pub struct CoverageIdsInfo {
     pub counters_seen: DenseBitSet<CounterId>,
     pub zero_expressions: DenseBitSet<ExpressionId>,
+
+    pub phys_counter_for_node: FxIndexMap<BasicCoverageBlock, CounterId>,
+    pub term_for_bcb: IndexVec<BasicCoverageBlock, Option<CovTerm>>,
+    pub expressions: IndexVec<ExpressionId, Expression>,
 }
 
 impl CoverageIdsInfo {
@@ -334,3 +327,28 @@ rustc_index::newtype_index! {
         const START_BCB = 0;
     }
 }
+
+/// Data representing a view of some underlying graph, in which each node's
+/// successors have been merged into a single "supernode".
+///
+/// The resulting supernodes have no obvious meaning on their own.
+/// However, merging successor nodes means that a node's out-edges can all
+/// be combined into a single out-edge, whose flow is the same as the flow
+/// (execution count) of its corresponding node in the original graph.
+///
+/// With all node flows now in the original graph now represented as edge flows
+/// in the merged graph, it becomes possible to analyze the original node flows
+/// using techniques for analyzing edge flows.
+#[derive(Clone, Debug)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
+pub struct NodeFlowData<Node: Idx> {
+    /// Maps each node to the supernode that contains it, indicated by some
+    /// arbitrary "root" node that is part of that supernode.
+    pub supernodes: IndexVec<Node, Node>,
+    /// For each node, stores the single supernode that all of its successors
+    /// have been merged into.
+    ///
+    /// (Note that each node in a supernode can potentially have a _different_
+    /// successor supernode from its peers.)
+    pub succ_supernodes: IndexVec<Node, Node>,
+}

compiler/rustc_middle/src/mir/pretty.rs

@@ -619,13 +619,9 @@ fn write_function_coverage_info(
     function_coverage_info: &coverage::FunctionCoverageInfo,
     w: &mut dyn io::Write,
 ) -> io::Result<()> {
-    let coverage::FunctionCoverageInfo { body_span, expressions, mappings, .. } =
-        function_coverage_info;
+    let coverage::FunctionCoverageInfo { body_span, mappings, .. } = function_coverage_info;
 
     writeln!(w, "{INDENT}coverage body span: {body_span:?}")?;
-    for (id, expression) in expressions.iter_enumerated() {
-        writeln!(w, "{INDENT}coverage {id:?} => {expression:?};")?;
-    }
     for coverage::Mapping { kind, span } in mappings {
         writeln!(w, "{INDENT}coverage {kind:?} => {span:?};")?;
     }

compiler/rustc_middle/src/query/mod.rs

@@ -615,9 +615,16 @@ rustc_queries! {
         feedable
     }
 
-    /// Summarizes coverage IDs inserted by the `InstrumentCoverage` MIR pass
-    /// (for compiler option `-Cinstrument-coverage`), after MIR optimizations
-    /// have had a chance to potentially remove some of them.
+    /// Scans through a function's MIR after MIR optimizations, to prepare the
+    /// information needed by codegen when `-Cinstrument-coverage` is active.
+    ///
+    /// This includes the details of where to insert `llvm.instrprof.increment`
+    /// intrinsics, and the expression tables to be embedded in the function's
+    /// coverage metadata.
+    ///
+    /// FIXME(Zalathar): This query's purpose has drifted a bit and should
+    /// probably be renamed, but that can wait until after the potential
+    /// follow-ups to #136053 have settled down.
     ///
     /// Returns `None` for functions that were not instrumented.
     query coverage_ids_info(key: ty::InstanceKind<'tcx>) -> Option<&'tcx mir::coverage::CoverageIdsInfo> {

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,7 +77,7 @@ 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>,
 ) -> CoverageCounters {
@@ -129,15 +132,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>,
+    pub(crate) phys_counter_for_node: FxIndexMap<BasicCoverageBlock, CounterId>,
     next_counter_id: CounterId,
 
     /// Coverage counters/expressions that are associated with individual BCBs.
     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 +191,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,30 +199,4 @@ impl CoverageCounters {
         );
         counter
     }
-
-    /// 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
-    }
 }