coverage: Completely overhaul counter assignment, using node-flow graphs

11 files changed, 733 insertions, 632 deletions

diff --git a/compiler/rustc_data_structures/src/graph/iterate/mod.rs b/compiler/rustc_data_structures/src/graph/iterate/mod.rs
index ecda7d3fba8..7b4573d7a84 100644
--- a/compiler/rustc_data_structures/src/graph/iterate/mod.rs
+++ b/compiler/rustc_data_structures/src/graph/iterate/mod.rs
@@ -125,6 +125,16 @@ where
     pub fn visited(&self, node: G::Node) -> bool {
         self.visited.contains(node)
     }
+
+    /// Returns a reference to the set of nodes that have been visited, with
+    /// the same caveats as [`Self::visited`].
+    ///
+    /// When incorporating the visited nodes into another bitset, using bulk
+    /// operations like `union` or `intersect` can be more efficient than
+    /// processing each node individually.
+    pub fn visited_set(&self) -> &DenseBitSet<G::Node> {
+        &self.visited
+    }
 }
 
 impl<G> std::fmt::Debug for DepthFirstSearch<G>
diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs
index a9111a5ef9b..1a9323329f6 100644
--- a/compiler/rustc_mir_transform/src/coverage/counters.rs
+++ b/compiler/rustc_mir_transform/src/coverage/counters.rs
@@ -1,18 +1,24 @@
 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::graph::DirectedGraph;
 use rustc_index::IndexVec;
 use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::mir::coverage::{CounterId, CovTerm, Expression, ExpressionId, Op};
-use tracing::{debug, debug_span, instrument};
 
-use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph, ReadyFirstTraversal};
+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::graph::{BasicCoverageBlock, CoverageGraph};
 
-#[cfg(test)]
-mod tests;
+mod balanced_flow;
+mod iter_nodes;
+mod node_flow;
+mod union_find;
 
 /// The coverage counter or counter expression associated with a particular
 /// BCB node or BCB edge.
@@ -48,10 +54,12 @@ struct BcbExpression {
 }
 
 /// 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 },
-    Edge { from_bcb: BasicCoverageBlock, to_bcb: BasicCoverageBlock },
 }
 
 /// Generates and stores coverage counter and coverage expression information
@@ -79,10 +87,38 @@ impl CoverageCounters {
         graph: &CoverageGraph,
         bcb_needs_counter: &DenseBitSet<BasicCoverageBlock>,
     ) -> Self {
-        let mut builder = CountersBuilder::new(graph, bcb_needs_counter);
-        builder.make_bcb_counters();
+        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(),
+        );
 
-        builder.into_coverage_counters()
+        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 {
@@ -182,321 +218,51 @@ impl CoverageCounters {
     }
 }
 
