Make will_cache_on_disk_for_key_fn optional in query vtables

Author: Zalathar

compiler/rustc_macros/src/query.rs

@@ -280,31 +280,21 @@ fn add_query_desc_cached_impl(
         let crate::query::Providers { #name: _, .. };
     };
 
-    // Find out if we should cache the query on disk
-    let cache = if let Some((args, expr)) = modifiers.cache.as_ref() {
+    // Generate a function to check whether we should cache the query to disk, for some key.
+    if let Some((args, expr)) = modifiers.cache.as_ref() {
         let tcx = args.as_ref().map(|t| quote! { #t }).unwrap_or_else(|| quote! { _ });
         // expr is a `Block`, meaning that `{ #expr }` gets expanded
         // to `{ { stmts... } }`, which triggers the `unused_braces` lint.
         // we're taking `key` by reference, but some rustc types usually prefer being passed by value
-        quote! {
+        cached.extend(quote! {
             #[allow(unused_variables, unused_braces, rustc::pass_by_value)]
             #[inline]
             pub fn #name<'tcx>(#tcx: TyCtxt<'tcx>, #key: &crate::query::queries::#name::Key<'tcx>) -> bool {
                 #ra_hint
                 #expr
             }
-        }
-    } else {
-        quote! {
-            // we're taking `key` by reference, but some rustc types usually prefer being passed by value
-            #[allow(rustc::pass_by_value)]
-            #[inline]
-            pub fn #name<'tcx>(_: TyCtxt<'tcx>, _: &crate::query::queries::#name::Key<'tcx>) -> bool {
-                #ra_hint
-                false
-            }
-        }
-    };
+        });
+    }
 
     let (tcx, desc) = &modifiers.desc;
     let tcx = tcx.as_ref().map_or_else(|| quote! { _ }, |t| quote! { #t });
@@ -322,10 +312,6 @@ fn add_query_desc_cached_impl(
     descs.extend(quote! {
         #desc
     });
-
-    cached.extend(quote! {
-        #cache
-    });
 }
 
 pub(super) fn rustc_queries(input: TokenStream) -> TokenStream {

compiler/rustc_middle/src/query/plumbing.rs

@@ -18,6 +18,8 @@ use crate::query::{
 };
 use crate::ty::TyCtxt;
 
+pub type WillCacheOnDiskForKeyFn<'tcx, Key> = fn(tcx: TyCtxt<'tcx>, key: &Key) -> bool;
+
 pub type TryLoadFromDiskFn<'tcx, Key, Value> = fn(
     tcx: TyCtxt<'tcx>,
     key: &Key,
@@ -41,7 +43,7 @@ pub struct QueryVTable<'tcx, C: QueryCache> {
     pub query_state: usize,
     // Offset of this query's cache field in the QueryCaches struct
     pub query_cache: usize,
-    pub cache_on_disk: fn(tcx: TyCtxt<'tcx>, key: &C::Key) -> bool,
+    pub will_cache_on_disk_for_key_fn: Option<WillCacheOnDiskForKeyFn<'tcx, C::Key>>,
     pub execute_query: fn(tcx: TyCtxt<'tcx>, k: C::Key) -> C::Value,
     pub compute: fn(tcx: TyCtxt<'tcx>, key: C::Key) -> C::Value,
     pub try_load_from_disk_fn: Option<TryLoadFromDiskFn<'tcx, C::Key, C::Value>>,

compiler/rustc_query_impl/src/lib.rs

@@ -81,8 +81,8 @@ where
     }
 
     #[inline(always)]
-    fn cache_on_disk(self, tcx: TyCtxt<'tcx>, key: &Self::Key) -> bool {
-        (self.vtable.cache_on_disk)(tcx, key)
+    fn will_cache_on_disk_for_key(self, tcx: TyCtxt<'tcx>, key: &Self::Key) -> bool {
+        self.vtable.will_cache_on_disk_for_key_fn.map_or(false, |f| f(tcx, key))
     }
 
     #[inline(always)]

compiler/rustc_query_impl/src/plumbing.rs

@@ -396,7 +396,7 @@ pub(crate) fn encode_query_results<'a, 'tcx, Q>(
     assert!(query.query_state(qcx).all_inactive());
     let cache = query.query_cache(qcx);
     cache.iter(&mut |key, value, dep_node| {
-        if query.cache_on_disk(qcx.tcx, key) {
+        if query.will_cache_on_disk_for_key(qcx.tcx, key) {
             let dep_node = SerializedDepNodeIndex::new(dep_node.index());
 
             // Record position of the cache entry.
@@ -445,7 +445,7 @@ where
     let key = Q::Key::recover(tcx, &dep_node).unwrap_or_else(|| {
         panic!("Failed to recover key for {:?} with hash {}", dep_node, dep_node.hash)
     });
-    if query.cache_on_disk(tcx, &key) {
+    if query.will_cache_on_disk_for_key(tcx, &key) {
         let _ = query.execute_query(tcx, key);
     }
 }
@@ -648,7 +648,11 @@ macro_rules! define_queries {
                     cycle_error_handling: cycle_error_handling!([$($modifiers)*]),
                     query_state: std::mem::offset_of!(QueryStates<'tcx>, $name),
                     query_cache: std::mem::offset_of!(QueryCaches<'tcx>, $name),
-                    cache_on_disk: |tcx, key| ::rustc_middle::query::cached::$name(tcx, key),
+                    will_cache_on_disk_for_key_fn: should_ever_cache_on_disk!([$($modifiers)*] {
+                        Some(::rustc_middle::query::cached::$name)
+                    } {
+                        None
+                    }),
                     execute_query: |tcx, key| erase::erase_val(tcx.$name(key)),
                     compute: |tcx, key| {
                         #[cfg(debug_assertions)]

compiler/rustc_query_system/src/query/dispatcher.rs

@@ -46,7 +46,7 @@ pub trait QueryDispatcher<'tcx>: Copy {
     // Don't use this method to access query results, instead use the methods on TyCtxt
     fn query_cache<'a>(self, tcx: Self::Qcx) -> &'a Self::Cache;
 
-    fn cache_on_disk(self, tcx: DepContextOf<'tcx, Self>, key: &Self::Key) -> bool;
+    fn will_cache_on_disk_for_key(self, tcx: DepContextOf<'tcx, Self>, key: &Self::Key) -> bool;
 
     // Don't use this method to compute query results, instead use the methods on TyCtxt
     fn execute_query(self, tcx: DepContextOf<'tcx, Self>, k: Self::Key) -> Self::Value;

compiler/rustc_query_system/src/query/plumbing.rs

@@ -623,7 +623,7 @@ where
     // We always expect to find a cached result for things that
     // can be forced from `DepNode`.
     debug_assert!(
-        !query.cache_on_disk(*qcx.dep_context(), key)
+        !query.will_cache_on_disk_for_key(*qcx.dep_context(), key)
             || !qcx.dep_context().fingerprint_style(dep_node.kind).reconstructible(),
         "missing on-disk cache entry for {dep_node:?}"
     );