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

Zalathar · Zalathar · commit f1cc40a91b6e · 2024-04-23T13:32:44.000+10:00
diff --git a/compiler/rustc_middle/src/mir/coverage.rs b/compiler/rustc_middle/src/mir/coverage.rs
@@ -279,17 +279,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
@@ -475,14 +475,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 { span, condition_info, true_marker, false_marker } in
@@ -502,8 +506,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
@@ -4,7 +4,7 @@ use std::collections::VecDeque;
 
 use rustc_data_structures::fx::FxHashMap;
 use rustc_middle::mir::coverage::{
-    BlockMarkerId, BranchSpan, ConditionId, ConditionInfo, CoverageKind, MCDCBranchSpan,
+    BlockMarkerId, BranchArm, ConditionId, ConditionInfo, CoverageKind, MCDCBranchSpan,
     MCDCDecisionSpan,
 };
 use rustc_middle::mir::{self, BasicBlock, SourceInfo, UnOp};
@@ -21,7 +21,7 @@ pub(crate) struct BranchInfoBuilder {
     nots: FxHashMap<ExprId, NotInfo>,
 
     num_block_markers: usize,
-    branch_spans: Vec<BranchSpan>,
+    branch_arm_lists: Vec<Vec<BranchArm>>,
 
     mcdc_branch_spans: Vec<MCDCBranchSpan>,
     mcdc_decision_spans: Vec<MCDCDecisionSpan>,
@@ -46,7 +46,7 @@ impl BranchInfoBuilder {
             Some(Self {
                 nots: FxHashMap::default(),
                 num_block_markers: 0,
-                branch_spans: vec![],
+                branch_arm_lists: vec![],
                 mcdc_branch_spans: vec![],
                 mcdc_decision_spans: vec![],
                 mcdc_state: MCDCState::new_if_enabled(tcx),
@@ -144,7 +144,12 @@ impl BranchInfoBuilder {
         let true_marker = self.inject_block_marker(cfg, source_info, true_block);
         let false_marker = self.inject_block_marker(cfg, source_info, false_block);
 
-        self.branch_spans.push(BranchSpan { span: source_info.span, true_marker, false_marker });
+        let arm = |marker| BranchArm {
+            span: source_info.span,
+            pre_guard_marker: marker,
+            arm_taken_marker: marker,
+        };
+        self.branch_arm_lists.push(vec![arm(true_marker), arm(false_marker)]);
     }
 
     fn next_block_marker_id(&mut self) -> BlockMarkerId {
@@ -174,20 +179,20 @@ impl BranchInfoBuilder {
         let Self {
             nots: _,
             num_block_markers,
-            branch_spans,
+            branch_arm_lists,
             mcdc_branch_spans,
             mcdc_decision_spans,
             mcdc_state: _,
         } = self;
 
         if num_block_markers == 0 {
-            assert!(branch_spans.is_empty());
+            assert!(branch_arm_lists.is_empty());
             return None;
         }
 
         Some(Box::new(mir::coverage::BranchInfo {
             num_block_markers,
-            branch_spans,
+            branch_arm_lists,
             mcdc_branch_spans,
             mcdc_decision_spans,
         }))
@@ -337,7 +342,7 @@ impl MCDCState {
 
 impl Builder<'_, '_> {
     /// 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 table of branches.
     pub(crate) fn visit_coverage_branch_condition(
         &mut self,
         mut expr_id: ExprId,
diff --git a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
@@ -2,7 +2,7 @@ use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_index::IndexVec;
 use rustc_middle::mir::coverage::{
-    BlockMarkerId, BranchSpan, CoverageKind, MCDCBranchSpan, MCDCDecisionSpan,
+    BlockMarkerId, BranchArm, CoverageKind, MCDCBranchSpan, MCDCDecisionSpan,
 };
 use rustc_middle::mir::{
     self, AggregateKind, BasicBlock, FakeReadCause, Rvalue, Statement, StatementKind, Terminator,
@@ -403,25 +403,38 @@ 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;
-            }
-            let (span, _) =
-                unexpand_into_body_span_with_visible_macro(raw_span, hir_info.body_span)?;
+        .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 &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 bcb_from_marker =
-                |marker: BlockMarkerId| basic_coverage_blocks.bcb_from_bb(block_markers[marker]?);
+                let (span, _) =
+                    unexpand_into_body_span_with_visible_macro(raw_span, hir_info.body_span)?;
 
-            let true_bcb = bcb_from_marker(true_marker)?;
-            let false_bcb = bcb_from_marker(false_marker)?;
+                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(BcbBranchArm { span, pre_guard_bcb, arm_taken_bcb });
+            }
+            assert_eq!(arms.len(), bcb_arms.len());
+
+            if bcb_arms.len() < 2 {
+                debug_assert!(false, "MIR building shouldn't create branches with <2 arms");
+                return None;
+            }
 
-            let arm = |bcb| BcbBranchArm { 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,10 +5,14 @@
       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) => $DIR/instrument_coverage_cleanup.rs:14:8: 14:36 (#0)
+          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 ExpressionId(1) => Expression { lhs: Counter(1), op: Add, rhs: Expression(0) };
+      coverage ExpressionId(2) => Expression { lhs: Expression(0), op: Add, rhs: Counter(1) };
       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;
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,10 +5,14 @@
       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) => $DIR/instrument_coverage_cleanup.rs:14:8: 14:36 (#0)
+          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 ExpressionId(1) => Expression { lhs: Counter(1), op: Add, rhs: Expression(0) };
++     coverage ExpressionId(2) => Expression { lhs: Expression(0), op: Add, rhs: Counter(1) };
 +     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;
