Stabilize query cache size by fixing encoding order of some query results

Author: saethlin

Cargo.lock

@@ -4435,6 +4435,7 @@ dependencies = [
 name = "rustc_query_system"
 version = "0.0.0"
 dependencies = [
+ "indexmap",
  "parking_lot",
  "rustc-rayon-core",
  "rustc_abi",

compiler/rustc_middle/src/query/keys.rs

@@ -3,7 +3,7 @@
 use rustc_hir::def_id::{CrateNum, DefId, LOCAL_CRATE, LocalDefId, LocalModDefId, ModDefId};
 use rustc_hir::hir_id::{HirId, OwnerId};
 use rustc_query_system::dep_graph::DepNodeIndex;
-use rustc_query_system::query::{DefIdCache, DefaultCache, SingleCache, VecCache};
+use rustc_query_system::query::{DefIdCache, DefaultCache, IndexCache, SingleCache, VecCache};
 use rustc_span::{DUMMY_SP, Ident, Span, Symbol};
 
 use crate::infer::canonical::CanonicalQueryInput;
@@ -471,7 +471,7 @@ impl<'tcx> Key for ty::ParamEnv<'tcx> {
 }
 
 impl<'tcx, T: Key> Key for ty::PseudoCanonicalInput<'tcx, T> {
-    type Cache<V> = DefaultCache<Self, V>;
+    type Cache<V> = IndexCache<Self, V>;
 
     fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
         self.value.default_span(tcx)

compiler/rustc_query_system/Cargo.toml

@@ -5,6 +5,7 @@ edition = "2021"
 
 [dependencies]
 # tidy-alphabetical-start
+indexmap = "2.4"
 parking_lot = "0.12"
 rustc-rayon-core = { version = "0.5.0" }
 rustc_abi = { path = "../rustc_abi" }

compiler/rustc_query_system/src/query/caches.rs

@@ -1,7 +1,7 @@
 use std::fmt::Debug;
 use std::hash::Hash;
 
-use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
 use rustc_data_structures::sharded::{self, Sharded};
 use rustc_data_structures::sync::OnceLock;
 pub use rustc_data_structures::vec_cache::VecCache;
@@ -187,3 +187,48 @@ where
         self.iter(f)
     }
 }
+
+pub struct IndexCache<K, V> {
+    cache: Sharded<FxIndexMap<K, (V, DepNodeIndex)>>,
+}
+
+impl<K, V> Default for IndexCache<K, V> {
+    fn default() -> Self {
+        IndexCache { cache: Default::default() }
+    }
+}
+
+impl<K, V> QueryCache for IndexCache<K, V>
+where
+    K: Eq + Hash + Copy + Debug,
+    V: Copy,
+{
+    type Key = K;
+    type Value = V;
+
+    #[inline(always)]
+    fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> {
+        use indexmap::map::raw_entry_v1::RawEntryApiV1;
+        let key_hash = sharded::make_hash(key);
+        let lock = self.cache.lock_shard_by_hash(key_hash);
+        let result = lock.raw_entry_v1().from_key_hashed_nocheck(key_hash, key);
+
+        if let Some((_, value)) = result { Some(*value) } else { None }
+    }
+
+    #[inline]
+    fn complete(&self, key: K, value: V, index: DepNodeIndex) {
+        let mut lock = self.cache.lock_shard_by_value(&key);
+        // We may be overwriting another value. This is all right, since the dep-graph
+        // will check that the fingerprint matches.
+        lock.insert(key, (value, index));
+    }
+
+    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
+        for shard in self.cache.lock_shards() {
+            for (k, v) in shard.iter() {
+                f(k, &v.0, v.1);
+            }
+        }
+    }
+}

compiler/rustc_query_system/src/query/mod.rs

@@ -8,7 +8,7 @@ pub use self::job::{
 };
 
 mod caches;
-pub use self::caches::{DefIdCache, DefaultCache, QueryCache, SingleCache, VecCache};
+pub use self::caches::{DefIdCache, DefaultCache, IndexCache, QueryCache, SingleCache, VecCache};
 
 mod config;
 use rustc_data_structures::sync::Lock;