[WIP] Introduce SingletonCache for use with ZST query keys

compiler/rustc_middle/src/query/mod.rs

@@ -21,6 +21,7 @@ rustc_queries! {
     /// To avoid this fate, do not call `tcx.hir().krate()`; instead,
     /// prefer wrappers like `tcx.visit_all_items_in_krate()`.
     query hir_crate(key: ()) -> &'tcx Crate<'tcx> {
+        storage(SingletonCacheSelector)
         eval_always
         no_hash
         desc { "get the crate HIR" }
@@ -29,6 +30,7 @@ rustc_queries! {
     /// The indexed HIR. This can be conveniently accessed by `tcx.hir()`.
     /// Avoid calling this query directly.
     query index_hir(_: ()) -> &'tcx crate::hir::IndexedHir<'tcx> {
+        storage(SingletonCacheSelector)
         eval_always
         no_hash
         desc { "index HIR" }
@@ -115,6 +117,7 @@ rustc_queries! {
     }
 
     query analysis(key: ()) -> Result<(), ErrorReported> {
+        storage(SingletonCacheSelector)
         eval_always
         desc { "running analysis passes on this crate" }
     }
@@ -701,6 +704,7 @@ rustc_queries! {
         }
 
     query typeck_item_bodies(_: ()) -> () {
+        storage(SingletonCacheSelector)
         desc { "type-checking all item bodies" }
     }
 
@@ -769,6 +773,7 @@ rustc_queries! {
     /// Not meant to be used directly outside of coherence.
     query crate_inherent_impls_overlap_check(_: ())
         -> () {
+        storage(SingletonCacheSelector)
         eval_always
         desc { "check for overlap between inherent impls defined in this crate" }
     }
@@ -862,10 +867,12 @@ rustc_queries! {
 
     /// Performs part of the privacy check and computes "access levels".
     query privacy_access_levels(_: ()) -> &'tcx AccessLevels {
+        storage(SingletonCacheSelector)
         eval_always
         desc { "privacy access levels" }
     }
     query check_private_in_public(_: ()) -> () {
+        storage(SingletonCacheSelector)
         eval_always
         desc { "checking for private elements in public interfaces" }
     }
@@ -981,6 +988,7 @@ rustc_queries! {
     ///
     /// [`LOCAL_CRATE`]: rustc_hir::def_id::LOCAL_CRATE
     query all_local_trait_impls(_: ()) -> &'tcx BTreeMap<DefId, Vec<LocalDefId>> {
+        storage(SingletonCacheSelector)
         desc { "local trait impls" }
     }
 
@@ -1084,6 +1092,7 @@ rustc_queries! {
     }
 
     query dependency_formats(_: ()) -> Lrc<crate::middle::dependency_format::Dependencies> {
+        storage(SingletonCacheSelector)
         desc { "get the linkage format of all dependencies" }
     }
 
@@ -1229,12 +1238,15 @@ rustc_queries! {
     /// Identifies the entry-point (e.g., the `main` function) for a given
     /// crate, returning `None` if there is no entry point (such as for library crates).
     query entry_fn(_: ()) -> Option<(DefId, EntryFnType)> {
+        storage(SingletonCacheSelector)
         desc { "looking up the entry function of a crate" }
     }
     query plugin_registrar_fn(_: ()) -> Option<LocalDefId> {
+        storage(SingletonCacheSelector)
         desc { "looking up the plugin registrar for a crate" }
     }
     query proc_macro_decls_static(_: ()) -> Option<LocalDefId> {
+        storage(SingletonCacheSelector)
         desc { "looking up the derive registrar for a crate" }
     }
     query crate_disambiguator(_: CrateNum) -> CrateDisambiguator {
@@ -1411,6 +1423,7 @@ rustc_queries! {
         desc { "looking at the source for a crate" }
     }
     query postorder_cnums(_: ()) -> &'tcx [CrateNum] {
+        storage(SingletonCacheSelector)
         eval_always
         desc { "generating a postorder list of CrateNums" }
     }
@@ -1430,6 +1443,7 @@ rustc_queries! {
         desc { |tcx| "maybe_unused_trait_import for `{}`", tcx.def_path_str(def_id.to_def_id()) }
     }
     query maybe_unused_extern_crates(_: ()) -> &'tcx [(LocalDefId, Span)] {
+        storage(SingletonCacheSelector)
         eval_always
         desc { "looking up all possibly unused extern crates" }
     }
@@ -1445,6 +1459,7 @@ rustc_queries! {
         desc { "calculating the stability index for the local crate" }
     }
     query all_crate_nums(_: ()) -> &'tcx [CrateNum] {
+        storage(SingletonCacheSelector)
         eval_always
         desc { "fetching all foreign CrateNum instances" }
     }
@@ -1453,6 +1468,7 @@ rustc_queries! {
     /// (i.e., including those from subcrates). This is used only for
     /// error reporting.
     query all_traits(_: ()) -> &'tcx [DefId] {
+        storage(SingletonCacheSelector)
         desc { "fetching all foreign and local traits" }
     }
 
@@ -1467,6 +1483,7 @@ rustc_queries! {
     }
 
     query collect_and_partition_mono_items(_: ()) -> (&'tcx DefIdSet, &'tcx [CodegenUnit<'tcx>]) {
+        storage(SingletonCacheSelector)
         eval_always
         desc { "collect_and_partition_mono_items" }
     }
@@ -1476,8 +1493,9 @@ rustc_queries! {
 
     /// All items participating in code generation together with items inlined into them.
     query codegened_and_inlined_items(_: ()) -> &'tcx DefIdSet {
+        storage(SingletonCacheSelector)
         eval_always
-       desc { "codegened_and_inlined_items" }
+        desc { "codegened_and_inlined_items" }
     }
 
     query codegen_unit(_: Symbol) -> &'tcx CodegenUnit<'tcx> {
@@ -1491,10 +1509,12 @@ rustc_queries! {
         }
     }
     query backend_optimization_level(_: ()) -> OptLevel {
+        storage(SingletonCacheSelector)
         desc { "optimization level used by backend" }
     }
 
     query output_filenames(_: ()) -> Arc<OutputFilenames> {
+        storage(SingletonCacheSelector)
         eval_always
         desc { "output_filenames" }
     }
@@ -1671,6 +1691,7 @@ rustc_queries! {
     }
 
     query features_query(_: ()) -> &'tcx rustc_feature::Features {
+        storage(SingletonCacheSelector)
         eval_always
         desc { "looking up enabled feature gates" }
     }

compiler/rustc_middle/src/ty/query/mod.rs

@@ -116,7 +116,7 @@ macro_rules! query_storage {
     ([storage($ty:ty) $($rest:tt)*][$K:ty, $V:ty]) => {
         <$ty as CacheSelector<$K, $V>>::Cache
     };
-    ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)*][$($args:tt)*]) => {
+    ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)?][$($args:tt)*]) => {
         query_storage!([$($($modifiers)*)*][$($args)*])
     };
 }

