Make some load-from-disk function pointers optional in query vtables

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,18 @@ 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,
+    prev_index: SerializedDepNodeIndex,
+    index: DepNodeIndex,
+) -> Option<Value>;
+
+pub type IsLoadableFromDiskFn<'tcx, Key> =
+    fn(tcx: TyCtxt<'tcx>, key: &Key, index: SerializedDepNodeIndex) -> bool;
+
 /// Stores function pointers and other metadata for a particular query.
 ///
 /// Used indirectly by query plumbing in `rustc_query_system`, via a trait.
@@ -31,18 +43,11 @@ 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 can_load_from_disk: bool,
-    pub try_load_from_disk: fn(
-        tcx: TyCtxt<'tcx>,
-        key: &C::Key,
-        prev_index: SerializedDepNodeIndex,
-        index: DepNodeIndex,
-    ) -> Option<C::Value>,
-    pub loadable_from_disk:
-        fn(tcx: TyCtxt<'tcx>, key: &C::Key, index: SerializedDepNodeIndex) -> bool,
+    pub try_load_from_disk_fn: Option<TryLoadFromDiskFn<'tcx, C::Key, C::Value>>,
+    pub is_loadable_from_disk_fn: Option<IsLoadableFromDiskFn<'tcx, C::Key>>,
     pub hash_result: HashResult<C::Value>,
     pub value_from_cycle_error:
         fn(tcx: TyCtxt<'tcx>, cycle_error: &CycleError, guar: ErrorGuaranteed) -> C::Value,

compiler/rustc_query_impl/src/lib.rs

@@ -82,8 +82,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)]
@@ -135,21 +135,18 @@ where
         prev_index: SerializedDepNodeIndex,
         index: DepNodeIndex,
     ) -> Option<Self::Value> {
-        if self.vtable.can_load_from_disk {
-            (self.vtable.try_load_from_disk)(qcx.tcx, key, prev_index, index)
-        } else {
-            None
-        }
+        // `?` will return None immediately for queries that never cache to disk.
+        self.vtable.try_load_from_disk_fn?(qcx.tcx, key, prev_index, index)
     }
 
     #[inline]
-    fn loadable_from_disk(
+    fn is_loadable_from_disk(
         self,
         qcx: QueryCtxt<'tcx>,
         key: &Self::Key,
         index: SerializedDepNodeIndex,
     ) -> bool {
-        (self.vtable.loadable_from_disk)(qcx.tcx, key, index)
+        self.vtable.is_loadable_from_disk_fn.map_or(false, |f| f(qcx.tcx, key, index))
     }
 
     fn value_from_cycle_error(

compiler/rustc_query_impl/src/plumbing.rs

@@ -400,7 +400,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.
@@ -449,7 +449,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);
     }
 }
@@ -652,7 +652,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(tcx.$name(key)),
                     compute: |tcx, key| {
                         #[cfg(debug_assertions)]
@@ -670,37 +674,34 @@ macro_rules! define_queries {
                             )
                         )
                     },
-                    can_load_from_disk: should_ever_cache_on_disk!([$($modifiers)*] true false),
-                    try_load_from_disk: should_ever_cache_on_disk!([$($modifiers)*] {
-                        |tcx, key, prev_index, index| {
-                            if ::rustc_middle::query::cached::$name(tcx, key) {
-                                let value = $crate::plumbing::try_load_from_disk::<
-                                    queries::$name::ProvidedValue<'tcx>
-                                >(
-                                    tcx,
-                                    prev_index,
-                                    index,
-                                );
-                                value.map(|value| queries::$name::provided_to_erased(tcx, value))
-                            } else {
-                                None
+                    try_load_from_disk_fn: should_ever_cache_on_disk!([$($modifiers)*] {
+                        Some(|tcx, key, prev_index, index| {
+                            // Check the `cache_on_disk_if` condition for this key.
+                            if !::rustc_middle::query::cached::$name(tcx, key) {
+                                return None;
                             }
-                        }
+
+                            let value: queries::$name::ProvidedValue<'tcx> =
+                                $crate::plumbing::try_load_from_disk(tcx, prev_index, index)?;
+
+                            // Arena-alloc the value if appropriate, and erase it.
+                            Some(queries::$name::provided_to_erased(tcx, value))
+                        })
                     } {
-                        |_tcx, _key, _prev_index, _index| None
+                        None
+                    }),
+                    is_loadable_from_disk_fn: should_ever_cache_on_disk!([$($modifiers)*] {
+                        Some(|tcx, key, index| -> bool {
+                            ::rustc_middle::query::cached::$name(tcx, key) &&
+                                $crate::plumbing::loadable_from_disk(tcx, index)
+                        })
+                    } {
+                        None
                     }),
                     value_from_cycle_error: |tcx, cycle, guar| {
                         let result: queries::$name::Value<'tcx> = Value::from_cycle_error(tcx, cycle, guar);
                         erase(result)
                     },
-                    loadable_from_disk: |_tcx, _key, _index| {
-                        should_ever_cache_on_disk!([$($modifiers)*] {
-                            ::rustc_middle::query::cached::$name(_tcx, _key) &&
-                                $crate::plumbing::loadable_from_disk(_tcx, _index)
-                        } {
-                            false
-                        })
-                    },
                     hash_result: hash_result!([$($modifiers)*][queries::$name::Value<'tcx>]),
                     format_value: |value| format!("{:?}", restore::<queries::$name::Value<'tcx>>(*value)),
                 }

compiler/rustc_query_system/src/query/dispatcher.rs

@@ -42,7 +42,7 @@ pub trait QueryDispatcher<Qcx: QueryContext>: Copy {
     where
         Qcx: 'a;
 
-    fn cache_on_disk(self, tcx: Qcx::DepContext, key: &Self::Key) -> bool;
+    fn will_cache_on_disk_for_key(self, tcx: Qcx::DepContext, key: &Self::Key) -> bool;
 
     // Don't use this method to compute query results, instead use the methods on TyCtxt
     fn execute_query(self, tcx: Qcx::DepContext, k: Self::Key) -> Self::Value;
@@ -57,7 +57,7 @@ pub trait QueryDispatcher<Qcx: QueryContext>: Copy {
         index: DepNodeIndex,
     ) -> Option<Self::Value>;
 
-    fn loadable_from_disk(self, qcx: Qcx, key: &Self::Key, idx: SerializedDepNodeIndex) -> bool;
+    fn is_loadable_from_disk(self, qcx: Qcx, key: &Self::Key, idx: SerializedDepNodeIndex) -> bool;
 
     /// Synthesize an error value to let compilation continue after a cycle.
     fn value_from_cycle_error(

compiler/rustc_query_system/src/query/plumbing.rs

@@ -630,15 +630,15 @@ 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:?}"
     );
 
     // Sanity check for the logic in `ensure`: if the node is green and the result loadable,
     // we should actually be able to load it.
     debug_assert!(
-        !query.loadable_from_disk(qcx, key, prev_dep_node_index),
+        !query.is_loadable_from_disk(qcx, key, prev_dep_node_index),
         "missing on-disk cache entry for loadable {dep_node:?}"
     );
 
@@ -806,7 +806,7 @@ where
         return (false, None);
     }
 
-    let loadable = query.loadable_from_disk(qcx, key, serialized_dep_node_index);
+    let loadable = query.is_loadable_from_disk(qcx, key, serialized_dep_node_index);
     (!loadable, Some(dep_node))
 }