specialize the query caches

Author: SparrowLii

compiler/rustc_data_structures/src/sharded.rs

@@ -1,5 +1,5 @@
 use crate::fx::{FxHashMap, FxHasher};
-use crate::sync::{CacheAligned, Lock, LockGuard};
+use crate::sync::{Lock, LockGuard};
 use std::borrow::Borrow;
 use std::collections::hash_map::RawEntryMut;
 use std::hash::{Hash, Hasher};

compiler/rustc_data_structures/src/sync.rs

@@ -40,7 +40,7 @@
 //! [^2] `MTLockRef` is a typedef.
 
 pub use crate::marker::*;
-use std::cell::{Cell, UnsafeCell};
+use std::cell::{Cell, RefCell, RefMut, UnsafeCell};
 use std::collections::HashMap;
 use std::fmt::{Debug, Formatter};
 use std::hash::{BuildHasher, Hash};
@@ -50,10 +50,10 @@ use std::ops::{Deref, DerefMut};
 use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe};
 
 mod worker_local;
-pub use worker_local::{Registry, WorkerLocal};
 
 pub use std::sync::atomic::Ordering;
 pub use std::sync::atomic::Ordering::SeqCst;
+use std::sync::{Mutex, MutexGuard};
 
 pub use vec::{AppendOnlyIndexVec, AppendOnlyVec};
 
@@ -559,8 +559,8 @@ impl<T> FromDyn<T> {
     }
 }
 
-#[derive(Default)]
-#[cfg_attr(parallel_compiler, repr(align(64)))]
+#[derive(Default, Debug)]
+#[repr(align(64))]
 pub struct CacheAligned<T>(pub T);
 
 pub trait HashMapExt<K, V> {
@@ -768,6 +768,103 @@ impl<'a, T> Drop for LockGuard<'a, T> {
     }
 }
 
+pub trait SLock: Copy {
+    type Lock<T>: LockLike<T>;
+}
+
+pub trait LockLike<T> {
+    type LockGuard<'a>: DerefMut<Target = T>
+    where
+        Self: 'a;
+
+    fn new(val: T) -> Self;
+
+    fn into_inner(self) -> T;
+
+    fn get_mut(&mut self) -> &mut T;
+
+    fn try_lock(&self) -> Option<Self::LockGuard<'_>>;
+
+    fn lock(&self) -> Self::LockGuard<'_>;
+}
+
+#[derive(Copy, Clone, Default)]
+pub struct SRefCell;
+
+impl SLock for SRefCell {
+    type Lock<T> = RefCell<T>;
+}
+
+impl<T> LockLike<T> for RefCell<T> {
+    type LockGuard<'a> = RefMut<'a, T> where T: 'a;
+
+    #[inline]
+    fn new(val: T) -> Self {
+        RefCell::new(val)
+    }
+
+    #[inline]
+    fn into_inner(self) -> T {
+        self.into_inner()
+    }
+
+    #[inline]
+    fn get_mut(&mut self) -> &mut T {
+        self.get_mut()
+    }
+
+    #[inline]
+    fn try_lock(&self) -> Option<RefMut<'_, T>> {
+        self.try_borrow_mut().ok()
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    fn lock(&self) -> RefMut<'_, T> {
+        self.borrow_mut()
+    }
+}
+
+#[derive(Copy, Clone, Default)]
+pub struct SMutex;
+
+impl SLock for SMutex {
+    type Lock<T> = Mutex<T>;
+}
+
+impl<T> LockLike<T> for Mutex<T> {
+    type LockGuard<'a> = MutexGuard<'a, T> where T: 'a;
+
+    #[inline]
+    fn new(val: T) -> Self {
+        Mutex::new(val)
+    }
+
+    #[inline]
+    fn into_inner(self) -> T {
+        self.into_inner().unwrap()
+    }
+
+    #[inline]
+    fn get_mut(&mut self) -> &mut T {
+        self.get_mut().unwrap()
+    }
+
+    #[inline]
+    fn try_lock(&self) -> Option<MutexGuard<'_, T>> {
+        self.try_lock().ok()
+    }
+
+    #[inline(always)]
+    #[track_caller]
+    fn lock(&self) -> MutexGuard<'_, T> {
+        self.lock().unwrap_or_else(|e| {
+            self.clear_poison();
+            e.into_inner()
+        })
+    }
+}
+
 pub struct MappedReadGuard<'a, T: ?Sized> {
     raw: &'a RwLockRaw,
     data: *const T,
