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.

4 files changed, 30 insertions, 28 deletions

diff --git a/compiler/rustc_middle/src/mir/coverage.rs b/compiler/rustc_middle/src/mir/coverage.rs
index 164a17ff77a..2587a6eae1a 100644
--- a/compiler/rustc_middle/src/mir/coverage.rs
+++ b/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>,
 }
diff --git a/compiler/rustc_middle/src/mir/pretty.rs b/compiler/rustc_middle/src/mir/pretty.rs
index c5e3ee575e1..11269e749f8 100644
--- a/compiler/rustc_middle/src/mir/pretty.rs
+++ b/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"),
diff --git a/compiler/rustc_middle/src/mir/syntax.rs b/compiler/rustc_middle/src/mir/syntax.rs
index c7d99648f1e..77b26e0a4cd 100644
--- a/compiler/rustc_middle/src/mir/syntax.rs
+++ b/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)]
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index ed0f73fcf2f..afe94d10752 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/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