coverage: Make `is_eligible_for_coverage` a hook method

This will allow MIR building to check whether a function is eligible for
coverage instrumentation, and avoid collecting branch coverage info if it is
not.

2 files changed, 41 insertions, 38 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index 4c5be0a3f4b..bde13583486 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -14,7 +14,6 @@ use self::spans::{BcbMapping, BcbMappingKind, CoverageSpans};
 use crate::MirPass;
 
 use rustc_middle::hir;
-use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::coverage::*;
 use rustc_middle::mir::{
     self, BasicBlock, BasicBlockData, Coverage, SourceInfo, Statement, StatementKind, Terminator,
@@ -44,7 +43,7 @@ impl<'tcx> MirPass<'tcx> for InstrumentCoverage {
 
         let def_id = mir_source.def_id().expect_local();
 
-        if !is_eligible_for_coverage(tcx, def_id) {
+        if !tcx.is_eligible_for_coverage(def_id) {
             trace!("InstrumentCoverage skipped for {def_id:?} (not eligible)");
             return;
         }
@@ -349,37 +348,6 @@ fn check_code_region(code_region: CodeRegion) -> Option<CodeRegion> {
     }
 }
 
-fn is_eligible_for_coverage(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
-    // Only instrument functions, methods, and closures (not constants since they are evaluated
-    // at compile time by Miri).
-    // FIXME(#73156): Handle source code coverage in const eval, but note, if and when const
-    // expressions get coverage spans, we will probably have to "carve out" space for const
-    // expressions from coverage spans in enclosing MIR's, like we do for closures. (That might
-    // be tricky if const expressions have no corresponding statements in the enclosing MIR.
-    // Closures are carved out by their initial `Assign` statement.)
-    if !tcx.def_kind(def_id).is_fn_like() {
-        trace!("InstrumentCoverage skipped for {def_id:?} (not an fn-like)");
-        return false;
-    }
-
-    // Don't instrument functions with `#[automatically_derived]` on their
-    // enclosing impl block, on the assumption that most users won't care about
-    // coverage for derived impls.
-    if let Some(impl_of) = tcx.impl_of_method(def_id.to_def_id())
-        && tcx.is_automatically_derived(impl_of)
-    {
-        trace!("InstrumentCoverage skipped for {def_id:?} (automatically derived)");
-        return false;
-    }
-
-    if tcx.codegen_fn_attrs(def_id).flags.contains(CodegenFnAttrFlags::NO_COVERAGE) {
-        trace!("InstrumentCoverage skipped for {def_id:?} (`#[coverage(off)]`)");
-        return false;
-    }
-
-    true
-}
-
 /// Function information extracted from HIR by the coverage instrumentor.
 #[derive(Debug)]
 struct ExtractedHirInfo {
diff --git a/compiler/rustc_mir_transform/src/coverage/query.rs b/compiler/rustc_mir_transform/src/coverage/query.rs
index 1354273e2c4..1de7b6f66a7 100644
--- a/compiler/rustc_mir_transform/src/coverage/query.rs
+++ b/compiler/rustc_mir_transform/src/coverage/query.rs
@@ -1,16 +1,51 @@
-use super::*;
-
 use rustc_data_structures::captures::Captures;
-use rustc_middle::mir::coverage::*;
-use rustc_middle::mir::{Body, CoverageIdsInfo};
-use rustc_middle::ty::{self};
+use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
+use rustc_middle::mir::coverage::{CounterId, CoverageKind};
+use rustc_middle::mir::{Body, Coverage, CoverageIdsInfo, Statement, StatementKind};
+use rustc_middle::query::TyCtxtAt;
+use rustc_middle::ty::{self, TyCtxt};
 use rustc_middle::util::Providers;
+use rustc_span::def_id::LocalDefId;
 
 /// Registers query/hook implementations related to coverage.
 pub(crate) fn provide(providers: &mut Providers) {
+    providers.hooks.is_eligible_for_coverage =
+        |TyCtxtAt { tcx, .. }, def_id| is_eligible_for_coverage(tcx, def_id);
     providers.queries.coverage_ids_info = coverage_ids_info;
 }
 
+/// Hook implementation for [`TyCtxt::is_eligible_for_coverage`].
+fn is_eligible_for_coverage(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
+    // Only instrument functions, methods, and closures (not constants since they are evaluated
+    // at compile time by Miri).
+    // FIXME(#73156): Handle source code coverage in const eval, but note, if and when const
+    // expressions get coverage spans, we will probably have to "carve out" space for const
+    // expressions from coverage spans in enclosing MIR's, like we do for closures. (That might
+    // be tricky if const expressions have no corresponding statements in the enclosing MIR.
+    // Closures are carved out by their initial `Assign` statement.)
+    if !tcx.def_kind(def_id).is_fn_like() {
+        trace!("InstrumentCoverage skipped for {def_id:?} (not an fn-like)");
+        return false;
+    }
+
+    // Don't instrument functions with `#[automatically_derived]` on their
+    // enclosing impl block, on the assumption that most users won't care about
+    // coverage for derived impls.
+    if let Some(impl_of) = tcx.impl_of_method(def_id.to_def_id())
+        && tcx.is_automatically_derived(impl_of)
+    {
+        trace!("InstrumentCoverage skipped for {def_id:?} (automatically derived)");
+        return false;
+    }
+
+    if tcx.codegen_fn_attrs(def_id).flags.contains(CodegenFnAttrFlags::NO_COVERAGE) {
+        trace!("InstrumentCoverage skipped for {def_id:?} (`#[coverage(off)]`)");
+        return false;
+    }
+
+    true
+}
+
 /// Query implementation for `coverage_ids_info`.
 fn coverage_ids_info<'tcx>(
     tcx: TyCtxt<'tcx>,