coverage: Collect HIR info during MIR building

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/coverage.rs

@@ -2,7 +2,7 @@
 
 use rustc_index::IndexVec;
 use rustc_macros::HashStable;
-use rustc_span::Symbol;
+use rustc_span::{Span, Symbol};
 
 use std::fmt::{self, Debug, Formatter};
 
@@ -184,3 +184,15 @@ pub struct FunctionCoverageInfo {
     pub expressions: IndexVec<ExpressionId, Expression>,
     pub mappings: Vec<Mapping>,
 }
+
+/// Coverage information captured from HIR/THIR during MIR building.
+///
+/// A MIR body that does not have this information cannot be instrumented for
+/// coverage.
+#[derive(Clone, Debug)]
+#[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)]
+pub struct HirInfo {
+    pub function_source_hash: u64,
+    pub fn_sig_span: Span,
+    pub body_span: Span,
+}

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,112 @@
+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 {
+    // 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() {
+        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<'_>) {
+    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 {
+    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::DefId;
 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
@@ -39,28 +36,15 @@ impl<'tcx> MirPass<'tcx> for InstrumentCoverage {
     fn run_pass(&self, tcx: TyCtxt<'tcx>, mir_body: &mut mir::Body<'tcx>) {
         let mir_source = mir_body.source;
 
-        // If the InstrumentCoverage pass is called on promoted MIRs, skip them.
-        // See: https://github.com/rust-lang/rust/pull/73011#discussion_r438317601
-        if mir_source.promoted.is_some() {
-            trace!(
-                "InstrumentCoverage skipped for {:?} (already promoted for Miri evaluation)",
-                mir_source.def_id()
-            );
-            return;
-        }
+        // This pass runs after MIR promotion, but before promoted MIR starts to
+        // be transformed, so it should never see promoted MIR.
+        assert!(mir_source.promoted.is_none());
 
-        let is_fn_like =
-            tcx.hir().get_by_def_id(mir_source.def_id().expect_local()).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 {
-            trace!("InstrumentCoverage skipped for {:?} (not an fn-like)", mir_source.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;
         }
 
@@ -72,14 +56,9 @@ impl<'tcx> MirPass<'tcx> for InstrumentCoverage {
             _ => {}
         }
 
-        let codegen_fn_attrs = tcx.codegen_fn_attrs(mir_source.def_id());
-        if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_COVERAGE) {
-            return;
-        }
-
-        trace!("InstrumentCoverage starting for {:?}", mir_source.def_id());
+        trace!("InstrumentCoverage starting for {def_id:?}");
         Instrumentor::new(tcx, mir_body).inject_counters();
-        trace!("InstrumentCoverage done for {:?}", mir_source.def_id());
+        trace!("InstrumentCoverage done for {def_id:?}");
     }
 }
 
@@ -97,29 +76,17 @@ struct Instrumentor<'a, 'tcx> {
 impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
     fn new(tcx: TyCtxt<'tcx>, mir_body: &'a mut mir::Body<'tcx>) -> Self {
         let source_map = tcx.sess.source_map();
-        let def_id = mir_body.source.def_id();
-        let (some_fn_sig, hir_body) = fn_sig_and_body(tcx, def_id);
 
-        let body_span = get_body_span(tcx, hir_body, mir_body);
+        let def_id = mir_body.source.def_id().expect_local();
+        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());
-        let fn_sig_span = match some_fn_sig.filter(|fn_sig| {
-            fn_sig.span.eq_ctxt(body_span)
-                && Lrc::ptr_eq(&source_file, &source_map.lookup_source_file(fn_sig.span.lo()))
-        }) {
-            Some(fn_sig) => fn_sig.span.with_hi(body_span.lo()),
-            None => body_span.shrink_to_lo(),
-        };
-
-        debug!(
-            "instrumenting {}: {:?}, fn sig span: {:?}, body span: {:?}",
-            if tcx.is_closure(def_id) { "closure" } else { "function" },
-            def_id,
-            fn_sig_span,
-            body_span
-        );
 
-        let function_source_hash = hash_mir_source(tcx, hir_body);
+        debug!(?fn_sig_span, ?body_span, "instrumenting {def_id:?}",);
+
         let basic_coverage_blocks = CoverageGraph::from_mir(mir_body);
         let coverage_counters = CoverageCounters::new(&basic_coverage_blocks);
 
@@ -324,55 +291,3 @@ fn make_code_region(
         end_col: end_col as u32,
     }
 }
-
-fn fn_sig_and_body(
-    tcx: TyCtxt<'_>,
-    def_id: DefId,
-) -> (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_if_local(def_id).expect("expected DefId is local");
-    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>,
-    mir_body: &mut mir::Body<'tcx>,
-) -> Span {
-    let mut body_span = hir_body.value.span;
-    let def_id = mir_body.source.def_id();
-
-    if tcx.is_closure(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()
-}