coverage: Separate branch pairs from other mapping kinds

This clears the way for larger changes to how branches are handled by the
coverage instrumentor, in order to support branch coverage for more language
constructs.

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -9,7 +9,7 @@ mod tests;
 
 use self::counters::{CounterIncrementSite, CoverageCounters};
 use self::graph::{BasicCoverageBlock, CoverageGraph};
-use self::spans::{BcbMapping, BcbMappingKind, CoverageSpans};
+use self::spans::{BcbBranchPair, BcbMapping, BcbMappingKind, CoverageSpans};
 
 use crate::MirPass;
 
@@ -141,14 +141,10 @@ fn create_mappings<'tcx>(
 
     let mut mappings = Vec::new();
 
-    mappings.extend(coverage_spans.all_bcb_mappings().filter_map(
+    mappings.extend(coverage_spans.mappings.iter().filter_map(
         |BcbMapping { kind: bcb_mapping_kind, span }| {
             let kind = match *bcb_mapping_kind {
                 BcbMappingKind::Code(bcb) => MappingKind::Code(term_for_bcb(bcb)),
-                BcbMappingKind::Branch { true_bcb, false_bcb } => MappingKind::Branch {
-                    true_term: term_for_bcb(true_bcb),
-                    false_term: term_for_bcb(false_bcb),
-                },
                 BcbMappingKind::MCDCBranch { true_bcb, false_bcb, condition_info: None } => {
                     MappingKind::Branch {
                         true_term: term_for_bcb(true_bcb),
@@ -173,6 +169,16 @@ fn create_mappings<'tcx>(
         },
     ));
 
+    mappings.extend(coverage_spans.branch_pairs.iter().filter_map(
+        |&BcbBranchPair { span, true_bcb, false_bcb }| {
+            let true_term = term_for_bcb(true_bcb);
+            let false_term = term_for_bcb(false_bcb);
+            let kind = MappingKind::Branch { true_term, false_term };
+            let code_region = make_code_region(source_map, file_name, span, body_span)?;
+            Some(Mapping { kind, code_region })
+        },
+    ));
+
     mappings
 }
 
@@ -241,7 +247,7 @@ fn inject_mcdc_statements<'tcx>(
 
     // Inject test vector update first because `inject_statement` always insert new statement at head.
     for (end_bcbs, bitmap_idx) in
-        coverage_spans.all_bcb_mappings().filter_map(|mapping| match &mapping.kind {
+        coverage_spans.mappings.iter().filter_map(|mapping| match &mapping.kind {
             BcbMappingKind::MCDCDecision { end_bcbs, bitmap_idx, .. } => {
                 Some((end_bcbs, *bitmap_idx))
             }
@@ -255,7 +261,7 @@ fn inject_mcdc_statements<'tcx>(
     }
 
     for (true_bcb, false_bcb, condition_id) in
-        coverage_spans.all_bcb_mappings().filter_map(|mapping| match mapping.kind {
+        coverage_spans.mappings.iter().filter_map(|mapping| match mapping.kind {
             BcbMappingKind::MCDCBranch { true_bcb, false_bcb, condition_info } => {
                 Some((true_bcb, false_bcb, condition_info?.condition_id))
             }

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -15,8 +15,10 @@ mod from_mir;
 pub(super) enum BcbMappingKind {
     /// Associates an ordinary executable code span with its corresponding BCB.
     Code(BasicCoverageBlock),
-    /// Associates a branch span with BCBs for its true and false arms.
-    Branch { true_bcb: BasicCoverageBlock, false_bcb: 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,
@@ -35,9 +37,20 @@ pub(super) struct BcbMapping {
     pub(super) span: Span,
 }
 
+/// This is separate from [`BcbMappingKind`] 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) span: Span,
+    pub(super) true_bcb: BasicCoverageBlock,
+    pub(super) false_bcb: BasicCoverageBlock,
+}
+
 pub(super) struct CoverageSpans {
     bcb_has_mappings: BitSet<BasicCoverageBlock>,
-    mappings: Vec<BcbMapping>,
+    pub(super) mappings: Vec<BcbMapping>,
+    pub(super) branch_pairs: Vec<BcbBranchPair>,
     test_vector_bitmap_bytes: u32,
 }
 
@@ -46,10 +59,6 @@ impl CoverageSpans {
         self.bcb_has_mappings.contains(bcb)
     }
 
-    pub(super) fn all_bcb_mappings(&self) -> impl Iterator<Item = &BcbMapping> {
-        self.mappings.iter()
-    }
-
     pub(super) fn test_vector_bitmap_bytes(&self) -> u32 {
         self.test_vector_bitmap_bytes
     }
@@ -65,6 +74,7 @@ pub(super) fn generate_coverage_spans(
     basic_coverage_blocks: &CoverageGraph,
 ) -> Option<CoverageSpans> {
     let mut mappings = vec![];
+    let mut branch_pairs = vec![];
 
     if hir_info.is_async_fn {
         // An async function desugars into a function that returns a future,
@@ -86,7 +96,7 @@ pub(super) fn generate_coverage_spans(
             BcbMapping { kind: BcbMappingKind::Code(bcb), span }
         }));
 
-        mappings.extend(from_mir::extract_branch_mappings(
+        branch_pairs.extend(from_mir::extract_branch_pairs(
             mir_body,
             hir_info,
             basic_coverage_blocks,
@@ -99,7 +109,7 @@ pub(super) fn generate_coverage_spans(
         ));
     }
 
-    if mappings.is_empty() {
+    if mappings.is_empty() && branch_pairs.is_empty() {
         return None;
     }
 
@@ -112,8 +122,7 @@ pub(super) fn generate_coverage_spans(
     for BcbMapping { kind, span: _ } in &mappings {
         match *kind {
             BcbMappingKind::Code(bcb) => insert(bcb),
-            BcbMappingKind::Branch { true_bcb, false_bcb }
-            | BcbMappingKind::MCDCBranch { true_bcb, false_bcb, .. } => {
+            BcbMappingKind::MCDCBranch { true_bcb, false_bcb, .. } => {
                 insert(true_bcb);
                 insert(false_bcb);
             }
@@ -126,8 +135,12 @@ pub(super) fn generate_coverage_spans(
             }
         }
     }
+    for &BcbBranchPair { true_bcb, false_bcb, .. } in &branch_pairs {
+        insert(true_bcb);
+        insert(false_bcb);
+    }
 
-    Some(CoverageSpans { bcb_has_mappings, mappings, test_vector_bitmap_bytes })
+    Some(CoverageSpans { bcb_has_mappings, mappings, branch_pairs, test_vector_bitmap_bytes })
 }
 
 #[derive(Debug)]

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

@@ -13,7 +13,7 @@ use rustc_span::{ExpnKind, MacroKind, Span, Symbol};
 use crate::coverage::graph::{
     BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph, START_BCB,
 };
-use crate::coverage::spans::{BcbMapping, BcbMappingKind};
+use crate::coverage::spans::{BcbBranchPair, BcbMapping, BcbMappingKind};
 use crate::coverage::ExtractedHirInfo;
 
 /// Traverses the MIR body to produce an initial collection of coverage-relevant
@@ -388,16 +388,16 @@ fn resolve_block_markers(
 }
 
 // FIXME: There is currently a lot of redundancy between
-// `extract_branch_mappings` and `extract_mcdc_mappings`. This is needed so
+// `extract_branch_pairs` and `extract_mcdc_mappings`. This is needed so
 // that they can each be modified without interfering with the other, but in
 // the long term we should try to bring them together again when branch coverage
 // and MC/DC coverage support are more mature.
 
-pub(super) fn extract_branch_mappings(
+pub(super) fn extract_branch_pairs(
     mir_body: &mir::Body<'_>,
     hir_info: &ExtractedHirInfo,
     basic_coverage_blocks: &CoverageGraph,
-) -> Vec<BcbMapping> {
+) -> Vec<BcbBranchPair> {
     let Some(branch_info) = mir_body.coverage_branch_info.as_deref() else { return vec![] };
 
     let block_markers = resolve_block_markers(branch_info, mir_body);
@@ -420,7 +420,7 @@ pub(super) fn extract_branch_mappings(
             let true_bcb = bcb_from_marker(true_marker)?;
             let false_bcb = bcb_from_marker(false_marker)?;
 
-            Some(BcbMapping { kind: BcbMappingKind::Branch { true_bcb, false_bcb }, span })
+            Some(BcbBranchPair { span, true_bcb, false_bcb })
         })
         .collect::<Vec<_>>()
 }