coverage: Represent branches as a list of arms during MIR building, too

Zalathar · Zalathar · commit 05390e7fe87d · 2024-06-25T11:51:34.000+10:00
diff --git a/compiler/rustc_middle/src/mir/coverage.rs b/compiler/rustc_middle/src/mir/coverage.rs
@@ -282,17 +282,21 @@ pub struct BranchInfo {
     /// injected into the MIR body. This makes it possible to allocate per-ID
     /// data structures without having to scan the entire body first.
     pub num_block_markers: usize,
-    pub branch_spans: Vec<BranchSpan>,
+    pub branch_arm_lists: Vec<Vec<BranchArm>>,
     pub mcdc_branch_spans: Vec<MCDCBranchSpan>,
     pub mcdc_decision_spans: Vec<MCDCDecisionSpan>,
 }
 
 #[derive(Clone, Debug)]
 #[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
-pub struct BranchSpan {
+pub struct BranchArm {
     pub span: Span,
-    pub true_marker: BlockMarkerId,
-    pub false_marker: BlockMarkerId,
+    /// Marks the block that is jumped to after this arm's pattern matches,
+    /// but before its guard is checked.
+    pub pre_guard_marker: BlockMarkerId,
+    /// Marks the block that is jumped to after this arm's guard succeeds.
+    /// If this is equal to `pre_guard_marker`, the arm has no guard.
+    pub arm_taken_marker: BlockMarkerId,
 }
 
 #[derive(Copy, Clone, Debug)]
diff --git a/compiler/rustc_middle/src/mir/pretty.rs b/compiler/rustc_middle/src/mir/pretty.rs
@@ -487,14 +487,18 @@ fn write_coverage_branch_info(
     branch_info: &coverage::BranchInfo,
     w: &mut dyn io::Write,
 ) -> io::Result<()> {
-    let coverage::BranchInfo { branch_spans, mcdc_branch_spans, mcdc_decision_spans, .. } =
+    let coverage::BranchInfo { branch_arm_lists, mcdc_branch_spans, mcdc_decision_spans, .. } =
         branch_info;
 
-    for coverage::BranchSpan { span, true_marker, false_marker } in branch_spans {
-        writeln!(
-            w,
-            "{INDENT}coverage branch {{ true: {true_marker:?}, false: {false_marker:?} }} => {span:?}",
-        )?;
+    for arms in branch_arm_lists {
+        writeln!(w, "{INDENT}coverage branches {{")?;
+        for coverage::BranchArm { span, pre_guard_marker, arm_taken_marker } in arms {
+            writeln!(w, "{INDENT}{INDENT}{pre_guard_marker:?}, {arm_taken_marker:?} => {span:?}")?;
+        }
+        writeln!(w, "{INDENT}}}")?;
+    }
+    if !branch_arm_lists.is_empty() {
+        writeln!(w)?;
     }
 
     for coverage::MCDCBranchSpan {
@@ -521,8 +525,7 @@ fn write_coverage_branch_info(
         )?;
     }
 
-    if !branch_spans.is_empty() || !mcdc_branch_spans.is_empty() || !mcdc_decision_spans.is_empty()
-    {
+    if !mcdc_branch_spans.is_empty() || !mcdc_decision_spans.is_empty() {
         writeln!(w)?;
     }
 
diff --git a/compiler/rustc_mir_build/src/build/coverageinfo.rs b/compiler/rustc_mir_build/src/build/coverageinfo.rs
@@ -2,7 +2,7 @@ use std::assert_matches::assert_matches;
 use std::collections::hash_map::Entry;
 
 use rustc_data_structures::fx::FxHashMap;
-use rustc_middle::mir::coverage::{BlockMarkerId, BranchSpan, CoverageKind};
+use rustc_middle::mir::coverage::{BlockMarkerId, BranchArm, CoverageKind};
 use rustc_middle::mir::{self, BasicBlock, SourceInfo, UnOp};
 use rustc_middle::thir::{ExprId, ExprKind, Pat, Thir};
 use rustc_middle::ty::TyCtxt;
@@ -18,7 +18,7 @@ pub(crate) struct BranchInfoBuilder {
     nots: FxHashMap<ExprId, NotInfo>,
 
     markers: BlockMarkerGen,
-    branch_spans: Vec<BranchSpan>,
+    branch_arm_lists: Vec<Vec<BranchArm>>,
 
     mcdc_info: Option<MCDCInfoBuilder>,
 }
@@ -70,7 +70,7 @@ impl BranchInfoBuilder {
             Some(Self {
                 nots: FxHashMap::default(),
                 markers: BlockMarkerGen::default(),
-                branch_spans: vec![],
+                branch_arm_lists: vec![],
                 mcdc_info: tcx.sess.instrument_coverage_mcdc().then(MCDCInfoBuilder::new),
             })
         } else {
@@ -141,24 +141,25 @@ impl BranchInfoBuilder {
             let true_marker = self.markers.inject_block_marker(cfg, source_info, true_block);
             let false_marker = self.markers.inject_block_marker(cfg, source_info, false_block);
 
-            self.branch_spans.push(BranchSpan {
+            let arm = |marker| BranchArm {
                 span: source_info.span,
-                true_marker,
-                false_marker,
-            });
+                pre_guard_marker: marker,
+                arm_taken_marker: marker,
+            };
+            self.branch_arm_lists.push(vec![arm(true_marker), arm(false_marker)]);
         }
     }
 
     pub(crate) fn into_done(self) -> Option<Box<mir::coverage::BranchInfo>> {
         let Self {
             nots: _,
             markers: BlockMarkerGen { num_block_markers },
-            branch_spans,
+            branch_arm_lists,
             mcdc_info,
         } = self;
 
         if num_block_markers == 0 {
-            assert!(branch_spans.is_empty());
+            assert!(branch_arm_lists.is_empty());
             return None;
         }
 
@@ -167,7 +168,7 @@ impl BranchInfoBuilder {
 
         Some(Box::new(mir::coverage::BranchInfo {
             num_block_markers,
-            branch_spans,
+            branch_arm_lists,
             mcdc_branch_spans,
             mcdc_decision_spans,
         }))
@@ -240,7 +241,7 @@ impl<'tcx> Builder<'_, 'tcx> {
     }
 
     /// If branch coverage is enabled, inject marker statements into `then_block`
-    /// and `else_block`, and record their IDs in the table of branch spans.
+    /// and `else_block`, and record their IDs in the branch table.
     pub(crate) fn visit_coverage_branch_condition(
         &mut self,
         mut expr_id: ExprId,
diff --git a/compiler/rustc_mir_transform/src/coverage/mappings.rs b/compiler/rustc_mir_transform/src/coverage/mappings.rs
@@ -3,7 +3,7 @@ use std::collections::BTreeSet;
 use rustc_data_structures::graph::DirectedGraph;
 use rustc_index::bit_set::BitSet;
 use rustc_index::IndexVec;
-use rustc_middle::mir::coverage::{BlockMarkerId, BranchSpan, ConditionInfo, CoverageKind};
+use rustc_middle::mir::coverage::{BlockMarkerId, ConditionInfo, CoverageKind};
 use rustc_middle::mir::{self, BasicBlock, StatementKind};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::Span;
@@ -193,25 +193,40 @@ pub(super) fn extract_branch_arm_lists(
     let block_markers = resolve_block_markers(branch_info, mir_body);
 
     branch_info
-        .branch_spans
+        .branch_arm_lists
         .iter()
-        .filter_map(|&BranchSpan { span: raw_span, true_marker, false_marker }| {
-            // For now, ignore any branch span that was introduced by
-            // expansion. This makes things like assert macros less noisy.
-            if !raw_span.ctxt().outer_expn_data().is_root() {
-                return None;
+        .filter_map(|arms| {
+            let mut bcb_arms = Vec::with_capacity(arms.len());
+
+            // If any arm can't be resolved, return None to skip the entire list
+            // of arms that contains it.
+            for &mir::coverage::BranchArm { span: raw_span, pre_guard_marker, arm_taken_marker } in
+                arms
+            {
+                // For now, ignore any branch span that was introduced by
+                // expansion. This makes things like assert macros less noisy.
+                if !raw_span.ctxt().outer_expn_data().is_root() {
+                    return None;
+                }
+
+                let (span, _) =
+                    unexpand_into_body_span_with_visible_macro(raw_span, hir_info.body_span)?;
+
+                let pre_guard_bcb =
+                    basic_coverage_blocks.bcb_from_bb(block_markers[pre_guard_marker]?)?;
+                let arm_taken_bcb =
+                    basic_coverage_blocks.bcb_from_bb(block_markers[arm_taken_marker]?)?;
+
+                bcb_arms.push(BranchArm { span, pre_guard_bcb, arm_taken_bcb });
             }
-            let (span, _) =
-                unexpand_into_body_span_with_visible_macro(raw_span, hir_info.body_span)?;
-
-            let bcb_from_marker =
-                |marker: BlockMarkerId| basic_coverage_blocks.bcb_from_bb(block_markers[marker]?);
+            assert_eq!(arms.len(), bcb_arms.len());
 
-            let true_bcb = bcb_from_marker(true_marker)?;
-            let false_bcb = bcb_from_marker(false_marker)?;
+            if bcb_arms.len() < 2 {
+                debug_assert!(false, "MIR building shouldn't create branches with <2 arms");
+                return None;
+            }
 
-            let arm = |bcb| BranchArm { span, pre_guard_bcb: bcb, arm_taken_bcb: bcb };
-            Some(vec![arm(true_bcb), arm(false_bcb)])
+            Some(bcb_arms)
         })
         .collect::<Vec<_>>()
 }
diff --git a/tests/mir-opt/coverage/instrument_coverage_cleanup.main.CleanupPostBorrowck.diff b/tests/mir-opt/coverage/instrument_coverage_cleanup.main.CleanupPostBorrowck.diff
@@ -5,14 +5,17 @@
       let mut _0: ();
       let mut _1: bool;
   
-      coverage branch { true: BlockMarkerId(0), false: BlockMarkerId(1) } => $DIR/instrument_coverage_cleanup.rs:14:8: 14:36 (#0)
+      coverage branches {
+          BlockMarkerId(0), BlockMarkerId(0) => $DIR/instrument_coverage_cleanup.rs:14:8: 14:36 (#0)
+          BlockMarkerId(1), BlockMarkerId(1) => $DIR/instrument_coverage_cleanup.rs:14:8: 14:36 (#0)
+      }
   
       coverage ExpressionId(0) => Expression { lhs: Counter(0), op: Subtract, rhs: Counter(1) };
+      coverage Branch { true_term: Expression(0), false_term: Counter(1) } => $DIR/instrument_coverage_cleanup.rs:14:8 - 14:36;
       coverage Code(Counter(0)) => $DIR/instrument_coverage_cleanup.rs:13:1 - 14:36;
       coverage Code(Expression(0)) => $DIR/instrument_coverage_cleanup.rs:14:37 - 14:39;
       coverage Code(Counter(1)) => $DIR/instrument_coverage_cleanup.rs:14:39 - 14:40;
       coverage Code(Counter(0)) => $DIR/instrument_coverage_cleanup.rs:15:1 - 15:2;
-      coverage Branch { true_term: Expression(0), false_term: Counter(1) } => $DIR/instrument_coverage_cleanup.rs:14:8 - 14:36;
   
       bb0: {
           Coverage::CounterIncrement(0);
diff --git a/tests/mir-opt/coverage/instrument_coverage_cleanup.main.InstrumentCoverage.diff b/tests/mir-opt/coverage/instrument_coverage_cleanup.main.InstrumentCoverage.diff
@@ -5,14 +5,17 @@
       let mut _0: ();
       let mut _1: bool;
   
-      coverage branch { true: BlockMarkerId(0), false: BlockMarkerId(1) } => $DIR/instrument_coverage_cleanup.rs:14:8: 14:36 (#0)
+      coverage branches {
+          BlockMarkerId(0), BlockMarkerId(0) => $DIR/instrument_coverage_cleanup.rs:14:8: 14:36 (#0)
+          BlockMarkerId(1), BlockMarkerId(1) => $DIR/instrument_coverage_cleanup.rs:14:8: 14:36 (#0)
+      }
   
 +     coverage ExpressionId(0) => Expression { lhs: Counter(0), op: Subtract, rhs: Counter(1) };
++     coverage Branch { true_term: Expression(0), false_term: Counter(1) } => $DIR/instrument_coverage_cleanup.rs:14:8 - 14:36;
 +     coverage Code(Counter(0)) => $DIR/instrument_coverage_cleanup.rs:13:1 - 14:36;
 +     coverage Code(Expression(0)) => $DIR/instrument_coverage_cleanup.rs:14:37 - 14:39;
 +     coverage Code(Counter(1)) => $DIR/instrument_coverage_cleanup.rs:14:39 - 14:40;
 +     coverage Code(Counter(0)) => $DIR/instrument_coverage_cleanup.rs:15:1 - 15:2;
-+     coverage Branch { true_term: Expression(0), false_term: Counter(1) } => $DIR/instrument_coverage_cleanup.rs:14:8 - 14:36;
 + 
       bb0: {
 +         Coverage::CounterIncrement(0);
