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

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -23,14 +23,13 @@ pub struct FunctionCoverage<'tcx> {
     function_coverage_info: &'tcx FunctionCoverageInfo,
     is_used: bool,
 
-    /// Tracks which counters have been seen, to avoid duplicate mappings
-    /// that might be introduced by MIR inlining.
+    /// Tracks which counters have been seen, so that we can identify mappings
+    /// to counters that were optimized out, and set them to zero.
     counters_seen: BitSet<CounterId>,
     expressions: IndexVec<ExpressionId, Option<Expression>>,
     /// Tracks which expressions are known to always have a value of zero.
     /// Only updated during the finalize step.
     zero_expressions: FxIndexSet<ExpressionId>,
-    mappings: Vec<Mapping>,
 }
 
 impl<'tcx> FunctionCoverage<'tcx> {
@@ -67,7 +66,6 @@ impl<'tcx> FunctionCoverage<'tcx> {
             counters_seen: BitSet::new_empty(num_counters),
             expressions: IndexVec::from_elem_n(None, num_expressions),
             zero_expressions: FxIndexSet::default(),
-            mappings: Vec::new(),
         }
     }
 
@@ -76,28 +74,20 @@ impl<'tcx> FunctionCoverage<'tcx> {
         self.is_used
     }
 
-    /// Adds code regions to be counted by an injected counter intrinsic.
+    /// Marks a counter ID as having been seen in a counter-increment statement.
     #[instrument(level = "debug", skip(self))]
