coverage: Collect HIR info during MIR building

The coverage instrumentor pass operates on MIR, but it needs access to some
source-related information that is normally not present in MIR.

Historically, that information was obtained by looking back at HIR during MIR
transformation. This patch arranges for the necessary information to be
collected ahead of time during MIR building instead.

Author: Zalathar

compiler/rustc_const_eval/src/transform/promote_consts.rs

@@ -971,6 +971,7 @@ pub fn promote_candidates<'tcx>(
             body.span,
             body.coroutine_kind(),
             body.tainted_by_errors,
+            None,
         );
         promoted.phase = MirPhase::Analysis(AnalysisPhase::Initial);
 

compiler/rustc_middle/src/mir/mod.rs

@@ -346,6 +346,13 @@ pub struct Body<'tcx> {
 
     pub tainted_by_errors: Option<ErrorGuaranteed>,
 
+    /// Extra information about this function's source code captured during MIR
+    /// building, for use by coverage instrumentation.
+    ///
+    /// If `-Cinstrument-coverage` is not active, or if an individual function
+    /// is not eligible for coverage, then this should always be `None`.
+    pub coverage_hir_info: Option<Box<coverage::HirInfo>>,
+
     /// Per-function coverage information added by the `InstrumentCoverage`
     /// pass, to be used in conjunction with the coverage statements injected
     /// into this body's blocks.
@@ -367,6 +374,7 @@ impl<'tcx> Body<'tcx> {
         span: Span,
         coroutine_kind: Option<CoroutineKind>,
         tainted_by_errors: Option<ErrorGuaranteed>,
+        coverage_hir_info: Option<Box<coverage::HirInfo>>,
     ) -> Self {
         // We need `arg_count` locals, and one for the return place.
         assert!(
@@ -400,6 +408,7 @@ impl<'tcx> Body<'tcx> {
             is_polymorphic: false,
             injection_phase: None,
             tainted_by_errors,
+            coverage_hir_info,
             function_coverage_info: None,
         };
         body.is_polymorphic = body.has_non_region_param();
@@ -429,6 +438,7 @@ impl<'tcx> Body<'tcx> {
             is_polymorphic: false,
             injection_phase: None,
             tainted_by_errors: None,
+            coverage_hir_info: None,
             function_coverage_info: None,
         };
         body.is_polymorphic = body.has_non_region_param();

compiler/rustc_mir_build/src/build/coverageinfo.rs

@@ -0,0 +1,117 @@
+use rustc_middle::hir;
+use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
+use rustc_middle::mir;
+use rustc_middle::ty::TyCtxt;
+use rustc_span::def_id::LocalDefId;
+use rustc_span::{ExpnKind, Span};
+
+/// If the given item is eligible for coverage instrumentation, collect relevant
+/// HIR information that will be needed by the instrumentor pass.
+pub(crate) fn make_coverage_hir_info_if_eligible(
+    tcx: TyCtxt<'_>,
+    def_id: LocalDefId,
+) -> Option<Box<mir::coverage::HirInfo>> {
+    assert!(tcx.sess.instrument_coverage());
+
+    is_eligible_for_coverage(tcx, def_id).then(|| Box::new(make_coverage_hir_info(tcx, def_id)))
+}
+
+fn is_eligible_for_coverage(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
+    let is_fn_like = tcx.hir().get_by_def_id(def_id).fn_kind().is_some();
+
+    // 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 !is_fn_like {
+        return false;
+    }
+
+    let codegen_fn_attrs = tcx.codegen_fn_attrs(def_id);
+    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_COVERAGE) {
+        return false;
+    }
+
+    true
+}
+
+fn make_coverage_hir_info(tcx: TyCtxt<'_>, def_id: LocalDefId) -> mir::coverage::HirInfo {
+    let (maybe_fn_sig, hir_body) = fn_sig_and_body(tcx, def_id);
+
+    let function_source_hash = hash_mir_source(tcx, hir_body);
+    let body_span = get_body_span(tcx, hir_body, def_id);
+
+    let spans_are_compatible = {
+        let source_map = tcx.sess.source_map();
+        |a: Span, b: Span| {
+            a.eq_ctxt(b)
+                && source_map.lookup_source_file_idx(a.lo())
+                    == source_map.lookup_source_file_idx(b.lo())
+        }
+    };
+
+    let fn_sig_span = if let Some(fn_sig) = maybe_fn_sig
+        && spans_are_compatible(fn_sig.span, body_span)
+        && fn_sig.span.lo() <= body_span.lo()
+    {
+        fn_sig.span.with_hi(body_span.lo())
+    } else {
+        body_span.shrink_to_lo()
+    };
+
+    mir::coverage::HirInfo { function_source_hash, fn_sig_span, body_span }
+}
+
+fn fn_sig_and_body(
+    tcx: TyCtxt<'_>,
+    def_id: LocalDefId,
+) -> (Option<&rustc_hir::FnSig<'_>>, &rustc_hir::Body<'_>) {
+    // FIXME(#79625): Consider improving MIR to provide the information needed, to avoid going back
+    // to HIR for it.
+    let hir_node = tcx.hir().get_by_def_id(def_id);
+    let (_, fn_body_id) =
+        hir::map::associated_body(hir_node).expect("HIR node is a function with body");
+    (hir_node.fn_sig(), tcx.hir().body(fn_body_id))
+}
+
+fn get_body_span<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    hir_body: &rustc_hir::Body<'tcx>,
+    def_id: LocalDefId,
+) -> Span {
+    let mut body_span = hir_body.value.span;
+
+    if tcx.is_closure(def_id.to_def_id()) {
+        // If the MIR function is a closure, and if the closure body span
+        // starts from a macro, but it's content is not in that macro, try
+        // to find a non-macro callsite, and instrument the spans there
+        // instead.
+        loop {
+            let expn_data = body_span.ctxt().outer_expn_data();
+            if expn_data.is_root() {
+                break;
+            }
+            if let ExpnKind::Macro { .. } = expn_data.kind {
+                body_span = expn_data.call_site;
+            } else {
+                break;
+            }
+        }
+    }
+
+    body_span
+}
+
+fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &'tcx rustc_hir::Body<'tcx>) -> u64 {
+    // FIXME(cjgillot) Stop hashing HIR manually here.
+    let owner = hir_body.id().hir_id.owner;
+    tcx.hir_owner_nodes(owner)
+        .unwrap()
+        .opt_hash_including_bodies
+        .unwrap()
+        .to_smaller_hash()
+        .as_u64()
+}

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,
+        coverage_hir_info: None,
         function_coverage_info: None,
     };
 

