use RefCell for query caches in single thread

rust-lang · Aug 25, 2023 · cd6850c · cd6850c
1 parent 16b9941
commit cd6850c
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 11 deletions.
diff --git a/compiler/rustc_middle/src/query/plumbing.rs b/compiler/rustc_middle/src/query/plumbing.rs
@@ -13,6 +13,8 @@ use field_offset::FieldOffset;
 use measureme::StringId;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sync::AtomicU64;
+#[cfg(parallel_compiler)]
+use rustc_data_structures::sync::DynSync;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_hir::hir_id::OwnerId;
@@ -97,6 +99,10 @@ pub struct QuerySystem<'tcx> {
     pub jobs: AtomicU64,
 }
 
+// It's thread safe since `QueryCaches` only used under single thread
+#[cfg(parallel_compiler)]
+unsafe impl<'tcx> DynSync for QuerySystem<'tcx> {}
+
 #[derive(Copy, Clone)]
 pub struct TyCtxtAt<'tcx> {
     pub tcx: TyCtxt<'tcx>,
@@ -428,18 +434,46 @@ macro_rules! define_callbacks {
             })*
         }
 
+        #[cfg(not(parallel_compiler))]
         impl<'tcx> TyCtxtAt<'tcx> {
             $($(#[$attr])*
+
             #[inline(always)]
             pub fn $name(self, key: query_helper_param_ty!($($K)*)) -> $V
             {
-                restore::<$V>(with_query_caches!(query_get_at(
+                restore::<$V>(query_get_at(
                     self.tcx,
                     self.tcx.query_system.fns.engine.$name,
-                    :self.tcx, $name,
+                    &self.tcx.query_system.caches.$name,
                     self.span,
                     key.into_query_param(),
-                )))
+                ))
+            })*
+        }
+
+        #[cfg(parallel_compiler)]
+        impl<'tcx> TyCtxtAt<'tcx> {
+            $($(#[$attr])*
+            #[inline(always)]
+            pub fn $name(self, key: query_helper_param_ty!($($K)*)) -> $V
+            {
+                if std::intrinsics::likely(self.tcx.query_system.single_thread) {
+                    restore::<$V>(query_get_at(
+                        self.tcx,
+                        self.tcx.query_system.fns.engine.$name,
+                        &self.tcx.query_system.caches.$name,
+                        self.span,
+                        key.into_query_param(),
+                    ))
+                } else {
+                    restore::<$V>(query_get_at(
+                        self.tcx,
+                        self.tcx.query_system.fns.engine.$name,
+                        &self.tcx.query_system.mt_caches.$name,
+                        self.span,
+                        key.into_query_param(),
+                    ))
+                }
             })*
         }
 

diff --git a/compiler/rustc_query_system/src/query/caches.rs b/compiler/rustc_query_system/src/query/caches.rs
@@ -1,4 +1,5 @@
 use crate::dep_graph::DepNodeIndex;
+use std::cell::RefCell;
 
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sharded::{self, Sharded};
@@ -41,7 +42,7 @@ impl<'tcx, K: Eq + Hash, V: 'tcx> CacheSelector<'tcx, V> for DefaultCacheSelecto
 }
 
 pub struct DefaultCache<K, V> {
-    cache: Lock<FxHashMap<K, (V, DepNodeIndex)>>,
+    cache: RefCell<FxHashMap<K, (V, DepNodeIndex)>>,
 }
 
 impl<K, V> Default for DefaultCache<K, V> {
@@ -61,22 +62,22 @@ where
     #[inline(always)]
     fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> {
         let key_hash = sharded::make_hash(key);
-        let lock = self.cache.lock();
+        let lock = self.cache.borrow_mut();
         let result = lock.raw_entry().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();
+        let mut lock = self.cache.borrow_mut();
         // 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)) {
-        let map = self.cache.lock();
+        let map = self.cache.borrow_mut();
         for (k, v) in map.iter() {
             f(k, &v.0, v.1);
         }
@@ -186,7 +187,7 @@ impl<'tcx, K: Idx, V: 'tcx> CacheSelector<'tcx, V> for VecCacheSelector<K> {
 }
 
 pub struct VecCache<K: Idx, V> {
-    cache: Lock<IndexVec<K, Option<(V, DepNodeIndex)>>>,
+    cache: RefCell<IndexVec<K, Option<(V, DepNodeIndex)>>>,
 }
 
 impl<K: Idx, V> Default for VecCache<K, V> {
@@ -205,18 +206,18 @@ where
 
     #[inline(always)]
     fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> {
-        let lock = self.cache.lock();
+        let lock = self.cache.borrow_mut();
         if let Some(Some(value)) = lock.get(*key) { Some(*value) } else { None }
     }
 
     #[inline]
     fn complete(&self, key: K, value: V, index: DepNodeIndex) {
-        let mut lock = self.cache.lock();
+        let mut lock = self.cache.borrow_mut();
         lock.insert(key, (value, index));
     }
 
     fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
-        let map = self.cache.lock();
+        let map = self.cache.borrow_mut();
         for (k, v) in map.iter_enumerated() {
             if let Some(v) = v {
                 f(&k, &v.0, v.1);