-/// Symbolic representation of the coverage counter to be used for a particular
-/// node or edge in the coverage graph. The same site counter can be used for
-/// multiple sites, if they have been determined to have the same count.
-#[derive(Clone, Copy, Debug)]
-enum SiteCounter {
-    /// A physical counter at some node/edge.
-    Phys { site: Site },
-    /// A counter expression for a node that takes the sum of all its in-edge
-    /// counters.
-    NodeSumExpr { bcb: BasicCoverageBlock },
-    /// A counter expression for an edge that takes the counter of its source
-    /// node, and subtracts the counters of all its sibling out-edges.
-    EdgeDiffExpr { from_bcb: BasicCoverageBlock, to_bcb: BasicCoverageBlock },
-}
-
-/// Yields the graph successors of `from_bcb` that aren't `to_bcb`. This is
-/// used when creating a counter expression for [`SiteCounter::EdgeDiffExpr`].
-///
-/// For example, in this diagram the sibling out-edge targets of edge `AC` are
-/// the nodes `B` and `D`.
-///
-/// ```text
-///    A
-///  / | \
-/// B  C  D
-/// ```
-fn sibling_out_edge_targets(
-    graph: &CoverageGraph,
-    from_bcb: BasicCoverageBlock,
-    to_bcb: BasicCoverageBlock,
-) -> impl Iterator<Item = BasicCoverageBlock> + Captures<'_> {
-    graph.successors[from_bcb].iter().copied().filter(move |&t| t != to_bcb)
-}
-
-/// Helper struct that allows counter creation to inspect the BCB graph, and
-/// the set of nodes that need counters.
-struct CountersBuilder<'a> {
-    graph: &'a CoverageGraph,
-    bcb_needs_counter: &'a DenseBitSet<BasicCoverageBlock>,
-
-    site_counters: FxHashMap<Site, SiteCounter>,
-}
-
-impl<'a> CountersBuilder<'a> {
-    fn new(
-        graph: &'a CoverageGraph,
-        bcb_needs_counter: &'a DenseBitSet<BasicCoverageBlock>,
-    ) -> Self {
-        assert_eq!(graph.num_nodes(), bcb_needs_counter.domain_size());
-        Self { graph, bcb_needs_counter, site_counters: FxHashMap::default() }
-    }
-
-    fn make_bcb_counters(&mut self) {
-        debug!("make_bcb_counters(): adding a counter or expression to each BasicCoverageBlock");
-
-        // Traverse the coverage graph, ensuring that every node that needs a
-        // coverage counter has one.
-        for bcb in ReadyFirstTraversal::new(self.graph) {
-            let _span = debug_span!("traversal", ?bcb).entered();
-            if self.bcb_needs_counter.contains(bcb) {
-                self.make_node_counter_and_out_edge_counters(bcb);
-            }
-        }
-    }
-
-    /// Make sure the given node has a node counter, and then make sure each of
-    /// its out-edges has an edge counter (if appropriate).
-    #[instrument(level = "debug", skip(self))]
-    fn make_node_counter_and_out_edge_counters(&mut self, from_bcb: BasicCoverageBlock) {
-        // First, ensure that this node has a counter of some kind.
-        // We might also use that counter to compute one of the out-edge counters.
-        self.get_or_make_node_counter(from_bcb);
-
-        // If this node's out-edges won't sum to the node's counter,
-        // then there's no reason to create edge counters here.
-        if !self.graph[from_bcb].is_out_summable {
-            return;
-        }
-
-        // When choosing which out-edge should be given a counter expression, ignore edges that
-        // already have counters, or could use the existing counter of their target node.
-        let out_edge_has_counter = |to_bcb| {
-            if self.site_counters.contains_key(&Site::Edge { from_bcb, to_bcb }) {
-                return true;
-            }
-            self.graph.sole_predecessor(to_bcb) == Some(from_bcb)
-                && self.site_counters.contains_key(&Site::Node { bcb: to_bcb })
-        };
-
-        // Determine the set of out-edges that could benefit from being given an expression.
-        let candidate_successors = self.graph.successors[from_bcb]
-            .iter()
-            .copied()
-            .filter(|&to_bcb| !out_edge_has_counter(to_bcb))
-            .collect::<Vec<_>>();
-        debug!(?candidate_successors);
-
-        // If there are out-edges without counters, choose one to be given an expression
-        // (computed from this node and the other out-edges) instead of a physical counter.
-        let Some(to_bcb) = self.choose_out_edge_for_expression(from_bcb, &candidate_successors)
-        else {
-            return;
-        };
-
-        // For each out-edge other than the one that was chosen to get an expression,
-        // ensure that it has a counter (existing counter/expression or a new counter).
-        for target in sibling_out_edge_targets(self.graph, from_bcb, to_bcb) {
-            self.get_or_make_edge_counter(from_bcb, target);
-        }
-
-        // Now create an expression for the chosen edge, by taking the counter
-        // for its source node and subtracting the sum of its sibling out-edges.
-        let counter = SiteCounter::EdgeDiffExpr { from_bcb, to_bcb };
-        self.site_counters.insert(Site::Edge { from_bcb, to_bcb }, counter);
-    }
-
-    #[instrument(level = "debug", skip(self))]
-    fn get_or_make_node_counter(&mut self, bcb: BasicCoverageBlock) -> SiteCounter {
-        // If the BCB already has a counter, return it.
-        if let Some(&counter) = self.site_counters.get(&Site::Node { bcb }) {
-            debug!("{bcb:?} already has a counter: {counter:?}");
-            return counter;
-        }
-
-        let counter = self.make_node_counter_inner(bcb);
-        self.site_counters.insert(Site::Node { bcb }, counter);
-        counter
-    }
-
-    fn make_node_counter_inner(&mut self, bcb: BasicCoverageBlock) -> SiteCounter {
-        // If the node's sole in-edge already has a counter, use that.
-        if let Some(sole_pred) = self.graph.sole_predecessor(bcb)
-            && let Some(&edge_counter) =
-                self.site_counters.get(&Site::Edge { from_bcb: sole_pred, to_bcb: bcb })
-        {
-            return edge_counter;
-        }
-
-        let predecessors = self.graph.predecessors[bcb].as_slice();
-
-        // Handle cases where we can't compute a node's count from its in-edges:
-        // - START_BCB has no in-edges, so taking the sum would panic (or be wrong).
-        // - For nodes with one in-edge, or that directly loop to themselves,
-        //   trying to get the in-edge counts would require this node's counter,
-        //   leading to infinite recursion.
-        if predecessors.len() <= 1 || predecessors.contains(&bcb) {
-            debug!(?bcb, ?predecessors, "node has <=1 predecessors or is its own predecessor");
-            let counter = SiteCounter::Phys { site: Site::Node { bcb } };
-            debug!(?bcb, ?counter, "node gets a physical counter");
-            return counter;
-        }
-
-        // A BCB with multiple incoming edges can compute its count by ensuring that counters
-        // exist for each of those edges, and then adding them up to get a total count.
-        for &from_bcb in predecessors {
-            self.get_or_make_edge_counter(from_bcb, bcb);
-        }
-        let sum_of_in_edges = SiteCounter::NodeSumExpr { bcb };
-
-        debug!("{bcb:?} gets a new counter (sum of predecessor counters): {sum_of_in_edges:?}");
-        sum_of_in_edges
-    }
-
-    #[instrument(level = "debug", skip(self))]
-    fn get_or_make_edge_counter(
-        &mut self,
-        from_bcb: BasicCoverageBlock,
-        to_bcb: BasicCoverageBlock,
-    ) -> SiteCounter {
-        // If the edge already has a counter, return it.
-        if let Some(&counter) = self.site_counters.get(&Site::Edge { from_bcb, to_bcb }) {
-            debug!("Edge {from_bcb:?}->{to_bcb:?} already has a counter: {counter:?}");
-            return counter;
-        }
-
-        let counter = self.make_edge_counter_inner(from_bcb, to_bcb);
-        self.site_counters.insert(Site::Edge { from_bcb, to_bcb }, counter);
-        counter
-    }
-
-    fn make_edge_counter_inner(
-        &mut self,
-        from_bcb: BasicCoverageBlock,
-        to_bcb: BasicCoverageBlock,
-    ) -> SiteCounter {
-        // If the target node has exactly one in-edge (i.e. this one), then just
-        // use the node's counter, since it will have the same value.
-        if let Some(sole_pred) = self.graph.sole_predecessor(to_bcb) {
-            assert_eq!(sole_pred, from_bcb);
-            // This call must take care not to invoke `get_or_make_edge` for
-            // this edge, since that would result in infinite recursion!
-            return self.get_or_make_node_counter(to_bcb);
-        }
-
-        // If the source node has exactly one out-edge (i.e. this one) and would have
-        // the same execution count as that edge, then just use the node's counter.
-        if let Some(simple_succ) = self.graph.simple_successor(from_bcb) {
-            assert_eq!(simple_succ, to_bcb);
-            return self.get_or_make_node_counter(from_bcb);
-        }
-
-        // Make a new counter to count this edge.
-        let counter = SiteCounter::Phys { site: Site::Edge { from_bcb, to_bcb } };
-        debug!(?from_bcb, ?to_bcb, ?counter, "edge gets a physical counter");
-        counter
-    }
-
-    /// Given a set of candidate out-edges (represented by their successor node),
-    /// choose one to be given a counter expression instead of a physical counter.
-    fn choose_out_edge_for_expression(
-        &self,
-        from_bcb: BasicCoverageBlock,
-        candidate_successors: &[BasicCoverageBlock],
-    ) -> Option<BasicCoverageBlock> {
-        // Try to find a candidate that leads back to the top of a loop,
-        // because reloop edges tend to be executed more times than loop-exit edges.
-        if let Some(reloop_target) = self.find_good_reloop_edge(from_bcb, &candidate_successors) {
-            debug!("Selecting reloop target {reloop_target:?} to get an expression");
-            return Some(reloop_target);
-        }
-
-        // We couldn't identify a "good" edge, so just choose an arbitrary one.
-        let arbitrary_target = candidate_successors.first().copied()?;
-        debug!(?arbitrary_target, "selecting arbitrary out-edge to get an expression");
-        Some(arbitrary_target)
-    }
-
-    /// Given a set of candidate out-edges (represented by their successor node),
-    /// tries to find one that leads back to the top of a loop.
-    ///
-    /// Reloop edges are good candidates for counter expressions, because they
-    /// will tend to be executed more times than a loop-exit edge, so it's nice
-    /// for them to be able to avoid a physical counter increment.
-    fn find_good_reloop_edge(
-        &self,
-        from_bcb: BasicCoverageBlock,
-        candidate_successors: &[BasicCoverageBlock],
-    ) -> Option<BasicCoverageBlock> {
-        // If there are no candidates, avoid iterating over the loop stack.
-        if candidate_successors.is_empty() {
-            return None;
-        }
-
-        // Consider each loop on the current traversal context stack, top-down.
-        for loop_header_node in self.graph.loop_headers_containing(from_bcb) {
-            // Try to find a candidate edge that doesn't exit this loop.
-            for &target_bcb in candidate_successors {
-                // An edge is a reloop edge if its target dominates any BCB that has
-                // an edge back to the loop header. (Otherwise it's an exit edge.)
-                let is_reloop_edge = self
-                    .graph
-                    .reloop_predecessors(loop_header_node)
-                    .any(|reloop_bcb| self.graph.dominates(target_bcb, reloop_bcb));
-                if is_reloop_edge {
-                    // We found a good out-edge to be given an expression.
-                    return Some(target_bcb);
-                }
-            }
-
-            // All of the candidate edges exit this loop, so keep looking
-            // for a good reloop edge for one of the outer loops.
-        }
-
-        None
-    }
-
-    fn into_coverage_counters(self) -> CoverageCounters {
-        Transcriber::new(&self).transcribe_counters()
-    }
-}
-
-/// Helper struct for converting `CountersBuilder` into a final `CoverageCounters`.
-struct Transcriber<'a> {
-    old: &'a CountersBuilder<'a>,
+struct Transcriber {
+    old: NodeCounters<BasicCoverageBlock>,
     new: CoverageCounters,
     phys_counter_for_site: FxHashMap<Site, BcbCounter>,
 }
 