compiler/rustc_mir_build/src/build/mod.rs

@@ -693,6 +693,7 @@ fn construct_error(tcx: TyCtxt<'_>, def_id: LocalDefId, guar: ErrorGuaranteed) -
         span,
         coroutine_kind,
         Some(guar),
+        None,
     );
 
     body.coroutine.as_mut().map(|gen| gen.yield_ty = yield_ty);
@@ -775,6 +776,12 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             }
         }
 
+        let coverage_hir_info = if self.tcx.sess.instrument_coverage() {
+            coverageinfo::make_coverage_hir_info_if_eligible(self.tcx, self.def_id)
+        } else {
+            None
+        };
+
         Body::new(
             MirSource::item(self.def_id.to_def_id()),
             self.cfg.basic_blocks,
@@ -786,6 +793,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             self.fn_span,
             self.coroutine_kind,
             None,
+            coverage_hir_info,
         )
     }
 
@@ -1056,6 +1064,7 @@ pub(crate) fn parse_float_into_scalar(
 
 mod block;
 mod cfg;
+mod coverageinfo;
 mod custom;
 mod expr;
 mod matches;

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -14,17 +14,14 @@ use self::spans::CoverageSpans;
 use crate::MirPass;
 
 use rustc_data_structures::sync::Lrc;
-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,
     TerminatorKind,
 };
 use rustc_middle::ty::TyCtxt;
-use rustc_span::def_id::LocalDefId;
 use rustc_span::source_map::SourceMap;
-use rustc_span::{ExpnKind, SourceFile, Span, Symbol};
+use rustc_span::{SourceFile, Span, Symbol};
 
 /// Inserts `StatementKind::Coverage` statements that either instrument the binary with injected
 /// counters, via intrinsic `llvm.instrprof.increment`, and/or inject metadata used during codegen
