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.)

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

@@ -11,7 +11,8 @@ use rustc_codegen_ssa::traits::{
     BaseTypeCodegenMethods, ConstCodegenMethods, StaticCodegenMethods,
 };
 use rustc_middle::mir::coverage::{
-    CovTerm, CoverageIdsInfo, Expression, FunctionCoverageInfo, Mapping, MappingKind, Op,
+    BasicCoverageBlock, CovTerm, CoverageIdsInfo, Expression, FunctionCoverageInfo, Mapping,
+    MappingKind, Op,
 };
 use rustc_middle::ty::{Instance, TyCtxt};
 use rustc_span::Span;
@@ -53,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);
 
     let mut covfun = CovfunRecord {
         mangled_function_name: tcx.symbol_name(instance).name,
@@ -75,26 +76,14 @@ 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> {
-    // If any counters or expressions were removed by MIR opts, replace their
-    // terms with zero.
-    let counter_for_term = |term| {
-        if !is_used || ids_info.is_zero_term(term) {
-            ffi::Counter::ZERO
-        } else {
-            ffi::Counter::from_term(term)
-        }
-    };
+fn prepare_expressions(ids_info: &CoverageIdsInfo) -> Vec<ffi::CounterExpression> {
+    let counter_for_term = ffi::Counter::from_term;
 
     // We know that LLVM will optimize out any unused expressions before
     // 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 {
@@ -136,11 +125,16 @@ fn fill_region_tables<'tcx>(
 
     // For each counter/region pair in this function+file, convert it to a
     // form suitable for FFI.
-    let is_zero_term = |term| !covfun.is_used || ids_info.is_zero_term(term);
     for &Mapping { ref kind, span } in &fn_cov_info.mappings {
-        // If the mapping refers to counters/expressions that were removed by
-        // MIR opts, replace those occurrences with zero.
-        let kind = kind.map_terms(|term| if is_zero_term(term) { CovTerm::Zero } else { term });
+        // If this function is unused, replace all counters with zero.
+        let counter_for_bcb = |bcb: BasicCoverageBlock| -> ffi::Counter {
+            let term = if covfun.is_used {
+                ids_info.term_for_bcb[bcb].expect("every BCB in a mapping was given a term")
+            } else {
+                CovTerm::Zero
+            };
+            ffi::Counter::from_term(term)
+        };
 
         // Convert the `Span` into coordinates that we can pass to LLVM, or
         // discard the span if conversion fails. In rare, cases _all_ of a
@@ -154,23 +148,22 @@ fn fill_region_tables<'tcx>(
             continue;
         }
 
-        match kind {
-            MappingKind::Code(term) => {
-                code_regions
-                    .push(ffi::CodeRegion { cov_span, counter: ffi::Counter::from_term(term) });
+        match *kind {
+            MappingKind::Code { bcb } => {
+                code_regions.push(ffi::CodeRegion { cov_span, counter: counter_for_bcb(bcb) });
             }
-            MappingKind::Branch { true_term, false_term } => {
+            MappingKind::Branch { true_bcb, false_bcb } => {
                 branch_regions.push(ffi::BranchRegion {
                     cov_span,
-                    true_counter: ffi::Counter::from_term(true_term),
-                    false_counter: ffi::Counter::from_term(false_term),
+                    true_counter: counter_for_bcb(true_bcb),
+                    false_counter: counter_for_bcb(false_bcb),
                 });
             }
-            MappingKind::MCDCBranch { true_term, false_term, mcdc_params } => {
+            MappingKind::MCDCBranch { true_bcb, false_bcb, mcdc_params } => {
                 mcdc_branch_regions.push(ffi::MCDCBranchRegion {
                     cov_span,
-                    true_counter: ffi::Counter::from_term(true_term),
-                    false_counter: ffi::Counter::from_term(false_term),
+                    true_counter: counter_for_bcb(true_bcb),
+                    false_counter: counter_for_bcb(false_bcb),
                     mcdc_branch_params: ffi::mcdc::BranchParameters::from(mcdc_params),
                 });
             }

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -160,32 +160,21 @@ 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.
-                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
-                );
-
+            CoverageKind::VirtualCounter { bcb }
+                if let Some(&id) = ids_info.phys_counter_for_node.get(&bcb) =>
+            {
                 let fn_name = bx.get_pgo_func_name_var(instance);
                 let hash = bx.const_u64(function_coverage_info.function_source_hash);
-                let num_counters = bx.const_u32(num_counters);
+                let num_counters = bx.const_u32(ids_info.num_counters);
                 let index = bx.const_u32(id.as_u32());
                 debug!(
                     "codegen intrinsic instrprof.increment(fn_name={:?}, hash={:?}, num_counters={:?}, index={:?})",
                     fn_name, hash, num_counters, index,
                 );
                 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,8 @@
 
 use std::fmt::{self, Debug, Formatter};
 
-use rustc_index::IndexVec;
-use rustc_index::bit_set::DenseBitSet;
+use rustc_data_structures::fx::FxIndexMap;
+use rustc_index::{Idx, IndexVec};
 use rustc_macros::{HashStable, TyDecodable, TyEncodable};
 use rustc_span::Span;
 
@@ -103,23 +103,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 +127,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)
             }
@@ -179,34 +167,19 @@ pub struct Expression {
 #[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub enum MappingKind {
     /// Associates a normal region of code with a counter/expression/zero.
-    Code(CovTerm),
+    Code { bcb: BasicCoverageBlock },
     /// Associates a branch region with separate counters for true and false.
-    Branch { true_term: CovTerm, false_term: CovTerm },
+    Branch { true_bcb: BasicCoverageBlock, false_bcb: BasicCoverageBlock },
     /// Associates a branch region with separate counters for true and false.
-    MCDCBranch { true_term: CovTerm, false_term: CovTerm, mcdc_params: ConditionInfo },
+    MCDCBranch {
+        true_bcb: BasicCoverageBlock,
+        false_bcb: BasicCoverageBlock,
+        mcdc_params: ConditionInfo,
+    },
     /// Associates a decision region with a bitmap and number of conditions.
     MCDCDecision(DecisionInfo),
 }
 
-impl MappingKind {
-    /// Returns a copy of this mapping kind, in which all coverage terms have
-    /// been replaced with ones returned by the given function.
-    pub fn map_terms(&self, map_fn: impl Fn(CovTerm) -> CovTerm) -> Self {
-        match *self {
-            Self::Code(term) => Self::Code(map_fn(term)),
-            Self::Branch { true_term, false_term } => {
-                Self::Branch { true_term: map_fn(true_term), false_term: map_fn(false_term) }
-            }
-            Self::MCDCBranch { true_term, false_term, mcdc_params } => Self::MCDCBranch {
-                true_term: map_fn(true_term),
-                false_term: map_fn(false_term),
-                mcdc_params,
-            },
-            Self::MCDCDecision(param) => Self::MCDCDecision(param),
-        }
-    }
-}
-
 #[derive(Clone, Debug)]
 #[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub struct Mapping {
@@ -222,10 +195,15 @@ pub struct Mapping {
 pub struct FunctionCoverageInfo {
     pub function_source_hash: u64,
     pub body_span: Span,
-    pub num_counters: usize,
-    pub mcdc_bitmap_bits: 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 mcdc_bitmap_bits: usize,
     /// The depth of the deepest decision is used to know how many
     /// temp condbitmaps should be allocated for the function.
     pub mcdc_num_condition_bitmaps: usize,
@@ -292,40 +270,55 @@ 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 num_counters: u32,
+    pub phys_counter_for_node: FxIndexMap<BasicCoverageBlock, CounterId>,
+    pub term_for_bcb: IndexVec<BasicCoverageBlock, Option<CovTerm>>,
+    pub expressions: IndexVec<ExpressionId, Expression>,
 }
 
-impl CoverageIdsInfo {
-    /// Coverage codegen needs to know how many coverage counters are ever
-    /// incremented within a function, so that it can set the `num-counters`
-    /// argument of the `llvm.instrprof.increment` intrinsic.
+rustc_index::newtype_index! {
+    /// During the `InstrumentCoverage` MIR pass, a BCB is a node in the
+    /// "coverage graph", which is a refinement of the MIR control-flow graph
+    /// that merges or omits some blocks that aren't relevant to coverage.
     ///
-    /// This may be less than the highest counter ID emitted by the
-    /// InstrumentCoverage MIR pass, if the highest-numbered counter increments
-    /// were removed by MIR optimizations.
-    pub fn num_counters_after_mir_opts(&self) -> u32 {
-        // FIXME(Zalathar): Currently this treats an unused counter as "used"
-        // if its ID is less than that of the highest counter that really is
-        // used. Fixing this would require adding a renumbering step somewhere.
-        self.counters_seen.last_set_in(..).map_or(0, |max| max.as_u32() + 1)
+    /// After that pass is complete, the coverage graph no longer exists, so a
+    /// BCB is effectively an opaque ID.
+    #[derive(HashStable)]
+    #[encodable]
+    #[orderable]
+    #[debug_format = "bcb{}"]
+    pub struct BasicCoverageBlock {
+        const START_BCB = 0;
     }
+}
 
-    /// Returns `true` if the given term is known to have a value of zero, taking
-    /// into account knowledge of which counters are unused and which expressions
-    /// are always zero.
-    pub fn is_zero_term(&self, term: CovTerm) -> bool {
-        match term {
-            CovTerm::Zero => true,
-            CovTerm::Counter(id) => !self.counters_seen.contains(id),
-            CovTerm::Expression(id) => self.zero_expressions.contains(id),
-        }
-    }
+/// 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

@@ -614,9 +614,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> {