@@ -1138,7 +1235,7 @@ impl<T: Clone> Clone for RwLock<T> {
 pub struct WorkerLocal<T> {
     single_thread: bool,
     inner: Option<T>,
-    mt_inner: Option<rayon_core::WorkerLocal<T>>,
+    mt_inner: Option<worker_local::WorkerLocal<T>>,
 }
 
 impl<T> WorkerLocal<T> {
@@ -1152,20 +1249,10 @@ impl<T> WorkerLocal<T> {
             WorkerLocal {
                 single_thread: false,
                 inner: None,
-                mt_inner: Some(rayon_core::WorkerLocal::new(f)),
+                mt_inner: Some(worker_local::WorkerLocal::new(f)),
             }
         }
     }
-
-    /// Returns the worker-local value for each thread
-    #[inline]
-    pub fn into_inner(self) -> Vec<T> {
-        if self.single_thread {
-            vec![self.inner.unwrap()]
-        } else {
-            self.mt_inner.unwrap().into_inner()
-        }
-    }
 }
 
 impl<T> Deref for WorkerLocal<T> {
@@ -1185,6 +1272,7 @@ impl<T> Deref for WorkerLocal<T> {
 unsafe impl<T: Send> std::marker::Sync for WorkerLocal<T> {}
 
 use std::thread;
+pub use worker_local::Registry;
 
 /// A type which only allows its inner value to be used in one thread.
 /// It will panic if it is used on multiple threads.

compiler/rustc_data_structures/src/sync/worker_local.rs

@@ -5,7 +5,6 @@ use std::ops::Deref;
 use std::ptr;
 use std::sync::Arc;
 
-#[cfg(parallel_compiler)]
 use {crate::cold_path, crate::sync::CacheAligned};
 
 /// A pointer to the `RegistryData` which uniquely identifies a registry.
@@ -21,7 +20,6 @@ impl RegistryId {
     ///
     /// Note that there's a race possible where the identifer in `THREAD_DATA` could be reused
     /// so this can succeed from a different registry.
-    #[cfg(parallel_compiler)]
     fn verify(self) -> usize {
         let (id, index) = THREAD_DATA.with(|data| (data.registry_id.get(), data.index.get()));
 
@@ -33,13 +31,14 @@ impl RegistryId {
     }
 }
 
+#[derive(Debug)]
 struct RegistryData {
     thread_limit: usize,
     threads: Lock<usize>,
 }
 
 /// Represents a list of threads which can access worker locals.
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct Registry(Arc<RegistryData>);
 
 thread_local! {
@@ -105,73 +104,35 @@ impl Registry {
 /// Holds worker local values for each possible thread in a registry. You can only access the
 /// worker local value through the `Deref` impl on the registry associated with the thread it was
 /// created on. It will panic otherwise.
-pub struct WorkerLocal<T> {
-    #[cfg(not(parallel_compiler))]
-    local: T,
-    #[cfg(parallel_compiler)]
+#[derive(Debug)]
+pub(crate) struct WorkerLocal<T> {
     locals: Box<[CacheAligned<T>]>,
-    #[cfg(parallel_compiler)]
     registry: Registry,
 }
 
 // This is safe because the `deref` call will return a reference to a `T` unique to each thread
 // or it will panic for threads without an associated local. So there isn't a need for `T` to do
 // it's own synchronization. The `verify` method on `RegistryId` has an issue where the the id
 // can be reused, but `WorkerLocal` has a reference to `Registry` which will prevent any reuse.
-#[cfg(parallel_compiler)]
 unsafe impl<T: Send> Sync for WorkerLocal<T> {}
 
 impl<T> WorkerLocal<T> {
     /// Creates a new worker local where the `initial` closure computes the
     /// value this worker local should take for each thread in the registry.
     #[inline]
     pub fn new<F: FnMut(usize) -> T>(mut initial: F) -> WorkerLocal<T> {
-        #[cfg(parallel_compiler)]
-        {
-            let registry = Registry::current();
-            WorkerLocal {
-                locals: (0..registry.0.thread_limit).map(|i| CacheAligned(initial(i))).collect(),
-                registry,
-            }
-        }
-        #[cfg(not(parallel_compiler))]
-        {
-            WorkerLocal { local: initial(0) }
-        }
-    }
-
-    /// Returns the worker-local values for each thread
-    #[inline]
-    pub fn into_inner(self) -> impl Iterator<Item = T> {
-        #[cfg(parallel_compiler)]
-        {
-            self.locals.into_vec().into_iter().map(|local| local.0)
-        }
-        #[cfg(not(parallel_compiler))]
-        {
-            std::iter::once(self.local)
+        let registry = Registry::current();
+        WorkerLocal {
+            locals: (0..registry.0.thread_limit).map(|i| CacheAligned(initial(i))).collect(),
+            registry,
         }
     }
 }
 
-impl<T> WorkerLocal<Vec<T>> {
-    /// Joins the elements of all the worker locals into one Vec
-    pub fn join(self) -> Vec<T> {
-        self.into_inner().into_iter().flat_map(|v| v).collect()
-    }
-}
-
 impl<T> Deref for WorkerLocal<T> {
     type Target = T;
 
     #[inline(always)]
-    #[cfg(not(parallel_compiler))]
-    fn deref(&self) -> &T {
-        &self.local
-    }
-
-    #[inline(always)]
-    #[cfg(parallel_compiler)]
     fn deref(&self) -> &T {
         // This is safe because `verify` will only return values less than
         // `self.registry.thread_limit` which is the size of the `self.locals` array.