From 80d265240bf1b884753bafc8f07e43fc8236679f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= <john.kare.alsaker@gmail.com>
Date: Fri, 3 Feb 2023 18:39:31 +0100
Subject: [PATCH] Create a single value cache for the () query key

---
 compiler/rustc_middle/src/query/keys.rs       |  4 +-
 .../rustc_query_system/src/query/caches.rs    | 47 ++++++++++++++++++-
 compiler/rustc_query_system/src/query/mod.rs  |  3 +-
 3 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/compiler/rustc_middle/src/query/keys.rs b/compiler/rustc_middle/src/query/keys.rs
index e4bb3ce3d5a..dc02fd53ed0 100644
--- a/compiler/rustc_middle/src/query/keys.rs
+++ b/compiler/rustc_middle/src/query/keys.rs
@@ -8,7 +8,7 @@ use crate::ty::subst::{GenericArg, SubstsRef};
 use crate::ty::{self, layout::TyAndLayout, Ty, TyCtxt};
 use rustc_hir::def_id::{CrateNum, DefId, LocalDefId, LOCAL_CRATE};
 use rustc_hir::hir_id::{HirId, OwnerId};
-use rustc_query_system::query::{DefaultCacheSelector, VecCacheSelector};
+use rustc_query_system::query::{DefaultCacheSelector, SingleCacheSelector, VecCacheSelector};
 use rustc_span::symbol::{Ident, Symbol};
 use rustc_span::{Span, DUMMY_SP};
 
@@ -45,7 +45,7 @@ pub trait Key: Sized {
 }
 
 impl Key for () {
-    type CacheSelector = DefaultCacheSelector<Self>;
+    type CacheSelector = SingleCacheSelector;
 
     #[inline(always)]
     fn query_crate_is_local(&self) -> bool {
diff --git a/compiler/rustc_query_system/src/query/caches.rs b/compiler/rustc_query_system/src/query/caches.rs
index 9f875b43731..c9dd75e4d55 100644
--- a/compiler/rustc_query_system/src/query/caches.rs
+++ b/compiler/rustc_query_system/src/query/caches.rs
@@ -5,7 +5,6 @@ use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sharded;
 #[cfg(parallel_compiler)]
 use rustc_data_structures::sharded::Sharded;
-#[cfg(not(parallel_compiler))]
 use rustc_data_structures::sync::Lock;
 use rustc_data_structures::sync::WorkerLocal;
 use rustc_index::vec::{Idx, IndexVec};
@@ -117,6 +116,52 @@ where
     }
 }
 
+pub struct SingleCacheSelector;
+
+impl<'tcx, V: 'tcx> CacheSelector<'tcx, V> for SingleCacheSelector {
+    type Cache = SingleCache<V>
+    where
+        V: Copy;
+    type ArenaCache = ArenaCache<'tcx, (), V>;
+}
+
+pub struct SingleCache<V> {
+    cache: Lock<Option<(V, DepNodeIndex)>>,
+}
+
+impl<V> Default for SingleCache<V> {
+    fn default() -> Self {
+        SingleCache { cache: Lock::new(None) }
+    }
+}
+
+impl<V: Copy + Debug> QueryStorage for SingleCache<V> {
+    type Value = V;
+    type Stored = V;
+}
+
+impl<V> QueryCache for SingleCache<V>
+where
+    V: Copy + Debug,
+{
+    type Key = ();
+
+    #[inline(always)]
+    fn lookup(&self, _key: &()) -> Option<(V, DepNodeIndex)> {
+        *self.cache.lock()
+    }
+
+    #[inline]
+    fn complete(&self, _key: (), value: V, index: DepNodeIndex) -> Self::Stored {
+        *self.cache.lock() = Some((value.clone(), index));
+        value
+    }
+
+    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
+        self.cache.lock().as_ref().map(|value| f(&(), &value.0, value.1));
+    }
+}
+
 pub struct ArenaCache<'tcx, K, V> {
     arena: WorkerLocal<TypedArena<(V, DepNodeIndex)>>,
     #[cfg(parallel_compiler)]
diff --git a/compiler/rustc_query_system/src/query/mod.rs b/compiler/rustc_query_system/src/query/mod.rs
index d308af19207..6c0ee2bc2f6 100644
--- a/compiler/rustc_query_system/src/query/mod.rs
+++ b/compiler/rustc_query_system/src/query/mod.rs
@@ -8,7 +8,8 @@ pub use self::job::{print_query_stack, QueryInfo, QueryJob, QueryJobId, QueryJob
 
 mod caches;
 pub use self::caches::{
-    CacheSelector, DefaultCacheSelector, QueryCache, QueryStorage, VecCacheSelector,
+    CacheSelector, DefaultCacheSelector, QueryCache, QueryStorage, SingleCacheSelector,
+    VecCacheSelector,
 };
 
 mod config;