Auto merge of #101307 - jyn514:simplify-storage, r=cjgillot

Simplify caching and storage for queries

I highly recommend reviewing commit-by-commit; each individual commit is quite small but it can be hard to see looking at the overall diff that the behavior is the same. Each commit depends on the previous.

r? `@cjgillot`

Author: bors

compiler/rustc_macros/src/query.rs

@@ -86,14 +86,11 @@ struct QueryModifiers {
     desc: (Option<Ident>, Punctuated<Expr, Token![,]>),
 
     /// Use this type for the in-memory cache.
-    storage: Option<Type>,
+    arena_cache: Option<Ident>,
 
     /// Cache the query to disk if the `Block` returns true.
     cache: Option<(Option<Pat>, Block)>,
 
-    /// Custom code to load the query from disk.
-    load_cached: Option<(Ident, Ident, Block)>,
-
     /// A cycle error for this query aborting the compilation with a fatal error.
     fatal_cycle: Option<Ident>,
 
@@ -120,8 +117,7 @@ struct QueryModifiers {
 }
 
 fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
-    let mut load_cached = None;
-    let mut storage = None;
+    let mut arena_cache = None;
     let mut cache = None;
     let mut desc = None;
     let mut fatal_cycle = None;
@@ -173,21 +169,8 @@ fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
             };
             let block = input.parse()?;
             try_insert!(cache = (args, block));
