Rollup merge of #135873 - Zalathar:be-prepared, r=oli-obk

coverage: Prepare for upcoming changes to counter creation

This is a collection of smaller changes to coverage instrumentation code that have been extracted from a larger PR that I'm still working on, in order to hopefully make review easier.

Each individual change should hopefully be mostly self-explanatory. One of the big goals of the upcoming PR will be to defer certain parts of counter-creation until codegen, via the query system, so that ends up being a recurring theme in these changes. Several of the changes are follow-ups to #135481.

There should be no observable change in compiler output.

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

@@ -51,7 +51,7 @@ pub(crate) fn prepare_covfun_record<'tcx>(
     is_used: bool,
 ) -> Option<CovfunRecord<'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 ids_info = tcx.coverage_ids_info(instance.def)?;
 
     let expressions = prepare_expressions(fn_cov_info, ids_info, is_used);
 

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -8,7 +8,6 @@ use rustc_codegen_ssa::traits::{
 use rustc_data_structures::fx::{FxHashMap, FxIndexSet};
 use rustc_middle::mir::coverage::CoverageKind;
 use rustc_middle::ty::Instance;
-use rustc_middle::ty::layout::HasTyCtxt;
 use tracing::{debug, instrument};
 
 use crate::builder::Builder;
@@ -147,6 +146,10 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
             debug!("function has a coverage statement but no coverage info");
             return;
         };
+        let Some(ids_info) = bx.tcx.coverage_ids_info(instance.def) else {
+            debug!("function has a coverage statement but no IDs info");
+            return;
+        };
 
         // Mark the instance as used in this CGU, for coverage purposes.
         // This includes functions that were not partitioned into this CGU,
@@ -162,8 +165,7 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
                 // 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 =
-                    bx.tcx().coverage_ids_info(instance.def).num_counters_after_mir_opts();
+                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 {}",

compiler/rustc_middle/src/mir/coverage.rs

@@ -4,7 +4,7 @@ use std::fmt::{self, Debug, Formatter};
 
 use rustc_index::IndexVec;
 use rustc_index::bit_set::DenseBitSet;
