coverage. Record branch blocks for match

Author: Lambdaris

compiler/rustc_mir_build/src/build/coverageinfo.rs

@@ -1,12 +1,14 @@
 use std::assert_matches::assert_matches;
 use std::collections::hash_map::Entry;
+use std::collections::BTreeMap;
 
 use rustc_data_structures::fx::FxHashMap;
 use rustc_middle::mir::coverage::{BlockMarkerId, BranchSpan, CoverageKind};
 use rustc_middle::mir::{self, BasicBlock, UnOp};
 use rustc_middle::thir::{ExprId, ExprKind, Thir};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::def_id::LocalDefId;
+use rustc_span::Span;
 
 use crate::build::Builder;
 
@@ -16,6 +18,7 @@ pub(crate) struct BranchInfoBuilder {
 
     num_block_markers: usize,
     branch_spans: Vec<BranchSpan>,
+    pattern_match_branches: BTreeMap<Span, (Vec<BasicBlock>, Vec<BasicBlock>)>,
 }
 
 #[derive(Clone, Copy)]
@@ -33,7 +36,12 @@ 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_spans: vec![],
+                pattern_match_branches: BTreeMap::new(),
+            })
         } else {
             None
         }
@@ -86,7 +94,7 @@ 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_spans, .. } = self;
 
         if num_block_markers == 0 {
             assert!(branch_spans.is_empty());
@@ -143,4 +151,29 @@ impl Builder<'_, '_> {
             false_marker,
         });
     }
+
+    #[allow(unused)]
+    pub(crate) fn visit_pattern_match_branches(
+        &mut self,
+        targets: impl Iterator<Item = (Span, BasicBlock)>,
+        otherwise_block: BasicBlock,
+    ) {
+        // TODO! Add get_block_marker_id here to transform BasicBlock to BlockMarkerId then `pattern_match_branches` could store BlockMarkerId.
+        let targets = targets.collect::<Vec<_>>();
+
+        let Some(branch_info) = self.coverage_branch_info.as_mut() else { return };
+        for (span, true_blk) in &targets {
+            let (true_blks, false_blks) =
+                branch_info.pattern_match_branches.entry(*span).or_insert_with(|| (vec![], vec![]));
+            // SomeEnum::A | SomeEnum::B would be lowered to something like switchInt(_1) -> [ 0: bb1, 1: bb3, otherwise: otherwise_block ]
+            // Thus bb3 and otherwise_block both are false blocks for SomeEnum::A.
+            true_blks.push(*true_blk);
+            false_blks.extend(
+                targets
+                    .iter()
+                    .filter_map(|(_, blk)| (blk != true_blk).then_some(*blk))
+                    .chain(std::iter::once(otherwise_block)),
+            );
+        }
+    }
 }

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

@@ -28,6 +28,7 @@ mod test;
 mod util;
 
 use std::borrow::Borrow;
+use std::collections::BTreeMap;
 use std::mem;
 
 /// Arguments to [`Builder::then_else_break_inner`] that are usually forwarded
@@ -1855,6 +1856,18 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         // Extract the match-pair from the highest priority candidate and build a test from it.
         let (match_place, test) = self.pick_test(candidates);
 
+        // Record spans of targets. `test.span` only represents `Enum::A` out of `Enum::A | Enum::B`
+        // Span in extra data to identify associated candidate later.
+        let mut coverage_targets: BTreeMap<_, _> = candidates
+            .iter()
+            .filter_map(|candidate| {
+                candidate
+                    .match_pairs
+                    .first()
+                    .map(|match_pair| (candidate.extra_data.span, (match_pair.pattern.span, None)))
+            })
+            .collect();
+
         // For each of the N possible test outcomes, build the vector of candidates that applies if
         // the test has that particular outcome.
         let (remaining_candidates, target_candidates) =
@@ -1881,6 +1894,12 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             .into_iter()
             .map(|(branch, mut candidates)| {
                 let candidate_start = self.cfg.start_new_block();
+                // After `match_candidates` above the candidate replace its `matched_pairs` with sub patterns.
+                // But with luck the extra_data.span is unchanged. So we can use it to find the associated target span and update its target blocks.
+                coverage_targets
+                    .get_mut(&candidates.first().expect("must be non-empty").extra_data.span)
+                    .expect("must exist")
+                    .1 = Some(candidate_start);
                 self.match_candidates(
                     span,
                     scrutinee_span,
@@ -1892,6 +1911,13 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             })
             .collect();
 
+        self.visit_pattern_match_branches(
+            coverage_targets
+                .into_values()
+                .filter_map(|(span, target)| target.map(|blk| (span, blk))),
+            remainder_start,
+        );
+
         // Perform the test, branching to one of N blocks.
         self.perform_test(
             span,