-impl<'a> Transcriber<'a> {
-    fn new(old: &'a CountersBuilder<'a>) -> Self {
+impl Transcriber {
+    fn new(num_nodes: usize, old: NodeCounters<BasicCoverageBlock>) -> Self {
         Self {
             old,
-            new: CoverageCounters::with_num_bcbs(old.graph.num_nodes()),
+            new: CoverageCounters::with_num_bcbs(num_nodes),
             phys_counter_for_site: FxHashMap::default(),
         }
     }
 
-    fn transcribe_counters(mut self) -> CoverageCounters {
-        for bcb in self.old.bcb_needs_counter.iter() {
+    fn transcribe_counters(
+        mut self,
+        bcb_needs_counter: &DenseBitSet<BasicCoverageBlock>,
+    ) -> CoverageCounters {
+        for bcb in bcb_needs_counter.iter() {
             let site = Site::Node { bcb };
-            let site_counter = self.site_counter(site);
-
-            // Resolve the site counter into flat lists of nodes/edges whose
-            // physical counts contribute to the counter for this node.
-            // Distinguish between counts that will be added vs subtracted.
-            let mut pos = vec![];
-            let mut neg = vec![];
-            self.push_resolved_sites(site_counter, &mut pos, &mut neg);
-
-            // Simplify by cancelling out sites that appear on both sides.
-            let (mut pos, mut neg) = sort_and_cancel(pos, neg);
+            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 after cancellation,
-                // fall back to creating a physical counter. There's no known way
-                // for this to happen, but it's hard to confidently rule it out.
+                // 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![Some(site)];
+                pos = vec![bcb];
                 neg = vec![];
             }
 
-            let mut new_counters_for_sites = |sites: Vec<Option<Site>>| {
+            pos.sort();
+            neg.sort();
+
+            let mut new_counters_for_sites = |sites: Vec<BasicCoverageBlock>| {
                 sites
                     .into_iter()
-                    .filter_map(|id| try { self.ensure_phys_counter(id?) })
+                    .map(|node| self.ensure_phys_counter(Site::Node { bcb: node }))
                     .collect::<Vec<_>>()
             };
             let mut pos = new_counters_for_sites(pos);
@@ -513,79 +279,7 @@ impl<'a> Transcriber<'a> {
         self.new
     }
 
-    fn site_counter(&self, site: Site) -> SiteCounter {
-        self.old.site_counters.get(&site).copied().unwrap_or_else(|| {
-            // We should have already created all necessary site counters.
-            // But if we somehow didn't, avoid crashing in release builds,
-            // and just use an extra physical counter instead.
-            debug_assert!(false, "{site:?} should have a counter");
-            SiteCounter::Phys { site }
-        })
-    }
-
     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))
     }
-
-    /// Resolves the given counter into flat lists of nodes/edges, whose counters
-    /// will then be added and subtracted to form a counter expression.
-    fn push_resolved_sites(&self, counter: SiteCounter, pos: &mut Vec<Site>, neg: &mut Vec<Site>) {
-        match counter {
-            SiteCounter::Phys { site } => pos.push(site),
-            SiteCounter::NodeSumExpr { bcb } => {
-                for &from_bcb in &self.old.graph.predecessors[bcb] {
-                    let edge_counter = self.site_counter(Site::Edge { from_bcb, to_bcb: bcb });
-                    self.push_resolved_sites(edge_counter, pos, neg);
-                }
-            }
-            SiteCounter::EdgeDiffExpr { from_bcb, to_bcb } => {
-                // First, add the count for `from_bcb`.
-                let node_counter = self.site_counter(Site::Node { bcb: from_bcb });
-                self.push_resolved_sites(node_counter, pos, neg);
-
-                // Then subtract the counts for the other out-edges.
-                for target in sibling_out_edge_targets(self.old.graph, from_bcb, to_bcb) {
-                    let edge_counter = self.site_counter(Site::Edge { from_bcb, to_bcb: target });
-                    // Swap `neg` and `pos` so that the counter is subtracted.
-                    self.push_resolved_sites(edge_counter, neg, pos);
-                }
-            }
-        }
-    }
-}
-
-/// Given two lists:
-/// - Sorts each list.
-/// - Converts each list to `Vec<Option<T>>`.
-/// - Scans for values that appear in both lists, and cancels them out by
-///   replacing matching pairs of values with `None`.
-fn sort_and_cancel<T: Ord>(mut pos: Vec<T>, mut neg: Vec<T>) -> (Vec<Option<T>>, Vec<Option<T>>) {
-    pos.sort();
-    neg.sort();
-
-    // Convert to `Vec<Option<T>>`. If `T` has a niche, this should be zero-cost.
-    let mut pos = pos.into_iter().map(Some).collect::<Vec<_>>();
-    let mut neg = neg.into_iter().map(Some).collect::<Vec<_>>();
-
-    // Scan through the lists using two cursors. When either cursor reaches the
-    // end of its list, there can be no more equal pairs, so stop.
-    let mut p = 0;
-    let mut n = 0;
-    while p < pos.len() && n < neg.len() {
-        // If the values are equal, remove them and advance both cursors.
-        // Otherwise, advance whichever cursor points to the lesser value.
-        // (Choosing which cursor to advance relies on both lists being sorted.)
-        match pos[p].cmp(&neg[n]) {
-            Ordering::Less => p += 1,
-            Ordering::Equal => {
-                pos[p] = None;
-                neg[n] = None;
-                p += 1;
-                n += 1;
-            }
-            Ordering::Greater => n += 1,
-        }
-    }
-
-    (pos, neg)
 }
diff --git a/compiler/rustc_mir_transform/src/coverage/counters/balanced_flow.rs b/compiler/rustc_mir_transform/src/coverage/counters/balanced_flow.rs
new file mode 100644
index 00000000000..c108f96a564
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/coverage/counters/balanced_flow.rs
@@ -0,0 +1,133 @@
+//! A control-flow graph can be said to have “balanced flow” if the flow
+//! (execution count) of each node is equal to the sum of its in-edge flows,
+//! and also equal to the sum of its out-edge flows.
+//!
+//! Control-flow graphs typically have one or more nodes that don't satisfy the
+//! balanced-flow property, e.g.:
+//! - The start node has out-edges, but no in-edges.
+//! - Return nodes have in-edges, but no out-edges.
+//! - `Yield` nodes can have an out-flow that is less than their in-flow.
+//! - Inescapable loops cause the in-flow/out-flow relationship to break down.
+//!
+//! Balanced-flow graphs are nevertheless useful for analysis, so this module
+//! provides a wrapper type ([`BalancedFlowGraph`]) that imposes balanced flow
+//! on an underlying graph. This is done by non-destructively adding synthetic
+//! nodes and edges as necessary.
+
+use rustc_data_structures::graph;
+use rustc_data_structures::graph::iterate::DepthFirstSearch;
+use rustc_data_structures::graph::reversed::ReversedGraph;
+use rustc_index::Idx;
+use rustc_index::bit_set::DenseBitSet;
+
+use crate::coverage::counters::iter_nodes::IterNodes;
+
+/// A view of an underlying graph that has been augmented to have “balanced flow”.
+/// This means that the flow (execution count) of each node is equal to the
+/// sum of its in-edge flows, and also equal to the sum of its out-edge flows.
+///
+/// To achieve this, a synthetic "sink" node is non-destructively added to the
+/// graph, with synthetic in-edges from these nodes:
+/// - Any node that has no out-edges.
+/// - Any node that explicitly requires a sink edge, as indicated by a
+///   caller-supplied `force_sink_edge` function.
+/// - Any node that would otherwise be unable to reach the sink, because it is
+///   part of an inescapable loop.
+///
+/// To make the graph fully balanced, there is also a synthetic edge from the
+/// sink node back to the start node.
+///
+/// ---
+/// The benefit of having a balanced-flow graph is that it can be subsequently
+/// transformed in ways that are guaranteed to preserve balanced flow
+/// (e.g. merging nodes together), which is useful for discovering relationships
+/// between the node flows of different nodes in the graph.
+pub(crate) struct BalancedFlowGraph<G: graph::DirectedGraph> {
+    graph: G,
+    sink_edge_nodes: DenseBitSet<G::Node>,
+    pub(crate) sink: G::Node,
+}
+
+impl<G: graph::DirectedGraph> BalancedFlowGraph<G> {
+    /// Creates a balanced view of an underlying graph, by adding a synthetic
+    /// sink node that has in-edges from nodes that need or request such an edge,
+    /// and a single out-edge to the start node.
+    ///
+    /// Assumes that all nodes in the underlying graph are reachable from the
+    /// start node.
+    pub(crate) fn for_graph(graph: G, force_sink_edge: impl Fn(G::Node) -> bool) -> Self
+    where
+        G: graph::ControlFlowGraph,
+    {
+        let mut sink_edge_nodes = DenseBitSet::new_empty(graph.num_nodes());
+        let mut dfs = DepthFirstSearch::new(ReversedGraph::new(&graph));
+
+        // First, determine the set of nodes that explicitly request or require
+        // an out-edge to the sink.
+        for node in graph.iter_nodes() {
+            if force_sink_edge(node) || graph.successors(node).next().is_none() {
+                sink_edge_nodes.insert(node);
+                dfs.push_start_node(node);
+            }
+        }
+
+        // Next, find all nodes that are currently not reverse-reachable from
+        // `sink_edge_nodes`, and add them to the set as well.
+        dfs.complete_search();
+        sink_edge_nodes.union_not(dfs.visited_set());
+
+        // The sink node is 1 higher than the highest real node.
+        let sink = G::Node::new(graph.num_nodes());
+
+        BalancedFlowGraph { graph, sink_edge_nodes, sink }
+    }
+}
+
+impl<G> graph::DirectedGraph for BalancedFlowGraph<G>
+where
+    G: graph::DirectedGraph,
+{
+    type Node = G::Node;
+
+    /// Returns the number of nodes in this balanced-flow graph, which is 1
+    /// more than the number of nodes in the underlying graph, to account for
+    /// the synthetic sink node.
+    fn num_nodes(&self) -> usize {
+        // The sink node's index is already the size of the underlying graph,
+        // so just add 1 to that instead.
+        self.sink.index() + 1
+    }
+}
+
+impl<G> graph::StartNode for BalancedFlowGraph<G>
+where
+    G: graph::StartNode,
+{
+    fn start_node(&self) -> Self::Node {
+        self.graph.start_node()
+    }
+}
+
+impl<G> graph::Successors for BalancedFlowGraph<G>
+where
+    G: graph::StartNode + graph::Successors,
+{
+    fn successors(&self, node: Self::Node) -> impl Iterator<Item = Self::Node> {
+        let real_edges;
+        let sink_edge;
+
+        if node == self.sink {
+            // The sink node has no real out-edges, and one synthetic out-edge
+            // to the start node.
+            real_edges = None;
+            sink_edge = Some(self.graph.start_node());
+        } else {
+            // Real nodes have their real out-edges, and possibly one synthetic
+            // out-edge to the sink node.
+            real_edges = Some(self.graph.successors(node));
+            sink_edge = self.sink_edge_nodes.contains(node).then_some(self.sink);
+        }
+
+        real_edges.into_iter().flatten().chain(sink_edge)
+    }
+}
diff --git a/compiler/rustc_mir_transform/src/coverage/counters/iter_nodes.rs b/compiler/rustc_mir_transform/src/coverage/counters/iter_nodes.rs
new file mode 100644
index 00000000000..9d87f7af1b0
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/coverage/counters/iter_nodes.rs
@@ -0,0 +1,16 @@
+use rustc_data_structures::graph;
+use rustc_index::Idx;
+
+pub(crate) trait IterNodes: graph::DirectedGraph {
+    /// Iterates over all nodes of a graph in ascending numeric order.
+    /// Assumes that nodes are densely numbered, i.e. every index in
+    /// `0..num_nodes` is a valid node.
+    ///
+    /// FIXME: Can this just be part of [`graph::DirectedGraph`]?
+    fn iter_nodes(
+        &self,
+    ) -> impl Iterator<Item = Self::Node> + DoubleEndedIterator + ExactSizeIterator {
+        (0..self.num_nodes()).map(<Self::Node as Idx>::new)
+    }
+}
+impl<G: graph::DirectedGraph> IterNodes for G {}
diff --git a/compiler/rustc_mir_transform/src/coverage/counters/node_flow.rs b/compiler/rustc_mir_transform/src/coverage/counters/node_flow.rs
new file mode 100644
index 00000000000..5e5d6624959
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/coverage/counters/node_flow.rs
@@ -0,0 +1,290 @@
+//! For each node in a control-flow graph, determines whether that node should
+//! have a physical counter, or a counter expression that is derived from the
+//! physical counters of other nodes.
+//!
+//! Based on the algorithm given in
+//! "Optimal measurement points for program frequency counts"
+//! (Knuth & Stevenson, 1973).
+
+use rustc_data_structures::graph;
+use rustc_index::bit_set::DenseBitSet;
+use rustc_index::{Idx, 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};
+
+#[cfg(test)]
+mod tests;
+
+/// 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(Debug)]
+pub(crate) struct MergedNodeFlowGraph<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>,
+    /// 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.)
+    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
+    where
+        G: graph::DirectedGraph<Node = Node> + graph::Successors,
+    {
+        let mut supernodes = UnionFind::<G::Node>::new(graph.num_nodes());
+
+        // For each node, merge its successors into a single supernode, and
+        // arbitrarily choose one of those successors to represent all of them.
+        let successors = graph
+            .iter_nodes()
+            .map(|node| {
+                graph
+                    .successors(node)
+                    .reduce(|a, b| supernodes.unify(a, b))
+                    .expect("each node in a balanced graph must have at least one out-edge")
+            })
+            .collect::<IndexVec<G::Node, G::Node>>();
+
+        // Now that unification is complete, freeze 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();
+
+        Self { 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
+    /// permutation of all nodes in the graph, to create physical counters and
+    /// counter expressions for each node in the underlying graph.
+    ///
+    /// The given list must contain exactly one copy of each node in the
+    /// underlying balanced-flow graph. The order of nodes is used as a hint to
+    /// 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);
+
+        for &node in all_nodes_permutation {
+            builder.visit_node(node);
+        }
+
+        NodeCounters { counter_exprs: 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
+    /// that counter as its sole element.
+    pub(crate) fn counter_expr(&self, this: Node) -> &[CounterTerm<Node>] {
+        self.counter_exprs[this].as_slice()
+    }
+}
+
+#[derive(Debug)]
+struct SpantreeEdge<Node> {
+    /// If true, this edge in the spantree has been reversed an odd number of
+    /// times, so all physical counters added to its node's counter expression
+    /// need to be negated.
+    is_reversed: bool,
+    /// Each spantree edge is "claimed" by the (regular) node that caused it to
+    /// be created. When a node with a physical counter traverses this edge,
+    /// that counter is added to the claiming node's counter expression.
+    claiming_node: Node,
+    /// Supernode at the other end of this spantree edge. Transitively points
+    /// to the "root" of this supernode's spantree component.
+    span_parent: Node,
+}
+
+/// Part of a node's counter expression, which is a sum of counter terms.
+#[derive(Debug)]
+pub(crate) struct CounterTerm<Node> {
+    /// Whether to add or subtract the value of the node's physical counter.
+    pub(crate) op: Op,
+    /// The node whose physical counter is represented by this term.
+    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>,
+    is_unvisited: DenseBitSet<Node>,
+    /// Links supernodes to each other, gradually forming a spanning tree of
+    /// the merged-flow graph.
+    ///
+    /// A supernode without a span edge is the root of its component of the
+    /// spantree. Nodes that aren't supernodes cannot have a spantree edge.
+    span_edges: IndexVec<Node, Option<SpantreeEdge<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>>,
+}
+
+impl<'a, Node: Idx> SpantreeBuilder<'a, Node> {
+    fn new(graph: &'a MergedNodeFlowGraph<Node>) -> Self {
+        let num_nodes = graph.num_nodes();
+        Self {
+            graph,
+            is_unvisited: DenseBitSet::new_filled(num_nodes),
+            span_edges: IndexVec::from_fn_n(|_| None, num_nodes),
+            counter_exprs: IndexVec::from_fn_n(|_| SmallVec::new(), num_nodes),
+        }
+    }
+
+    /// 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));
+
+        match self.span_edges[this] {
+            None => this,
+            Some(SpantreeEdge { span_parent, .. }) => self.spantree_root(span_parent),
+        }
+    }
+
+    /// 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));
+
+        // Temporarily remove this supernode (any any spantree-children) from its
+        // spantree component, by disconnecting the edge to its spantree-parent.
+        let Some(SpantreeEdge { is_reversed, claiming_node, span_parent }) =
+            self.span_edges[this].take()
+        else {
+            // This supernode has no spantree-parent edge, so it is already the
+            // root of its spantree component.
+            return;
+        };
+
+        // Recursively make our immediate spantree-parent the root of what's
+        // left of its component, so that only one more edge rotation is needed.
+        self.yank_to_spantree_root(span_parent);
+
+        // Recreate the removed edge, but in the opposite direction.
+        // Now `this` is the root of its spantree component.
+        self.span_edges[span_parent] =
+            Some(SpantreeEdge { is_reversed: !is_reversed, claiming_node, span_parent: this });
+    }
+
+    /// Must be called exactly once for each node in the balanced-flow graph.
+    fn visit_node(&mut self, this: Node) {
+        // Assert that this node was unvisited, and mark it visited.
+        assert!(self.is_unvisited.remove(this), "node has already been visited: {this:?}");
+
+        // Get the supernode containing `this`, and make it the root of its
+        // component of the spantree.
+        let this_supernode = self.graph.supernodes.find(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));
+
+        // If two supernodes are already connected in the spantree, they will
+        // have the same spantree root. (Each supernode is connected to itself.)
+        if this_supernode != self.spantree_root(succ_supernode) {
+            // Adding this node's flow edge to the spantree would cause two
+            // previously-disconnected supernodes to become connected, so add
+            // it. That spantree-edge is now "claimed" by this node.
+            //
+            // Claiming a spantree-edge means that this node will get a counter
+            // expression instead of a physical counter. That expression is
+            // currently empty, but will be built incrementally as the other
+            // nodes are visited.
+            self.span_edges[this_supernode] = Some(SpantreeEdge {
+                is_reversed: false,
+                claiming_node: this,
+                span_parent: succ_supernode,
+            });
+        } else {
+            // This node's flow edge would join two supernodes that are already
+            // connected in the spantree (or are the same supernode). That would
+            // create a cycle in the spantree, so don't add an edge.
+            //
+            // Instead, create a physical counter for this node, and add that
+            // counter to all expressions on the path from `succ_supernode` to
+            // `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 });
+
+            // Walk the spantree from `this.successor` back to `this`. For each
+            // spantree edge along the way, add this node's physical counter to
+            // the counter expression of the node that claimed the spantree edge.
+            let mut curr = succ_supernode;
+            while curr != this_supernode {
+                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 });
+
+                curr = span_parent;
+            }
+        }
+    }
+
+    /// 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, counter_exprs } = 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) }),
+            "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_exprs
+    }
+}
diff --git a/compiler/rustc_mir_transform/src/coverage/counters/node_flow/tests.rs b/compiler/rustc_mir_transform/src/coverage/counters/node_flow/tests.rs
new file mode 100644
index 00000000000..9e7f754523d
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/coverage/counters/node_flow/tests.rs
@@ -0,0 +1,64 @@
+use itertools::Itertools;
+use rustc_data_structures::graph;
+use rustc_data_structures::graph::vec_graph::VecGraph;
+use rustc_index::Idx;
+use rustc_middle::mir::coverage::Op;
+
+use super::{CounterTerm, MergedNodeFlowGraph, NodeCounters};
+
+fn merged_node_flow_graph<G: graph::Successors>(graph: G) -> MergedNodeFlowGraph<G::Node> {
+    MergedNodeFlowGraph::for_balanced_graph(graph)
+}
+
+fn make_graph<Node: Idx + Ord>(num_nodes: usize, edge_pairs: Vec<(Node, Node)>) -> VecGraph<Node> {
+    VecGraph::new(num_nodes, edge_pairs)
+}
+
+/// Example used in "Optimal Measurement Points for Program Frequency Counts"
+/// (Knuth & Stevenson, 1973), but with 0-based node IDs.
+#[test]
+fn example_driver() {
+    let graph = make_graph::<u32>(5, vec![
+        (0, 1),
+        (0, 3),
+        (1, 0),
+        (1, 2),
+        (2, 1),
+        (2, 4),
+        (3, 3),
+        (3, 4),
+        (4, 0),
+    ]);
+
+    let merged = merged_node_flow_graph(&graph);
+    let counters = merged.make_node_counters(&[3, 1, 2, 0, 4]);
+
+    assert_eq!(format_counter_expressions(&counters), &[
+        // (comment to force vertical formatting for clarity)
+        "[0]: +c0",
+        "[1]: +c0 +c2 -c4",
+        "[2]: +c2",
+        "[3]: +c3",
+        "[4]: +c4",
+    ]);
+}
+
+fn format_counter_expressions<Node: Idx>(counters: &NodeCounters<Node>) -> Vec<String> {
+    let format_item = |&CounterTerm { node, op }| {
+        let op = match op {
+            Op::Subtract => '-',
+            Op::Add => '+',
+        };
+        format!("{op}c{node:?}")
+    };
+
+    counters
+        .counter_exprs
+        .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(" "))
+        })
+        .collect()
+}
diff --git a/compiler/rustc_mir_transform/src/coverage/counters/tests.rs b/compiler/rustc_mir_transform/src/coverage/counters/tests.rs
deleted file mode 100644
index 794d4358f82..00000000000
--- a/compiler/rustc_mir_transform/src/coverage/counters/tests.rs
+++ /dev/null
@@ -1,41 +0,0 @@
-use std::fmt::Debug;
-
-use super::sort_and_cancel;
-
-fn flatten<T>(input: Vec<Option<T>>) -> Vec<T> {
-    input.into_iter().flatten().collect()
-}
-
-fn sort_and_cancel_and_flatten<T: Clone + Ord>(pos: Vec<T>, neg: Vec<T>) -> (Vec<T>, Vec<T>) {
-    let (pos_actual, neg_actual) = sort_and_cancel(pos, neg);
-    (flatten(pos_actual), flatten(neg_actual))
-}
-
-#[track_caller]
-fn check_test_case<T: Clone + Debug + Ord>(
-    pos: Vec<T>,
-    neg: Vec<T>,
-    pos_expected: Vec<T>,
-    neg_expected: Vec<T>,
-) {
-    eprintln!("pos = {pos:?}; neg = {neg:?}");
-    let output = sort_and_cancel_and_flatten(pos, neg);
-    assert_eq!(output, (pos_expected, neg_expected));
-}
-
-#[test]
-fn cancellation() {
-    let cases: &[(Vec<u32>, Vec<u32>, Vec<u32>, Vec<u32>)] = &[
-        (vec![], vec![], vec![], vec![]),
-        (vec![4, 2, 1, 5, 3], vec![], vec![1, 2, 3, 4, 5], vec![]),
-        (vec![5, 5, 5, 5, 5], vec![5], vec![5, 5, 5, 5], vec![]),
-        (vec![1, 1, 2, 2, 3, 3], vec![1, 2, 3], vec![1, 2, 3], vec![]),
-        (vec![1, 1, 2, 2, 3, 3], vec![2, 4, 2], vec![1, 1, 3, 3], vec![4]),
-    ];
-
-    for (pos, neg, pos_expected, neg_expected) in cases {
-        check_test_case(pos.to_vec(), neg.to_vec(), pos_expected.to_vec(), neg_expected.to_vec());
-        // Same test case, but with its inputs flipped and its outputs flipped.
-        check_test_case(neg.to_vec(), pos.to_vec(), neg_expected.to_vec(), pos_expected.to_vec());
-    }
-}
diff --git a/compiler/rustc_mir_transform/src/coverage/counters/union_find.rs b/compiler/rustc_mir_transform/src/coverage/counters/union_find.rs
new file mode 100644
index 00000000000..2da4f5f5fce
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/coverage/counters/union_find.rs
@@ -0,0 +1,116 @@
+use std::cmp::Ordering;
+use std::mem;
+
+use rustc_index::{Idx, IndexVec};
+
+#[cfg(test)]
+mod tests;
+
+/// Simple implementation of a union-find data structure, i.e. a disjoint-set
+/// forest.
+#[derive(Debug)]
+pub(crate) struct UnionFind<Key: Idx> {
+    table: IndexVec<Key, UnionFindEntry<Key>>,
+}
+
+#[derive(Debug)]
+struct UnionFindEntry<Key> {
+    /// Transitively points towards the "root" of the set containing this key.
+    ///
+    /// Invariant: A root key is its own parent.
+    parent: Key,
+    /// When merging two "root" keys, their ranks determine which key becomes
+    /// the new root, to prevent the parent tree from becoming unnecessarily
+    /// tall. See [`UnionFind::unify`] for details.
+    rank: u32,
+}
+
+impl<Key: Idx> UnionFind<Key> {
+    /// Creates a new disjoint-set forest containing the keys `0..num_keys`.
+    /// Initially, every key is part of its own one-element set.
+    pub(crate) fn new(num_keys: usize) -> Self {
+        // Initially, every key is the root of its own set, so its parent is itself.
+        Self { table: IndexVec::from_fn_n(|key| UnionFindEntry { parent: key, rank: 0 }, num_keys) }
+    }
+
+    /// 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.
+    ///
+    /// Also updates internal data structures to make subsequent `find`
+    /// operations faster.
+    pub(crate) fn find(&mut self, key: Key) -> Key {
+        // Loop until we find a key that is its own parent.
+        let mut curr = key;
+        while let parent = self.table[curr].parent
+            && curr != parent
+        {
+            // Perform "path compression" by peeking one layer ahead, and
+            // setting the current key's parent to that value.
+            // (This works even when `parent` is the root of its set, because
+            // of the invariant that a root is its own parent.)
+            let parent_parent = self.table[parent].parent;
+            self.table[curr].parent = parent_parent;
+
+            // Advance by one step and continue.
+            curr = parent;
+        }
+        curr
+    }
+
+    /// Merges the set containing `a` and the set containing `b` into one set.
+    ///
+    /// Returns the common root of both keys, after the merge.
+    pub(crate) fn unify(&mut self, a: Key, b: Key) -> Key {
+        let mut a = self.find(a);
+        let mut b = self.find(b);
+
+        // If both keys have the same root, they're already in the same set,
+        // so there's nothing more to do.
+        if a == b {
+            return a;
+        };
+
+        // Ensure that `a` has strictly greater rank, swapping if necessary.
+        // If both keys have the same rank, increment the rank of `a` so that
+        // future unifications will also prefer `a`, leading to flatter trees.
+        match Ord::cmp(&self.table[a].rank, &self.table[b].rank) {
+            Ordering::Less => mem::swap(&mut a, &mut b),
+            Ordering::Equal => self.table[a].rank += 1,
+            Ordering::Greater => {}
+        }
+
+        debug_assert!(self.table[a].rank > self.table[b].rank);
+        debug_assert_eq!(self.table[b].parent, b);
+
+        // Make `a` the parent of `b`.
+        self.table[b].parent = a;
+
+        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]
+    }
+}
diff --git a/compiler/rustc_mir_transform/src/coverage/counters/union_find/tests.rs b/compiler/rustc_mir_transform/src/coverage/counters/union_find/tests.rs
new file mode 100644
index 00000000000..34a4e4f8e6e
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/coverage/counters/union_find/tests.rs
@@ -0,0 +1,32 @@
+use super::UnionFind;
+
+#[test]
+fn empty() {
+    let mut sets = UnionFind::<u32>::new(10);
+
+    for i in 1..10 {
+        assert_eq!(sets.find(i), i);
+    }
+}
+
+#[test]
+fn transitive() {
+    let mut sets = UnionFind::<u32>::new(10);
+
+    sets.unify(3, 7);
+    sets.unify(4, 2);
+
+    assert_eq!(sets.find(7), sets.find(3));
+    assert_eq!(sets.find(2), sets.find(4));
+    assert_ne!(sets.find(3), sets.find(4));
+
+    sets.unify(7, 4);
+
+    assert_eq!(sets.find(7), sets.find(3));
+    assert_eq!(sets.find(2), sets.find(4));
+    assert_eq!(sets.find(3), sets.find(4));
+
+    for i in [0, 1, 5, 6, 8, 9] {
+        assert_eq!(sets.find(i), i);
+    }
+}
diff --git a/compiler/rustc_mir_transform/src/coverage/graph.rs b/compiler/rustc_mir_transform/src/coverage/graph.rs
index 3fa8b063fa7..25dc7f31227 100644
--- a/compiler/rustc_mir_transform/src/coverage/graph.rs
+++ b/compiler/rustc_mir_transform/src/coverage/graph.rs
@@ -1,7 +1,6 @@
 use std::cmp::Ordering;
