Rollup merge of rust-lang#124615 - Zalathar:extracted-mappings, r=davidtwco

coverage: Further simplify extraction of mapping info from MIR

This is another round of rearrangement and simplification that builds on top of the changes made to mapping-extraction by rust-lang#124603.

The overall theme is to take the computation of `bcb_has_mappings` and `test_vector_bitmap_bytes` out of the main body of `generate_coverage_spans`, which then lets us perform a few other small changes that had previously been held up by the need to work around those computations.

Author: matthiaskrgr

compiler/rustc_mir_transform/src/coverage/mappings.rs

@@ -52,104 +52,104 @@ pub(super) struct MCDCDecision {
     pub(super) decision_depth: u16,
 }
 
-pub(super) struct CoverageSpans {
-    bcb_has_mappings: BitSet<BasicCoverageBlock>,
+#[derive(Default)]
+pub(super) struct ExtractedMappings {
     pub(super) code_mappings: Vec<CodeMapping>,
     pub(super) branch_pairs: Vec<BranchPair>,
-    test_vector_bitmap_bytes: u32,
+    pub(super) mcdc_bitmap_bytes: u32,
     pub(super) mcdc_branches: Vec<MCDCBranch>,
     pub(super) mcdc_decisions: Vec<MCDCDecision>,
 }
 
-impl CoverageSpans {
-    pub(super) fn bcb_has_coverage_spans(&self, bcb: BasicCoverageBlock) -> bool {
-        self.bcb_has_mappings.contains(bcb)
-    }
-
-    pub(super) fn test_vector_bitmap_bytes(&self) -> u32 {
-        self.test_vector_bitmap_bytes
-    }
-}
-
 /// Extracts coverage-relevant spans from MIR, and associates them with
 /// their corresponding BCBs.
-///
-/// Returns `None` if no coverage-relevant spans could be extracted.
-pub(super) fn generate_coverage_spans(
+pub(super) fn extract_all_mapping_info_from_mir(
     mir_body: &mir::Body<'_>,
     hir_info: &ExtractedHirInfo,
     basic_coverage_blocks: &CoverageGraph,
-) -> Option<CoverageSpans> {
-    let mut code_mappings = vec![];
-    let mut branch_pairs = vec![];
-    let mut mcdc_branches = vec![];
-    let mut mcdc_decisions = vec![];
-
+) -> ExtractedMappings {
     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.
+        let mut mappings = ExtractedMappings::default();
         if let Some(span) = hir_info.fn_sig_span_extended {
-            code_mappings.push(CodeMapping { span, bcb: START_BCB });
+            mappings.code_mappings.push(CodeMapping { span, bcb: START_BCB });
         }
-    } else {
-        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));
-
-        extract_mcdc_mappings(
-            mir_body,
-            hir_info.body_span,
-            basic_coverage_blocks,
-            &mut mcdc_branches,
-            &mut mcdc_decisions,
-        );
+        return mappings;
     }
 
-    if code_mappings.is_empty()
-        && branch_pairs.is_empty()
-        && mcdc_branches.is_empty()
-        && mcdc_decisions.is_empty()
-    {
-        return None;
-    }
+    let mut code_mappings = vec![];
+    let mut branch_pairs = vec![];
+    let mut mcdc_bitmap_bytes = 0;
+    let mut mcdc_branches = vec![];
+    let mut mcdc_decisions = vec![];
 
-    // Identify which BCBs have one or more mappings.
-    let mut bcb_has_mappings = BitSet::new_empty(basic_coverage_blocks.num_nodes());
-    let mut insert = |bcb| {
-        bcb_has_mappings.insert(bcb);
-    };
+    extract_refined_covspans(mir_body, hir_info, basic_coverage_blocks, &mut code_mappings);
 
-    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 &MCDCBranch { true_bcb, false_bcb, .. } in &mcdc_branches {
-        insert(true_bcb);
-        insert(false_bcb);
-    }
+    branch_pairs.extend(extract_branch_pairs(mir_body, hir_info, basic_coverage_blocks));
 
