coverage: Store all of a function's mappings in function coverage info

Previously, mappings were attached to individual coverage statements in MIR.
That necessitated special handling in MIR optimizations to avoid deleting those
statements, since otherwise codegen would be unable to reassemble the original
list of mappings.

With this change, a function's list of mappings is now attached to its MIR
body, and survives intact even if individual statements are deleted by
optimizations.

2 files changed, 15 insertions, 41 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index e1b656bd981..078123a7b7f 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -104,6 +104,7 @@ struct Instrumentor<'a, 'tcx> {
     function_source_hash: u64,
     basic_coverage_blocks: CoverageGraph,
     coverage_counters: CoverageCounters,
+    mappings: Vec<Mapping>,
 }
 
 impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
@@ -144,6 +145,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
             function_source_hash,
             basic_coverage_blocks,
             coverage_counters,
+            mappings: Vec::new(),
         }
     }
 
@@ -216,6 +218,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
             function_source_hash: self.function_source_hash,
             num_counters: self.coverage_counters.num_counters(),
             num_expressions: self.coverage_counters.num_expressions(),
+            mappings: std::mem::take(&mut self.mappings),
         }));
     }
 
@@ -232,18 +235,16 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
                 bug!("Every BasicCoverageBlock should have a Counter or Expression");
             });
 
-            // Convert the coverage spans into a vector of code regions to be
-            // associated with this BCB's coverage statement.
-            let code_regions = spans
-                .iter()
-                .map(|&span| make_code_region(source_map, file_name, span, body_span))
-                .collect::<Vec<_>>();
+            let term = counter_kind.as_term();
+            self.mappings.extend(spans.iter().map(|&span| {
+                let code_region = make_code_region(source_map, file_name, span, body_span);
+                Mapping { code_region, term }
+            }));
 
             inject_statement(
                 self.mir_body,
                 self.make_mir_coverage_kind(&counter_kind),
                 self.bcb_leader_bb(bcb),
-                code_regions,
             );
         }
     }
@@ -301,7 +302,6 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
                         self.mir_body,
                         self.make_mir_coverage_kind(&counter_kind),
                         inject_to_bb,
-                        Vec::new(),
                     );
                 }
                 BcbCounter::Expression { .. } => inject_intermediate_expression(
@@ -329,7 +329,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
 
     fn make_mir_coverage_kind(&self, counter_kind: &BcbCounter) -> CoverageKind {
         match *counter_kind {
-            BcbCounter::Counter { id } => CoverageKind::Counter { id },
+            BcbCounter::Counter { id } => CoverageKind::CounterIncrement { id },
             BcbCounter::Expression { id, lhs, op, rhs } => {
                 CoverageKind::Expression { id, lhs, op, rhs }
             }
@@ -360,18 +360,13 @@ fn inject_edge_counter_basic_block(
     new_bb
 }
 
-fn inject_statement(
-    mir_body: &mut mir::Body<'_>,
-    counter_kind: CoverageKind,
-    bb: BasicBlock,
-    code_regions: Vec<CodeRegion>,
-) {
-    debug!("  injecting statement {counter_kind:?} for {bb:?} at code regions: {code_regions:?}");
+fn inject_statement(mir_body: &mut mir::Body<'_>, counter_kind: CoverageKind, bb: BasicBlock) {
+    debug!("  injecting statement {counter_kind:?} for {bb:?}");
     let data = &mut mir_body[bb];
     let source_info = data.terminator().source_info;
     let statement = Statement {
         source_info,
-        kind: StatementKind::Coverage(Box::new(Coverage { kind: counter_kind, code_regions })),
+        kind: StatementKind::Coverage(Box::new(Coverage { kind: counter_kind })),
     };
     data.statements.insert(0, statement);
 }
@@ -385,10 +380,7 @@ fn inject_intermediate_expression(mir_body: &mut mir::Body<'_>, expression: Cove
     let source_info = data.terminator().source_info;
     let statement = Statement {
         source_info,
-        kind: StatementKind::Coverage(Box::new(Coverage {
-            kind: expression,
-            code_regions: Vec::new(),
-        })),
+        kind: StatementKind::Coverage(Box::new(Coverage { kind: expression })),
     };
     data.statements.push(statement);
 }
diff --git a/compiler/rustc_mir_transform/src/coverage/query.rs b/compiler/rustc_mir_transform/src/coverage/query.rs
index b1f46c8e05c..809407f897d 100644
--- a/compiler/rustc_mir_transform/src/coverage/query.rs
+++ b/compiler/rustc_mir_transform/src/coverage/query.rs
@@ -2,15 +2,13 @@ use super::*;
 
 use rustc_data_structures::captures::Captures;
 use rustc_middle::mir::coverage::*;
-use rustc_middle::mir::{self, Body, Coverage, CoverageIdsInfo};
+use rustc_middle::mir::{Body, Coverage, CoverageIdsInfo};
 use rustc_middle::query::Providers;
 use rustc_middle::ty::{self, TyCtxt};
-use rustc_span::def_id::DefId;
 
 /// A `query` provider for retrieving coverage information injected into MIR.
 pub(crate) fn provide(providers: &mut Providers) {
     providers.coverage_ids_info = |tcx, def_id| coverage_ids_info(tcx, def_id);
-    providers.covered_code_regions = |tcx, def_id| covered_code_regions(tcx, def_id);
 }
 
 /// Query implementation for `coverage_ids_info`.
@@ -22,7 +20,7 @@ fn coverage_ids_info<'tcx>(
 
     let max_counter_id = all_coverage_in_mir_body(mir_body)
         .filter_map(|coverage| match coverage.kind {
-            CoverageKind::Counter { id } => Some(id),
+            CoverageKind::CounterIncrement { id } => Some(id),
             _ => None,
         })
         .max()
@@ -31,14 +29,6 @@ fn coverage_ids_info<'tcx>(
     CoverageIdsInfo { max_counter_id }
 }
 
-fn covered_code_regions(tcx: TyCtxt<'_>, def_id: DefId) -> Vec<&CodeRegion> {
-    let body = mir_body(tcx, def_id);
-    all_coverage_in_mir_body(body)
-        // Coverage statements have a list of code regions (possibly empty).
-        .flat_map(|coverage| coverage.code_regions.as_slice())
-        .collect()
-}
-
 fn all_coverage_in_mir_body<'a, 'tcx>(
     body: &'a Body<'tcx>,
 ) -> impl Iterator<Item = &'a Coverage> + Captures<'tcx> {
@@ -56,11 +46,3 @@ fn is_inlined(body: &Body<'_>, statement: &Statement<'_>) -> bool {
     let scope_data = &body.source_scopes[statement.source_info.scope];
     scope_data.inlined.is_some() || scope_data.inlined_parent_scope.is_some()
 }
-
-/// This function ensures we obtain the correct MIR for the given item irrespective of
-/// whether that means const mir or runtime mir. For `const fn` this opts for runtime
-/// mir.
-fn mir_body(tcx: TyCtxt<'_>, def_id: DefId) -> &mir::Body<'_> {
-    let def = ty::InstanceDef::Item(def_id);
-    tcx.instance_mir(def)
-}