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.

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 [`ZeroExpressions`]. Mappings that refer to a zero expression
     /// can instead become mappings to a constant zero value.
@@ -235,7 +199,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)
     }
@@ -280,9 +244,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

@@ -88,13 +88,12 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
     /// For used/called functions, the coverageinfo was already added to the
     /// `function_coverage_map` (keyed by function `Instance`) during codegen.
     /// But in this case, since the unused function was _not_ previously
-    /// codegenned, collect the coverage `CodeRegion`s from the MIR and add
-    /// them. Since the function is never called, all of its `CodeRegion`s can be
-    /// added as `unreachable_region`s.
+    /// codegenned, collect the function coverage info from MIR and add an
+    /// "unused" entry to the function coverage map.
     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 +115,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);
@@ -147,10 +146,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);
             }
         }
     }
@@ -213,16 +209,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,13 @@ 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,
-    },
+    /// Marks the point in MIR control flow represented by a coverage counter.
+    ///
+    /// This is eventually lowered to `llvm.instrprof.increment` in LLVM IR.
+    ///
+    /// If this statement does not survive MIR optimizations, any mappings that
+    /// refer to this counter can have those references simplified to zero.
+    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 +76,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 +94,6 @@ impl Debug for CoverageKind {
                 },
                 rhs,
             ),
-            Unreachable => write!(fmt, "Unreachable"),
         }
     }
 }
@@ -158,4 +158,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

@@ -493,6 +493,24 @@ pub fn write_mir_intro<'tcx>(
     // Add an empty line before the first block is printed.
     writeln!(w)?;
 
+    if let Some(function_coverage_info) = &body.function_coverage_info {
+        write_function_coverage_info(function_coverage_info, w)?;
+    }
+
+    Ok(())
+}
+
+fn write_function_coverage_info(
+    function_coverage_info: &coverage::FunctionCoverageInfo,
+    w: &mut dyn io::Write,
+) -> io::Result<()> {
+    let coverage::FunctionCoverageInfo { mappings, .. } = function_coverage_info;
+
+    for coverage::Mapping { term, code_region } in mappings {
+        writeln!(w, "{INDENT}coverage {term:?} => {code_region:?};")?;
+    }
+    writeln!(w)?;
+
     Ok(())
 }
 
@@ -685,13 +703,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;
@@ -514,7 +514,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

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