-    // 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);
+    extract_mcdc_mappings(
+        mir_body,
+        hir_info.body_span,
+        basic_coverage_blocks,
+        &mut mcdc_bitmap_bytes,
+        &mut mcdc_branches,
+        &mut mcdc_decisions,
+    );
 
-    Some(CoverageSpans {
-        bcb_has_mappings,
+    ExtractedMappings {
         code_mappings,
         branch_pairs,
-        test_vector_bitmap_bytes,
+        mcdc_bitmap_bytes,
         mcdc_branches,
         mcdc_decisions,
-    })
+    }
+}
+
+impl ExtractedMappings {
+    pub(super) fn all_bcbs_with_counter_mappings(
+        &self,
+        basic_coverage_blocks: &CoverageGraph, // Only used for allocating a correctly-sized set
+    ) -> BitSet<BasicCoverageBlock> {
+        // Fully destructure self to make sure we don't miss any fields that have mappings.
+        let Self {
+            code_mappings,
+            branch_pairs,
+            mcdc_bitmap_bytes: _,
+            mcdc_branches,
+            mcdc_decisions,
+        } = self;
+
+        // Identify which BCBs have one or more mappings.
+        let mut bcbs_with_counter_mappings = BitSet::new_empty(basic_coverage_blocks.num_nodes());
+        let mut insert = |bcb| {
+            bcbs_with_counter_mappings.insert(bcb);
+        };
+
+        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 &MCDCBranch { true_bcb, false_bcb, .. } in mcdc_branches {
+            insert(true_bcb);
+            insert(false_bcb);
+        }
+
+        // MC/DC decisions refer to BCBs, but don't require those BCBs to have counters.
+        if bcbs_with_counter_mappings.is_empty() {
+            debug_assert!(
+                mcdc_decisions.is_empty(),
+                "A function with no counter mappings shouldn't have any decisions: {mcdc_decisions:?}",
+            );
+        }
+
+        bcbs_with_counter_mappings
+    }
 }
 
 fn resolve_block_markers(
@@ -215,6 +215,7 @@ pub(super) fn extract_mcdc_mappings(
     mir_body: &mir::Body<'_>,
     body_span: Span,
     basic_coverage_blocks: &CoverageGraph,
+    mcdc_bitmap_bytes: &mut u32,
     mcdc_branches: &mut impl Extend<MCDCBranch>,
     mcdc_decisions: &mut impl Extend<MCDCDecision>,
 ) {
@@ -253,8 +254,6 @@ pub(super) fn extract_mcdc_mappings(
         },
     ));
 
-    let mut next_bitmap_idx = 0;
-
     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)?;
@@ -265,8 +264,11 @@ pub(super) fn extract_mcdc_mappings(
                 .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);
