coverage: Split out MC/DC mappings from BcbMappingKind

compiler/rustc_mir_transform/src/coverage/mappings.rs

@@ -2,9 +2,8 @@ use std::collections::BTreeSet;
 
 use rustc_data_structures::graph::DirectedGraph;
 use rustc_index::bit_set::BitSet;
-use rustc_middle::mir::coverage::{
-    BlockMarkerId, BranchSpan, ConditionInfo, CoverageKind, MCDCBranchSpan, MCDCDecisionSpan,
-};
+use rustc_index::IndexVec;
+use rustc_middle::mir::coverage::{BlockMarkerId, BranchSpan, ConditionInfo, CoverageKind};
 use rustc_middle::mir::{self, BasicBlock, StatementKind};
 use rustc_span::Span;
 
@@ -13,55 +12,53 @@ use crate::coverage::spans::{
     extract_refined_covspans, unexpand_into_body_span_with_visible_macro,
 };
 use crate::coverage::ExtractedHirInfo;
-use rustc_index::IndexVec;
-
-#[derive(Clone, Debug)]
-pub(super) enum BcbMappingKind {
-    /// Associates an ordinary executable code span with its corresponding BCB.
-    Code(BasicCoverageBlock),
-
-    // Ordinary branch mappings are stored separately, so they don't have a
-    // variant in this enum.
-    //
-    /// Associates a mcdc branch span with condition info besides fields for normal branch.
-    MCDCBranch {
-        true_bcb: BasicCoverageBlock,
-        false_bcb: BasicCoverageBlock,
-        /// If `None`, this actually represents a normal branch mapping inserted
-        /// for code that was too complex for MC/DC.
-        condition_info: Option<ConditionInfo>,
-        decision_depth: u16,
-    },
-    /// Associates a mcdc decision with its join BCB.
-    MCDCDecision {
-        end_bcbs: BTreeSet<BasicCoverageBlock>,
-        bitmap_idx: u32,
-        conditions_num: u16,
-        decision_depth: u16,
-    },
-}
 
+/// Associates an ordinary executable code span with its corresponding BCB.
 #[derive(Debug)]
