coverage: Support match statements in branch coverage

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -66,7 +66,14 @@ 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 mapping in function_coverage_info.mappings.iter() {
+        for mapping 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 { .. }))
+        {
             // Currently we only worry about ordinary code mappings.
             // For branch and MC/DC mappings, expressions might not correspond
             // to any particular point in the control-flow graph.

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

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)]
@@ -42,6 +43,17 @@ struct NotInfo {
     is_flipped: bool,
 }
 
+pub(crate) struct MatchArm {
+    pub(crate) sub_branches: Vec<MatchArmSubBranch>,
+}
+
+#[derive(Debug)]
+pub(crate) struct MatchArmSubBranch {
+    pub(crate) source_info: SourceInfo,
+    pub(crate) pre_binding_block: BasicBlock,
+    pub(crate) branch_taken_block: BasicBlock,
+}
+
 #[derive(Default)]
 struct BlockMarkerGen {
     num_block_markers: usize,
@@ -152,37 +164,87 @@ impl CoverageInfoBuilder {
                 false_block,
                 inject_block_marker,
             );
-            return;
+        } else {
+            // Bail out if branch coverage is not enabled.
+            let Some(branch_info) = self.branch_info.as_mut() else { return };
+
+            // Avoid duplicates coverage markers.
+            // When lowering match sub-branches (like or-patterns), `if` guards will
+            // be added multiple times for each sub-branch
+            // FIXME: This feels dirty. It would be nice to find a smarter way to avoid duplicate
+            // coverage markers.
+            for arms in &branch_info.branch_arm_lists {
+                for arm in arms {
+                    if arm.span == source_info.span {
+                        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);
+
+            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 };
+    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;
+        }
 
-        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 Some(branch_info) = self.branch_info.as_mut() else {
+            return;
+        };
 
-        branch_info.branch_spans.push(BranchSpan {
-            span: source_info.span,
-            true_marker,
-            false_marker,
-        });
+        let branch_arms = arms
+            .iter()
+            .flat_map(|MatchArm { sub_branches }| {
+                sub_branches
+                    .iter()
+                    .map(|sub_branch| {
+                        let pre_guard_marker = self.markers.inject_block_marker(
+                            cfg,
+                            sub_branch.source_info,
+                            sub_branch.pre_binding_block,
+                        );
+                        let arm_taken_marker = self.markers.inject_block_marker(
+                            cfg,
+                            sub_branch.source_info,
+                            sub_branch.branch_taken_block,
+                        );
+
+                        BranchArm {
+                            span: sub_branch.source_info.span,
+                            pre_guard_marker,
+                            arm_taken_marker,
+                        }
+                    })
+                    .collect::<Vec<_>>()
+            })
+            .collect::<Vec<_>>();
+
+        branch_info.branch_arm_lists.push(branch_arms);
     }
 
     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 +317,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

@@ -22,7 +22,7 @@ use crate::build::ForGuard::{self, OutsideGuard, RefWithinGuard};
 use crate::build::expr::as_place::PlaceBuilder;
 use crate::build::scope::DropKind;
 use crate::build::{
-    BlockAnd, BlockAndExtension, Builder, GuardFrame, GuardFrameLocal, LocalsForNode,
+    BlockAnd, BlockAndExtension, Builder, GuardFrame, GuardFrameLocal, LocalsForNode, coverageinfo,
 };
 
 // helper functions, broken out by category:
@@ -416,7 +416,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     where
         'tcx: 'pat,
     {
-        let arm_end_blocks: Vec<BasicBlock> = arms
+        let mut coverage_match_arms = self.coverage_info.is_some().then_some(vec![]);
+
+        let arm_end_blocks: Vec<_> = arms
             .into_iter()
             .zip(built_match_tree.branches)
             .map(|(arm, branch)| {
@@ -451,6 +453,16 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         opt_scrutinee_place,
                     );
 
+                    let mut sub_branches: Vec<_> = branch
+                        .sub_branches
+                        .iter()
+                        .map(|b| coverageinfo::MatchArmSubBranch {
+                            source_info: this.source_info(b.span),
+                            pre_binding_block: b.success_block,
+                            branch_taken_block: b.success_block,
+                        })
+                        .collect();
+
                     let arm_block = this.bind_pattern(
                         outer_source_info,
                         branch,
@@ -460,6 +472,18 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         EmitStorageLive::Yes,
                     );
 
+                    // If the match arm has a guard, change the branch_taken_block for all of the
+                    // sub-branches to be the guard block.
+                    if arm.guard.is_some() {
+                        for sub_branch in sub_branches.iter_mut() {
+                            sub_branch.branch_taken_block = arm_block;
+                        }
+                    }
+
+                    if let Some(coverage_match_arms) = coverage_match_arms.as_mut() {
+                        coverage_match_arms.push(coverageinfo::MatchArm { sub_branches });
+                    }
+
                     this.fixed_temps_scope = old_dedup_scope;
 
                     if let Some(source_scope) = scope {
@@ -472,6 +496,13 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             })
             .collect();
 
+        if let Some(coverage_match_arms) = coverage_match_arms {
+            self.coverage_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();
 
@@ -1840,7 +1871,10 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         });
         for candidate in candidates_to_expand.iter_mut() {
             if !candidate.subcandidates.is_empty() {
-                self.merge_trivial_subcandidates(candidate);
+                // FIXME: Support merging trival candidates in branch coverage instrumentation
+                if self.coverage_info.is_none() {
+                    self.merge_trivial_subcandidates(candidate);
+                }
                 self.remove_never_subcandidates(candidate);
             }
         }

compiler/rustc_mir_transform/src/coverage/counters.rs

@@ -122,7 +122,12 @@ impl CoverageCounters {
         self.set_bcb_edge_counter(from_bcb, to_bcb, counter)
     }
 
-    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