-use std::collections::VecDeque;
 use std::ops::{Index, IndexMut};
-use std::{iter, mem, slice};
+use std::{mem, slice};
 
 use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::FxHashSet;
@@ -211,54 +210,6 @@ impl CoverageGraph {
         self.dominator_order_rank[a].cmp(&self.dominator_order_rank[b])
     }
 
-    /// Returns the source of this node's sole in-edge, if it has exactly one.
-    /// That edge can be assumed to have the same execution count as the node
-    /// itself (in the absence of panics).
-    pub(crate) fn sole_predecessor(
-        &self,
-        to_bcb: BasicCoverageBlock,
-    ) -> Option<BasicCoverageBlock> {
-        // Unlike `simple_successor`, there is no need for extra checks here.
-        if let &[from_bcb] = self.predecessors[to_bcb].as_slice() { Some(from_bcb) } else { None }
-    }
-
-    /// Returns the target of this node's sole out-edge, if it has exactly
-    /// one, but only if that edge can be assumed to have the same execution
-    /// count as the node itself (in the absence of panics).
-    pub(crate) fn simple_successor(
-        &self,
-        from_bcb: BasicCoverageBlock,
-    ) -> Option<BasicCoverageBlock> {
-        // If a node's count is the sum of its out-edges, and it has exactly
-        // one out-edge, then that edge has the same count as the node.
-        if self.bcbs[from_bcb].is_out_summable
-            && let &[to_bcb] = self.successors[from_bcb].as_slice()
-        {
-            Some(to_bcb)
-        } else {
-            None
-        }
-    }
-
-    /// For each loop that contains the given node, yields the "loop header"
-    /// node representing that loop, from innermost to outermost. If the given
-    /// node is itself a loop header, it is yielded first.
-    pub(crate) fn loop_headers_containing(
-        &self,
-        bcb: BasicCoverageBlock,
-    ) -> impl Iterator<Item = BasicCoverageBlock> + Captures<'_> {
-        let self_if_loop_header = self.is_loop_header.contains(bcb).then_some(bcb).into_iter();
-
-        let mut curr = Some(bcb);
-        let strictly_enclosing = iter::from_fn(move || {
-            let enclosing = self.enclosing_loop_header[curr?];
-            curr = enclosing;
-            enclosing
-        });
-
-        self_if_loop_header.chain(strictly_enclosing)
-    }
-
     /// For the given node, yields the subset of its predecessor nodes that
     /// it dominates. If that subset is non-empty, the node is a "loop header",
     /// and each of those predecessors represents an in-edge that jumps back to
