Auto merge of #153685 - Zalathar:for-each-query, r=<try>

bors · bors · commit 5b8248425951 · 2026-03-11T08:16:46.000Z
Introduce `for_each_query_vtable!` to move more code out of query macros
diff --git a/compiler/rustc_query_impl/src/execution.rs b/compiler/rustc_query_impl/src/execution.rs
@@ -16,8 +16,8 @@ use rustc_middle::ty::TyCtxt;
 use rustc_middle::verify_ich::incremental_verify_ich;
 use rustc_span::{DUMMY_SP, Span};
 
-use crate::collect_active_jobs_from_all_queries;
 use crate::dep_graph::{DepNode, DepNodeIndex};
+use crate::for_each_query_vtable;
 use crate::job::{QueryJobInfo, QueryJobMap, find_cycle_in_stack, report_cycle};
 use crate::plumbing::{current_query_job, next_job_id, start_query};
 
@@ -30,20 +30,49 @@ pub(crate) fn all_inactive<'tcx, K>(state: &QueryState<'tcx, K>) -> bool {
     state.active.lock_shards().all(|shard| shard.is_empty())
 }
 
+/// Returns a map of currently active query jobs, collected from all queries.
+///
+/// If `require_complete` is `true`, this function locks all shards of the
+/// query results to produce a complete map, which always returns `Ok`.
+/// Otherwise, it may return an incomplete map as an error if any shard
+/// lock cannot be acquired.
+///
+/// Prefer passing `false` to `require_complete` to avoid potential deadlocks,
+/// especially when called from within a deadlock handler, unless a
+/// complete map is needed and no deadlock is possible at this call site.
+pub fn collect_active_jobs_from_all_queries<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    require_complete: bool,
+) -> Result<QueryJobMap<'tcx>, QueryJobMap<'tcx>> {
+    let mut job_map_out = QueryJobMap::default();
+    let mut complete = true;
+
+    for_each_query_vtable!(
+        gather_active_jobs,
+        tcx,
+        ALL_QUERIES,
+        require_complete,
+        &mut job_map_out,
+        &mut complete,
+    );
+
+    if complete { Ok(job_map_out) } else { Err(job_map_out) }
+}
+
 /// Internal plumbing for collecting the set of active jobs for this query.
 ///
 /// Should only be called from `collect_active_jobs_from_all_queries`.
 ///
 /// (We arbitrarily use the word "gather" when collecting the jobs for
 /// each individual query, so that we have distinct function names to
 /// grep for.)
-pub(crate) fn gather_active_jobs<'tcx, C>(
-    query: &'tcx QueryVTable<'tcx, C>,
+fn gather_active_jobs<'tcx, C>(
     tcx: TyCtxt<'tcx>,
+    query: &'tcx QueryVTable<'tcx, C>,
     require_complete: bool,
     job_map_out: &mut QueryJobMap<'tcx>, // Out-param; job info is gathered into this map
-) -> Option<()>
-where
+    complete: &mut bool, // Out-param; set to false if `try_lock` fails for any shard
+) where
     C: QueryCache<Key: QueryKey + DynSend + DynSync>,
     QueryVTable<'tcx, C>: DynSync,
 {
@@ -76,7 +105,8 @@ where
                         "Failed to collect active jobs for query with name `{}`!",
                         query.name
                     );
-                    return None;
+                    *complete = false;
+                    return;
                 }
                 Some(shard) => gather_shard_jobs(&shard),
             }
@@ -89,8 +119,6 @@ where
         let frame = crate::plumbing::create_deferred_query_stack_frame(tcx, query, key);
         job_map_out.insert(job.id, QueryJobInfo { frame, job });
     }
-
-    Some(())
 }
 
 #[cold]
diff --git a/compiler/rustc_query_impl/src/job.rs b/compiler/rustc_query_impl/src/job.rs
@@ -14,7 +14,7 @@ use rustc_middle::ty::TyCtxt;
 use rustc_session::Session;
 use rustc_span::{DUMMY_SP, Span};
 
-use crate::collect_active_jobs_from_all_queries;
+use crate::execution::collect_active_jobs_from_all_queries;
 
 /// Map from query job IDs to job information collected by
 /// `collect_active_jobs_from_all_queries`.
diff --git a/compiler/rustc_query_impl/src/lib.rs b/compiler/rustc_query_impl/src/lib.rs
@@ -3,6 +3,7 @@
 // tidy-alphabetical-start
 #![allow(internal_features)]
 #![feature(core_intrinsics)]
+#![feature(macro_metavar_expr)]
 #![feature(min_specialization)]
 #![feature(rustc_attrs)]
 #![feature(try_blocks)]