-    pub(crate) fn add_counter(&mut self, id: CounterId, code_regions: &[CodeRegion]) {
-        if self.counters_seen.insert(id) {
-            self.add_mappings(CovTerm::Counter(id), code_regions);
-        }
+    pub(crate) fn add_counter(&mut self, id: CounterId) {
+        self.counters_seen.insert(id);
     }
 
-    /// Adds information about a coverage expression, along with zero or more
-    /// code regions mapped to that expression.
-    ///
-    /// Both counters and "counter expressions" (or simply, "expressions") can be operands in other
-    /// expressions. These are tracked as separate variants of `CovTerm`, so there is no ambiguity
-    /// between operands that are counter IDs and operands that are expression IDs.
+    /// Adds information about a coverage expression.
     #[instrument(level = "debug", skip(self))]
     pub(crate) fn add_counter_expression(
         &mut self,
         expression_id: ExpressionId,
         lhs: CovTerm,
         op: Op,
         rhs: CovTerm,
-        code_regions: &[CodeRegion],
     ) {
         debug_assert!(
             expression_id.as_usize() < self.expressions.len(),
@@ -111,10 +101,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
         let expression = Expression { lhs, op, rhs };
         let slot = &mut self.expressions[expression_id];
         match slot {
-            None => {
-                *slot = Some(expression);
-                self.add_mappings(CovTerm::Expression(expression_id), code_regions);
-            }
+            None => *slot = Some(expression),
             // If this expression ID slot has already been filled, it should
             // contain identical information.
             Some(ref previous_expression) => assert_eq!(
@@ -124,25 +111,8 @@ impl<'tcx> FunctionCoverage<'tcx> {
         }
     }
 
-    /// Adds regions that will be marked as "unreachable", with a constant "zero counter".
-    #[instrument(level = "debug", skip(self))]
-    pub(crate) fn add_unreachable_regions(&mut self, code_regions: &[CodeRegion]) {
-        assert!(!code_regions.is_empty(), "unreachable regions always have code regions");
-        self.add_mappings(CovTerm::Zero, code_regions);
-    }
-
-    #[instrument(level = "debug", skip(self))]
-    fn add_mappings(&mut self, term: CovTerm, code_regions: &[CodeRegion]) {
-        self.mappings
-            .extend(code_regions.iter().cloned().map(|code_region| Mapping { term, code_region }));
-    }
-
     pub(crate) fn finalize(&mut self) {
         self.update_zero_expressions();
-
-        // Sort all of the collected mappings into a predictable order.
-        // (Mappings have a total order, so an unstable sort should be fine.)
-        self.mappings.sort_unstable();
     }
 
     /// Identify expressions that will always have a value of zero, and note
@@ -249,8 +219,17 @@ impl<'tcx> FunctionCoverage<'tcx> {
     /// Converts this function's coverage mappings into an intermediate form
     /// that will be used by `mapgen` when preparing for FFI.
     pub(crate) fn mappings_for_ffi(&self) -> impl Iterator<Item = (Counter, &CodeRegion)> {
-        self.mappings.iter().map(|&Mapping { term, ref code_region }| {
-            let counter = Counter::from_term(term);
+        self.function_coverage_info.mappings.iter().map(|&Mapping { term, ref code_region }| {
+            // Historically, mappings were stored directly in counter/expression
+            // statements in MIR, and MIR optimizations would sometimes remove them.
+            // That's mostly no longer true, so now we detect cases where that would
+            // have happened, and zero out the corresponding mappings here instead.
+            let force_to_zero = match term {
+                CovTerm::Counter(id) => !self.counters_seen.contains(id),
+                CovTerm::Expression(id) => self.zero_expressions.contains(&id),
+                CovTerm::Zero => false,
+            };
+            let counter = if force_to_zero { Counter::ZERO } else { Counter::from_term(term) };
             (counter, code_region)
         })
     }

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -96,7 +96,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
     fn define_unused_fn(&self, def_id: DefId, function_coverage_info: &'tcx FunctionCoverageInfo) {
         let instance = declare_unused_fn(self, def_id);
         codegen_unused_fn_and_counter(self, instance);
-        add_unused_function_coverage(self, instance, def_id, function_coverage_info);
+        add_unused_function_coverage(self, instance, function_coverage_info);
     }
 }
 
@@ -116,10 +116,10 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
             .entry(instance)
             .or_insert_with(|| FunctionCoverage::new(instance, function_coverage_info));
 
-        let Coverage { kind, code_regions } = coverage;
+        let Coverage { kind } = coverage;
         match *kind {
             CoverageKind::Counter { id } => {
-                func_coverage.add_counter(id, code_regions);
+                func_coverage.add_counter(id);
                 // We need to explicitly drop the `RefMut` before calling into `instrprof_increment`,
                 // as that needs an exclusive borrow.
                 drop(coverage_map);
@@ -148,10 +148,7 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
                 bx.instrprof_increment(fn_name, hash, num_counters, index);
             }
             CoverageKind::Expression { id, lhs, op, rhs } => {
-                func_coverage.add_counter_expression(id, lhs, op, rhs, code_regions);
-            }
-            CoverageKind::Unreachable => {
-                func_coverage.add_unreachable_regions(code_regions);
+                func_coverage.add_counter_expression(id, lhs, op, rhs);
             }
         }
     }
@@ -214,16 +211,9 @@ fn codegen_unused_fn_and_counter<'tcx>(cx: &CodegenCx<'_, 'tcx>, instance: Insta
 fn add_unused_function_coverage<'tcx>(
     cx: &CodegenCx<'_, 'tcx>,
     instance: Instance<'tcx>,
-    def_id: DefId,
     function_coverage_info: &'tcx FunctionCoverageInfo,
 ) {
-    let tcx = cx.tcx;
-
-    let mut function_coverage = FunctionCoverage::unused(instance, function_coverage_info);
-    for &code_region in tcx.covered_code_regions(def_id) {
-        let code_region = std::slice::from_ref(code_region);
-        function_coverage.add_unreachable_regions(code_region);
-    }
+    let function_coverage = FunctionCoverage::unused(instance, function_coverage_info);
 
     if let Some(coverage_context) = cx.coverage_context() {
         coverage_context.function_coverage_map.borrow_mut().insert(instance, function_coverage);

compiler/rustc_middle/src/mir/coverage.rs

@@ -44,10 +44,9 @@ impl ExpressionId {
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
 #[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
 pub enum CovTerm {
-    // Coverage codegen relies on this variant order when it sorts a function's mappings.
+    Zero,
     Counter(CounterId),
     Expression(ExpressionId),
-    Zero,
 }
 
 impl Debug for CovTerm {
@@ -75,7 +74,6 @@ pub enum CoverageKind {
         op: Op,
         rhs: CovTerm,
     },
-    Unreachable,
 }
 
 impl Debug for CoverageKind {
@@ -94,7 +92,6 @@ impl Debug for CoverageKind {
                 },
                 rhs,
             ),
-            Unreachable => write!(fmt, "Unreachable"),
         }
     }
 }
@@ -139,6 +136,7 @@ impl Op {
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
 pub struct Mapping {
     // Coverage codegen relies on this field order when it sorts a function's mappings.
     pub code_region: CodeRegion,
@@ -161,4 +159,6 @@ pub struct FunctionCoverageInfo {
     pub function_source_hash: u64,
     pub num_counters: usize,
     pub num_expressions: usize,
+
+    pub mappings: Vec<Mapping>,
 }

compiler/rustc_middle/src/mir/mod.rs

@@ -1484,12 +1484,8 @@ impl Debug for Statement<'_> {
             AscribeUserType(box (ref place, ref c_ty), ref variance) => {
                 write!(fmt, "AscribeUserType({place:?}, {variance:?}, {c_ty:?})")
             }
-            Coverage(box self::Coverage { ref kind, ref code_regions }) => {
-                if code_regions.is_empty() {
-                    write!(fmt, "Coverage::{kind:?}")
-                } else {
-                    write!(fmt, "Coverage::{kind:?} for {code_regions:?}")
-                }
+            Coverage(box self::Coverage { ref kind }) => {
+                write!(fmt, "Coverage::{kind:?}")
             }
             Intrinsic(box ref intrinsic) => write!(fmt, "{intrinsic}"),
             ConstEvalCounter => write!(fmt, "ConstEvalCounter"),

compiler/rustc_middle/src/mir/syntax.rs

@@ -5,7 +5,7 @@
 
 use super::{BasicBlock, Constant, Local, SwitchTargets, UserTypeProjection};
 
-use crate::mir::coverage::{CodeRegion, CoverageKind};
+use crate::mir::coverage::CoverageKind;
 use crate::traits::Reveal;
 use crate::ty::adjustment::PointerCoercion;
 use crate::ty::GenericArgsRef;
@@ -523,7 +523,6 @@ pub enum FakeReadCause {
 #[derive(TypeFoldable, TypeVisitable)]
 pub struct Coverage {
     pub kind: CoverageKind,
-    pub code_regions: Vec<CodeRegion>,
 }
 
 #[derive(Clone, Debug, PartialEq, TyEncodable, TyDecodable, Hash, HashStable)]

compiler/rustc_middle/src/query/mod.rs

@@ -575,17 +575,6 @@ rustc_queries! {
         arena_cache
     }
 
-    /// Returns the `CodeRegions` for a function that has instrumented coverage, in case the
-    /// function was optimized out before codegen, and before being added to the Coverage Map.
-    query covered_code_regions(key: DefId) -> &'tcx Vec<&'tcx mir::coverage::CodeRegion> {
-        desc {
-            |tcx| "retrieving the covered `CodeRegion`s, if instrumented, for `{}`",
-            tcx.def_path_str(key)
-        }
-        arena_cache
-        cache_on_disk_if { key.is_local() }
-    }
-
     /// The `DefId` is the `DefId` of the containing MIR body. Promoteds do not have their own
     /// `DefId`. This function returns all promoteds in the specified body. The body references
     /// promoteds by the `DefId` and the `mir::Promoted` index. This is necessary, because

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -107,6 +107,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> {
@@ -148,6 +149,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
             function_source_hash,
             basic_coverage_blocks,
             coverage_counters,
+            mappings: Vec::new(),
         }
     }
 
@@ -281,6 +283,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),
         }));
     }
 
@@ -311,20 +314,17 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
             debug_used_expressions.add_expression_operands(&counter_kind);
             graphviz_data.add_bcb_coverage_spans_with_counter(bcb, bcb_spans, &counter_kind);
 
-            // Convert the coverage spans into a vector of code regions to be
-            // associated with this BCB's coverage statement.
-            let code_regions = bcb_spans
-                .iter()
-                .map(|coverage_span| {
-                    make_code_region(source_map, file_name, coverage_span.span, body_span)
-                })
-                .collect::<Vec<_>>();
+            let term = counter_kind.as_term();
+            self.mappings.extend(bcb_spans.iter().map(|coverage_span| {
+                let code_region =
+                    make_code_region(source_map, file_name, coverage_span.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,
             );
         }
     }
@@ -407,7 +407,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(
@@ -471,18 +470,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);
 }
@@ -496,10 +490,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);
 }