@@ -489,126 +440,3 @@ impl<'a, 'tcx> graph::Successors for CoverageRelevantSubgraph<'a, 'tcx> {
         self.coverage_successors(bb).into_iter()
     }
 }
-
-/// State of a node in the coverage graph during ready-first traversal.
-#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
-enum ReadyState {
-    /// This node has not yet been added to the fallback queue or ready queue.
-    Unqueued,
-    /// This node is currently in the fallback queue.
-    InFallbackQueue,
-    /// This node's predecessors have all been visited, so it is in the ready queue.
-    /// (It might also have a stale entry in the fallback queue.)
-    InReadyQueue,
-    /// This node has been visited.
-    /// (It might also have a stale entry in the fallback queue.)
-    Visited,
-}
-
-/// Iterator that visits nodes in the coverage graph, in an order that always
-/// prefers "ready" nodes whose predecessors have already been visited.
-pub(crate) struct ReadyFirstTraversal<'a> {
-    graph: &'a CoverageGraph,
-
-    /// For each node, the number of its predecessor nodes that haven't been visited yet.
-    n_unvisited_preds: IndexVec<BasicCoverageBlock, u32>,
-    /// Indicates whether a node has been visited, or which queue it is in.
-    state: IndexVec<BasicCoverageBlock, ReadyState>,
-
-    /// Holds unvisited nodes whose predecessors have all been visited.
-    ready_queue: VecDeque<BasicCoverageBlock>,
-    /// Holds unvisited nodes with some unvisited predecessors.
-    /// Also contains stale entries for nodes that were upgraded to ready.
-    fallback_queue: VecDeque<BasicCoverageBlock>,
-}
-
-impl<'a> ReadyFirstTraversal<'a> {
-    pub(crate) fn new(graph: &'a CoverageGraph) -> Self {
-        let num_nodes = graph.num_nodes();
-
-        let n_unvisited_preds =
-            IndexVec::from_fn_n(|node| graph.predecessors[node].len() as u32, num_nodes);
-        let mut state = IndexVec::from_elem_n(ReadyState::Unqueued, num_nodes);
-
-        // We know from coverage graph construction that the start node is the
-        // only node with no predecessors.
-        debug_assert!(
-            n_unvisited_preds.iter_enumerated().all(|(node, &n)| (node == START_BCB) == (n == 0))
-        );
-        let ready_queue = VecDeque::from(vec![START_BCB]);
-        state[START_BCB] = ReadyState::InReadyQueue;
-
-        Self { graph, state, n_unvisited_preds, ready_queue, fallback_queue: VecDeque::new() }
-    }
-
-    /// Returns the next node from the ready queue, or else the next unvisited
-    /// node from the fallback queue.
-    fn next_inner(&mut self) -> Option<BasicCoverageBlock> {
-        // Always prefer to yield a ready node if possible.
-        if let Some(node) = self.ready_queue.pop_front() {
-            assert_eq!(self.state[node], ReadyState::InReadyQueue);
-            return Some(node);
-        }
-
-        while let Some(node) = self.fallback_queue.pop_front() {
-            match self.state[node] {
-                // This entry in the fallback queue is not stale, so yield it.
-                ReadyState::InFallbackQueue => return Some(node),
-                // This node was added to the fallback queue, but later became
-                // ready and was visited via the ready queue. Ignore it here.
-                ReadyState::Visited => {}
-                // Unqueued nodes can't be in the fallback queue, by definition.
-                // We know that the ready queue is empty at this point.
-                ReadyState::Unqueued | ReadyState::InReadyQueue => unreachable!(
-                    "unexpected state for {node:?} in the fallback queue: {:?}",
-                    self.state[node]
-                ),
-            }
-        }
-
-        None
-    }
-
-    fn mark_visited_and_enqueue_successors(&mut self, node: BasicCoverageBlock) {
-        assert!(self.state[node] < ReadyState::Visited);
-        self.state[node] = ReadyState::Visited;
-
-        // For each of this node's successors, decrease the successor's
-        // "unvisited predecessors" count, and enqueue it if appropriate.
-        for &succ in &self.graph.successors[node] {
-            let is_unqueued = match self.state[succ] {
-                ReadyState::Unqueued => true,
-                ReadyState::InFallbackQueue => false,
-                ReadyState::InReadyQueue => {
-                    unreachable!("nodes in the ready queue have no unvisited predecessors")
-                }
-                // The successor was already visited via one of its other predecessors.
-                ReadyState::Visited => continue,
-            };
-
-            self.n_unvisited_preds[succ] -= 1;
-            if self.n_unvisited_preds[succ] == 0 {
-                // This node's predecessors have all been visited, so add it to
-                // the ready queue. If it's already in the fallback queue, that
-                // fallback entry will be ignored later.
-                self.state[succ] = ReadyState::InReadyQueue;
-                self.ready_queue.push_back(succ);
-            } else if is_unqueued {
-                // This node has unvisited predecessors, so add it to the
-                // fallback queue in case we run out of ready nodes later.
-                self.state[succ] = ReadyState::InFallbackQueue;
-                self.fallback_queue.push_back(succ);
-            }
-        }
-    }
-}
-
-impl<'a> Iterator for ReadyFirstTraversal<'a> {
-    type Item = BasicCoverageBlock;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        let node = self.next_inner()?;
-        self.mark_visited_and_enqueue_successors(node);
-        Some(node)
-    }
-}
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index 57956448414..b1b609595b7 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -15,10 +15,7 @@ use rustc_middle::hir::nested_filter;
 use rustc_middle::mir::coverage::{
     CoverageKind, DecisionInfo, FunctionCoverageInfo, Mapping, MappingKind,
 };