@@ -11,16 +12,16 @@
 use rustc_data_structures::sync::AtomicU64;
 use rustc_middle::dep_graph;
 use rustc_middle::queries::{self, ExternProviders, Providers};
-use rustc_middle::query::on_disk_cache::{CacheEncoder, EncodedDepNodeIndex, OnDiskCache};
+use rustc_middle::query::on_disk_cache::OnDiskCache;
 use rustc_middle::query::plumbing::{QuerySystem, QueryVTable};
 use rustc_middle::query::{AsLocalQueryKey, QueryCache, QueryMode};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::Span;
 
 pub use crate::dep_kind_vtables::make_dep_kind_vtables;
+pub use crate::execution::collect_active_jobs_from_all_queries;
 use crate::from_cycle_error::FromCycleError;
 pub use crate::job::{QueryJobMap, break_query_cycles, print_query_stack};
-use crate::profiling_support::QueryKeyStringCache;
 
 #[macro_use]
 mod plumbing;
@@ -65,7 +66,8 @@ pub fn query_system<'tcx>(
 rustc_middle::rustc_with_all_queries! { define_queries! }
 
 pub fn provide(providers: &mut rustc_middle::util::Providers) {
-    providers.hooks.alloc_self_profile_query_strings = alloc_self_profile_query_strings;
-    providers.hooks.query_key_hash_verify_all = query_key_hash_verify_all;
-    providers.hooks.encode_all_query_results = encode_all_query_results;
+    providers.hooks.alloc_self_profile_query_strings =
+        profiling_support::alloc_self_profile_query_strings;
+    providers.hooks.query_key_hash_verify_all = plumbing::query_key_hash_verify_all;
+    providers.hooks.encode_all_query_results = plumbing::encode_all_query_results;
 }
diff --git a/compiler/rustc_query_impl/src/plumbing.rs b/compiler/rustc_query_impl/src/plumbing.rs
@@ -33,7 +33,7 @@ use rustc_span::def_id::LOCAL_CRATE;
 use crate::error::{QueryOverflow, QueryOverflowNote};
 use crate::execution::{all_inactive, force_query};
 use crate::job::find_dep_kind_root;
-use crate::{GetQueryVTable, collect_active_jobs_from_all_queries};
+use crate::{GetQueryVTable, collect_active_jobs_from_all_queries, for_each_query_vtable};
 
 fn depth_limit_error<'tcx>(tcx: TyCtxt<'tcx>, job: QueryJobId) {
     let job_map =
@@ -146,7 +146,21 @@ where
     QueryStackFrame::new(info, kind, def_id, def_id_for_ty_in_cycle)
 }
 
-pub(crate) fn encode_query_results<'a, 'tcx, C, V>(
+pub(crate) fn encode_all_query_results<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    encoder: &mut CacheEncoder<'_, 'tcx>,
+    query_result_index: &mut EncodedDepNodeIndex,
+) {
+    for_each_query_vtable!(
+        encode_query_results,
+        tcx,
+        CACHE_ON_DISK_QUERIES,
+        encoder,
+        query_result_index,
+    );
+}
+
+fn encode_query_results<'a, 'tcx, C, V>(
     tcx: TyCtxt<'tcx>,
     query: &'tcx QueryVTable<'tcx, C>,
     encoder: &mut CacheEncoder<'a, 'tcx>,
@@ -172,9 +186,17 @@ pub(crate) fn encode_query_results<'a, 'tcx, C, V>(
     });
 }
 
-pub(crate) fn query_key_hash_verify<'tcx, C: QueryCache>(
-    query: &'tcx QueryVTable<'tcx, C>,
+pub(crate) fn query_key_hash_verify_all<'tcx>(tcx: TyCtxt<'tcx>) {
+    if tcx.sess.opts.unstable_opts.incremental_verify_ich || cfg!(debug_assertions) {
+        tcx.sess.time("query_key_hash_verify_all", || {
+            for_each_query_vtable!(query_key_hash_verify, tcx, ALL_QUERIES)
+        });
+    }
+}
+
+fn query_key_hash_verify<'tcx, C: QueryCache>(
     tcx: TyCtxt<'tcx>,
+    query: &'tcx QueryVTable<'tcx, C>,
 ) {
     let _timer = tcx.prof.generic_activity_with_arg("query_key_hash_verify_for", query.name);
 
@@ -512,95 +534,55 @@ macro_rules! define_queries {
             }
         }
 
-        /// Returns a map of currently active query jobs, collected from all queries.
+        /// Given a function name, a `tcx`, and a filter condition
+        /// (e.g. `ALL_QUERIES` or `CACHE_ON_DISK_QUERIES`),
+        /// this macro calls that function with each query vtable that satisfies
+        /// the filter condition.
         ///
