coverage: Replace `bcb_has_multiple_in_edges` with `sole_predecessor`

This does a better job of expressing the special cases that occur when a node
in the coverage graph has exactly one in-edge.

2 files changed, 30 insertions, 35 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs
index 6f3814f4552..52c3a6a12e5 100644
--- a/compiler/rustc_mir_transform/src/coverage/counters.rs
+++ b/compiler/rustc_mir_transform/src/coverage/counters.rs
@@ -364,10 +364,13 @@ impl<'a> MakeBcbCounters<'a> {
         );
 
         debug!("{expression_to_bcb:?} gets an expression: {expression:?}");
-        if self.basic_coverage_blocks.bcb_has_multiple_in_edges(expression_to_bcb) {
-            self.coverage_counters.set_bcb_edge_counter(from_bcb, expression_to_bcb, expression);
-        } else {
+        if let Some(sole_pred) = self.basic_coverage_blocks.sole_predecessor(expression_to_bcb) {
+            // This edge normally wouldn't get its own counter, so attach the expression
+            // to its target node instead, so that `edge_has_no_counter` can see it.
+            assert_eq!(sole_pred, from_bcb);
             self.coverage_counters.set_bcb_counter(expression_to_bcb, expression);
+        } else {
+            self.coverage_counters.set_bcb_edge_counter(from_bcb, expression_to_bcb, expression);
         }
     }
 
@@ -413,10 +416,12 @@ impl<'a> MakeBcbCounters<'a> {
         from_bcb: BasicCoverageBlock,
         to_bcb: BasicCoverageBlock,
     ) -> BcbCounter {
-        // If the target BCB has only one in-edge (i.e. this one), then create
-        // a node counter instead, since it will have the same value.
-        if !self.basic_coverage_blocks.bcb_has_multiple_in_edges(to_bcb) {
-            assert_eq!([from_bcb].as_slice(), self.basic_coverage_blocks.predecessors[to_bcb]);
+        // 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.basic_coverage_blocks.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);
         }
 
@@ -508,18 +513,14 @@ impl<'a> MakeBcbCounters<'a> {
         from_bcb: BasicCoverageBlock,
         to_bcb: BasicCoverageBlock,
     ) -> bool {
-        self.edge_counter(from_bcb, to_bcb).is_none()
-    }
-
-    fn edge_counter(
-        &self,
-        from_bcb: BasicCoverageBlock,
-        to_bcb: BasicCoverageBlock,
-    ) -> Option<&BcbCounter> {
-        if self.basic_coverage_blocks.bcb_has_multiple_in_edges(to_bcb) {
-            self.coverage_counters.bcb_edge_counters.get(&(from_bcb, to_bcb))
-        } else {
-            self.coverage_counters.bcb_counters[to_bcb].as_ref()
-        }
+        let edge_counter =
+            if let Some(sole_pred) = self.basic_coverage_blocks.sole_predecessor(to_bcb) {
+                assert_eq!(sole_pred, from_bcb);
+                self.coverage_counters.bcb_counters[to_bcb]
+            } else {
+                self.coverage_counters.bcb_edge_counters.get(&(from_bcb, to_bcb)).copied()
+            };
+
+        edge_counter.is_none()
     }
 }
diff --git a/compiler/rustc_mir_transform/src/coverage/graph.rs b/compiler/rustc_mir_transform/src/coverage/graph.rs
index bf57a99224d..743aa679058 100644
--- a/compiler/rustc_mir_transform/src/coverage/graph.rs
+++ b/compiler/rustc_mir_transform/src/coverage/graph.rs
@@ -165,21 +165,15 @@ impl CoverageGraph {
         self.dominators.as_ref().unwrap().cmp_in_dominator_order(a, b)
     }
 
-    /// Returns true if the given node has 2 or more in-edges, i.e. 2 or more
-    /// predecessors.
-    ///
-    /// This property is interesting to code that assigns counters to nodes and
-    /// edges, because if a node _doesn't_ have multiple in-edges, then there's
-    /// no benefit in having a separate counter for its in-edge, because it
-    /// would have the same value as the node's own counter.
-    #[inline(always)]
-    pub(crate) fn bcb_has_multiple_in_edges(&self, bcb: BasicCoverageBlock) -> bool {
-        // Even though bcb0 conceptually has an extra virtual in-edge due to
-        // being the entry point, we've already asserted that it has no _other_
-        // in-edges, so there's no possibility of it having _multiple_ in-edges.
-        // (And since its virtual in-edge doesn't exist in the graph, that edge
-        // can't have a separate counter anyway.)
-        self.predecessors[bcb].len() > 1
+    /// 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