Register queries with self profiler in rustc_interface

4 files changed, 13 insertions, 2 deletions

diff --git a/src/librustc/session/mod.rs b/src/librustc/session/mod.rs
index bba377392b4..cedac74d774 100644
--- a/src/librustc/session/mod.rs
+++ b/src/librustc/session/mod.rs
@@ -1099,7 +1099,6 @@ fn build_session_(
             );
             match profiler {
                 Ok(profiler) => {
-                    crate::ty::query::QueryName::register_with_profiler(&profiler);
                     Some(Arc::new(profiler))
                 },
                 Err(e) => {
diff --git a/src/librustc/ty/query/plumbing.rs b/src/librustc/ty/query/plumbing.rs
index 0f77a637b4d..1c15b7d5f35 100644
--- a/src/librustc/ty/query/plumbing.rs
+++ b/src/librustc/ty/query/plumbing.rs
@@ -827,7 +827,9 @@ macro_rules! define_queries_inner {
         }
 
         impl QueryName {
-            pub fn register_with_profiler(profiler: &rustc_data_structures::profiling::SelfProfiler) {
+            pub fn register_with_profiler(
+                profiler: &rustc_data_structures::profiling::SelfProfiler,
+            ) {
                 $(profiler.register_query_name(QueryName::$name);)*
             }
 
diff --git a/src/librustc_data_structures/profiling.rs b/src/librustc_data_structures/profiling.rs
index b89170cccdb..86f59bfabe6 100644
--- a/src/librustc_data_structures/profiling.rs
+++ b/src/librustc_data_structures/profiling.rs
@@ -205,6 +205,12 @@ impl SelfProfilerRef {
             TimingGuard::none()
         }));
     }
+
+    pub fn register_queries(&self, f: impl FnOnce(&SelfProfiler)) {
+        if let Some(profiler) = &self.profiler {
+            f(&profiler)
+        }
+    }
 }
 
 pub struct SelfProfiler {
diff --git a/src/librustc_interface/util.rs b/src/librustc_interface/util.rs
index 115345955ae..6b9d708cefa 100644
--- a/src/librustc_interface/util.rs
+++ b/src/librustc_interface/util.rs
@@ -108,6 +108,10 @@ pub fn create_session(
         process_configure_mod,
     );
 
+    sess.prof.register_queries(|profiler| {
+        rustc::ty::query::QueryName::register_with_profiler(&profiler);
+    });
+
     let codegen_backend = get_codegen_backend(&sess);
 
     let mut cfg = config::build_configuration(&sess, config::to_crate_config(cfg));