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

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

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

compiler/rustc_middle/src/mir/pretty.rs

@@ -486,14 +486,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
@@ -513,8 +517,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

@@ -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>,
@@ -38,6 +38,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,
+}
+
 impl BranchInfoBuilder {
     /// Creates a new branch info builder, but only if branch coverage instrumentation
     /// is enabled and `def_id` represents a function that is eligible for coverage.
@@ -46,7 +52,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 +150,37 @@ 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)]);
+    }
+
+    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.inject_block_marker(cfg, source_info, pre_binding_block.unwrap());
+                let arm_taken_marker = self.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);
     }
 
     fn next_block_marker_id(&mut self) -> BlockMarkerId {
@@ -174,20 +210,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 +373,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,

compiler/rustc_mir_build/src/build/matches/mod.rs

@@ -5,6 +5,7 @@
 //! This also includes code for pattern bindings in `let` statements and
 //! function parameters.
 
+use crate::build::coverageinfo;
 use crate::build::expr::as_place::PlaceBuilder;
 use crate::build::scope::DropKind;
 use crate::build::ForGuard::{self, OutsideGuard, RefWithinGuard};
@@ -461,6 +462,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)| {
@@ -495,6 +498,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,
@@ -504,6 +509,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 {
@@ -515,6 +528,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

@@ -10,7 +10,7 @@ use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph, TraverseCoverage
 
 /// The coverage counter or counter expression associated with a particular
 /// BCB node or BCB edge.
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, PartialEq, Eq)]
 pub(super) enum BcbCounter {
     Counter { id: CounterId },
     Expression { id: ExpressionId },
@@ -34,6 +34,13 @@ impl Debug for BcbCounter {
     }
 }
 
+#[derive(Debug)]
+struct BcbExpression {
+    lhs: BcbCounter,
+    op: Op,
+    rhs: BcbCounter,
+}
+
 #[derive(Debug)]
 pub(super) enum CounterIncrementSite {
     Node { bcb: BasicCoverageBlock },
@@ -57,7 +64,7 @@ pub(super) struct CoverageCounters {
     bcb_edge_counters: FxHashMap<(BasicCoverageBlock, BasicCoverageBlock), BcbCounter>,
     /// Table of expression data, associating each expression ID with its
     /// corresponding operator (+ or -) and its LHS/RHS operands.
-    expressions: IndexVec<ExpressionId, Expression>,
+    expressions: IndexVec<ExpressionId, BcbExpression>,
 }
 
 impl CoverageCounters {
@@ -88,9 +95,23 @@ impl CoverageCounters {
         BcbCounter::Counter { id }
     }
 
-    fn make_expression(&mut self, lhs: BcbCounter, op: Op, rhs: BcbCounter) -> BcbCounter {
-        let expression = Expression { lhs: lhs.as_term(), op, rhs: rhs.as_term() };
-        let id = self.expressions.push(expression);
+    pub(super) fn make_expression(
+        &mut self,
+        lhs: BcbCounter,
+        op: Op,
+        rhs: BcbCounter,
+    ) -> BcbCounter {
+        // Replace `(a + b) - b` with `a`, since this happens often.
+        if op == Op::Subtract
+            && let BcbCounter::Expression { id } = lhs
+            && let lhs_expr = &self.expressions[id]
+            && lhs_expr.op == Op::Add
+            && lhs_expr.rhs == rhs
+        {
+            return lhs_expr.lhs;
+        }
+
+        let id = self.expressions.push(BcbExpression { lhs, op, rhs });
         BcbCounter::Expression { id }
     }
 
@@ -165,7 +186,21 @@ impl CoverageCounters {
     }
 
     pub(super) fn into_expressions(self) -> IndexVec<ExpressionId, Expression> {
-        self.expressions
+        let old_len = self.expressions.len();
+        let expressions = self
+            .expressions
+            .into_iter()
+            .map(|BcbExpression { lhs, op, rhs }| Expression {
+                lhs: lhs.as_term(),
+                op,
+                rhs: rhs.as_term(),
+            })
+            .collect::<IndexVec<ExpressionId, _>>();
+
+        // Expression IDs are indexes into this vector, so make sure we didn't
+        // accidentally invalidate them by changing its length.
+        assert_eq!(old_len, expressions.len());
+        expressions
     }
 }