-pub(super) struct BcbMapping {
-    pub(super) kind: BcbMappingKind,
+pub(super) struct CodeMapping {
     pub(super) span: Span,
+    pub(super) bcb: BasicCoverageBlock,
 }
 
-/// This is separate from [`BcbMappingKind`] to help prepare for larger changes
+/// This is separate from [`MCDCBranch`] to help prepare for larger changes
 /// that will be needed for improved branch coverage in the future.
 /// (See <https://github.com/rust-lang/rust/pull/124217>.)
 #[derive(Debug)]
-pub(super) struct BcbBranchPair {
+pub(super) struct BranchPair {
     pub(super) span: Span,
     pub(super) true_bcb: BasicCoverageBlock,
     pub(super) false_bcb: BasicCoverageBlock,
 }
 
+/// Associates an MC/DC branch span with condition info besides fields for normal branch.
+#[derive(Debug)]
+pub(super) struct MCDCBranch {
+    pub(super) span: Span,
+    pub(super) true_bcb: BasicCoverageBlock,
+    pub(super) false_bcb: BasicCoverageBlock,
+    /// If `None`, this actually represents a normal branch mapping inserted
+    /// for code that was too complex for MC/DC.
+    pub(super) condition_info: Option<ConditionInfo>,
+    pub(super) decision_depth: u16,
+}
+
+/// Associates an MC/DC decision with its join BCBs.
+#[derive(Debug)]
+pub(super) struct MCDCDecision {
+    pub(super) span: Span,
+    pub(super) end_bcbs: BTreeSet<BasicCoverageBlock>,
+    pub(super) bitmap_idx: u32,
+    pub(super) conditions_num: u16,
+    pub(super) decision_depth: u16,
+}
+
 pub(super) struct CoverageSpans {
     bcb_has_mappings: BitSet<BasicCoverageBlock>,
-    pub(super) mappings: Vec<BcbMapping>,
-    pub(super) branch_pairs: Vec<BcbBranchPair>,
+    pub(super) code_mappings: Vec<CodeMapping>,
+    pub(super) branch_pairs: Vec<BranchPair>,
     test_vector_bitmap_bytes: u32,
+    pub(super) mcdc_branches: Vec<MCDCBranch>,
+    pub(super) mcdc_decisions: Vec<MCDCDecision>,
 }
 
 impl CoverageSpans {
@@ -83,26 +80,38 @@ pub(super) fn generate_coverage_spans(
     hir_info: &ExtractedHirInfo,
     basic_coverage_blocks: &CoverageGraph,
 ) -> Option<CoverageSpans> {
-    let mut mappings = vec![];
+    let mut code_mappings = vec![];
     let mut branch_pairs = vec![];
+    let mut mcdc_branches = vec![];
+    let mut mcdc_decisions = vec![];
 
     if hir_info.is_async_fn {
         // An async function desugars into a function that returns a future,
         // with the user code wrapped in a closure. Any spans in the desugared
         // outer function will be unhelpful, so just keep the signature span
         // and ignore all of the spans in the MIR body.
         if let Some(span) = hir_info.fn_sig_span_extended {
-            mappings.push(BcbMapping { kind: BcbMappingKind::Code(START_BCB), span });
+            code_mappings.push(CodeMapping { span, bcb: START_BCB });
         }
     } else {
-        extract_refined_covspans(mir_body, hir_info, basic_coverage_blocks, &mut mappings);
+        extract_refined_covspans(mir_body, hir_info, basic_coverage_blocks, &mut code_mappings);
 
         branch_pairs.extend(extract_branch_pairs(mir_body, hir_info, basic_coverage_blocks));
 
-        mappings.extend(extract_mcdc_mappings(mir_body, hir_info.body_span, basic_coverage_blocks));
+        extract_mcdc_mappings(
+            mir_body,
+            hir_info.body_span,
+            basic_coverage_blocks,
+            &mut mcdc_branches,
+            &mut mcdc_decisions,
+        );
     }
 
-    if mappings.is_empty() && branch_pairs.is_empty() {
+    if code_mappings.is_empty()
+        && branch_pairs.is_empty()
+        && mcdc_branches.is_empty()
+        && mcdc_decisions.is_empty()
+    {
         return None;
     }
 
@@ -111,29 +120,36 @@ pub(super) fn generate_coverage_spans(
     let mut insert = |bcb| {
         bcb_has_mappings.insert(bcb);
     };
-    let mut test_vector_bitmap_bytes = 0;
-    for BcbMapping { kind, span: _ } in &mappings {
-        match *kind {
-            BcbMappingKind::Code(bcb) => insert(bcb),
-            BcbMappingKind::MCDCBranch { true_bcb, false_bcb, .. } => {
-                insert(true_bcb);
-                insert(false_bcb);
-            }
-            BcbMappingKind::MCDCDecision { bitmap_idx, conditions_num, .. } => {
-                // `bcb_has_mappings` is used for inject coverage counters
-                // but they are not needed for decision BCBs.
-                // While the length of test vector bitmap should be calculated here.
-                test_vector_bitmap_bytes = test_vector_bitmap_bytes
-                    .max(bitmap_idx + (1_u32 << conditions_num as u32).div_ceil(8));
-            }
-        }
+
+    for &CodeMapping { span: _, bcb } in &code_mappings {
+        insert(bcb);
+    }
+    for &BranchPair { true_bcb, false_bcb, .. } in &branch_pairs {
+        insert(true_bcb);
+        insert(false_bcb);
     }
-    for &BcbBranchPair { true_bcb, false_bcb, .. } in &branch_pairs {
+    for &MCDCBranch { true_bcb, false_bcb, .. } in &mcdc_branches {
         insert(true_bcb);
         insert(false_bcb);
     }
 
-    Some(CoverageSpans { bcb_has_mappings, mappings, branch_pairs, test_vector_bitmap_bytes })
+    // Determine the length of the test vector bitmap.
+    let test_vector_bitmap_bytes = mcdc_decisions
+        .iter()
+        .map(|&MCDCDecision { bitmap_idx, conditions_num, .. }| {
+            bitmap_idx + (1_u32 << u32::from(conditions_num)).div_ceil(8)
+        })
+        .max()
+        .unwrap_or(0);
+
+    Some(CoverageSpans {
+        bcb_has_mappings,
+        code_mappings,
+        branch_pairs,
+        test_vector_bitmap_bytes,
+        mcdc_branches,
+        mcdc_decisions,
+    })
 }
 
 fn resolve_block_markers(
@@ -167,7 +183,7 @@ pub(super) fn extract_branch_pairs(
     mir_body: &mir::Body<'_>,
     hir_info: &ExtractedHirInfo,
     basic_coverage_blocks: &CoverageGraph,
-) -> Vec<BcbBranchPair> {
+) -> Vec<BranchPair> {
     let Some(branch_info) = mir_body.coverage_branch_info.as_deref() else { return vec![] };
 
     let block_markers = resolve_block_markers(branch_info, mir_body);
@@ -190,7 +206,7 @@ pub(super) fn extract_branch_pairs(
             let true_bcb = bcb_from_marker(true_marker)?;
             let false_bcb = bcb_from_marker(false_marker)?;
 
-            Some(BcbBranchPair { span, true_bcb, false_bcb })
+            Some(BranchPair { span, true_bcb, false_bcb })
         })
         .collect::<Vec<_>>()
 }
@@ -199,10 +215,10 @@ pub(super) fn extract_mcdc_mappings(
     mir_body: &mir::Body<'_>,
     body_span: Span,
     basic_coverage_blocks: &CoverageGraph,
-) -> Vec<BcbMapping> {
-    let Some(branch_info) = mir_body.coverage_branch_info.as_deref() else {
-        return vec![];
-    };
+    mcdc_branches: &mut impl Extend<MCDCBranch>,
+    mcdc_decisions: &mut impl Extend<MCDCDecision>,
+) {
+    let Some(branch_info) = mir_body.coverage_branch_info.as_deref() else { return };
 
     let block_markers = resolve_block_markers(branch_info, mir_body);
 
@@ -223,53 +239,42 @@ pub(super) fn extract_mcdc_mappings(
             Some((span, true_bcb, false_bcb))
         };
 
-    let mcdc_branch_filter_map = |&MCDCBranchSpan {
-                                      span: raw_span,
-                                      true_marker,
-                                      false_marker,
-                                      condition_info,
-                                      decision_depth,
-                                  }| {
-        check_branch_bcb(raw_span, true_marker, false_marker).map(|(span, true_bcb, false_bcb)| {
-            BcbMapping {
-                kind: BcbMappingKind::MCDCBranch {
-                    true_bcb,
-                    false_bcb,
-                    condition_info,
-                    decision_depth,
-                },
-                span,
-            }
-        })
-    };
+    mcdc_branches.extend(branch_info.mcdc_branch_spans.iter().filter_map(
+        |&mir::coverage::MCDCBranchSpan {
+             span: raw_span,
+             condition_info,
+             true_marker,
+             false_marker,
+             decision_depth,
+         }| {
+            let (span, true_bcb, false_bcb) =
+                check_branch_bcb(raw_span, true_marker, false_marker)?;
+            Some(MCDCBranch { span, true_bcb, false_bcb, condition_info, decision_depth })
+        },
+    ));
 
     let mut next_bitmap_idx = 0;
 
-    let decision_filter_map = |decision: &MCDCDecisionSpan| {
-        let (span, _) = unexpand_into_body_span_with_visible_macro(decision.span, body_span)?;
+    mcdc_decisions.extend(branch_info.mcdc_decision_spans.iter().filter_map(
+        |decision: &mir::coverage::MCDCDecisionSpan| {
+            let (span, _) = unexpand_into_body_span_with_visible_macro(decision.span, body_span)?;
 
-        let end_bcbs = decision
-            .end_markers
-            .iter()
-            .map(|&marker| bcb_from_marker(marker))
-            .collect::<Option<_>>()?;
+            let end_bcbs = decision
+                .end_markers
+                .iter()
+                .map(|&marker| bcb_from_marker(marker))
+                .collect::<Option<_>>()?;
 
-        let bitmap_idx = next_bitmap_idx;
-        next_bitmap_idx += (1_u32 << decision.conditions_num).div_ceil(8);
+            let bitmap_idx = next_bitmap_idx;
+            next_bitmap_idx += (1_u32 << decision.conditions_num).div_ceil(8);
 
-        Some(BcbMapping {
-            kind: BcbMappingKind::MCDCDecision {
+            Some(MCDCDecision {
+                span,
                 end_bcbs,
                 bitmap_idx,
                 conditions_num: decision.conditions_num as u16,
                 decision_depth: decision.decision_depth,
-            },
-            span,
-        })
-    };
-
-    std::iter::empty()
-        .chain(branch_info.mcdc_branch_spans.iter().filter_map(mcdc_branch_filter_map))
-        .chain(branch_info.mcdc_decision_spans.iter().filter_map(decision_filter_map))
-        .collect::<Vec<_>>()
+            })
+        },
+    ));
 }