-use rustc_macros::{HashStable, TyDecodable, TyEncodable, TypeFoldable, TypeVisitable};
+use rustc_macros::{HashStable, TyDecodable, TyEncodable};
 use rustc_span::Span;
 
 rustc_index::newtype_index! {
@@ -72,7 +72,7 @@ impl ConditionId {
 /// Enum that can hold a constant zero value, the ID of an physical coverage
 /// counter, or the ID of a coverage-counter expression.
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
-#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub enum CovTerm {
     Zero,
     Counter(CounterId),
@@ -89,7 +89,7 @@ impl Debug for CovTerm {
     }
 }
 
-#[derive(Clone, PartialEq, TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(Clone, PartialEq, TyEncodable, TyDecodable, Hash, HashStable)]
 pub enum CoverageKind {
     /// Marks a span that might otherwise not be represented in MIR, so that
     /// coverage instrumentation can associate it with its enclosing block/BCB.
@@ -151,7 +151,7 @@ impl Debug for CoverageKind {
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, HashStable)]
-#[derive(TyEncodable, TyDecodable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable)]
 pub enum Op {
     Subtract,
     Add,
@@ -168,7 +168,7 @@ impl Op {
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
-#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub struct Expression {
     pub lhs: CovTerm,
     pub op: Op,
@@ -176,7 +176,7 @@ pub struct Expression {
 }
 
 #[derive(Clone, Debug)]
-#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub enum MappingKind {
     /// Associates a normal region of code with a counter/expression/zero.
     Code(CovTerm),
@@ -208,7 +208,7 @@ impl MappingKind {
 }
 
 #[derive(Clone, Debug)]
-#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub struct Mapping {
     pub kind: MappingKind,
     pub span: Span,
@@ -218,7 +218,7 @@ pub struct Mapping {
 /// to be used in conjunction with the individual coverage statements injected
 /// into the function's basic blocks.
 #[derive(Clone, Debug)]
-#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub struct FunctionCoverageInfo {
     pub function_source_hash: u64,
     pub body_span: Span,
@@ -238,7 +238,7 @@ pub struct FunctionCoverageInfo {
 /// ("Hi" indicates that this is "high-level" information collected at the
 /// THIR/MIR boundary, before the MIR-based coverage instrumentation pass.)
 #[derive(Clone, Debug)]
-#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub struct CoverageInfoHi {
     /// 1 more than the highest-numbered [`CoverageKind::BlockMarker`] that was
     /// injected into the MIR body. This makes it possible to allocate per-ID
@@ -252,7 +252,7 @@ pub struct CoverageInfoHi {
 }
 
 #[derive(Clone, Debug)]
-#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub struct BranchSpan {
     pub span: Span,
     pub true_marker: BlockMarkerId,
@@ -260,7 +260,7 @@ pub struct BranchSpan {
 }
 
 #[derive(Copy, Clone, Debug)]
-#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub struct ConditionInfo {
     pub condition_id: ConditionId,
     pub true_next_id: Option<ConditionId>,
@@ -268,7 +268,7 @@ pub struct ConditionInfo {
 }
 
 #[derive(Clone, Debug)]
-#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub struct MCDCBranchSpan {
     pub span: Span,
     pub condition_info: ConditionInfo,
@@ -277,14 +277,14 @@ pub struct MCDCBranchSpan {
 }
 
 #[derive(Copy, Clone, Debug)]
-#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub struct DecisionInfo {
     pub bitmap_idx: u32,
     pub num_conditions: u16,
 }
 
 #[derive(Clone, Debug)]
-#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub struct MCDCDecisionSpan {
     pub span: Span,
     pub end_markers: Vec<BlockMarkerId>,

compiler/rustc_middle/src/mir/mod.rs

@@ -358,6 +358,8 @@ pub struct Body<'tcx> {
     ///
     /// Only present if coverage is enabled and this function is eligible.
     /// Boxed to limit space overhead in non-coverage builds.
+    #[type_foldable(identity)]
+    #[type_visitable(ignore)]
     pub coverage_info_hi: Option<Box<coverage::CoverageInfoHi>>,
 
     /// Per-function coverage information added by the `InstrumentCoverage`
@@ -366,6 +368,8 @@ pub struct Body<'tcx> {
     ///
     /// If `-Cinstrument-coverage` is not active, or if an individual function
     /// is not eligible for coverage, then this should always be `None`.
+    #[type_foldable(identity)]
+    #[type_visitable(ignore)]
     pub function_coverage_info: Option<Box<coverage::FunctionCoverageInfo>>,
 }
 

compiler/rustc_middle/src/mir/syntax.rs

@@ -417,7 +417,14 @@ pub enum StatementKind<'tcx> {
     ///
     /// Interpreters and codegen backends that don't support coverage instrumentation
     /// can usually treat this as a no-op.
-    Coverage(CoverageKind),
+    Coverage(
+        // Coverage statements are unlikely to ever contain type information in
+        // the foreseeable future, so excluding them from TypeFoldable/TypeVisitable
+        // avoids some unhelpful derive boilerplate.
+        #[type_foldable(identity)]
+        #[type_visitable(ignore)]
+        CoverageKind,
+    ),
 
     /// Denotes a call to an intrinsic that does not require an unwind path and always returns.
     /// This avoids adding a new block and a terminator for simple intrinsics.

compiler/rustc_middle/src/query/mod.rs

@@ -618,7 +618,9 @@ rustc_queries! {
     /// 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.
-    query coverage_ids_info(key: ty::InstanceKind<'tcx>) -> &'tcx mir::coverage::CoverageIdsInfo {
+    ///
+    /// Returns `None` for functions that were not instrumented.
+    query coverage_ids_info(key: ty::InstanceKind<'tcx>) -> Option<&'tcx mir::coverage::CoverageIdsInfo> {
         desc { |tcx| "retrieving coverage IDs info from MIR for `{}`", tcx.def_path_str(key.def_id()) }
         arena_cache
     }

compiler/rustc_mir_transform/src/coverage/counters.rs

@@ -11,7 +11,9 @@ use rustc_middle::mir::coverage::{CounterId, CovTerm, Expression, ExpressionId,
 
 use crate::coverage::counters::balanced_flow::BalancedFlowGraph;
 use crate::coverage::counters::iter_nodes::IterNodes;
-use crate::coverage::counters::node_flow::{CounterTerm, MergedNodeFlowGraph, NodeCounters};
+use crate::coverage::counters::node_flow::{
+    CounterTerm, NodeCounters, make_node_counters, node_flow_data_for_balanced_graph,
+};
 use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
 
 mod balanced_flow;
@@ -27,12 +29,12 @@ pub(super) fn make_bcb_counters(
 ) -> 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);
+    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.
-    let nodes = make_node_counter_priority_list(graph, balanced_graph);
-    let node_counters = merged_graph.make_node_counters(&nodes);
+    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)
@@ -40,7 +42,7 @@ pub(super) fn make_bcb_counters(
 
 /// 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(
+fn make_node_flow_priority_list(
     graph: &CoverageGraph,
     balanced_graph: BalancedFlowGraph<&CoverageGraph>,
 ) -> Vec<BasicCoverageBlock> {
@@ -81,11 +83,11 @@ fn transcribe_counters(
     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
+        // 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_expr(bcb).iter().partition_map(|&CounterTerm { node, op }| match op {
+            old.counter_terms[bcb].iter().partition_map(|&CounterTerm { node, op }| match op {
                 Op::Add => Either::Left(node),
                 Op::Subtract => Either::Right(node),
             });

compiler/rustc_mir_transform/src/coverage/counters/node_flow.rs

@@ -8,18 +8,17 @@
 
 use rustc_data_structures::graph;
 use rustc_index::bit_set::DenseBitSet;
-use rustc_index::{Idx, IndexVec};
+use rustc_index::{Idx, IndexSlice, IndexVec};
 use rustc_middle::mir::coverage::Op;
-use smallvec::SmallVec;
 
 use crate::coverage::counters::iter_nodes::IterNodes;
-use crate::coverage::counters::union_find::{FrozenUnionFind, UnionFind};
+use crate::coverage::counters::union_find::UnionFind;
 
 #[cfg(test)]
 mod tests;
 
-/// View of some underlying graph, in which each node's successors have been
-/// merged into a single "supernode".
+/// 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
@@ -30,10 +29,10 @@ mod tests;
 /// in the merged graph, it becomes possible to analyze the original node flows
 /// using techniques for analyzing edge flows.
 #[derive(Debug)]
-pub(crate) struct MergedNodeFlowGraph<Node: Idx> {
+pub(crate) 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.
-    supernodes: FrozenUnionFind<Node>,
+    supernodes: IndexVec<Node, Node>,
     /// For each node, stores the single supernode that all of its successors
     /// have been merged into.
     ///
@@ -42,16 +41,15 @@ pub(crate) struct MergedNodeFlowGraph<Node: Idx> {
     succ_supernodes: IndexVec<Node, Node>,
 }
 
-impl<Node: Idx> MergedNodeFlowGraph<Node> {
 /// Creates a "merged" view of an underlying graph.
 ///
 /// The given graph is assumed to have [“balanced flow”](balanced-flow),
 /// though it does not necessarily have to be a `BalancedFlowGraph`.
 ///
 /// [balanced-flow]: `crate::coverage::counters::balanced_flow::BalancedFlowGraph`.
-    pub(crate) fn for_balanced_graph<G>(graph: G) -> Self
+pub(crate) fn node_flow_data_for_balanced_graph<G>(graph: G) -> NodeFlowData<G::Node>
 where
-        G: graph::DirectedGraph<Node = Node> + graph::Successors,
+    G: graph::Successors,
 {
     let mut supernodes = UnionFind::<G::Node>::new(graph.num_nodes());
 
@@ -67,24 +65,16 @@ impl<Node: Idx> MergedNodeFlowGraph<Node> {
         })
         .collect::<IndexVec<G::Node, G::Node>>();
 
-        // Now that unification is complete, freeze the supernode forest,
+    // Now that unification is complete, take a snapshot of the supernode forest,
     // and resolve each arbitrarily-chosen successor to its canonical root.
     // (This avoids having to explicitly resolve them later.)
-        let supernodes = supernodes.freeze();
-        let succ_supernodes = successors.into_iter().map(|succ| supernodes.find(succ)).collect();
+    let supernodes = supernodes.snapshot();
+    let succ_supernodes = successors.into_iter().map(|succ| supernodes[succ]).collect();
 
-        Self { supernodes, succ_supernodes }
+    NodeFlowData { supernodes, succ_supernodes }
 }
 
-    fn num_nodes(&self) -> usize {
-        self.succ_supernodes.len()
-    }
-
-    fn is_supernode(&self, node: Node) -> bool {
-        self.supernodes.find(node) == node
-    }
-
-    /// Using the information in this merged graph, together with a given
+/// Uses the graph information in `node_flow_data`, together with a given
 /// permutation of all nodes in the graph, to create physical counters and
 /// counter expressions for each node in the underlying graph.
 ///
@@ -93,33 +83,29 @@ impl<Node: Idx> MergedNodeFlowGraph<Node> {
 /// influence counter allocation:
 /// - Earlier nodes are more likely to receive counter expressions.
 /// - Later nodes are more likely to receive physical counters.
-    pub(crate) fn make_node_counters(&self, all_nodes_permutation: &[Node]) -> NodeCounters<Node> {
-        let mut builder = SpantreeBuilder::new(self);
+pub(crate) fn make_node_counters<Node: Idx>(
+    node_flow_data: &NodeFlowData<Node>,
+    priority_list: &[Node],
+) -> NodeCounters<Node> {
+    let mut builder = SpantreeBuilder::new(node_flow_data);
 
-        for &node in all_nodes_permutation {
+    for &node in priority_list {
         builder.visit_node(node);
     }
 
-        NodeCounters { counter_exprs: builder.finish() }
-    }
+    NodeCounters { counter_terms: builder.finish() }
 }
 
 /// End result of allocating physical counters and counter expressions for the
 /// nodes of a graph.
 #[derive(Debug)]
 pub(crate) struct NodeCounters<Node: Idx> {
-    counter_exprs: IndexVec<Node, CounterExprVec<Node>>,
-}
-
-impl<Node: Idx> NodeCounters<Node> {
     /// For the given node, returns the finished list of terms that represent
     /// its physical counter or counter expression. Always non-empty.
     ///
-    /// If a node was given a physical counter, its "expression" will contain
+    /// If a node was given a physical counter, the term list will contain
     /// that counter as its sole element.
-    pub(crate) fn counter_expr(&self, this: Node) -> &[CounterTerm<Node>] {
-        self.counter_exprs[this].as_slice()
-    }
+    pub(crate) counter_terms: IndexVec<Node, Vec<CounterTerm<Node>>>,
 }
 
 #[derive(Debug)]
@@ -146,12 +132,11 @@ pub(crate) struct CounterTerm<Node> {
     pub(crate) node: Node,
 }
 
-/// Stores the list of counter terms that make up a node's counter expression.
-type CounterExprVec<Node> = SmallVec<[CounterTerm<Node>; 2]>;
-
 #[derive(Debug)]
 struct SpantreeBuilder<'a, Node: Idx> {
-    graph: &'a MergedNodeFlowGraph<Node>,
+    supernodes: &'a IndexSlice<Node, Node>,
+    succ_supernodes: &'a IndexSlice<Node, Node>,
+
     is_unvisited: DenseBitSet<Node>,
     /// Links supernodes to each other, gradually forming a spanning tree of
     /// the merged-flow graph.
@@ -163,26 +148,32 @@ struct SpantreeBuilder<'a, Node: Idx> {
     yank_buffer: Vec<Node>,
     /// An in-progress counter expression for each node. Each expression is
     /// initially empty, and will be filled in as relevant nodes are visited.
-    counter_exprs: IndexVec<Node, CounterExprVec<Node>>,
+    counter_terms: IndexVec<Node, Vec<CounterTerm<Node>>>,
 }
 
 impl<'a, Node: Idx> SpantreeBuilder<'a, Node> {
-    fn new(graph: &'a MergedNodeFlowGraph<Node>) -> Self {
-        let num_nodes = graph.num_nodes();
+    fn new(node_flow_data: &'a NodeFlowData<Node>) -> Self {
+        let NodeFlowData { supernodes, succ_supernodes } = node_flow_data;
+        let num_nodes = supernodes.len();
         Self {
-            graph,
+            supernodes,
+            succ_supernodes,
             is_unvisited: DenseBitSet::new_filled(num_nodes),
             span_edges: IndexVec::from_fn_n(|_| None, num_nodes),
             yank_buffer: vec![],
-            counter_exprs: IndexVec::from_fn_n(|_| SmallVec::new(), num_nodes),
+            counter_terms: IndexVec::from_fn_n(|_| vec![], num_nodes),
         }
     }
 
+    fn is_supernode(&self, node: Node) -> bool {
+        self.supernodes[node] == node
+    }
+
     /// Given a supernode, finds the supernode that is the "root" of its
     /// spantree component. Two nodes that have the same spantree root are
     /// connected in the spantree.
     fn spantree_root(&self, this: Node) -> Node {
-        debug_assert!(self.graph.is_supernode(this));
+        debug_assert!(self.is_supernode(this));
 
         match self.span_edges[this] {
             None => this,
@@ -193,7 +184,7 @@ impl<'a, Node: Idx> SpantreeBuilder<'a, Node> {
     /// Rotates edges in the spantree so that `this` is the root of its
     /// spantree component.
     fn yank_to_spantree_root(&mut self, this: Node) {
-        debug_assert!(self.graph.is_supernode(this));
+        debug_assert!(self.is_supernode(this));
 
         // The rotation is done iteratively, by first traversing from `this` to
         // its root and storing the path in a buffer, and then traversing the
@@ -235,12 +226,12 @@ impl<'a, Node: Idx> SpantreeBuilder<'a, Node> {
 
         // Get the supernode containing `this`, and make it the root of its
         // component of the spantree.
-        let this_supernode = self.graph.supernodes.find(this);
+        let this_supernode = self.supernodes[this];
         self.yank_to_spantree_root(this_supernode);
 
         // Get the supernode containing all of this's successors.
-        let succ_supernode = self.graph.succ_supernodes[this];
-        debug_assert!(self.graph.is_supernode(succ_supernode));
+        let succ_supernode = self.succ_supernodes[this];
+        debug_assert!(self.is_supernode(succ_supernode));
 
         // If two supernodes are already connected in the spantree, they will
         // have the same spantree root. (Each supernode is connected to itself.)
@@ -268,8 +259,8 @@ impl<'a, Node: Idx> SpantreeBuilder<'a, Node> {
             // `this_supernode`.
 
             // Instead of setting `this.measure = true` as in the original paper,
-            // we just add the node's ID to its own "expression".
-            self.counter_exprs[this].push(CounterTerm { node: this, op: Op::Add });
+            // we just add the node's ID to its own list of terms.
+            self.counter_terms[this].push(CounterTerm { node: this, op: Op::Add });
 
             // Walk the spantree from `this.successor` back to `this`. For each
             // spantree edge along the way, add this node's physical counter to
@@ -279,7 +270,7 @@ impl<'a, Node: Idx> SpantreeBuilder<'a, Node> {
                 let &SpantreeEdge { is_reversed, claiming_node, span_parent } =
                     self.span_edges[curr].as_ref().unwrap();
                 let op = if is_reversed { Op::Subtract } else { Op::Add };
-                self.counter_exprs[claiming_node].push(CounterTerm { node: this, op });
+                self.counter_terms[claiming_node].push(CounterTerm { node: this, op });
 
                 curr = span_parent;
             }
@@ -288,19 +279,20 @@ impl<'a, Node: Idx> SpantreeBuilder<'a, Node> {
 
     /// Asserts that all nodes have been visited, and returns the computed
     /// counter expressions (made up of physical counters) for each node.
-    fn finish(self) -> IndexVec<Node, CounterExprVec<Node>> {
-        let Self { graph, is_unvisited, span_edges, yank_buffer: _, counter_exprs } = self;
+    fn finish(self) -> IndexVec<Node, Vec<CounterTerm<Node>>> {
+        let Self { ref span_edges, ref is_unvisited, ref counter_terms, .. } = self;
         assert!(is_unvisited.is_empty(), "some nodes were never visited: {is_unvisited:?}");
         debug_assert!(
             span_edges
                 .iter_enumerated()
-                .all(|(node, span_edge)| { span_edge.is_some() <= graph.is_supernode(node) }),
+                .all(|(node, span_edge)| { span_edge.is_some() <= self.is_supernode(node) }),
             "only supernodes can have a span edge",
         );
         debug_assert!(
-            counter_exprs.iter().all(|expr| !expr.is_empty()),
-            "after visiting all nodes, every node should have a non-empty expression",
+            counter_terms.iter().all(|terms| !terms.is_empty()),
+            "after visiting all nodes, every node should have at least one term",
         );
-        counter_exprs
+
+        self.counter_terms
     }
 }

compiler/rustc_mir_transform/src/coverage/counters/node_flow/tests.rs

@@ -4,10 +4,12 @@ use rustc_data_structures::graph::vec_graph::VecGraph;
 use rustc_index::Idx;
 use rustc_middle::mir::coverage::Op;
 
-use super::{CounterTerm, MergedNodeFlowGraph, NodeCounters};
+use crate::coverage::counters::node_flow::{
+    CounterTerm, NodeCounters, NodeFlowData, make_node_counters, node_flow_data_for_balanced_graph,
+};
 
-fn merged_node_flow_graph<G: graph::Successors>(graph: G) -> MergedNodeFlowGraph<G::Node> {
-    MergedNodeFlowGraph::for_balanced_graph(graph)
+fn node_flow_data<G: graph::Successors>(graph: G) -> NodeFlowData<G::Node> {
+    node_flow_data_for_balanced_graph(graph)
 }
 
 fn make_graph<Node: Idx + Ord>(num_nodes: usize, edge_pairs: Vec<(Node, Node)>) -> VecGraph<Node> {
@@ -30,8 +32,8 @@ fn example_driver() {
         (4, 0),
     ]);
 
-    let merged = merged_node_flow_graph(&graph);
-    let counters = merged.make_node_counters(&[3, 1, 2, 0, 4]);
+    let node_flow_data = node_flow_data(&graph);
+    let counters = make_node_counters(&node_flow_data, &[3, 1, 2, 0, 4]);
 
     assert_eq!(format_counter_expressions(&counters), &[
         // (comment to force vertical formatting for clarity)
@@ -53,12 +55,12 @@ fn format_counter_expressions<Node: Idx>(counters: &NodeCounters<Node>) -> Vec<S
     };
 
     counters
-        .counter_exprs
+        .counter_terms
         .indices()
         .map(|node| {
-            let mut expr = counters.counter_expr(node).iter().collect::<Vec<_>>();
-            expr.sort_by_key(|item| item.node.index());
-            format!("[{node:?}]: {}", expr.into_iter().map(format_item).join(" "))
+            let mut terms = counters.counter_terms[node].iter().collect::<Vec<_>>();
+            terms.sort_by_key(|item| item.node.index());
+            format!("[{node:?}]: {}", terms.into_iter().map(format_item).join(" "))
         })
         .collect()
 }

compiler/rustc_mir_transform/src/coverage/counters/union_find.rs

@@ -88,29 +88,9 @@ impl<Key: Idx> UnionFind<Key> {
         a
     }
 
-    /// Creates a snapshot of this disjoint-set forest that can no longer be
-    /// mutated, but can be queried without mutation.
-    pub(crate) fn freeze(&mut self) -> FrozenUnionFind<Key> {
-        // Just resolve each key to its actual root.
-        let roots = self.table.indices().map(|key| self.find(key)).collect();
-        FrozenUnionFind { roots }
-    }
-}
-
-/// Snapshot of a disjoint-set forest that can no longer be mutated, but can be
-/// queried in O(1) time without mutation.
-///
-/// This is really just a wrapper around a direct mapping from keys to roots,
-/// but with a [`Self::find`] method that resembles [`UnionFind::find`].
-#[derive(Debug)]
-pub(crate) struct FrozenUnionFind<Key: Idx> {
-    roots: IndexVec<Key, Key>,
-}
-
-impl<Key: Idx> FrozenUnionFind<Key> {
-    /// Returns the "root" key of the disjoint-set containing the given key.
-    /// If two keys have the same root, they belong to the same set.
-    pub(crate) fn find(&self, key: Key) -> Key {
-        self.roots[key]
+    /// Takes a "snapshot" of the current state of this disjoint-set forest, in
+    /// the form of a vector that directly maps each key to its current root.
+    pub(crate) fn snapshot(&mut self) -> IndexVec<Key, Key> {
+        self.table.indices().map(|key| self.find(key)).collect()
     }
 }

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -180,7 +180,12 @@ fn create_mappings(
     ));
 
     for (decision, branches) in mcdc_mappings {
-        let num_conditions = branches.len() as u16;
+        // FIXME(#134497): Previously it was possible for some of these branch
+        // conversions to fail, in which case the remaining branches in the
+        // decision would be degraded to plain `MappingKind::Branch`.
+        // The changes in #134497 made that failure impossible, because the
+        // fallible step was deferred to codegen. But the corresponding code
+        // in codegen wasn't updated to detect the need for a degrade step.
         let conditions = branches
             .into_iter()
             .map(
@@ -206,24 +211,13 @@ fn create_mappings(
             )
             .collect::<Vec<_>>();
 
-        if conditions.len() == num_conditions as usize {
         // LLVM requires end index for counter mapping regions.
         let kind = MappingKind::MCDCDecision(DecisionInfo {
             bitmap_idx: (decision.bitmap_idx + decision.num_test_vectors) as u32,
-                num_conditions,
+            num_conditions: u16::try_from(conditions.len()).unwrap(),
         });
         let span = decision.span;
         mappings.extend(std::iter::once(Mapping { kind, span }).chain(conditions.into_iter()));
-        } else {
-            mappings.extend(conditions.into_iter().map(|mapping| {
-                let MappingKind::MCDCBranch { true_term, false_term, mcdc_params: _ } =
-                    mapping.kind
-                else {
-                    unreachable!("all mappings here are MCDCBranch as shown above");
-                };
-                Mapping { kind: MappingKind::Branch { true_term, false_term }, span: mapping.span }
-            }))
-        }
     }
 
     mappings

compiler/rustc_mir_transform/src/coverage/query.rs

@@ -87,15 +87,9 @@ fn coverage_attr_on(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
 fn coverage_ids_info<'tcx>(
     tcx: TyCtxt<'tcx>,
     instance_def: ty::InstanceKind<'tcx>,
-) -> CoverageIdsInfo {
+) -> Option<CoverageIdsInfo> {
     let mir_body = tcx.instance_mir(instance_def);
-
-    let Some(fn_cov_info) = mir_body.function_coverage_info.as_deref() else {
-        return CoverageIdsInfo {
-            counters_seen: DenseBitSet::new_empty(0),
-            zero_expressions: DenseBitSet::new_empty(0),
-        };
-    };
+    let fn_cov_info = mir_body.function_coverage_info.as_deref()?;
 
     let mut counters_seen = DenseBitSet::new_empty(fn_cov_info.num_counters);
     let mut expressions_seen = DenseBitSet::new_filled(fn_cov_info.expressions.len());
@@ -129,7 +123,7 @@ fn coverage_ids_info<'tcx>(
     let zero_expressions =
         identify_zero_expressions(fn_cov_info, &counters_seen, &expressions_seen);
 
-    CoverageIdsInfo { counters_seen, zero_expressions }
+    Some(CoverageIdsInfo { counters_seen, zero_expressions })
 }
 
 fn all_coverage_in_mir_body<'a, 'tcx>(

tests/mir-opt/coverage/instrument_coverage.bar.InstrumentCoverage.diff

@@ -4,8 +4,8 @@
   fn bar() -> bool {
       let mut _0: bool;
   
-+     coverage body span: $DIR/instrument_coverage.rs:19:18: 21:2 (#0)
-+     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:19:1: 21:2 (#0);
++     coverage body span: $DIR/instrument_coverage.rs:29:18: 31:2 (#0)
++     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:29:1: 31:2 (#0);
 + 
       bb0: {
 +         Coverage::CounterIncrement(0);

tests/mir-opt/coverage/instrument_coverage.main.InstrumentCoverage.diff

@@ -7,13 +7,13 @@
       let mut _2: bool;
       let mut _3: !;
   
-+     coverage body span: $DIR/instrument_coverage.rs:10:11: 16:2 (#0)
++     coverage body span: $DIR/instrument_coverage.rs:14:11: 20:2 (#0)
 +     coverage ExpressionId(0) => Expression { lhs: Counter(1), op: Subtract, rhs: Counter(0) };
-+     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:10:1: 10:11 (#0);
-+     coverage Code(Counter(1)) => $DIR/instrument_coverage.rs:12:12: 12:17 (#0);
-+     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:13:13: 13:18 (#0);
-+     coverage Code(Expression(0)) => $DIR/instrument_coverage.rs:14:10: 14:10 (#0);
-+     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:16:2: 16:2 (#0);
++     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:14:1: 14:11 (#0);
++     coverage Code(Counter(1)) => $DIR/instrument_coverage.rs:16:12: 16:17 (#0);
++     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:17:13: 17:18 (#0);
++     coverage Code(Expression(0)) => $DIR/instrument_coverage.rs:18:10: 18:10 (#0);
++     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:20:2: 20:2 (#0);
 + 
       bb0: {
 +         Coverage::CounterIncrement(0);

tests/mir-opt/coverage/instrument_coverage.rs

@@ -6,7 +6,11 @@
 //@ compile-flags: -Cinstrument-coverage -Zno-profiler-runtime
 
 // EMIT_MIR instrument_coverage.main.InstrumentCoverage.diff
-// EMIT_MIR instrument_coverage.bar.InstrumentCoverage.diff
+// CHECK-LABEL: fn main()
+// CHECK: coverage body span:
+// CHECK: coverage Code(Counter({{[0-9]+}})) =>
+// CHECK: bb0:
+// CHECK: Coverage::CounterIncrement
 fn main() {
     loop {
         if bar() {
@@ -15,14 +19,13 @@ fn main() {
     }
 }
 
+// EMIT_MIR instrument_coverage.bar.InstrumentCoverage.diff
+// CHECK-LABEL: fn bar()
+// CHECK: coverage body span:
+// CHECK: coverage Code(Counter({{[0-9]+}})) =>
+// CHECK: bb0:
+// CHECK: Coverage::CounterIncrement
 #[inline(never)]
 fn bar() -> bool {
     true
 }
-
-// CHECK:     coverage ExpressionId({{[0-9]+}}) =>
-// CHECK-DAG: coverage Code(Counter({{[0-9]+}})) =>
-// CHECK-DAG: coverage Code(Expression({{[0-9]+}})) =>
-// CHECK:     bb0:
-// CHECK-DAG: Coverage::ExpressionUsed({{[0-9]+}})
-// CHECK-DAG: Coverage::CounterIncrement({{[0-9]+}})