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

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -65,7 +65,15 @@ impl<'tcx> FunctionCoverageCollector<'tcx> {
         // For each expression ID that is directly used by one or more mappings,
         // mark it as not-yet-seen. This indicates that we expect to see a
         // corresponding `ExpressionUsed` statement during MIR traversal.
-        for term in function_coverage_info.mappings.iter().flat_map(|m| m.kind.terms()) {
+        for term in function_coverage_info
+            .mappings
+            .iter()
+            // For many-armed branches, some branch mappings will have expressions
+            // that don't correspond to any node in the control-flow graph, so don't
+            // expect to see `ExpressionUsed` statements for them.
+            .filter(|m| !matches!(m.kind, MappingKind::Branch { .. }))
+            .flat_map(|m| m.kind.terms())
+        {
             if let CovTerm::Expression(id) = term {
                 expressions_seen.remove(id);
             }

compiler/rustc_middle/src/mir/coverage.rs

@@ -225,13 +225,12 @@ 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>>,
 }
 
 #[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,
+    pub marker: BlockMarkerId,
 }

compiler/rustc_middle/src/mir/pretty.rs

@@ -475,15 +475,16 @@ fn write_coverage_branch_info(
     branch_info: &coverage::BranchInfo,
     w: &mut dyn io::Write,
 ) -> io::Result<()> {
-    let coverage::BranchInfo { branch_spans, .. } = branch_info;
+    let coverage::BranchInfo { branch_arm_lists, .. } = 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, marker } in arms {
+            writeln!(w, "{INDENT}{INDENT}{marker:?} => {span:?}")?;
+        }
+        writeln!(w, "{INDENT}}}")?;
     }
-    if !branch_spans.is_empty() {
+    if !branch_arm_lists.is_empty() {
         writeln!(w)?;
     }
 

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, SourceScope, UnOp};
 use rustc_middle::thir::{ExprId, ExprKind, Thir};
 use rustc_middle::ty::TyCtxt;
@@ -15,7 +15,7 @@ pub(crate) struct BranchInfoBuilder {
     nots: FxHashMap<ExprId, NotInfo>,
 
     num_block_markers: usize,
-    branch_spans: Vec<BranchSpan>,
+    branch_arm_lists: Vec<Vec<BranchArm>>,
 }
 
 #[derive(Clone, Copy)]
@@ -33,7 +33,11 @@ impl BranchInfoBuilder {
     /// is enabled and `def_id` represents a function that is eligible for coverage.
     pub(crate) fn new_if_enabled(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<Self> {
         if tcx.sess.instrument_coverage_branch() && tcx.is_eligible_for_coverage(def_id) {
-            Some(Self { nots: FxHashMap::default(), num_block_markers: 0, branch_spans: vec![] })
+            Some(Self {
+                nots: FxHashMap::default(),
+                num_block_markers: 0,
+                branch_arm_lists: vec![],
+            })
         } else {
             None
         }
@@ -102,7 +106,10 @@ impl BranchInfoBuilder {
         let true_marker = self.inject_block_marker(cfg, source_info, then_block);
         let false_marker = self.inject_block_marker(cfg, source_info, else_block);
 
-        self.branch_spans.push(BranchSpan { span, true_marker, false_marker });
+        self.branch_arm_lists.push(vec![
+            BranchArm { span, marker: true_marker },
+            BranchArm { span, marker: false_marker },
+        ]);
     }
 
     fn next_block_marker_id(&mut self) -> BlockMarkerId {
@@ -129,20 +136,20 @@ impl BranchInfoBuilder {
     }
 
     pub(crate) fn into_done(self) -> Option<Box<mir::coverage::BranchInfo>> {
-        let Self { nots: _, num_block_markers, branch_spans } = self;
+        let Self { nots: _, num_block_markers, branch_arm_lists } = 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 }))
+        Some(Box::new(mir::coverage::BranchInfo { num_block_markers, branch_arm_lists }))
     }
 }
 
 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,
         expr_id: ExprId,

compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs

@@ -1,7 +1,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};
+use rustc_middle::mir::coverage::{BlockMarkerId, BranchArm, CoverageKind};
 use rustc_middle::mir::{
     self, AggregateKind, BasicBlock, FakeReadCause, Rvalue, Statement, StatementKind, Terminator,
     TerminatorKind,
@@ -382,24 +382,26 @@ pub(super) fn extract_branch_arm_lists(
     }
 
     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![];
+            for &BranchArm { span: raw_span, 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 bcb = basic_coverage_blocks.bcb_from_bb(block_markers[marker]?)?;
+
+                bcb_arms.push(BcbBranchArm { span, bcb });
+            }
 
-            Some(vec![BcbBranchArm { span, bcb: true_bcb }, BcbBranchArm { span, bcb: false_bcb }])
+            (bcb_arms.len() >= 2).then_some(bcb_arms)
         })
         .collect::<Vec<_>>()
 }

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;

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;