Trace tcx.layout_of by overriding it in query providers

Author: Stypox

compiler/rustc_ty_utils/src/layout.rs

@@ -31,7 +31,7 @@ pub(crate) fn provide(providers: &mut Providers) {
 }
 
 #[instrument(skip(tcx, query), level = "debug")]
-fn layout_of<'tcx>(
+pub fn layout_of<'tcx>(
     tcx: TyCtxt<'tcx>,
     query: ty::PseudoCanonicalInput<'tcx, Ty<'tcx>>,
 ) -> Result<TyAndLayout<'tcx>, &'tcx LayoutError<'tcx>> {

compiler/rustc_ty_utils/src/lib.rs

@@ -35,6 +35,8 @@ pub mod sig_types;
 mod structural_match;
 mod ty;
 
+pub use layout::layout_of;
+
 rustc_fluent_macro::fluent_messages! { "../messages.ftl" }
 
 pub fn provide(providers: &mut Providers) {

src/tools/miri/src/bin/miri.rs

@@ -23,6 +23,7 @@ extern crate rustc_metadata;
 extern crate rustc_middle;
 extern crate rustc_session;
 extern crate rustc_span;
+extern crate rustc_ty_utils;
 
 use std::env::{self, VarError};
 use std::num::NonZero;
@@ -50,7 +51,8 @@ use rustc_middle::middle::exported_symbols::{
 };
 use rustc_middle::query::LocalCrate;
 use rustc_middle::traits::{ObligationCause, ObligationCauseCode};
-use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_middle::ty::layout::{LayoutError, TyAndLayout};
+use rustc_middle::ty::{self, PseudoCanonicalInput, Ty, TyCtxt};
 use rustc_middle::util::Providers;
 use rustc_session::config::{CrateType, ErrorOutputType, OptLevel};
 use rustc_session::search_paths::PathKind;
@@ -143,6 +145,17 @@ impl rustc_driver::Callbacks for MiriCompilerCalls {
                 Arc::make_mut(&mut crate_source).rlib = Some((PathBuf::new(), PathKind::All));
                 crate_source
             };
+
+            // HACK: in order to trace how much time is spent in `layout_of` queries, we override
+            // the query with a wrapper around the original `layout_of`. The `assert` is to ensure
+            // we are really wrapping `rustc_ty_utils::layout_of`. Note that this required making
+            // `rustc_ty_utils::layout_of` public.
+            assert!(std::ptr::fn_addr_eq(providers.layout_of, rustc_ty_utils::layout_of as
+                for<'tcx> fn(TyCtxt<'tcx>, PseudoCanonicalInput<'tcx, Ty<'tcx>>) -> Result<TyAndLayout<'tcx>, &'tcx LayoutError<'tcx>>));
+            providers.layout_of = |tcx, query| {
+                let _span = tracing::info_span!("tcx.layout_of", "query = {:?}", query.value.kind()).entered();
+                rustc_ty_utils::layout_of(tcx, query)
+            };
         });
     }