coverage: Treat each match arm as a "branch" for branch coverage

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -66,7 +66,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

@@ -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)]

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

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>,
 }
@@ -33,6 +33,12 @@ struct NotInfo {
     is_flipped: bool,
 }
 
+pub(crate) struct MatchArm {
+    pub(crate) source_info: SourceInfo,
+    pub(crate) pre_binding_block: Option<BasicBlock>,
+    pub(crate) arm_block: BasicBlock,
+}
+
 #[derive(Default)]
 struct BlockMarkerGen {
     num_block_markers: usize,
@@ -70,7 +76,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 +147,51 @@ 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 add_match_arms(&mut self, cfg: &mut CFG<'_>, arms: &[MatchArm]) {
+        // Match expressions with 0-1 arms don't have any branches for their arms.
+        if arms.len() < 2 {
+            return;
         }
+
+        // FIXME(#124118) The current implementation of branch coverage for
+        // match arms can't handle or-patterns.
+        if arms.iter().any(|arm| arm.pre_binding_block.is_none()) {
+            return;
+        }
+
+        let branch_arms = arms
+            .iter()
+            .map(|&MatchArm { source_info, pre_binding_block, arm_block }| {
+                let pre_guard_marker =
+                    self.markers.inject_block_marker(cfg, source_info, pre_binding_block.unwrap());
+                let arm_taken_marker =
+                    self.markers.inject_block_marker(cfg, source_info, arm_block);
+                BranchArm { span: source_info.span, pre_guard_marker, arm_taken_marker }
+            })
+            .collect::<Vec<_>>();
+
+        self.branch_arm_lists.push(branch_arms);
     }
 
     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 +200,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 +273,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_build/src/build/matches/mod.rs

@@ -8,7 +8,7 @@
 use crate::build::expr::as_place::PlaceBuilder;
 use crate::build::scope::DropKind;
 use crate::build::ForGuard::{self, OutsideGuard, RefWithinGuard};
-use crate::build::{BlockAnd, BlockAndExtension, Builder};
+use crate::build::{coverageinfo, BlockAnd, BlockAndExtension, Builder};
 use crate::build::{GuardFrame, GuardFrameLocal, LocalsForNode};
 use rustc_data_structures::{fx::FxIndexMap, stack::ensure_sufficient_stack};
 use rustc_hir::{BindingMode, ByRef};
@@ -473,6 +473,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         outer_source_info: SourceInfo,
         fake_borrow_temps: Vec<(Place<'tcx>, Local, FakeBorrowKind)>,
     ) -> BlockAnd<()> {
+        let mut coverage_match_arms = self.coverage_branch_info.is_some().then_some(vec![]);
+
         let arm_end_blocks: Vec<_> = arm_candidates
             .into_iter()
             .map(|(arm, candidate)| {
@@ -507,6 +509,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         opt_scrutinee_place,
                     );
 
+                    let pre_binding_block = candidate.pre_binding_block;
+
                     let arm_block = this.bind_pattern(
                         outer_source_info,
                         candidate,
@@ -516,6 +520,14 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         false,
                     );
 
+                    if let Some(coverage_match_arms) = coverage_match_arms.as_mut() {
+                        coverage_match_arms.push(coverageinfo::MatchArm {
+                            source_info: this.source_info(arm.pattern.span),
+                            pre_binding_block,
+                            arm_block,
+                        })
+                    }
+
                     this.fixed_temps_scope = old_dedup_scope;
 
                     if let Some(source_scope) = scope {
@@ -527,6 +539,13 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             })
             .collect();
 
+        if let Some(coverage_match_arms) = coverage_match_arms {
+            self.coverage_branch_info
+                .as_mut()
+                .expect("checked when creating `coverage_match_arms`")
+                .add_match_arms(&mut self.cfg, &coverage_match_arms);
+        }
+
         // all the arm blocks will rejoin here
         let end_block = self.cfg.start_new_block();
 

compiler/rustc_mir_transform/src/coverage/counters.rs

@@ -101,7 +101,12 @@ impl CoverageCounters {
         BcbCounter::Counter { id }
     }
 
-    fn make_expression(&mut self, lhs: BcbCounter, op: Op, rhs: BcbCounter) -> BcbCounter {
+    pub(super) fn make_expression(
+        &mut self,
+        lhs: BcbCounter,
+        op: Op,
+        rhs: BcbCounter,
+    ) -> BcbCounter {
         let new_expr = BcbExpression { lhs, op, rhs };
         *self
             .expressions_memo