@@ -43,8 +40,10 @@ impl<'tcx> MirPass<'tcx> for InstrumentCoverage {
         // be transformed, so it should never see promoted MIR.
         assert!(mir_source.promoted.is_none());
 
-        let def_id = mir_source.def_id().expect_local();
-        if !is_eligible_for_coverage(tcx, def_id) {
+        let def_id = mir_source.def_id();
+        if mir_body.coverage_hir_info.is_none() {
+            // If we didn't capture HIR info during MIR building, this MIR
+            // wasn't eligible for coverage instrumentation, so skip it.
             trace!("InstrumentCoverage skipped for {def_id:?} (not eligible)");
             return;
         }
@@ -79,8 +78,10 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
         let source_map = tcx.sess.source_map();
 
         let def_id = mir_body.source.def_id().expect_local();
-        let mir::coverage::HirInfo { function_source_hash, fn_sig_span, body_span, .. } =
-            make_coverage_hir_info(tcx, def_id);
+        let &mir::coverage::HirInfo { function_source_hash, fn_sig_span, body_span, .. } = mir_body
+            .coverage_hir_info
+            .as_deref()
+            .expect("functions without HIR info have already been skipped");
 
         let source_file = source_map.lookup_source_file(body_span.lo());
 
@@ -290,103 +291,3 @@ fn make_code_region(
         end_col: end_col as u32,
     }
 }
-
-fn is_eligible_for_coverage(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
-    let is_fn_like = tcx.hir().get_by_def_id(def_id).fn_kind().is_some();
-
-    // 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 !is_fn_like {
-        return false;
-    }
-
-    let codegen_fn_attrs = tcx.codegen_fn_attrs(def_id);
-    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_COVERAGE) {
-        return false;
-    }
-
-    true
-}
-
-fn make_coverage_hir_info(tcx: TyCtxt<'_>, def_id: LocalDefId) -> mir::coverage::HirInfo {
-    let (maybe_fn_sig, hir_body) = fn_sig_and_body(tcx, def_id);
-
-    let function_source_hash = hash_mir_source(tcx, hir_body);
-    let body_span = get_body_span(tcx, hir_body, def_id);
-
-    let spans_are_compatible = {
-        let source_map = tcx.sess.source_map();
-        |a: Span, b: Span| {
-            a.eq_ctxt(b)
-                && source_map.lookup_source_file_idx(a.lo())
-                    == source_map.lookup_source_file_idx(b.lo())
-        }
-    };
-
-    let fn_sig_span = if let Some(fn_sig) = maybe_fn_sig
-        && spans_are_compatible(fn_sig.span, body_span)
-        && fn_sig.span.lo() <= body_span.lo()
-    {
-        fn_sig.span.with_hi(body_span.lo())
-    } else {
-        body_span.shrink_to_lo()
-    };
-
-    mir::coverage::HirInfo { function_source_hash, fn_sig_span, body_span }
-}
-
-fn fn_sig_and_body(
-    tcx: TyCtxt<'_>,
-    def_id: LocalDefId,
-) -> (Option<&rustc_hir::FnSig<'_>>, &rustc_hir::Body<'_>) {
-    // FIXME(#79625): Consider improving MIR to provide the information needed, to avoid going back
-    // to HIR for it.
-    let hir_node = tcx.hir().get_by_def_id(def_id);
-    let (_, fn_body_id) =
-        hir::map::associated_body(hir_node).expect("HIR node is a function with body");
-    (hir_node.fn_sig(), tcx.hir().body(fn_body_id))
-}
-
-fn get_body_span<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    hir_body: &rustc_hir::Body<'tcx>,
-    def_id: LocalDefId,
-) -> Span {
-    let mut body_span = hir_body.value.span;
-
-    if tcx.is_closure(def_id.to_def_id()) {
-        // If the MIR function is a closure, and if the closure body span
-        // starts from a macro, but it's content is not in that macro, try
-        // to find a non-macro callsite, and instrument the spans there
-        // instead.
-        loop {
-            let expn_data = body_span.ctxt().outer_expn_data();
-            if expn_data.is_root() {
-                break;
-            }
-            if let ExpnKind::Macro { .. } = expn_data.kind {
-                body_span = expn_data.call_site;
-            } else {
-                break;
-            }
-        }
-    }
-
-    body_span
-}
-
-fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &'tcx rustc_hir::Body<'tcx>) -> u64 {
-    // FIXME(cjgillot) Stop hashing HIR manually here.
-    let owner = hir_body.id().hir_id.owner;
-    tcx.hir_owner_nodes(owner)
-        .unwrap()
-        .opt_hash_including_bodies
-        .unwrap()
-        .to_smaller_hash()
-        .as_u64()
-}

compiler/rustc_mir_transform/src/shim.rs

@@ -281,6 +281,7 @@ fn new_body<'tcx>(
         None,
         // FIXME(compiler-errors): is this correct?
         None,
+        None,
     )
 }