compiler/rustc_query_impl/src/plumbing.rs

@@ -228,7 +228,7 @@ macro_rules! handle_cycle_error {
         $error.delay_as_bug();
         Value::from_cycle_error($tcx)
     }};
-    ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)*][$($args:tt)*]) => {
+    ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)?][$($args:tt)*]) => {
         handle_cycle_error!([$($($modifiers)*)*][$($args)*])
     };
 }
@@ -240,7 +240,7 @@ macro_rules! is_anon {
     ([anon $($rest:tt)*]) => {{
         true
     }};
-    ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)*]) => {
+    ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)?]) => {
         is_anon!([$($($modifiers)*)*])
     };
 }
@@ -252,7 +252,7 @@ macro_rules! is_eval_always {
     ([eval_always $($rest:tt)*]) => {{
         true
     }};
-    ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)*]) => {
+    ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)?]) => {
         is_eval_always!([$($($modifiers)*)*])
     };
 }
@@ -264,7 +264,7 @@ macro_rules! hash_result {
     ([no_hash $($rest:tt)*][$hcx:expr, $result:expr]) => {{
         None
     }};
-    ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)*][$($args:tt)*]) => {
+    ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)?][$($args:tt)*]) => {
         hash_result!([$($($modifiers)*)*][$($args)*])
     };
 }