-use rustc_middle::mir::{
-    self, BasicBlock, BasicBlockData, SourceInfo, Statement, StatementKind, Terminator,
-    TerminatorKind,
-};
+use rustc_middle::mir::{self, BasicBlock, Statement, StatementKind, TerminatorKind};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::Span;
 use rustc_span::def_id::LocalDefId;
@@ -248,19 +245,6 @@ fn inject_coverage_statements<'tcx>(
         // to create a new block between the two BCBs, and inject into that.
         let target_bb = match site {
             Site::Node { bcb } => graph[bcb].leader_bb(),
-            Site::Edge { from_bcb, to_bcb } => {
-                // Create a new block between the last block of `from_bcb` and
-                // the first block of `to_bcb`.
-                let from_bb = graph[from_bcb].last_bb();
-                let to_bb = graph[to_bcb].leader_bb();
-
-                let new_bb = inject_edge_counter_basic_block(mir_body, from_bb, to_bb);
-                debug!(
-                    "Edge {from_bcb:?} (last {from_bb:?}) -> {to_bcb:?} (leader {to_bb:?}) \
-                    requires a new MIR BasicBlock {new_bb:?} for counter increment {id:?}",
-                );
-                new_bb
-            }
         };
 
         inject_statement(mir_body, CoverageKind::CounterIncrement { id }, target_bb);
