coverage: Split out counter increment sites from BCB node/edge counters

Zalathar · Zalathar · commit 2e212b79e09b · 2024-02-02T10:50:05.000+11:00
This makes it possible for two nodes/edges in the coverage graph to share the
same counter, without causing the instrumentor to inject unwanted duplicate
counter-increment statements.
diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs
@@ -1,4 +1,5 @@
-use rustc_data_structures::fx::FxIndexMap;
+use rustc_data_structures::captures::Captures;
+use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::graph::WithNumNodes;
 use rustc_index::bit_set::BitSet;
 use rustc_index::IndexVec;
@@ -38,19 +39,27 @@ impl Debug for BcbCounter {
     }
 }
 
+#[derive(Debug)]
+pub(super) enum CounterIncrementSite {
+    Node { bcb: BasicCoverageBlock },
+    Edge { from_bcb: BasicCoverageBlock, to_bcb: BasicCoverageBlock },
+}
+
 /// Generates and stores coverage counter and coverage expression information
 /// associated with nodes/edges in the BCB graph.
 pub(super) struct CoverageCounters {
-    next_counter_id: CounterId,
+    /// List of places where a counter-increment statement should be injected
+    /// into MIR, each with its corresponding counter ID.
+    counter_increment_sites: IndexVec<CounterId, CounterIncrementSite>,
 
     /// Coverage counters/expressions that are associated with individual BCBs.
     bcb_counters: IndexVec<BasicCoverageBlock, Option<BcbCounter>>,
     /// Coverage counters/expressions that are associated with the control-flow
     /// edge between two BCBs.
     ///
-    /// The iteration order of this map can affect the precise contents of MIR,
-    /// so we use `FxIndexMap` to avoid query stability hazards.
-    bcb_edge_counters: FxIndexMap<(BasicCoverageBlock, BasicCoverageBlock), BcbCounter>,
+    /// We currently don't iterate over this map, but if we do in the future,
+    /// switch it back to `FxIndexMap` to avoid query stability hazards.
+    bcb_edge_counters: FxHashMap<(BasicCoverageBlock, BasicCoverageBlock), BcbCounter>,
     /// Tracks which BCBs have a counter associated with some incoming edge.
     /// Only used by assertions, to verify that BCBs with incoming edge
     /// counters do not have their own physical counters (expressions are allowed).
@@ -71,9 +80,9 @@ impl CoverageCounters {
         let num_bcbs = basic_coverage_blocks.num_nodes();
 
         let mut this = Self {
-            next_counter_id: CounterId::START,
+            counter_increment_sites: IndexVec::new(),
             bcb_counters: IndexVec::from_elem_n(None, num_bcbs),
-            bcb_edge_counters: FxIndexMap::default(),
+            bcb_edge_counters: FxHashMap::default(),
             bcb_has_incoming_edge_counters: BitSet::new_empty(num_bcbs),
             expressions: IndexVec::new(),
         };
@@ -84,8 +93,8 @@ impl CoverageCounters {
         this
     }
 
-    fn make_counter(&mut self) -> BcbCounter {
-        let id = self.next_counter();
+    fn make_counter(&mut self, site: CounterIncrementSite) -> BcbCounter {
+        let id = self.counter_increment_sites.push(site);
         BcbCounter::Counter { id }
     }
 
@@ -103,15 +112,8 @@ impl CoverageCounters {
         self.make_expression(lhs, Op::Add, rhs)
     }
 
-    /// Counter IDs start from one and go up.
-    fn next_counter(&mut self) -> CounterId {
-        let next = self.next_counter_id;
-        self.next_counter_id = self.next_counter_id + 1;
-        next
-    }
-
     pub(super) fn num_counters(&self) -> usize {
-        self.next_counter_id.as_usize()
+        self.counter_increment_sites.len()
     }
 
     #[cfg(test)]
@@ -171,22 +173,26 @@ impl CoverageCounters {
         self.bcb_counters[bcb]
     }
 
-    pub(super) fn bcb_node_counters(
+    /// Returns an iterator over all the nodes/edges in the coverage graph that
+    /// should have a counter-increment statement injected into MIR, along with
+    /// each site's corresponding counter ID.
+    pub(super) fn counter_increment_sites(
         &self,
-    ) -> impl Iterator<Item = (BasicCoverageBlock, &BcbCounter)> {
-        self.bcb_counters
-            .iter_enumerated()
-            .filter_map(|(bcb, counter_kind)| Some((bcb, counter_kind.as_ref()?)))
+    ) -> impl Iterator<Item = (CounterId, &CounterIncrementSite)> {
+        self.counter_increment_sites.iter_enumerated()
     }
 
-    /// For each edge in the BCB graph that has an associated counter, yields
-    /// that edge's *from* and *to* nodes, and its counter.
-    pub(super) fn bcb_edge_counters(
+    /// Returns an iterator over the subset of BCB nodes that have been associated
+    /// with a counter *expression*, along with the ID of that expression.
+    pub(super) fn bcb_nodes_with_coverage_expressions(
         &self,
-    ) -> impl Iterator<Item = (BasicCoverageBlock, BasicCoverageBlock, &BcbCounter)> {
-        self.bcb_edge_counters
-            .iter()
-            .map(|(&(from_bcb, to_bcb), counter_kind)| (from_bcb, to_bcb, counter_kind))
+    ) -> impl Iterator<Item = (BasicCoverageBlock, ExpressionId)> + Captures<'_> {
+        self.bcb_counters.iter_enumerated().filter_map(|(bcb, &counter_kind)| match counter_kind {
+            // Yield the BCB along with its associated expression ID.
+            Some(BcbCounter::Expression { id }) => Some((bcb, id)),
+            // This BCB is associated with a counter or nothing, so skip it.
+            Some(BcbCounter::Counter { .. }) | None => None,
+        })
     }
 
     pub(super) fn into_expressions(self) -> IndexVec<ExpressionId, Expression> {
@@ -339,7 +345,8 @@ impl<'a> MakeBcbCounters<'a> {
         // program results in a tight infinite loop, but it should still compile.
         let one_path_to_target = !self.basic_coverage_blocks.bcb_has_multiple_in_edges(bcb);
         if one_path_to_target || self.bcb_predecessors(bcb).contains(&bcb) {
-            let counter_kind = self.coverage_counters.make_counter();
+            let counter_kind =
+                self.coverage_counters.make_counter(CounterIncrementSite::Node { bcb });
             if one_path_to_target {
                 debug!("{bcb:?} gets a new counter: {counter_kind:?}");
             } else {
@@ -401,7 +408,8 @@ impl<'a> MakeBcbCounters<'a> {
         }
 
         // Make a new counter to count this edge.
-        let counter_kind = self.coverage_counters.make_counter();
+        let counter_kind =
+            self.coverage_counters.make_counter(CounterIncrementSite::Edge { from_bcb, to_bcb });
         debug!("Edge {from_bcb:?}->{to_bcb:?} gets a new counter: {counter_kind:?}");
         self.coverage_counters.set_bcb_edge_counter(from_bcb, to_bcb, counter_kind)
     }
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -7,7 +7,7 @@ mod spans;
 #[cfg(test)]
 mod tests;
 
-use self::counters::{BcbCounter, CoverageCounters};
+use self::counters::{CounterIncrementSite, CoverageCounters};
 use self::graph::{BasicCoverageBlock, CoverageGraph};
 use self::spans::{BcbMapping, BcbMappingKind, CoverageSpans};
 
@@ -155,61 +155,52 @@ fn inject_coverage_statements<'tcx>(
     bcb_has_coverage_spans: impl Fn(BasicCoverageBlock) -> bool,
     coverage_counters: &CoverageCounters,
 ) {
-    // Process the counters associated with BCB nodes.
-    for (bcb, counter_kind) in coverage_counters.bcb_node_counters() {
-        let do_inject = match counter_kind {
-            // Counter-increment statements always need to be injected.
-            BcbCounter::Counter { .. } => true,
-            // The only purpose of expression-used statements is to detect
-            // when a mapping is unreachable, so we only inject them for
-            // expressions with one or more mappings.
-            BcbCounter::Expression { .. } => bcb_has_coverage_spans(bcb),
-        };
-        if do_inject {
-            inject_statement(
-                mir_body,
-                make_mir_coverage_kind(counter_kind),
-                basic_coverage_blocks[bcb].leader_bb(),
-            );
-        }
-    }
-
-    // Process the counters associated with BCB edges.
-    for (from_bcb, to_bcb, counter_kind) in coverage_counters.bcb_edge_counters() {
-        let do_inject = match counter_kind {
-            // Counter-increment statements always need to be injected.
-            BcbCounter::Counter { .. } => true,
-            // BCB-edge expressions never have mappings, so they never need
-            // a corresponding statement.
-            BcbCounter::Expression { .. } => false,
+    // Inject counter-increment statements into MIR.
+    for (id, counter_increment_site) in coverage_counters.counter_increment_sites() {
+        // Determine the block to inject a counter-increment statement into.
+        // For BCB nodes this is just their first block, but for edges we need
+        // to create a new block between the two BCBs, and inject into that.
+        let target_bb = match *counter_increment_site {
+            CounterIncrementSite::Node { bcb } => basic_coverage_blocks[bcb].leader_bb(),
+            CounterIncrementSite::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 = basic_coverage_blocks[from_bcb].last_bb();
+                let to_bb = basic_coverage_blocks[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
+            }
         };
-        if !do_inject {
-            continue;
-        }
-
-        // We need to inject a coverage statement into a new BB between the
-        // last BB of `from_bcb` and the first BB of `to_bcb`.
-        let from_bb = basic_coverage_blocks[from_bcb].last_bb();
-        let to_bb = basic_coverage_blocks[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 edge counter {counter_kind:?}",
-        );
-
-        // Inject a counter into the newly-created BB.
-        inject_statement(mir_body, make_mir_coverage_kind(counter_kind), new_bb);
+        inject_statement(mir_body, CoverageKind::CounterIncrement { id }, target_bb);
     }
-}
 
-fn make_mir_coverage_kind(counter_kind: &BcbCounter) -> CoverageKind {
-    match *counter_kind {
-        BcbCounter::Counter { id } => CoverageKind::CounterIncrement { id },
-        BcbCounter::Expression { id } => CoverageKind::ExpressionUsed { id },
+    // For each counter expression that is directly associated with at least one
+    // span, we inject an "expression-used" statement, so that coverage codegen
+    // can check whether the injected statement survived MIR optimization.
+    // (BCB edges can't have spans, so we only need to process BCB nodes here.)
+    //
+    // See the code in `rustc_codegen_llvm::coverageinfo::map_data` that deals
+    // with "expressions seen" and "zero terms".
+    for (bcb, expression_id) in coverage_counters
+        .bcb_nodes_with_coverage_expressions()
+        .filter(|&(bcb, _)| bcb_has_coverage_spans(bcb))
+    {
+        inject_statement(
+            mir_body,
+            CoverageKind::ExpressionUsed { id: expression_id },
+            basic_coverage_blocks[bcb].leader_bb(),
+        );
     }
 }
 
+/// 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,
