Auto merge of #101785 - jyn514:query-struct-fn-ptrs, r=cjgillot

Use function pointers instead of macro-unrolled loops in rustc_query_impl

By making these standalone functions, we
a) allow making them extensible in the future with a new `QueryStruct`
b) greatly decrease the amount of code in each individual function, avoiding exponential blowup in llvm

Helps with https://github.com/rust-lang/rust/issues/96524. Based on https://github.com/rust-lang/rust/pull/101173; only the last commit is relevant.

r? `@cjgillot`

compiler/rustc_query_impl/src/lib.rs

@@ -1,6 +1,8 @@
 //! Support for serializing the dep-graph and reloading it.
 
 #![doc(html_root_url = "https://doc.rust-lang.org/nightly/nightly-rustc/")]
+// this shouldn't be necessary, but the check for `&mut _` is too naive and denies returning a function pointer that takes a mut ref
+#![feature(const_mut_refs)]
 #![feature(min_specialization)]
 #![feature(never_type)]
 #![feature(once_cell)]

compiler/rustc_query_impl/src/plumbing.rs

@@ -3,7 +3,8 @@
 //! manage the caches, and so forth.
 
 use crate::keys::Key;
-use crate::on_disk_cache::CacheDecoder;
+use crate::on_disk_cache::{CacheDecoder, CacheEncoder, EncodedDepNodeIndex};
+use crate::profiling_support::QueryKeyStringCache;
 use crate::{on_disk_cache, Queries};
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::{AtomicU64, Lock};
@@ -173,34 +174,14 @@ impl<'tcx> QueryCtxt<'tcx> {
 
     pub(super) fn encode_query_results(
         self,
-        encoder: &mut on_disk_cache::CacheEncoder<'_, 'tcx>,
+        encoder: &mut CacheEncoder<'_, 'tcx>,
-        query_result_index: &mut on_disk_cache::EncodedDepNodeIndex,
+        query_result_index: &mut EncodedDepNodeIndex,
     ) {
-        macro_rules! expand_if_cached {
+        for query in &self.queries.query_structs {
-            ([] $encode:expr) => {};
+            if let Some(encode) = query.encode_query_results {
-            ([(cache) $($rest:tt)*] $encode:expr) => {
+                encode(self, encoder, query_result_index);
-                $encode
-            };
-            ([$other:tt $($modifiers:tt)*] $encode:expr) => {
-                expand_if_cached!([$($modifiers)*] $encode)
-            };
-        }
-
-        macro_rules! encode_queries {
-            (
-            $($(#[$attr:meta])*
-                [$($modifiers:tt)*] fn $query:ident($($K:tt)*) -> $V:ty,)*) => {
-                $(
-                    expand_if_cached!([$($modifiers)*] on_disk_cache::encode_query_results::<_, super::queries::$query<'_>>(
-                        self,
-                        encoder,
-                        query_result_index
-                    ));
-                )*
             }
         }
-
-        rustc_query_append!(encode_queries!);
     }
 
     pub fn try_print_query_stack(
@@ -213,6 +194,14 @@ impl<'tcx> QueryCtxt<'tcx> {
     }
 }
 
+#[derive(Clone, Copy)]
+pub(crate) struct QueryStruct<'tcx> {
+    pub try_collect_active_jobs: fn(QueryCtxt<'tcx>, &mut QueryMap) -> Option<()>,
+    pub alloc_self_profile_query_strings: fn(TyCtxt<'tcx>, &mut QueryKeyStringCache),
+    pub encode_query_results:
+        Option<fn(QueryCtxt<'tcx>, &mut CacheEncoder<'_, 'tcx>, &mut EncodedDepNodeIndex)>,
+}
+
 macro_rules! handle_cycle_error {
     ([]) => {{
         rustc_query_system::HandleCycleError::Error
@@ -438,6 +427,18 @@ where
     }
 }
 
+macro_rules! expand_if_cached {
+    ([], $tokens:expr) => {{
+        None
+    }};
+    ([(cache) $($rest:tt)*], $tokens:expr) => {{
+        Some($tokens)
+    }};
+    ([$other:tt $($modifiers:tt)*], $tokens:expr) => {
+        expand_if_cached!([$($modifiers)*], $tokens)
+    };
+}
+
 // NOTE: `$V` isn't used here, but we still need to match on it so it can be passed to other macros
 // invoked by `rustc_query_append`.
 macro_rules! define_queries {
@@ -571,6 +572,59 @@ macro_rules! define_queries {
             })*
         }
 
+        mod query_structs {
+            use rustc_middle::ty::TyCtxt;
+            use $crate::plumbing::{QueryStruct, QueryCtxt};
+            use $crate::profiling_support::QueryKeyStringCache;
+            use rustc_query_system::query::{QueryDescription, QueryMap};
+
+            pub(super) const fn dummy_query_struct<'tcx>() -> QueryStruct<'tcx> {
+                fn noop_try_collect_active_jobs(_: QueryCtxt<'_>, _: &mut QueryMap) -> Option<()> {
+                    None
+                }
+                fn noop_alloc_self_profile_query_strings(_: TyCtxt<'_>, _: &mut QueryKeyStringCache) {}
+
+                QueryStruct {
+                    try_collect_active_jobs: noop_try_collect_active_jobs,
+                    alloc_self_profile_query_strings: noop_alloc_self_profile_query_strings,
+                    encode_query_results: None,
+                }
+            }
+
+            pub(super) use dummy_query_struct as Null;
+            pub(super) use dummy_query_struct as Red;
+            pub(super) use dummy_query_struct as TraitSelect;
+            pub(super) use dummy_query_struct as CompileCodegenUnit;
+            pub(super) use dummy_query_struct as CompileMonoItem;
+
+            $(
+            pub(super) const fn $name<'tcx>() -> QueryStruct<'tcx> { QueryStruct {
+                try_collect_active_jobs: |tcx, qmap| {
+                    let make_query = |tcx, key| {
+                        let kind = rustc_middle::dep_graph::DepKind::$name;
+                        let name = stringify!($name);
+                        $crate::plumbing::create_query_frame(tcx, super::queries::$name::describe, key, kind, name)
+                    };
+                    tcx.queries.$name.try_collect_active_jobs(
+                        tcx,
+                        make_query,
+                        qmap,
+                    )
+                },
+                alloc_self_profile_query_strings: |tcx, string_cache| {
+                    $crate::profiling_support::alloc_self_profile_query_strings_for_query_cache(
+                        tcx,
+                        stringify!($name),
+                        &tcx.query_caches.$name,
+                        string_cache,
+                    )
+                },
+                encode_query_results: expand_if_cached!([$($modifiers)*], |tcx, encoder, query_result_index|
+                    $crate::on_disk_cache::encode_query_results::<_, super::queries::$name<'_>>(tcx, encoder, query_result_index)
+                ),
+            }})*
+        }
+
         pub fn query_callbacks<'tcx>(arena: &'tcx Arena<'tcx>) -> &'tcx [DepKindStruct<'tcx>] {
             arena.alloc_from_iter(make_dep_kind_array!(query_callbacks))
         }
@@ -585,9 +639,11 @@ impl<'tcx> Queries<'tcx> {
         extern_providers: ExternProviders,
         on_disk_cache: Option<OnDiskCache<'tcx>>,
     ) -> Self {
+        use crate::query_structs;
         Queries {
             local_providers: Box::new(local_providers),
             extern_providers: Box::new(extern_providers),
+            query_structs: make_dep_kind_array!(query_structs).to_vec(),
             on_disk_cache,
             jobs: AtomicU64::new(1),
             ..Queries::default()
@@ -602,6 +658,7 @@ macro_rules! define_queries_struct {
         pub struct Queries<'tcx> {
             local_providers: Box<Providers>,
             extern_providers: Box<ExternProviders>,
+            query_structs: Vec<$crate::plumbing::QueryStruct<'tcx>>,
 
             pub on_disk_cache: Option<OnDiskCache<'tcx>>,
 
@@ -618,18 +675,9 @@ macro_rules! define_queries_struct {
                 let tcx = QueryCtxt { tcx, queries: self };
                 let mut jobs = QueryMap::default();
 
-                $(
+                for query in &self.query_structs {
-                    let make_query = |tcx, key| {
+                    (query.try_collect_active_jobs)(tcx, &mut jobs);
-                        let kind = dep_graph::DepKind::$name;
+                }
-                        let name = stringify!($name);
-                        $crate::plumbing::create_query_frame(tcx, queries::$name::describe, key, kind, name)
-                    };
-                    self.$name.try_collect_active_jobs(
-                        tcx,
-                        make_query,
-                        &mut jobs,
-                    )?;
-                )*
 
                 Some(jobs)
             }

compiler/rustc_query_impl/src/profiling_support.rs

@@ -1,3 +1,4 @@
+use crate::QueryCtxt;
 use measureme::{StringComponent, StringId};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::profiling::SelfProfiler;
@@ -8,7 +9,7 @@ use rustc_query_system::query::QueryCache;
 use std::fmt::Debug;
 use std::io::Write;
 
-struct QueryKeyStringCache {
+pub(crate) struct QueryKeyStringCache {
     def_id_cache: FxHashMap<DefId, StringId>,
 }
 
@@ -226,7 +227,7 @@ where
 /// Allocate the self-profiling query strings for a single query cache. This
 /// method is called from `alloc_self_profile_query_strings` which knows all
 /// the queries via macro magic.
-fn alloc_self_profile_query_strings_for_query_cache<'tcx, C>(
+pub(crate) fn alloc_self_profile_query_strings_for_query_cache<'tcx, C>(
     tcx: TyCtxt<'tcx>,
     query_name: &'static str,
     query_cache: &C,
@@ -298,27 +299,15 @@ fn alloc_self_profile_query_strings_for_query_cache<'tcx, C>(
 /// If we are recording only summary data, the ids will point to
 /// just the query names. If we are recording query keys too, we
 /// allocate the corresponding strings here.
-pub fn alloc_self_profile_query_strings(tcx: TyCtxt<'_>) {
+pub fn alloc_self_profile_query_strings<'tcx>(tcx: TyCtxt<'tcx>) {
     if !tcx.prof.enabled() {
         return;
     }
 
     let mut string_cache = QueryKeyStringCache::new();
+    let queries = QueryCtxt::from_tcx(tcx);
 
-    macro_rules! alloc_once {
+    for query in &queries.queries.query_structs {
-        (
+        (query.alloc_self_profile_query_strings)(tcx, &mut string_cache);
-        $($(#[$attr:meta])*
-            [$($modifiers:tt)*] fn $name:ident($($K:tt)*) -> $V:ty,)*) => {
-            $(
-                alloc_self_profile_query_strings_for_query_cache(
-                    tcx,
-                    stringify!($name),
-                    &tcx.query_caches.$name,
-                    &mut string_cache,
-                );
-            )+
-        }
     }
-
-    rustc_query_append! { alloc_once! }
 }