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

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -23,11 +23,10 @@ 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>>,
-    mappings: Vec<Mapping>,
 }
 
 impl<'tcx> FunctionCoverage<'tcx> {
@@ -63,7 +62,6 @@ impl<'tcx> FunctionCoverage<'tcx> {
             is_used,
             counters_seen: BitSet::new_empty(num_counters),
             expressions: IndexVec::from_elem_n(None, num_expressions),
-            mappings: Vec::new(),
         }
     }
 
@@ -72,28 +70,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 mark_counter_id_seen(&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(),
@@ -107,10 +97,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!(
@@ -120,29 +107,6 @@ 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) {
-        // Reorder the collected mappings so that counter mappings are first and
-        // zero mappings are last, matching the historical order.
-        self.mappings.sort_by_key(|mapping| match mapping.term {
-            CovTerm::Counter(_) => 0,
-            CovTerm::Expression(_) => 1,
-            CovTerm::Zero => u8::MAX,
-        });
-    }
-
     /// Identify expressions that will always have a value of zero, and note
     /// their IDs in `zero_expressions`. Mappings that refer to a zero expression
     /// can instead become mappings to a constant zero value.
@@ -222,7 +186,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
         // the two vectors should correspond 1:1.
         assert_eq!(self.expressions.len(), counter_expressions.len());
 
-        let counter_regions = self.counter_regions();
+        let counter_regions = self.counter_regions(zero_expressions);
 
         (counter_expressions, counter_regions)
     }
@@ -267,9 +231,26 @@ impl<'tcx> FunctionCoverage<'tcx> {
 
     /// Converts this function's coverage mappings into an intermediate form
     /// that will be used by `mapgen` when preparing for FFI.
-    fn counter_regions(&self) -> impl Iterator<Item = (Counter, &CodeRegion)> {
-        self.mappings.iter().map(|&Mapping { term, ref code_region }| {
-            let counter = Counter::from_term(term);
+    fn counter_regions(
+        &self,
+        zero_expressions: ZeroExpressions,
+    ) -> impl Iterator<Item = (Counter, &CodeRegion)> {
+        // 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 counter_for_term = move |term: CovTerm| {
+            let force_to_zero = match term {
+                CovTerm::Counter(id) => !self.counters_seen.contains(id),
+                CovTerm::Expression(id) => zero_expressions.contains(id),
+                CovTerm::Zero => false,
+            };
+            if force_to_zero { Counter::ZERO } else { Counter::from_term(term) }
+        };
+
+        self.function_coverage_info.mappings.iter().map(move |mapping| {
+            let &Mapping { term, ref code_region } = mapping;
+            let counter = counter_for_term(term);
             (counter, code_region)
         })
     }

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -59,10 +59,8 @@ pub fn finalize(cx: &CodegenCx<'_, '_>) {
 
     // Encode coverage mappings and generate function records
     let mut function_data = Vec::new();
-    for (instance, mut function_coverage) in function_coverage_map {
+    for (instance, function_coverage) in function_coverage_map {
         debug!("Generate function coverage for {}, {:?}", cx.codegen_unit.name(), instance);
-        function_coverage.finalize();
-        let function_coverage = function_coverage;
 
         let mangled_function_name = tcx.symbol_name(instance).name;
         let source_hash = function_coverage.source_hash();

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);
+            CoverageKind::CounterIncrement { id } => {
+                func_coverage.mark_counter_id_seen(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,11 @@ 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);
-    }
+    // An unused function's mappings will automatically be rewritten to map to
+    // zero, because none of its counters/expressions are marked as seen.
+    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

@@ -61,11 +61,11 @@ impl Debug for CovTerm {
 
 #[derive(Clone, PartialEq, TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
 pub enum CoverageKind {
-    Counter {
-        /// ID of this counter within its enclosing function.
-        /// Expressions in the same function can refer to it as an operand.
-        id: CounterId,
-    },
+    /// A counter increment statement serves two purposes:
+    /// - It is lowered to `llvm.instrprof.increment` in LLVM IR.
+    /// - Coverage mappings that refer to this counter ID should *not* be
+    ///   simplified to zero, because it survived MIR optimizations.
+    CounterIncrement { id: CounterId },
     Expression {
         /// ID of this coverage-counter expression within its enclosing function.
         /// Other expressions in the same function can refer to it as an operand.
@@ -74,14 +74,13 @@ pub enum CoverageKind {
         op: Op,
         rhs: CovTerm,
     },
-    Unreachable,
 }
 
 impl Debug for CoverageKind {
     fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
         use CoverageKind::*;
         match self {
-            Counter { id } => write!(fmt, "Counter({:?})", id.index()),
+            CounterIncrement { id } => write!(fmt, "CounterIncrement({:?})", id.index()),
             Expression { id, lhs, op, rhs } => write!(
                 fmt,
                 "Expression({:?}) = {:?} {} {:?}",
@@ -93,7 +92,6 @@ impl Debug for CoverageKind {
                 },
                 rhs,
             ),
-            Unreachable => write!(fmt, "Unreachable"),
         }
     }
 }
@@ -158,4 +156,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/pretty.rs

@@ -685,13 +685,7 @@ impl Debug for Statement<'_> {
             AscribeUserType(box (ref place, ref c_ty), ref variance) => {
                 write!(fmt, "AscribeUserType({place:?}, {variance:?}, {c_ty:?})")
             }
-            Coverage(box mir::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 mir::Coverage { ref kind }) => write!(fmt, "Coverage::{kind:?}"),
             Intrinsic(box ref intrinsic) => write!(fmt, "{intrinsic}"),
             ConstEvalCounter => write!(fmt, "ConstEvalCounter"),
             Nop => write!(fmt, "nop"),

compiler/rustc_middle/src/mir/syntax.rs

@@ -5,7 +5,7 @@
 
 use super::{BasicBlock, Const, Local, 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;
@@ -513,7 +513,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

@@ -573,17 +573,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

@@ -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);
 }