@@ -335,31 +319,6 @@ fn inject_mcdc_statements<'tcx>(
     }
 }
 
-/// Given two basic blocks that have a control-flow edge between them, creates
-/// and returns a new block that sits between those blocks.
-fn inject_edge_counter_basic_block(
-    mir_body: &mut mir::Body<'_>,
-    from_bb: BasicBlock,
-    to_bb: BasicBlock,
-) -> BasicBlock {
-    let span = mir_body[from_bb].terminator().source_info.span.shrink_to_hi();
-    let new_bb = mir_body.basic_blocks_mut().push(BasicBlockData {
-        statements: vec![], // counter will be injected here
-        terminator: Some(Terminator {
-            source_info: SourceInfo::outermost(span),
-            kind: TerminatorKind::Goto { target: to_bb },
-        }),
-        is_cleanup: false,
-    });
-    let edge_ref = mir_body[from_bb]
-        .terminator_mut()
-        .successors_mut()
-        .find(|successor| **successor == to_bb)
-        .expect("from_bb should have a successor for to_bb");
-    *edge_ref = new_bb;
-    new_bb
-}
-
 fn inject_statement(mir_body: &mut mir::Body<'_>, counter_kind: CoverageKind, bb: BasicBlock) {
     debug!("  injecting statement {counter_kind:?} for {bb:?}");
     let data = &mut mir_body[bb];