-        /// If `require_complete` is `true`, this function locks all shards of the
-        /// query results to produce a complete map, which always returns `Ok`.
-        /// Otherwise, it may return an incomplete map as an error if any shard
-        /// lock cannot be acquired.
+        /// The arguments passed to each function call are:
+        /// - `tcx`
+        /// - A query vtable that satisfies the filter condition
+        /// - All other arguments given after the filter condition
         ///
-        /// Prefer passing `false` to `require_complete` to avoid potential deadlocks,
-        /// especially when called from within a deadlock handler, unless a
-        /// complete map is needed and no deadlock is possible at this call site.
-        pub fn collect_active_jobs_from_all_queries<'tcx>(
-            tcx: TyCtxt<'tcx>,
-            require_complete: bool,
-        ) -> Result<QueryJobMap<'tcx>, QueryJobMap<'tcx>> {
-            let mut job_map_out = QueryJobMap::default();
-            let mut complete = true;
-
-            $(
-                let res = crate::execution::gather_active_jobs(
-                    &tcx.query_system.query_vtables.$name,
-                    tcx,
-                    require_complete,
-                    &mut job_map_out,
-                );
-                if res.is_none() {
-                    complete = false;
-                }
-            )*
-
-            if complete { Ok(job_map_out) } else { Err(job_map_out) }
-        }
-
-        /// All self-profiling events generated by the query engine use
-        /// virtual `StringId`s for their `event_id`. This method makes all
-        /// those virtual `StringId`s point to actual strings.
+        /// This needs to be a macro, because the vtables can have different
+        /// key/value/cache types for different queries.
         ///
-        /// 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<'_>) {
-            if !tcx.prof.enabled() {
-                return;
-            }
-
-            let _prof_timer = tcx.sess.prof.generic_activity("self_profile_alloc_query_strings");
-
-            let mut string_cache = QueryKeyStringCache::new();
-
-            $(
-                $crate::profiling_support::alloc_self_profile_query_strings_for_query_cache(
-                    tcx,
-                    stringify!($name),
-                    &tcx.query_system.query_vtables.$name.cache,
-                    &mut string_cache,
-                );
-            )*
-
-            tcx.sess.prof.store_query_cache_hits();
-        }
-
-        fn encode_all_query_results<'tcx>(
-            tcx: TyCtxt<'tcx>,
-            encoder: &mut CacheEncoder<'_, 'tcx>,
-            query_result_index: &mut EncodedDepNodeIndex,
-        ) {
-            $(
-                #[cfg($cache_on_disk)]
-                {
-                    $crate::plumbing::encode_query_results(
+        /// This macro's argument syntax is specifically intended to look like
+        /// plain Rust code, so that `for_each_query_vtable!(..)` calls will be
+        /// formatted by rustfmt.
+        ///
+        /// To avoid too much nested-macro complication, filter conditions are
+        /// implemented by hand as needed.
+        macro_rules! for_each_query_vtable {
+            // Call with all queries.
+            ($inner_fn:expr, $tcx:expr, ALL_QUERIES $$(,$args:expr)* $$(,)?) => {{
+                let tcx: rustc_middle::ty::TyCtxt<'_> = $tcx;
+                $(
+                    $inner_fn(
                         tcx,
                         &tcx.query_system.query_vtables.$name,
-                        encoder,
-                        query_result_index,
-                    )
-                }
-            )*
+                        $$($args),*
+                    );
+                )*
+            }};
+
+            // Only call with queries that can potentially cache to disk.
+            //
+            // This allows the use of trait bounds that only need to be satisfied
+            // by the subset of queries that actually cache to disk.
+            ($inner_fn:expr, $tcx:expr, CACHE_ON_DISK_QUERIES $$(,$args:expr)* $$(,)?) => {{
+                let tcx: rustc_middle::ty::TyCtxt<'_> = $tcx;
+                $(
+                    #[cfg($cache_on_disk)]
+                    $inner_fn(
+                        tcx,
+                        &tcx.query_system.query_vtables.$name,
+                        $$($args),*
+                    );
+                )*
+            }};
         }
 
-        pub fn query_key_hash_verify_all<'tcx>(tcx: TyCtxt<'tcx>) {
-            if tcx.sess.opts.unstable_opts.incremental_verify_ich || cfg!(debug_assertions) {
-                tcx.sess.time("query_key_hash_verify_all", || {
-                    $(
-                        $crate::plumbing::query_key_hash_verify(
-                            &tcx.query_system.query_vtables.$name,
-                            tcx
-                        );
-                    )*
-                })
-            }
-        }
+        pub(crate) use for_each_query_vtable;
     }
 }
diff --git a/compiler/rustc_query_impl/src/profiling_support.rs b/compiler/rustc_query_impl/src/profiling_support.rs
