Auto merge of #119977 - Mark-Simulacrum:defid-cache, r=<try>

Cache DefId-keyed queries without hashing

Not yet ready for review:

* My guess is that this will be a significant memory footprint hit for sparser queries and require some more logic.
* Likely merits some further consideration for parallel rustc, though as noted in a separate comment the existing IndexVec sharding looks useless to me (likely always selecting the same shard today in 99% of cases).

Perf notes:

* #119977 (comment) evaluated a `IndexVec<CrateNum, IndexVec<DefIndex, Option<(V, DepNodeIndex)>>` scheme. This showed poor performance on incremental scenarios as the `iter()` callbacks are slower when walking the sparse vecs. In `full` scenarios this was a win for many primary benchmarks (~1-6% instructions, ~1-10% cycles), but did show significant memory overhead (+50% on many benchmarks). Next attempt will (a) skip hashing for local storage (expected to be denser) and retains the hashing for foreign storage (expected to be sparse) and (b) keep a present Vec to speed up `iter()` callbacks.

cc #45275

r? `@ghost`

Author: bors

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

@@ -3,7 +3,10 @@ use crate::dep_graph::DepNodeIndex;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sharded::{self, Sharded};
 use rustc_data_structures::sync::OnceLock;
+use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_index::{Idx, IndexVec};
+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,70 @@ where
         }
     }
 }
+
+pub struct DefIdCacheSelector;
+
+impl<'tcx, V: 'tcx> CacheSelector<'tcx, V> for DefIdCacheSelector {
+    type Cache = DefIdCache<V>
+    where
+        V: Copy;
+}
+
+pub struct DefIdCache<V> {
+    local: rustc_data_structures::sync::Lock<(
+        IndexVec<DefIndex, Option<(V, DepNodeIndex)>>,
+        Vec<DefIndex>,
+    )>,
+    foreign: DefaultCache<DefId, V>,
+}
+
+impl<V> Default for DefIdCache<V> {
+    fn default() -> Self {
+        DefIdCache { local: Default::default(), foreign: 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)> {
+        if key.krate == LOCAL_CRATE {
+            let cache = self.local.lock();
+            cache.0.get(key.index).and_then(|v| *v)
+        } else {
+            self.foreign.lookup(key)
+        }
+    }
+
+    #[inline]
+    fn complete(&self, key: DefId, value: V, index: DepNodeIndex) {
+        if key.krate == LOCAL_CRATE {
+            let mut cache = self.local.lock();
+            let (cache, present) = &mut *cache;
+            let slot = cache.ensure_contains_elem(key.index, Default::default);
+            if slot.is_none() {
+                // FIXME: Only store the present set when running in incremental mode. `iter` is not
+                // used outside of saving caches to disk and self-profile.
+                present.push(key.index);
+            }
+            *slot = Some((value, index));
+        } else {
+            self.foreign.complete(key, value, index)
+        }
+    }
+
+    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
+        let guard = self.local.lock();
+        for &idx in guard.1.iter() {
+            // `present` only contains slots we filled via `complete`.
+            let value = guard.0[idx].unwrap();
+            f(&DefId { krate: LOCAL_CRATE, index: idx }, &value.0, value.1);
+        }
+        self.foreign.iter(f);
+    }
+}

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;