compiler/rustc_query_system/src/lib.rs

@@ -1,4 +1,5 @@
 #![feature(bool_to_option)]
+#![feature(const_panic)]
 #![feature(core_intrinsics)]
 #![feature(hash_raw_entry)]
 #![feature(iter_zip)]

compiler/rustc_query_system/src/query/caches.rs

@@ -221,3 +221,88 @@ where
         }
     }
 }
+
+pub struct SingletonCacheSelector;
+
+impl<K: Eq + Hash, V: Clone> CacheSelector<K, V> for SingletonCacheSelector {
+    type Cache = SingletonCache<K, V>;
+}
+
+pub struct SingletonCache<K, V>(PhantomData<(K, V)>);
+
+impl<K, V> Default for SingletonCache<K, V> {
+    fn default() -> Self {
+        SingletonCache(PhantomData)
+    }
+}
+
+impl<K, V> SingletonCache<K, V> {
+    const KEY_IS_ZST: () = {
+        if std::mem::size_of::<K>() != 0 {
+            panic!("Key must be a ZST");
+        }
+    };
+}
+
+impl<K: Eq + Hash, V: Clone + Debug> QueryStorage for SingletonCache<K, V> {
+    type Value = V;
+    type Stored = V;
+
+    #[inline]
+    fn store_nocache(&self, value: Self::Value) -> Self::Stored {
+        // We have no dedicated storage
+        value
+    }
+}
+
+impl<K, V> QueryCache for SingletonCache<K, V>
+where
+    K: Eq + Hash + Copy + Clone + Debug,
+    V: Clone + Debug,
+{
+    type Key = K;
+    type Sharded = Option<(K, V, DepNodeIndex)>;
+
+    #[inline(always)]
+    fn lookup<'s, R, OnHit>(
+        &self,
+        state: &'s QueryCacheStore<Self>,
+        _key: &K,
+        on_hit: OnHit,
+    ) -> Result<R, QueryLookup>
+    where
+        OnHit: FnOnce(&V, DepNodeIndex) -> R,
+    {
+        let _ = Self::KEY_IS_ZST;
+
+        let lock = state.shards.get_shard_by_index(0).lock();
+        if let Some((_key, val, dep_node)) = &*lock {
+            Ok(on_hit(&val, *dep_node))
+        } else {
+            Err(QueryLookup { key_hash: 0, shard: 0 })
+        }
+    }
+
+    #[inline]
+    fn complete(
+        &self,
+        lock_sharded_storage: &mut Self::Sharded,
+        key: K,
+        value: V,
+        index: DepNodeIndex,
+    ) -> Self::Stored {
+        *lock_sharded_storage = Some((key, value.clone(), index));
+        value
+    }
+
+    fn iter(
+        &self,
+        shards: &Sharded<Self::Sharded>,
+        f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex),
+    ) {
+        let lock = shards.get_shard_by_index(0).lock();
+        if let Some((key, val, dep_node)) = &*lock {
+            f(&key, &val, *dep_node);
+        }
+    }
+}

compiler/rustc_query_system/src/query/mod.rs

@@ -9,6 +9,7 @@ pub use self::job::{print_query_stack, QueryInfo, QueryJob, QueryJobId, QueryJob
 mod caches;
 pub use self::caches::{
     ArenaCacheSelector, CacheSelector, DefaultCacheSelector, QueryCache, QueryStorage,
+    SingletonCacheSelector,
 };
 
 mod config;

compiler/rustc_query_system/src/query/plumbing.rs

@@ -31,7 +31,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 
 pub struct QueryCacheStore<C: QueryCache> {
     cache: C,
-    shards: Sharded<C::Sharded>,
+    pub(super) shards: Sharded<C::Sharded>,
     #[cfg(debug_assertions)]
     pub cache_hits: AtomicUsize,
 }
@@ -50,7 +50,7 @@ impl<C: QueryCache + Default> Default for QueryCacheStore<C> {
 /// Values used when checking a query cache which can be reused on a cache-miss to execute the query.
 pub struct QueryLookup {
     pub(super) key_hash: u64,
-    shard: usize,
+    pub(super) shard: usize,
 }
 
 // We compute the key's hash once and then use it for both the