compiler/rustc_middle/src/mir/coverage.rs

@@ -288,7 +288,12 @@ pub struct BranchInfo {
 #[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
 pub struct BranchArm {
     pub span: Span,
-    pub 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

@@ -480,8 +480,8 @@ fn write_coverage_branch_info(
 
     for arms in branch_arm_lists {
         writeln!(w, "{INDENT}coverage branches {{")?;
-        for coverage::BranchArm { span, marker } in arms {
-            writeln!(w, "{INDENT}{INDENT}{marker:?} => {span:?}")?;
+        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}}}")?;
     }

compiler/rustc_mir_build/src/build/coverageinfo.rs

@@ -40,8 +40,8 @@ struct NotInfo {
 
 pub(crate) struct MatchArm {
     pub(crate) source_info: SourceInfo,
-    pub(crate) has_guard: bool,
     pub(crate) pre_binding_block: Option<BasicBlock>,
+    pub(crate) arm_block: BasicBlock,
 }
 
 impl BranchInfoBuilder {
@@ -150,10 +150,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_arm_lists.push(vec![
-            BranchArm { span: source_info.span, marker: true_marker },
-            BranchArm { span: source_info.span, 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: Vec<MatchArm>) {
@@ -163,16 +165,18 @@ impl BranchInfoBuilder {
         }
 
         // FIXME(#124118) The current implementation of branch coverage for
-        // match arms can't handle or-patterns or guards.
-        if arms.iter().any(|arm| arm.has_guard || arm.pre_binding_block.is_none()) {
+        // 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, .. }| {
-                let marker = self.inject_block_marker(cfg, source_info, pre_binding_block.unwrap());
-                BranchArm { span: source_info.span, marker }
+            .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<_>>();
 

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

@@ -318,10 +318,6 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             &mut candidates,
         );
 
-        // Record branch coverage info for this match.
-        // (Does nothing if branch coverage is not enabled.)
-        self.visit_coverage_match_expr(&candidates);
-
         self.lower_match_arms(
             destination,
             scrutinee_place,
@@ -371,26 +367,6 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             .collect()
     }
 
-    /// If branch coverage is enabled, inject marker statements into each match
-    /// arm, and record their IDs in the table of branches.
-    ///
-    /// Unlike some of the other branch coverage visitor methods, this one needs
-    /// access to private details of [`Candidate`], so having it here is easier.
-    fn visit_coverage_match_expr(&mut self, candidates: &[&mut Candidate<'_, 'tcx>]) {
-        // Bail out if branch coverage is not enabled for this function.
-        let Some(branch_info) = self.coverage_branch_info.as_mut() else { return };
-
-        let arms = candidates
-            .iter()
-            .map(|c| coverageinfo::MatchArm {
-                source_info: SourceInfo { span: c.extra_data.span, scope: self.source_scope },
-                has_guard: c.has_guard,
-                pre_binding_block: c.pre_binding_block,
-            })
-            .collect::<Vec<_>>();
-        branch_info.add_match_arms(&mut self.cfg, arms);
-    }
-
     /// Create the decision tree for the match expression, starting from `block`.
     ///
     /// Modifies `candidates` to store the bindings and type ascriptions for
@@ -486,6 +462,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         outer_source_info: SourceInfo,
         fake_borrow_temps: Vec<(Place<'tcx>, Local)>,
     ) -> 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)| {
@@ -520,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,
@@ -529,6 +509,15 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         false,
                     );
 
+                    if let Some(coverage_match_arms) = &mut coverage_match_arms {
+                        let source_info = this.source_info(arm.pattern.span);
+                        coverage_match_arms.push(coverageinfo::MatchArm {
+                            source_info,
+                            pre_binding_block,
+                            arm_block,
+                        })
+                    }
+
                     this.fixed_temps_scope = old_dedup_scope;
 
                     if let Some(source_scope) = scope {
@@ -540,6 +529,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 {
@@ -94,8 +101,17 @@ impl CoverageCounters {
         op: Op,
         rhs: BcbCounter,
     ) -> BcbCounter {
-        let expression = Expression { lhs: lhs.as_term(), op, rhs: rhs.as_term() };
-        let id = self.expressions.push(expression);
+        // 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 }
     }
 
@@ -170,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
     }
 }
 

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -175,24 +175,37 @@ fn create_mappings<'tcx>(
         let mut rest_counter = {
             // The last arm's span is ignored, because its BCB is only used as the
             // false branch of the second-last arm; it's not a branch of its own.
-            let Some(&BcbBranchArm { span: _, bcb }) = arms_rev.next() else { continue };
-            coverage_counters.bcb_counter(bcb).expect("all relevant BCBs have counters")
+            let Some(&BcbBranchArm { span: _, pre_guard_bcb, arm_taken_bcb }) = arms_rev.next()
+            else {
+                continue;
+            };
+            debug_assert_eq!(pre_guard_bcb, arm_taken_bcb, "last arm should not have a guard");
+            coverage_counters.bcb_counter(pre_guard_bcb).expect("all relevant BCBs have counters")
         };
 
-        for &BcbBranchArm { span, bcb } in arms_rev {
-            let true_counter =
-                coverage_counters.bcb_counter(bcb).expect("all relevant BCBs have counters");
+        // All relevant BCBs should have counters, so we can `.unwrap()` them.
+        for &BcbBranchArm { span, pre_guard_bcb, arm_taken_bcb } in arms_rev {
+            // Number of times the pattern matched.
+            let matched_counter = coverage_counters.bcb_counter(pre_guard_bcb).unwrap();
+            // Number of times the pattern matched and the guard succeeded.
+            let arm_taken_counter = coverage_counters.bcb_counter(arm_taken_bcb).unwrap();
+            // Total number of times execution logically reached this pattern.
+            let reached_counter =
+                coverage_counters.make_expression(rest_counter, Op::Add, arm_taken_counter);
+            // Number of times execution reached this pattern, but didn't match it.
+            let unmatched_counter =
+                coverage_counters.make_expression(reached_counter, Op::Subtract, matched_counter);
+
             let kind = MappingKind::Branch {
-                true_term: true_counter.as_term(),
-                false_term: rest_counter.as_term(),
+                true_term: matched_counter.as_term(),
+                false_term: unmatched_counter.as_term(),
             };
 
             if let Some(code_region) = make_code_region(source_map, file_name, span, body_span) {
                 mappings.push(Mapping { kind, code_region });
             }
 
-            // FIXME: Avoid creating an unused expression on the last iteration.
-            rest_counter = coverage_counters.make_expression(true_counter, Op::Add, rest_counter);
+            rest_counter = reached_counter;
         }
     }
 

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -40,7 +40,8 @@ pub(super) struct BcbMapping {
 #[derive(Debug)]
 pub(super) struct BcbBranchArm {
     pub(super) span: Span,
-    pub(super) bcb: BasicCoverageBlock,
+    pub(super) pre_guard_bcb: BasicCoverageBlock,
+    pub(super) arm_taken_bcb: BasicCoverageBlock,
 }
 
 pub(super) struct CoverageSpans {
@@ -131,8 +132,10 @@ pub(super) fn generate_coverage_spans(
             }
         }
     }
-    for &BcbBranchArm { bcb, .. } in branch_arm_lists.iter().flatten() {
-        insert(bcb);
+    for &BcbBranchArm { span: _, pre_guard_bcb, arm_taken_bcb } in branch_arm_lists.iter().flatten()
+    {
+        insert(pre_guard_bcb);
+        insert(arm_taken_bcb);
     }
 
     Some(CoverageSpans { bcb_has_mappings, mappings, branch_arm_lists, test_vector_bitmap_bytes })

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

@@ -407,7 +407,7 @@ pub(super) fn extract_branch_arm_lists(
         .iter()
         .filter_map(|arms| {
             let mut bcb_arms = vec![];
-            for &BranchArm { span: raw_span, marker } in arms {
+            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() {
@@ -417,9 +417,12 @@ pub(super) fn extract_branch_arm_lists(
                 let (span, _) =
                     unexpand_into_body_span_with_visible_macro(raw_span, hir_info.body_span)?;
 
-                let bcb = basic_coverage_blocks.bcb_from_bb(block_markers[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, bcb });
+                bcb_arms.push(BcbBranchArm { span, pre_guard_bcb, arm_taken_bcb });
             }
 
             (bcb_arms.len() >= 2).then_some(bcb_arms)

tests/coverage/async_block.cov-map

@@ -1,19 +1,17 @@
 Function name: async_block::main
-Raw bytes (38): 0x[01, 01, 02, 01, 05, 03, 05, 06, 01, 05, 01, 00, 0b, 05, 01, 09, 00, 0a, 03, 00, 0e, 00, 13, 05, 00, 14, 01, 16, 05, 07, 0a, 02, 06, 06, 03, 01, 00, 02]
+Raw bytes (36): 0x[01, 01, 01, 01, 05, 06, 01, 05, 01, 00, 0b, 05, 01, 09, 00, 0a, 03, 00, 0e, 00, 13, 05, 00, 14, 01, 16, 05, 07, 0a, 02, 06, 01, 03, 01, 00, 02]
 Number of files: 1
 - file 0 => global file 1
-Number of expressions: 2
+Number of expressions: 1
 - expression 0 operands: lhs = Counter(0), rhs = Counter(1)
-- expression 1 operands: lhs = Expression(0, Add), rhs = Counter(1)
 Number of file 0 mappings: 6
 - Code(Counter(0)) at (prev + 5, 1) to (start + 0, 11)
 - Code(Counter(1)) at (prev + 1, 9) to (start + 0, 10)
 - Code(Expression(0, Add)) at (prev + 0, 14) to (start + 0, 19)
     = (c0 + c1)
 - Code(Counter(1)) at (prev + 0, 20) to (start + 1, 22)
 - Code(Counter(1)) at (prev + 7, 10) to (start + 2, 6)
-- Code(Expression(1, Sub)) at (prev + 3, 1) to (start + 0, 2)
-    = ((c0 + c1) - c1)
+- Code(Counter(0)) at (prev + 3, 1) to (start + 0, 2)
 
 Function name: async_block::main::{closure#0}
 Raw bytes (28): 0x[01, 01, 02, 01, 05, 05, 02, 04, 01, 07, 1c, 01, 17, 05, 01, 18, 02, 0e, 02, 02, 14, 02, 0e, 07, 03, 09, 00, 0a]

tests/coverage/branch/guard-simple.cov-map

@@ -0,0 +1,35 @@
+Function name: guard_simple::never_taken
+Raw bytes (80): 0x[01, 01, 09, 07, 01, 17, 05, 0d, 09, 01, 05, 17, 11, 0d, 09, 11, 09, 23, 0d, 05, 09, 0a, 01, 08, 01, 02, 1e, 20, 01, 02, 02, 09, 00, 0a, 20, 05, 0e, 00, 0e, 00, 1e, 05, 00, 22, 00, 24, 20, 11, 12, 01, 09, 00, 0a, 11, 00, 0e, 00, 1e, 20, 09, 1a, 00, 0e, 00, 1e, 09, 00, 22, 00, 24, 0d, 01, 0e, 00, 10, 1f, 02, 01, 00, 02]
+Number of files: 1
+- file 0 => global file 1
+Number of expressions: 9
+- expression 0 operands: lhs = Expression(1, Add), rhs = Counter(0)
+- expression 1 operands: lhs = Expression(5, Add), rhs = Counter(1)
+- expression 2 operands: lhs = Counter(3), rhs = Counter(2)
+- expression 3 operands: lhs = Counter(0), rhs = Counter(1)
+- expression 4 operands: lhs = Expression(5, Add), rhs = Counter(4)
+- expression 5 operands: lhs = Counter(3), rhs = Counter(2)
+- expression 6 operands: lhs = Counter(4), rhs = Counter(2)
+- expression 7 operands: lhs = Expression(8, Add), rhs = Counter(3)
+- expression 8 operands: lhs = Counter(1), rhs = Counter(2)
+Number of file 0 mappings: 10
+- Code(Counter(0)) at (prev + 8, 1) to (start + 2, 30)
+- Branch { true: Counter(0), false: Expression(0, Sub) } at (prev + 2, 9) to (start + 0, 10)
+    true  = c0
+    false = (((c3 + c2) + c1) - c0)
+- Branch { true: Counter(1), false: Expression(3, Sub) } at (prev + 0, 14) to (start + 0, 30)
+    true  = c1
+    false = (c0 - c1)
+- Code(Counter(1)) at (prev + 0, 34) to (start + 0, 36)
+- Branch { true: Counter(4), false: Expression(4, Sub) } at (prev + 1, 9) to (start + 0, 10)
+    true  = c4
+    false = ((c3 + c2) - c4)
+- Code(Counter(4)) at (prev + 0, 14) to (start + 0, 30)
+- Branch { true: Counter(2), false: Expression(6, Sub) } at (prev + 0, 14) to (start + 0, 30)
+    true  = c2
+    false = (c4 - c2)
+- Code(Counter(2)) at (prev + 0, 34) to (start + 0, 36)
+- Code(Counter(3)) at (prev + 1, 14) to (start + 0, 16)
+- Code(Expression(7, Add)) at (prev + 2, 1) to (start + 0, 2)
+    = ((c1 + c2) + c3)
+