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

Author: Zalathar

compiler/rustc_middle/src/mir/coverage.rs

@@ -274,17 +274,21 @@ pub struct CoverageInfoHi {
     /// 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)]

compiler/rustc_middle/src/mir/pretty.rs

@@ -537,19 +537,20 @@ fn write_coverage_info_hi(
 ) -> io::Result<()> {
     let coverage::CoverageInfoHi {
         num_block_markers: _,
-        branch_spans,
+        branch_arm_lists,
         mcdc_branch_spans,
         mcdc_decision_spans,
     } = coverage_info_hi;
 
     // Only add an extra trailing newline if we printed at least one thing.
     let mut did_print = false;
 
-    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}}}")?;
         did_print = true;
     }
 
@@ -1739,12 +1740,13 @@ pub fn write_allocation_bytes<'tcx, Prov: Provenance, Extra, Bytes: AllocBytes>(
             }
         } else if let Some(prov) = alloc.provenance().get(i, &tcx) {
             // Memory with provenance must be defined
-            assert!(
-                alloc.init_mask().is_range_initialized(alloc_range(i, Size::from_bytes(1))).is_ok()
-            );
+            assert!(alloc
+                .init_mask()
+                .is_range_initialized(alloc_range(i, Size::from_bytes(1)))
+                .is_ok());
             ascii.push('━'); // HEAVY HORIZONTAL
-            // We have two characters to display this, which is obviously not enough.
-            // Format is similar to "oversized" above.
+                             // We have two characters to display this, which is obviously not enough.
+                             // Format is similar to "oversized" above.
             let j = i.bytes_usize();
             let c = alloc.inspect_with_uninit_and_ptr_outside_interpreter(j..j + 1)[0];
             write!(w, "╾{c:02x}{prov:#?} (1 ptr byte)╼")?;

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, CoverageInfoHi, CoverageKind};
+use rustc_middle::mir::coverage::{BlockMarkerId, BranchArm, CoverageInfoHi, CoverageKind};
 use rustc_middle::mir::{self, BasicBlock, SourceInfo, UnOp};
 use rustc_middle::thir::{ExprId, ExprKind, Pat, Thir};
 use rustc_middle::ty::TyCtxt;
@@ -23,13 +23,14 @@ pub(crate) struct CoverageInfoBuilder {
 
     /// Present if branch coverage is enabled.
     branch_info: Option<BranchInfo>,
+
     /// Present if MC/DC coverage is enabled.
     mcdc_info: Option<MCDCInfoBuilder>,
 }
 
 #[derive(Default)]
 struct BranchInfo {
-    branch_spans: Vec<BranchSpan>,
+    branch_arm_lists: Vec<Vec<BranchArm>>,
 }
 
 #[derive(Clone, Copy)]
@@ -141,6 +142,9 @@ impl CoverageInfoBuilder {
         true_block: BasicBlock,
         false_block: BasicBlock,
     ) {
+        // Bail out if branch coverage is not enabled.
+        let Some(branch_info) = self.branch_info.as_mut() else { return };
+
         // Separate path for handling branches when MC/DC is enabled.
         if let Some(mcdc_info) = self.mcdc_info.as_mut() {
             let inject_block_marker =
@@ -152,37 +156,31 @@ impl CoverageInfoBuilder {
                 false_block,
                 inject_block_marker,
             );
-            return;
+        } else {
+            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);
+
+            let arm = |marker| BranchArm {
+                span: source_info.span,
+                pre_guard_marker: marker,
+                arm_taken_marker: marker,
+            };
+            branch_info.branch_arm_lists.push(vec![arm(true_marker), arm(false_marker)]);
         }
-
-        // Bail out if branch coverage is not enabled.
-        let Some(branch_info) = self.branch_info.as_mut() else { return };
-
-        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);
-
-        branch_info.branch_spans.push(BranchSpan {
-            span: source_info.span,
-            true_marker,
-            false_marker,
-        });
     }
 
     pub(crate) fn into_done(self) -> Box<CoverageInfoHi> {
         let Self { nots: _, markers: BlockMarkerGen { num_block_markers }, branch_info, mcdc_info } =
             self;
 
-        let branch_spans =
-            branch_info.map(|branch_info| branch_info.branch_spans).unwrap_or_default();
-
         let (mcdc_decision_spans, mcdc_branch_spans) =
             mcdc_info.map(MCDCInfoBuilder::into_done).unwrap_or_default();
 
         // For simplicity, always return an info struct (without Option), even
         // if there's nothing interesting in it.
         Box::new(CoverageInfoHi {
             num_block_markers,
-            branch_spans,
+            branch_arm_lists: branch_info.map(|i| i.branch_arm_lists).unwrap_or_default(),
             mcdc_branch_spans,
             mcdc_decision_spans,
         })
@@ -255,7 +253,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,

compiler/rustc_mir_transform/src/coverage/mappings.rs

@@ -4,7 +4,7 @@ use rustc_data_structures::graph::DirectedGraph;
 use rustc_index::IndexVec;
 use rustc_index::bit_set::BitSet;
 use rustc_middle::mir::coverage::{
-    BlockMarkerId, BranchSpan, ConditionInfo, CoverageInfoHi, CoverageKind,
+    BlockMarkerId, ConditionInfo, CoverageInfoHi, CoverageKind,
 };
 use rustc_middle::mir::{self, BasicBlock, StatementKind};
 use rustc_middle::ty::TyCtxt;
@@ -206,24 +206,39 @@ pub(super) fn extract_branch_arm_lists(
     let block_markers = resolve_block_markers(coverage_info_hi, mir_body);
 
     coverage_info_hi
-        .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(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(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<_>>()
 }

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);

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);