Cache DefId-keyed queries without hashing

Author: Mark-Simulacrum

compiler/rustc_middle/src/query/keys.rs

@@ -9,6 +9,7 @@ use crate::ty::{self, Ty, TyCtxt};
 use crate::ty::{GenericArg, GenericArgsRef};
 use rustc_hir::def_id::{CrateNum, DefId, LocalDefId, LocalModDefId, ModDefId, LOCAL_CRATE};
 use rustc_hir::hir_id::{HirId, OwnerId};
+use rustc_query_system::query::DefIdCacheSelector;
 use rustc_query_system::query::{DefaultCacheSelector, SingleCacheSelector, VecCacheSelector};
 use rustc_span::symbol::{Ident, Symbol};
 use rustc_span::{Span, DUMMY_SP};
@@ -152,7 +153,7 @@ impl Key for LocalDefId {
 }
 
 impl Key for DefId {
-    type CacheSelector = DefaultCacheSelector<Self>;
+    type CacheSelector = DefIdCacheSelector;
 
     fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
         tcx.def_span(*self)

compiler/rustc_query_system/src/query/caches.rs

@@ -4,6 +4,9 @@ use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sharded::{self, Sharded};
 use rustc_data_structures::sync::OnceLock;
 use rustc_index::{Idx, IndexVec};
+use rustc_span::def_id::CrateNum;
+use rustc_span::def_id::DefId;
+use rustc_span::def_id::DefIndex;
 use std::fmt::Debug;
 use std::hash::Hash;
 use std::marker::PhantomData;
@@ -148,6 +151,8 @@ where
 
     #[inline(always)]
     fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> {
+        // BUG: the "hash" passed here is a terrible hash and sharding is basically useless, I
+        // think?
         let lock = self.cache.lock_shard_by_hash(key.index() as u64);
         if let Some(Some(value)) = lock.get(*key) { Some(*value) } else { None }
     }
@@ -168,3 +173,56 @@ where
         }
     }
 }
+
+pub struct DefIdCacheSelector;
+
+impl<'tcx, V: 'tcx> CacheSelector<'tcx, V> for DefIdCacheSelector {
+    type Cache = DefIdCache<V>
+    where
+        V: Copy;
+}
+
+pub struct DefIdCache<V> {
+    cache: rustc_data_structures::sync::Lock<
+        IndexVec<CrateNum, IndexVec<DefIndex, Option<(V, DepNodeIndex)>>>,
+    >,
+}
+
+impl<V> Default for DefIdCache<V> {
+    fn default() -> Self {
+        DefIdCache { cache: Default::default() }
+    }
+}
+
+impl<V> QueryCache for DefIdCache<V>
+where
+    V: Copy,
+{
+    type Key = DefId;
+    type Value = V;
+
+    #[inline(always)]
+    fn lookup(&self, key: &DefId) -> Option<(V, DepNodeIndex)> {
+        let cache = self.cache.lock();
+        let krate_cache = cache.get(key.krate)?;
+        krate_cache.get(key.index).and_then(|v| *v)
+    }
+
+    #[inline]
+    fn complete(&self, key: DefId, value: V, index: DepNodeIndex) {
+        let mut cache = self.cache.lock();
+        let krate_cache = cache.ensure_contains_elem(key.krate, Default::default);
+        *krate_cache.ensure_contains_elem(key.index, Default::default) = Some((value, index));
+    }
+
+    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
+        let guard = self.cache.lock();
+        for (krate, v) in guard.iter_enumerated() {
+            for (idx, v) in v.iter_enumerated() {
+                if let Some(v) = v {
+                    f(&DefId { krate, index: idx }, &v.0, v.1);
+                }
+            }
+        }
+    }
+}

compiler/rustc_query_system/src/query/mod.rs

@@ -10,7 +10,8 @@ pub use self::job::{
 
 mod caches;
 pub use self::caches::{
-    CacheSelector, DefaultCacheSelector, QueryCache, SingleCacheSelector, VecCacheSelector,
+    CacheSelector, DefIdCacheSelector, DefaultCacheSelector, QueryCache, SingleCacheSelector,
+    VecCacheSelector,
 };
 
 mod config;