moves coverage data computation from pass to query

Author: richkadel

src/librustc_codegen_llvm/intrinsic.rs

@@ -140,18 +140,15 @@ impl IntrinsicCallMethods<'tcx> for Builder<'a, 'll, 'tcx> {
                 self.call(llfn, &[], None)
             }
             "count_code_region" => {
-                let coverage_data = tcx
-                    .coverage_data(caller_instance.def_id())
-                    .as_ref()
-                    .expect("LLVM intrinsic count_code_region call has associated coverage_data");
+                let coverage_data = tcx.coverage_data(caller_instance.def_id());
                 let mangled_fn = tcx.symbol_name(caller_instance);
                 let (mangled_fn_name, _len_val) = self.const_str(mangled_fn.name);
                 let hash = self.const_u64(coverage_data.hash);
-                let index = args[0].immediate();
                 let num_counters = self.const_u32(coverage_data.num_counters);
+                let index = args[0].immediate();
                 debug!(
                     "count_code_region to LLVM intrinsic instrprof.increment(fn_name={}, hash={:?}, num_counters={:?}, index={:?})",
-                    mangled_fn.name, hash, index, num_counters
+                    mangled_fn.name, hash, num_counters, index
                 );
                 self.instrprof_increment(mangled_fn_name, hash, num_counters, index)
             }

src/librustc_middle/mir/mod.rs

@@ -88,19 +88,6 @@ impl MirPhase {
     }
 }
 
-/// Coverage data computed by the `InstrumentCoverage` MIR pass, when compiling with
-/// `-Zinstrument_coverage`.
-#[derive(Clone, RustcEncodable, RustcDecodable, Debug, HashStable, TypeFoldable)]
-pub struct CoverageData {
-    /// A hash value that can be used by the consumer of the coverage profile data to detect
-    /// changes to the instrumented source of the associated MIR body (typically, for an
-    /// individual function).
-    pub hash: u64,
-
-    /// The total number of coverage region counters added to this MIR Body.
-    pub num_counters: u32,
-}
-
 /// The lowered representation of a single function.
 #[derive(Clone, RustcEncodable, RustcDecodable, Debug, HashStable, TypeFoldable)]
 pub struct Body<'tcx> {
@@ -184,10 +171,6 @@ pub struct Body<'tcx> {
     /// FIXME(oli-obk): rewrite the promoted during promotion to eliminate the cell components.
     pub ignore_interior_mut_in_const_validation: bool,
 
-    /// If compiling with `-Zinstrument_coverage`, the `InstrumentCoverage` pass stores summary
-    /// information associated with the MIR, used in code generation of the coverage counters.
-    pub coverage_data: Option<CoverageData>,
-
     predecessor_cache: PredecessorCache,
 }
 
@@ -228,7 +211,6 @@ impl<'tcx> Body<'tcx> {
             required_consts: Vec::new(),
             ignore_interior_mut_in_const_validation: false,
             control_flow_destroyed,
-            coverage_data: None,
             predecessor_cache: PredecessorCache::new(),
         }
     }
@@ -256,7 +238,6 @@ impl<'tcx> Body<'tcx> {
             generator_kind: None,
             var_debug_info: Vec::new(),
             ignore_interior_mut_in_const_validation: false,
-            coverage_data: None,
             predecessor_cache: PredecessorCache::new(),
         }
     }
@@ -2938,3 +2919,18 @@ impl Location {
         }
     }
 }
+
+/// Coverage data associated with each function (MIR) instrumented with coverage counters, when
+/// compiled with `-Zinstrument_coverage`. The query `tcx.coverage_data(DefId)` computes these
+/// values on demand (during code generation). This query is only valid after executing the MIR pass
+/// `InstrumentCoverage`.
+#[derive(Clone, RustcEncodable, RustcDecodable, Debug, HashStable)]
+pub struct CoverageData {
+    /// A hash value that can be used by the consumer of the coverage profile data to detect
+    /// changes to the instrumented source of the associated MIR body (typically, for an
+    /// individual function).
+    pub hash: u64,
+
+    /// The total number of coverage region counters added to the MIR `Body`.
+    pub num_counters: u32,
+}

src/librustc_middle/query/mod.rs

@@ -214,7 +214,7 @@ rustc_queries! {
             cache_on_disk_if { key.is_local() }
         }
 
