coverage: Attach an optional `FunctionCoverageInfo` to `mir::Body`

This allows coverage information to be attached to the function as a whole when
appropriate, instead of being smuggled through coverage statements in the
function's basic blocks.

As an example, this patch moves the `function_source_hash` value out of
individual `CoverageKind::Counter` statements and into the per-function info.

When synthesizing unused functions for coverage purposes, the absence of this
info is taken to indicate that a function was not eligible for coverage and
should not be synthesized.

7 files changed, 74 insertions, 54 deletions

diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
index 55f43aa5341..f9d76108302 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
@@ -2,7 +2,9 @@ use crate::coverageinfo::ffi::{Counter, CounterExpression, ExprKind};
 
 use rustc_data_structures::fx::FxIndexSet;
 use rustc_index::IndexVec;
-use rustc_middle::mir::coverage::{CodeRegion, CounterId, ExpressionId, Op, Operand};
+use rustc_middle::mir::coverage::{
+    CodeRegion, CounterId, ExpressionId, FunctionCoverageInfo, Op, Operand,
+};
 use rustc_middle::ty::Instance;
 use rustc_middle::ty::TyCtxt;
 
@@ -27,8 +29,8 @@ pub struct Expression {
 /// line."
 #[derive(Debug)]
 pub struct FunctionCoverage<'tcx> {
-    instance: Instance<'tcx>,
-    source_hash: u64,
+    /// Coverage info that was attached to this function by the instrumentor.
+    function_coverage_info: &'tcx FunctionCoverageInfo,
     is_used: bool,
     counters: IndexVec<CounterId, Option<Vec<CodeRegion>>>,
     expressions: IndexVec<ExpressionId, Option<Expression>>,
@@ -37,24 +39,36 @@ pub struct FunctionCoverage<'tcx> {
 
 impl<'tcx> FunctionCoverage<'tcx> {
     /// Creates a new set of coverage data for a used (called) function.
-    pub fn new(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> Self {
-        Self::create(tcx, instance, true)
+    pub fn new(
+        tcx: TyCtxt<'tcx>,
+        instance: Instance<'tcx>,
+        function_coverage_info: &'tcx FunctionCoverageInfo,
+    ) -> Self {
+        Self::create(tcx, instance, function_coverage_info, true)
     }
 
     /// Creates a new set of coverage data for an unused (never called) function.
-    pub fn unused(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> Self {
-        Self::create(tcx, instance, false)
+    pub fn unused(
+        tcx: TyCtxt<'tcx>,
+        instance: Instance<'tcx>,
+        function_coverage_info: &'tcx FunctionCoverageInfo,
+    ) -> Self {
+        Self::create(tcx, instance, function_coverage_info, false)
     }
 
-    fn create(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>, is_used: bool) -> Self {
+    fn create(
+        tcx: TyCtxt<'tcx>,
+        instance: Instance<'tcx>,
+        function_coverage_info: &'tcx FunctionCoverageInfo,
+        is_used: bool,
+    ) -> Self {
         let coverageinfo = tcx.coverageinfo(instance.def);
         debug!(
             "FunctionCoverage::create(instance={:?}) has coverageinfo={:?}. is_used={}",
             instance, coverageinfo, is_used
         );
         Self {
-            instance,
-            source_hash: 0, // will be set with the first `add_counter()`
+            function_coverage_info,
             is_used,
             counters: IndexVec::from_elem_n(None, coverageinfo.num_counters as usize),
             expressions: IndexVec::from_elem_n(None, coverageinfo.num_expressions as usize),
@@ -67,16 +81,6 @@ impl<'tcx> FunctionCoverage<'tcx> {
         self.is_used
     }
 
-    /// Sets the function source hash value. If called multiple times for the same function, all
-    /// calls should have the same hash value.
-    pub fn set_function_source_hash(&mut self, source_hash: u64) {
-        if self.source_hash == 0 {
-            self.source_hash = source_hash;
-        } else {
-            debug_assert_eq!(source_hash, self.source_hash);
-        }
-    }
-
     /// Adds code regions to be counted by an injected counter intrinsic.
     #[instrument(level = "debug", skip(self))]
     pub(crate) fn add_counter(&mut self, id: CounterId, code_regions: &[CodeRegion]) {
@@ -195,7 +199,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
     /// Return the source hash, generated from the HIR node structure, and used to indicate whether
     /// or not the source code structure changed between different compilations.
     pub fn source_hash(&self) -> u64 {
-        self.source_hash
+        if self.is_used { self.function_coverage_info.function_source_hash } else { 0 }
     }
 
     /// Generate an array of CounterExpressions, and an iterator over all `Counter`s and their
@@ -204,12 +208,6 @@ impl<'tcx> FunctionCoverage<'tcx> {
     pub fn get_expressions_and_counter_regions(
         &self,
     ) -> (Vec<CounterExpression>, impl Iterator<Item = (Counter, &CodeRegion)>) {
-        assert!(
-            self.source_hash != 0 || !self.is_used,
-            "No counters provided the source_hash for used function: {:?}",
-            self.instance
-        );
-
         let counter_expressions = self.counter_expressions();
         // Expression IDs are indices into `self.expressions`, and on the LLVM
         // side they will be treated as indices into `counter_expressions`, so
diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
index d4e77525698..a2bf53e605d 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
@@ -10,9 +10,8 @@ use rustc_hir::def::DefKind;
 use rustc_hir::def_id::DefId;
 use rustc_index::IndexVec;
 use rustc_middle::bug;
-use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::coverage::CodeRegion;
-use rustc_middle::ty::TyCtxt;
+use rustc_middle::ty::{self, TyCtxt};
 use rustc_span::Symbol;
 
 /// Generates and exports the Coverage Map.
@@ -331,16 +330,14 @@ fn add_unused_functions(cx: &CodegenCx<'_, '_>) {
     for non_codegenned_def_id in
         eligible_def_ids.into_iter().filter(|id| !codegenned_def_ids.contains(id))
     {
-        let codegen_fn_attrs = tcx.codegen_fn_attrs(non_codegenned_def_id);
-
-        // If a function is marked `#[coverage(off)]`, then skip generating a
-        // dead code stub for it.
-        if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_COVERAGE) {
-            debug!("skipping unused fn marked #[coverage(off)]: {:?}", non_codegenned_def_id);
+        // Skip any function that didn't have coverage data added to it by the
+        // coverage instrumentor.
+        let body = tcx.instance_mir(ty::InstanceDef::Item(non_codegenned_def_id));
+        let Some(function_coverage_info) = body.function_coverage_info.as_deref() else {
             continue;
-        }
+        };
 
         debug!("generating unused fn: {:?}", non_codegenned_def_id);
-        cx.define_unused_fn(non_codegenned_def_id);
+        cx.define_unused_fn(non_codegenned_def_id, function_coverage_info);
     }
 }
diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
index dd2ce9b525b..0d9d3c7ad41 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
@@ -16,7 +16,7 @@ use rustc_hir as hir;
 use rustc_hir::def_id::DefId;
 use rustc_llvm::RustString;
 use rustc_middle::bug;
-use rustc_middle::mir::coverage::{CounterId, CoverageKind};
+use rustc_middle::mir::coverage::{CounterId, CoverageKind, FunctionCoverageInfo};
 use rustc_middle::mir::Coverage;
 use rustc_middle::ty;
 use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt};
@@ -91,31 +91,34 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
     /// 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.
-    fn define_unused_fn(&self, def_id: DefId) {
+    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);
+        add_unused_function_coverage(self, instance, def_id, function_coverage_info);
     }
 }
 
 impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
+    #[instrument(level = "debug", skip(self))]
     fn add_coverage(&mut self, instance: Instance<'tcx>, coverage: &Coverage) {
         let bx = self;
 
+        let Some(function_coverage_info) =
+            bx.tcx.instance_mir(instance.def).function_coverage_info.as_deref()
+        else {
+            debug!("function has a coverage statement but no coverage info");
+            return;
+        };
+
         let Some(coverage_context) = bx.coverage_context() else { return };
         let mut coverage_map = coverage_context.function_coverage_map.borrow_mut();
         let func_coverage = coverage_map
             .entry(instance)
-            .or_insert_with(|| FunctionCoverage::new(bx.tcx(), instance));
+            .or_insert_with(|| FunctionCoverage::new(bx.tcx(), instance, function_coverage_info));
 
         let Coverage { kind, code_regions } = coverage;
         match *kind {
-            CoverageKind::Counter { function_source_hash, id } => {
-                debug!(
-                    "ensuring function source hash is set for instance={:?}; function_source_hash={}",
-                    instance, function_source_hash,
-                );
-                func_coverage.set_function_source_hash(function_source_hash);
+            CoverageKind::Counter { id } => {
                 func_coverage.add_counter(id, code_regions);
                 // We need to explicitly drop the `RefMut` before calling into `instrprof_increment`,
                 // as that needs an exclusive borrow.
@@ -124,7 +127,7 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
                 let coverageinfo = bx.tcx().coverageinfo(instance.def);
 
                 let fn_name = bx.get_pgo_func_name_var(instance);
-                let hash = bx.const_u64(function_source_hash);
+                let hash = bx.const_u64(function_coverage_info.function_source_hash);
                 let num_counters = bx.const_u32(coverageinfo.num_counters);
                 let index = bx.const_u32(id.as_u32());
                 debug!(
@@ -201,10 +204,11 @@ 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(tcx, instance);
+    let mut function_coverage = FunctionCoverage::unused(tcx, 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);
diff --git a/compiler/rustc_middle/src/mir/coverage.rs b/compiler/rustc_middle/src/mir/coverage.rs
index a6d6f6f5df4..e1434e71a47 100644
--- a/compiler/rustc_middle/src/mir/coverage.rs
+++ b/compiler/rustc_middle/src/mir/coverage.rs
@@ -60,7 +60,6 @@ impl Debug for Operand {
 #[derive(Clone, PartialEq, TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
 pub enum CoverageKind {
     Counter {
-        function_source_hash: u64,
         /// ID of this counter within its enclosing function.
         /// Expressions in the same function can refer to it as an operand.
         id: CounterId,
@@ -80,7 +79,7 @@ impl Debug for CoverageKind {
     fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
         use CoverageKind::*;
         match self {
-            Counter { id, .. } => write!(fmt, "Counter({:?})", id.index()),
+            Counter { id } => write!(fmt, "Counter({:?})", id.index()),
             Expression { id, lhs, op, rhs } => write!(
                 fmt,
                 "Expression({:?}) = {:?} {} {:?}",
@@ -133,3 +132,12 @@ impl Op {
         matches!(self, Self::Subtract)
     }
 }
+
+/// Stores per-function coverage information attached to a `mir::Body`,
+/// to be used in conjunction with the individual coverage statements injected
+/// into the function's basic blocks.
+#[derive(Clone, Debug)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+pub struct FunctionCoverageInfo {
+    pub function_source_hash: u64,
+}
diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
index f979f736b15..3a5ff4dc91f 100644
--- a/compiler/rustc_middle/src/mir/mod.rs
+++ b/compiler/rustc_middle/src/mir/mod.rs
@@ -345,6 +345,14 @@ pub struct Body<'tcx> {
     pub injection_phase: Option<MirPhase>,
 
     pub tainted_by_errors: Option<ErrorGuaranteed>,
+
+    /// Per-function coverage information added by the `InstrumentCoverage`
+    /// pass, to be used in conjunction with the coverage statements injected
+    /// into this body's blocks.
+    ///
+    /// If `-Cinstrument-coverage` is not active, or if an individual function
+    /// is not eligible for coverage, then this should always be `None`.
+    pub function_coverage_info: Option<Box<coverage::FunctionCoverageInfo>>,
 }
 
 impl<'tcx> Body<'tcx> {
@@ -392,6 +400,7 @@ impl<'tcx> Body<'tcx> {
             is_polymorphic: false,
             injection_phase: None,
             tainted_by_errors,
+            function_coverage_info: None,
         };
         body.is_polymorphic = body.has_non_region_param();
         body
@@ -420,6 +429,7 @@ impl<'tcx> Body<'tcx> {
             is_polymorphic: false,
             injection_phase: None,
             tainted_by_errors: None,
+            function_coverage_info: None,
         };
         body.is_polymorphic = body.has_non_region_param();
         body
diff --git a/compiler/rustc_mir_build/src/build/custom/mod.rs b/compiler/rustc_mir_build/src/build/custom/mod.rs
index e5c2cc6c7bb..a81f70e3346 100644
--- a/compiler/rustc_mir_build/src/build/custom/mod.rs
+++ b/compiler/rustc_mir_build/src/build/custom/mod.rs
@@ -60,6 +60,7 @@ pub(super) fn build_custom_mir<'tcx>(
         tainted_by_errors: None,
         injection_phase: None,
         pass_count: 0,
+        function_coverage_info: None,
     };
 
     body.local_decls.push(LocalDecl::new(return_ty, return_ty_span));
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index abf13519e9e..1ce6ecc3bb5 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -211,6 +211,10 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
                 self.make_mir_coverage_kind(intermediate_expression),
             );
         }
+
+        self.mir_body.function_coverage_info = Some(Box::new(FunctionCoverageInfo {
+            function_source_hash: self.function_source_hash,
+        }));
     }
 
     /// Injects a single [`StatementKind::Coverage`] for each BCB that has one
@@ -323,9 +327,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 { function_source_hash: self.function_source_hash, id }
-            }
+            BcbCounter::Counter { id } => CoverageKind::Counter { id },
             BcbCounter::Expression { id, lhs, op, rhs } => {
                 CoverageKind::Expression { id, lhs, op, rhs }
             }