-        } else if modifier == "load_cached" {
-            // Parse a load_cached modifier like:
-            // `load_cached(tcx, id) { tcx.on_disk_cache.try_load_query_result(tcx, id) }`
-            let args;
-            parenthesized!(args in input);
-            let tcx = args.parse()?;
-            args.parse::<Token![,]>()?;
-            let id = args.parse()?;
-            let block = input.parse()?;
-            try_insert!(load_cached = (tcx, id, block));
-        } else if modifier == "storage" {
-            let args;
-            parenthesized!(args in input);
-            let ty = args.parse()?;
-            try_insert!(storage = ty);
+        } else if modifier == "arena_cache" {
+            try_insert!(arena_cache = modifier);
         } else if modifier == "fatal_cycle" {
             try_insert!(fatal_cycle = modifier);
         } else if modifier == "cycle_delay_bug" {
@@ -212,8 +195,7 @@ fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
         return Err(input.error("no description provided"));
     };
     Ok(QueryModifiers {
-        load_cached,
-        storage,
+        arena_cache,
         cache,
         desc,
         fatal_cycle,
@@ -262,20 +244,6 @@ fn add_query_description_impl(query: &Query, impls: &mut proc_macro2::TokenStrea
 
     // Find out if we should cache the query on disk
     let cache = if let Some((args, expr)) = modifiers.cache.as_ref() {
-        let try_load_from_disk = if let Some((tcx, id, block)) = modifiers.load_cached.as_ref() {
-            // Use custom code to load the query from disk
-            quote! {
-                const TRY_LOAD_FROM_DISK: Option<fn(QueryCtxt<'tcx>, SerializedDepNodeIndex) -> Option<Self::Value>>
-                    = Some(|#tcx, #id| { #block });
-            }
-        } else {
-            // Use the default code to load the query from disk
-            quote! {
-                const TRY_LOAD_FROM_DISK: Option<fn(QueryCtxt<'tcx>, SerializedDepNodeIndex) -> Option<Self::Value>>
-                    = Some(|tcx, id| tcx.on_disk_cache().as_ref()?.try_load_query_result(*tcx, id));
-            }
-        };
-
         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.
@@ -285,20 +253,13 @@ fn add_query_description_impl(query: &Query, impls: &mut proc_macro2::TokenStrea
             fn cache_on_disk(#tcx: TyCtxt<'tcx>, #key: &Self::Key) -> bool {
                 #expr
             }
-
-            #try_load_from_disk
         }
     } else {
-        if modifiers.load_cached.is_some() {
-            panic!("load_cached modifier on query `{}` without a cache modifier", name);
-        }
         quote! {
             #[inline]
             fn cache_on_disk(_: TyCtxt<'tcx>, _: &Self::Key) -> bool {
                 false
             }
-
-            const TRY_LOAD_FROM_DISK: Option<fn(QueryCtxt<'tcx>, SerializedDepNodeIndex) -> Option<Self::Value>> = None;
         }
     };
 
@@ -347,42 +308,28 @@ pub fn rustc_queries(input: TokenStream) -> TokenStream {
 
         let mut attributes = Vec::new();
 
-        // Pass on the fatal_cycle modifier
-        if let Some(fatal_cycle) = &modifiers.fatal_cycle {
-            attributes.push(quote! { (#fatal_cycle) });
-        };
-        // Pass on the storage modifier
-        if let Some(ref ty) = modifiers.storage {
-            let span = ty.span();
-            attributes.push(quote_spanned! {span=> (storage #ty) });
-        };
-        // Pass on the cycle_delay_bug modifier
-        if let Some(cycle_delay_bug) = &modifiers.cycle_delay_bug {
-            attributes.push(quote! { (#cycle_delay_bug) });
-        };
-        // Pass on the no_hash modifier
-        if let Some(no_hash) = &modifiers.no_hash {
-            attributes.push(quote! { (#no_hash) });
-        };
-        // Pass on the anon modifier
-        if let Some(anon) = &modifiers.anon {
-            attributes.push(quote! { (#anon) });
-        };
-        // Pass on the eval_always modifier
-        if let Some(eval_always) = &modifiers.eval_always {
-            attributes.push(quote! { (#eval_always) });
-        };
-        // Pass on the depth_limit modifier
-        if let Some(depth_limit) = &modifiers.depth_limit {
-            attributes.push(quote! { (#depth_limit) });
-        };
-        // Pass on the separate_provide_extern modifier
-        if let Some(separate_provide_extern) = &modifiers.separate_provide_extern {
-            attributes.push(quote! { (#separate_provide_extern) });
+        macro_rules! passthrough {
+            ( $( $modifier:ident ),+ $(,)? ) => {
+                $( if let Some($modifier) = &modifiers.$modifier {
+                    attributes.push(quote! { (#$modifier) });
+                }; )+
+            }
         }
-        // Pass on the remap_env_constness modifier
-        if let Some(remap_env_constness) = &modifiers.remap_env_constness {
-            attributes.push(quote! { (#remap_env_constness) });
+
+        passthrough!(
+            fatal_cycle,
+            arena_cache,
+            cycle_delay_bug,
+            no_hash,
+            anon,
+            eval_always,
+            depth_limit,
+            separate_provide_extern,
+            remap_env_constness,
+        );
+
+        if modifiers.cache.is_some() {
+            attributes.push(quote! { (cache) });
         }
 
         // This uses the span of the query definition for the commas,

compiler/rustc_middle/src/query/mod.rs

@@ -121,8 +121,8 @@ macro_rules! query_storage {
     ([][$K:ty, $V:ty]) => {
         <DefaultCacheSelector as CacheSelector<$K, $V>>::Cache
     };
-    ([(storage $ty:ty) $($rest:tt)*][$K:ty, $V:ty]) => {
-        <$ty as CacheSelector<$K, $V>>::Cache
+    ([(arena_cache) $($rest:tt)*][$K:ty, $V:ty]) => {
+        <ArenaCacheSelector<'tcx> as CacheSelector<$K, $V>>::Cache
     };
     ([$other:tt $($modifiers:tt)*][$($args:tt)*]) => {
         query_storage!([$($modifiers)*][$($args)*])

compiler/rustc_middle/src/ty/query.rs

@@ -17,7 +17,7 @@ extern crate rustc_middle;
 
 use rustc_data_structures::sync::AtomicU64;
 use rustc_middle::arena::Arena;
-use rustc_middle::dep_graph::{self, DepKindStruct, SerializedDepNodeIndex};
+use rustc_middle::dep_graph::{self, DepKindStruct};
 use rustc_middle::ty::query::{query_keys, query_storage, query_stored, query_values};
 use rustc_middle::ty::query::{ExternProviders, Providers, QueryEngine};
 use rustc_middle::ty::{self, TyCtxt};

compiler/rustc_query_impl/src/lib.rs

@@ -3,6 +3,7 @@
 //! manage the caches, and so forth.
 
 use crate::keys::Key;
+use crate::on_disk_cache::CacheDecoder;
 use crate::{on_disk_cache, Queries};
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::{AtomicU64, Lock};
@@ -19,6 +20,7 @@ use rustc_query_system::query::{
     QuerySideEffects, QueryStackFrame,
 };
 use rustc_query_system::Value;
+use rustc_serialize::Decodable;
 use std::any::Any;
 use std::num::NonZeroU64;
 use thin_vec::ThinVec;
@@ -253,6 +255,18 @@ macro_rules! get_provider {
     };
 }
 
+macro_rules! should_ever_cache_on_disk {
+    ([]) => {{
+        None
+    }};
+    ([(cache) $($rest:tt)*]) => {{
+        Some($crate::plumbing::try_load_from_disk::<Self::Value>)
+    }};
+    ([$other:tt $($modifiers:tt)*]) => {
+        should_ever_cache_on_disk!([$($modifiers)*])
+    };
+}
+
 pub(crate) fn create_query_frame<
     'tcx,
     K: Copy + Key + for<'a> HashStable<StableHashingContext<'a>>,
@@ -313,6 +327,16 @@ where
     }
 }
 
+pub(crate) fn try_load_from_disk<'tcx, V>(
+    tcx: QueryCtxt<'tcx>,
+    id: SerializedDepNodeIndex,
+) -> Option<V>
+where
+    V: for<'a> Decodable<CacheDecoder<'a, 'tcx>>,
+{
+    tcx.on_disk_cache().as_ref()?.try_load_query_result(*tcx, id)
+}
+
 fn force_from_dep_node<'tcx, Q>(tcx: TyCtxt<'tcx>, dep_node: DepNode) -> bool
 where
     Q: QueryDescription<QueryCtxt<'tcx>>,
@@ -418,8 +442,7 @@ macro_rules! define_queries {
                     hash_result: hash_result!([$($modifiers)*]),
                     handle_cycle_error: handle_cycle_error!([$($modifiers)*]),
                     compute,
-                    cache_on_disk,
-                    try_load_from_disk: Self::TRY_LOAD_FROM_DISK,
+                    try_load_from_disk: if cache_on_disk { should_ever_cache_on_disk!([$($modifiers)*]) } else { None },
                 }
             }
 

compiler/rustc_query_impl/src/plumbing.rs

@@ -25,11 +25,11 @@ pub struct QueryVTable<CTX: QueryContext, K, V> {
     pub dep_kind: CTX::DepKind,
     pub eval_always: bool,
     pub depth_limit: bool,
-    pub cache_on_disk: bool,
 
     pub compute: fn(CTX::DepContext, K) -> V,
     pub hash_result: Option<fn(&mut StableHashingContext<'_>, &V) -> Fingerprint>,
     pub handle_cycle_error: HandleCycleError,
+    // NOTE: this is also `None` if `cache_on_disk()` returns false, not just if it's unsupported by the query
     pub try_load_from_disk: Option<fn(CTX, SerializedDepNodeIndex) -> Option<V>>,
 }
 
@@ -44,18 +44,9 @@ impl<CTX: QueryContext, K, V> QueryVTable<CTX, K, V> {
     pub(crate) fn compute(&self, tcx: CTX::DepContext, key: K) -> V {
         (self.compute)(tcx, key)
     }
-
-    pub(crate) fn try_load_from_disk(&self, tcx: CTX, index: SerializedDepNodeIndex) -> Option<V> {
-        self.try_load_from_disk
-            .expect("QueryDescription::load_from_disk() called for an unsupported query.")(
-            tcx, index,
-        )
-    }
 }
 
 pub trait QueryDescription<CTX: QueryContext>: QueryConfig {
-    const TRY_LOAD_FROM_DISK: Option<fn(CTX, SerializedDepNodeIndex) -> Option<Self::Value>>;
-
     type Cache: QueryCache<Key = Self::Key, Stored = Self::Stored, Value = Self::Value>;
 
     fn describe(tcx: CTX, key: Self::Key) -> String;

compiler/rustc_query_system/src/query/config.rs

@@ -488,14 +488,14 @@ where
 
     // First we try to load the result from the on-disk cache.
     // Some things are never cached on disk.
-    if query.cache_on_disk {
+    if let Some(try_load_from_disk) = query.try_load_from_disk {
         let prof_timer = tcx.dep_context().profiler().incr_cache_loading();
 
         // The call to `with_query_deserialization` enforces that no new `DepNodes`
         // are created during deserialization. See the docs of that method for more
         // details.
-        let result = dep_graph
-            .with_query_deserialization(|| query.try_load_from_disk(tcx, prev_dep_node_index));
+        let result =
+            dep_graph.with_query_deserialization(|| try_load_from_disk(tcx, prev_dep_node_index));
 
         prof_timer.finish_with_query_invocation_id(dep_node_index.into());