-        query coverage_data(key: DefId) -> Option<mir::CoverageData> {
+        query coverage_data(key: DefId) -> mir::CoverageData {
             desc { |tcx| "retrieving coverage data from MIR for `{}`", tcx.def_path_str(key) }
             storage(ArenaCacheSelector<'tcx>)
             cache_on_disk_if { key.is_local() }

src/librustc_mir/transform/instrument_coverage.rs

@@ -7,10 +7,12 @@ use rustc_middle::hir;
 use rustc_middle::ich::StableHashingContext;
 use rustc_middle::mir::interpret::Scalar;
 use rustc_middle::mir::{
-    self, BasicBlock, BasicBlockData, CoverageData, Operand, Place, SourceInfo, StatementKind,
-    Terminator, TerminatorKind, START_BLOCK,
+    self, traversal, BasicBlock, BasicBlockData, CoverageData, Operand, Place, SourceInfo,
+    StatementKind, Terminator, TerminatorKind, START_BLOCK,
 };
 use rustc_middle::ty;
+use rustc_middle::ty::query::Providers;
+use rustc_middle::ty::FnDef;
 use rustc_middle::ty::TyCtxt;
 use rustc_span::def_id::DefId;
 use rustc_span::Span;
@@ -19,6 +21,31 @@ use rustc_span::Span;
 /// the intrinsic llvm.instrprof.increment.
 pub struct InstrumentCoverage;
 
+/// The `query` provider for `CoverageData`, requested by `codegen_intrinsic_call()` when
+/// constructing the arguments for `llvm.instrprof.increment`.
+pub(crate) fn provide(providers: &mut Providers<'_>) {
+    providers.coverage_data = |tcx, def_id| {
+        let body = tcx.optimized_mir(def_id);
+        let count_code_region_fn =
+            tcx.require_lang_item(lang_items::CountCodeRegionFnLangItem, None);
+        let mut num_counters: u32 = 0;
+        for (_, data) in traversal::preorder(body) {
+            if let Some(terminator) = &data.terminator {
+                if let TerminatorKind::Call { func: Operand::Constant(func), .. } = &terminator.kind
+                {
+                    if let FnDef(called_fn_def_id, _) = func.literal.ty.kind {
+                        if called_fn_def_id == count_code_region_fn {
+                            num_counters += 1;
+                        }
+                    }
+                }
+            }
+        }
+        let hash = if num_counters > 0 { hash_mir_source(tcx, def_id) } else { 0 };
+        CoverageData { num_counters, hash }
+    };
+}
+
 struct Instrumentor<'tcx> {
     tcx: TyCtxt<'tcx>,
     num_counters: u32,
@@ -30,20 +57,12 @@ impl<'tcx> MirPass<'tcx> for InstrumentCoverage {
             // If the InstrumentCoverage pass is called on promoted MIRs, skip them.
             // See: https://github.com/rust-lang/rust/pull/73011#discussion_r438317601
             if src.promoted.is_none() {
-                assert!(mir_body.coverage_data.is_none());
-
-                let hash = hash_mir_source(tcx, &src);
-
                 debug!(
-                    "instrumenting {:?}, hash: {}, span: {}",
+                    "instrumenting {:?}, span: {}",
                     src.def_id(),
-                    hash,
                     tcx.sess.source_map().span_to_string(mir_body.span)
                 );
-
-                let num_counters = Instrumentor::new(tcx).inject_counters(mir_body);
-
-                mir_body.coverage_data = Some(CoverageData { hash, num_counters });
+                Instrumentor::new(tcx).inject_counters(mir_body);
             }
         }
     }
@@ -60,15 +79,13 @@ impl<'tcx> Instrumentor<'tcx> {
         next
     }
 
-    fn inject_counters(&mut self, mir_body: &mut mir::Body<'tcx>) -> u32 {
+    fn inject_counters(&mut self, mir_body: &mut mir::Body<'tcx>) {
         // FIXME(richkadel): As a first step, counters are only injected at the top of each
         // function. The complete solution will inject counters at each conditional code branch.
         let top_of_function = START_BLOCK;
         let entire_function = mir_body.span;
 
         self.inject_counter(mir_body, top_of_function, entire_function);
-
-        self.num_counters
     }
 
     fn inject_counter(
@@ -138,14 +155,9 @@ fn placeholder_block(span: Span) -> BasicBlockData<'tcx> {
     }
 }
 
-fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, src: &MirSource<'tcx>) -> u64 {
-    let fn_body_id = match tcx.hir().get_if_local(src.def_id()) {
-        Some(node) => match hir::map::associated_body(node) {
-            Some(body_id) => body_id,
-            _ => bug!("instrumented MirSource does not include a function body: {:?}", node),
-        },
-        None => bug!("instrumented MirSource is not local: {:?}", src),
-    };
+fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> u64 {
+    let hir_node = tcx.hir().get_if_local(def_id).expect("DefId is local");
+    let fn_body_id = hir::map::associated_body(hir_node).expect("HIR node is a function with body");
     let hir_body = tcx.hir().body(fn_body_id);
     let mut hcx = tcx.create_no_span_stable_hashing_context();
     hash(&mut hcx, &hir_body.value).to_smaller_hash()

src/librustc_mir/transform/mod.rs

@@ -6,7 +6,7 @@ use rustc_hir::def_id::{CrateNum, DefId, LocalDefId, LOCAL_CRATE};
 use rustc_hir::intravisit::{self, NestedVisitorMap, Visitor};
 use rustc_index::vec::IndexVec;
 use rustc_middle::mir::visit::Visitor as _;
-use rustc_middle::mir::{traversal, Body, ConstQualifs, CoverageData, MirPhase, Promoted};
+use rustc_middle::mir::{traversal, Body, ConstQualifs, MirPhase, Promoted};
 use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::steal::Steal;
 use rustc_middle::ty::{InstanceDef, TyCtxt, TypeFoldable};
@@ -53,10 +53,10 @@ pub(crate) fn provide(providers: &mut Providers<'_>) {
         mir_drops_elaborated_and_const_checked,
         optimized_mir,
         is_mir_available,
-        coverage_data,
         promoted_mir,
         ..*providers
     };
+    instrument_coverage::provide(providers);
 }
 
 fn is_mir_available(tcx: TyCtxt<'_>, def_id: DefId) -> bool {
@@ -423,11 +423,6 @@ fn run_optimization_passes<'tcx>(
     );
 }
 
-fn coverage_data(tcx: TyCtxt<'_>, def_id: DefId) -> Option<CoverageData> {
-    let body = tcx.optimized_mir(def_id);
-    body.coverage_data.clone()
-}
-
 fn optimized_mir(tcx: TyCtxt<'_>, def_id: DefId) -> Body<'_> {
     if tcx.is_constructor(def_id) {
         // There's no reason to run all of the MIR passes on constructors when