+            // Each decision containing N conditions needs 2^N bits of space in
+            // the bitmap, rounded up to a whole number of bytes.
+            // The decision's "bitmap index" points to its first byte in the bitmap.
+            let bitmap_idx = *mcdc_bitmap_bytes;
+            *mcdc_bitmap_bytes += (1_u32 << decision.conditions_num).div_ceil(8);
 
             Some(MCDCDecision {
                 span,

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -7,13 +7,9 @@ mod spans;
 #[cfg(test)]
 mod tests;
 
-use self::counters::{CounterIncrementSite, CoverageCounters};
-use self::graph::{BasicCoverageBlock, CoverageGraph};
-use self::mappings::CoverageSpans;
-
-use crate::MirPass;
-
-use rustc_middle::mir::coverage::*;
+use rustc_middle::mir::coverage::{
+    CodeRegion, CoverageKind, DecisionInfo, FunctionCoverageInfo, Mapping, MappingKind,
+};
 use rustc_middle::mir::{
     self, BasicBlock, BasicBlockData, SourceInfo, Statement, StatementKind, Terminator,
     TerminatorKind,
@@ -23,6 +19,11 @@ use rustc_span::def_id::LocalDefId;
 use rustc_span::source_map::SourceMap;
 use rustc_span::{BytePos, Pos, RelativeBytePos, Span, Symbol};
 
+use crate::coverage::counters::{CounterIncrementSite, CoverageCounters};
+use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
+use crate::coverage::mappings::ExtractedMappings;
+use crate::MirPass;
+
 /// Inserts `StatementKind::Coverage` statements that either instrument the binary with injected
 /// counters, via intrinsic `llvm.instrprof.increment`, and/or inject metadata used during codegen
 /// to construct the coverage map.
@@ -69,24 +70,27 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
     let basic_coverage_blocks = CoverageGraph::from_mir(mir_body);
 
     ////////////////////////////////////////////////////
-    // Compute coverage spans from the `CoverageGraph`.
-    let Some(coverage_spans) =
-        mappings::generate_coverage_spans(mir_body, &hir_info, &basic_coverage_blocks)
-    else {
-        // No relevant spans were found in MIR, so skip instrumenting this function.
-        return;
-    };
+    // Extract coverage spans and other mapping info from MIR.
+    let extracted_mappings =
+        mappings::extract_all_mapping_info_from_mir(mir_body, &hir_info, &basic_coverage_blocks);
 
     ////////////////////////////////////////////////////
     // Create an optimized mix of `Counter`s and `Expression`s for the `CoverageGraph`. Ensure
     // every coverage span has a `Counter` or `Expression` assigned to its `BasicCoverageBlock`
     // and all `Expression` dependencies (operands) are also generated, for any other
     // `BasicCoverageBlock`s not already associated with a coverage span.
-    let bcb_has_coverage_spans = |bcb| coverage_spans.bcb_has_coverage_spans(bcb);
+    let bcbs_with_counter_mappings =
+        extracted_mappings.all_bcbs_with_counter_mappings(&basic_coverage_blocks);
+    if bcbs_with_counter_mappings.is_empty() {
+        // No relevant spans were found in MIR, so skip instrumenting this function.
+        return;
+    }
+
+    let bcb_has_counter_mappings = |bcb| bcbs_with_counter_mappings.contains(bcb);
     let coverage_counters =
-        CoverageCounters::make_bcb_counters(&basic_coverage_blocks, bcb_has_coverage_spans);
+        CoverageCounters::make_bcb_counters(&basic_coverage_blocks, bcb_has_counter_mappings);
 
-    let mappings = create_mappings(tcx, &hir_info, &coverage_spans, &coverage_counters);
+    let mappings = create_mappings(tcx, &hir_info, &extracted_mappings, &coverage_counters);
     if mappings.is_empty() {
         // No spans could be converted into valid mappings, so skip this function.
         debug!("no spans could be converted into valid mappings; skipping");
@@ -96,13 +100,13 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
     inject_coverage_statements(
         mir_body,
         &basic_coverage_blocks,
-        bcb_has_coverage_spans,
+        bcb_has_counter_mappings,
         &coverage_counters,
     );
 
-    inject_mcdc_statements(mir_body, &basic_coverage_blocks, &coverage_spans);
+    inject_mcdc_statements(mir_body, &basic_coverage_blocks, &extracted_mappings);
 
-    let mcdc_num_condition_bitmaps = coverage_spans
+    let mcdc_num_condition_bitmaps = extracted_mappings
         .mcdc_decisions
         .iter()
         .map(|&mappings::MCDCDecision { decision_depth, .. }| decision_depth)
@@ -112,7 +116,7 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
     mir_body.function_coverage_info = Some(Box::new(FunctionCoverageInfo {
         function_source_hash: hir_info.function_source_hash,
         num_counters: coverage_counters.num_counters(),
-        mcdc_bitmap_bytes: coverage_spans.test_vector_bitmap_bytes(),
+        mcdc_bitmap_bytes: extracted_mappings.mcdc_bitmap_bytes,
         expressions: coverage_counters.into_expressions(),
         mappings,
         mcdc_num_condition_bitmaps,
@@ -127,7 +131,7 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
 fn create_mappings<'tcx>(
     tcx: TyCtxt<'tcx>,
     hir_info: &ExtractedHirInfo,
-    coverage_spans: &CoverageSpans,
+    extracted_mappings: &ExtractedMappings,
     coverage_counters: &CoverageCounters,
 ) -> Vec<Mapping> {
     let source_map = tcx.sess.source_map();
@@ -148,17 +152,26 @@ fn create_mappings<'tcx>(
     };
     let region_for_span = |span: Span| make_code_region(source_map, file_name, span, body_span);
 
+    // Fully destructure the mappings struct to make sure we don't miss any kinds.
+    let ExtractedMappings {
+        code_mappings,
+        branch_pairs,
+        mcdc_bitmap_bytes: _,
+        mcdc_branches,
+        mcdc_decisions,
+    } = extracted_mappings;
     let mut mappings = Vec::new();
 
-    mappings.extend(coverage_spans.code_mappings.iter().filter_map(
+    mappings.extend(code_mappings.iter().filter_map(
+        // Ordinary code mappings are the simplest kind.
         |&mappings::CodeMapping { span, bcb }| {
             let code_region = region_for_span(span)?;
             let kind = MappingKind::Code(term_for_bcb(bcb));
             Some(Mapping { kind, code_region })
         },
     ));
 
-    mappings.extend(coverage_spans.branch_pairs.iter().filter_map(
+    mappings.extend(branch_pairs.iter().filter_map(
         |&mappings::BranchPair { span, true_bcb, false_bcb }| {
             let true_term = term_for_bcb(true_bcb);
             let false_term = term_for_bcb(false_bcb);
@@ -168,7 +181,7 @@ fn create_mappings<'tcx>(
         },
     ));
 
-    mappings.extend(coverage_spans.mcdc_branches.iter().filter_map(
+    mappings.extend(mcdc_branches.iter().filter_map(
         |&mappings::MCDCBranch { span, true_bcb, false_bcb, condition_info, decision_depth: _ }| {
             let code_region = region_for_span(span)?;
             let true_term = term_for_bcb(true_bcb);
@@ -181,7 +194,7 @@ fn create_mappings<'tcx>(
         },
     ));
 
-    mappings.extend(coverage_spans.mcdc_decisions.iter().filter_map(
+    mappings.extend(mcdc_decisions.iter().filter_map(
         |&mappings::MCDCDecision { span, bitmap_idx, conditions_num, .. }| {
             let code_region = region_for_span(span)?;
             let kind = MappingKind::MCDCDecision(DecisionInfo { bitmap_idx, conditions_num });
@@ -249,20 +262,16 @@ fn inject_coverage_statements<'tcx>(
 fn inject_mcdc_statements<'tcx>(
     mir_body: &mut mir::Body<'tcx>,
     basic_coverage_blocks: &CoverageGraph,
-    coverage_spans: &CoverageSpans,
+    extracted_mappings: &ExtractedMappings,
 ) {
-    if coverage_spans.test_vector_bitmap_bytes() == 0 {
-        return;
-    }
-
     // Inject test vector update first because `inject_statement` always insert new statement at head.
     for &mappings::MCDCDecision {
         span: _,
         ref end_bcbs,
         bitmap_idx,
         conditions_num: _,
         decision_depth,
-    } in &coverage_spans.mcdc_decisions
+    } in &extracted_mappings.mcdc_decisions
     {
         for end in end_bcbs {
             let end_bb = basic_coverage_blocks[*end].leader_bb();
@@ -275,7 +284,7 @@ fn inject_mcdc_statements<'tcx>(
     }
 
     for &mappings::MCDCBranch { span: _, true_bcb, false_bcb, condition_info, decision_depth } in
-        &coverage_spans.mcdc_branches
+        &extracted_mappings.mcdc_branches
     {
         let Some(condition_info) = condition_info else { continue };
         let id = condition_info.condition_id;