Add optimized lock methods for Sharded and refactor Lock

compiler/rustc_data_structures/src/sharded.rs

@@ -1,7 +1,7 @@
 use crate::fx::{FxHashMap, FxHasher};
 #[cfg(parallel_compiler)]
 use crate::sync::{is_dyn_thread_safe, CacheAligned};
-use crate::sync::{Lock, LockGuard};
+use crate::sync::{Lock, LockGuard, Mode};
 #[cfg(parallel_compiler)]
 use itertools::Either;
 use std::borrow::Borrow;
@@ -73,6 +73,56 @@ impl<T> Sharded<T> {
         }
     }
 
+    /// The shard is selected by hashing `val` with `FxHasher`.
+    #[inline]
+    #[track_caller]
+    pub fn lock_shard_by_value<K: Hash + ?Sized>(&self, _val: &K) -> LockGuard<'_, T> {
+        match self {
+            Self::Single(single) => {
+                // Syncronization is disabled so use the `lock_assume_no_sync` method optimized
+                // for that case.
+
+                // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus
+                // `might_be_dyn_thread_safe` was also false.
+                unsafe { single.lock_assume(Mode::NoSync) }
+            }
+            #[cfg(parallel_compiler)]
+            Self::Shards(..) => self.lock_shard_by_hash(make_hash(_val)),
+        }
+    }
+
+    #[inline]
+    #[track_caller]
+    pub fn lock_shard_by_hash(&self, hash: u64) -> LockGuard<'_, T> {
+        self.lock_shard_by_index(get_shard_hash(hash))
+    }
+
+    #[inline]
+    #[track_caller]
+    pub fn lock_shard_by_index(&self, _i: usize) -> LockGuard<'_, T> {
+        match self {
+            Self::Single(single) => {
+                // Syncronization is disabled so use the `lock_assume_no_sync` method optimized
+                // for that case.
+
+                // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus
+                // `might_be_dyn_thread_safe` was also false.
+                unsafe { single.lock_assume(Mode::NoSync) }
+            }
+            #[cfg(parallel_compiler)]
+            Self::Shards(shards) => {
+                // Syncronization is enabled so use the `lock_assume_sync` method optimized
+                // for that case.
+
+                // SAFETY (get_unchecked): The index gets ANDed with the shard mask, ensuring it is
+                // always inbounds.
+                // SAFETY (lock_assume_sync): We know `is_dyn_thread_safe` was true when creating
+                // the lock thus `might_be_dyn_thread_safe` was also true.
+                unsafe { shards.get_unchecked(_i & (SHARDS - 1)).0.lock_assume(Mode::Sync) }
+            }
+        }
+    }
+
     #[inline]
     pub fn lock_shards(&self) -> impl Iterator<Item = LockGuard<'_, T>> {
         match self {
@@ -124,7 +174,7 @@ impl<K: Eq + Hash + Copy> ShardedHashMap<K, ()> {
         Q: Hash + Eq,
     {
         let hash = make_hash(value);
-        let mut shard = self.get_shard_by_hash(hash).lock();
+        let mut shard = self.lock_shard_by_hash(hash);
         let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, value);
 
         match entry {
@@ -144,7 +194,7 @@ impl<K: Eq + Hash + Copy> ShardedHashMap<K, ()> {
         Q: Hash + Eq,
     {
         let hash = make_hash(&value);
-        let mut shard = self.get_shard_by_hash(hash).lock();
+        let mut shard = self.lock_shard_by_hash(hash);
         let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, &value);
 
         match entry {
@@ -166,7 +216,7 @@ pub trait IntoPointer {
 impl<K: Eq + Hash + Copy + IntoPointer> ShardedHashMap<K, ()> {
     pub fn contains_pointer_to<T: Hash + IntoPointer>(&self, value: &T) -> bool {
         let hash = make_hash(&value);
-        let shard = self.get_shard_by_hash(hash).lock();
+        let shard = self.lock_shard_by_hash(hash);
         let value = value.into_pointer();
         shard.raw_entry().from_hash(hash, |entry| entry.into_pointer() == value).is_some()
     }

compiler/rustc_data_structures/src/sync.rs

@@ -49,7 +49,7 @@ use std::ops::{Deref, DerefMut};
 use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe};
 
 mod lock;
-pub use lock::{Lock, LockGuard};
+pub use lock::{Lock, LockGuard, Mode};
 
 mod worker_local;
 pub use worker_local::{Registry, WorkerLocal};
@@ -86,7 +86,6 @@ mod mode {
 
     // Whether thread safety might be enabled.
     #[inline]
-    #[cfg(parallel_compiler)]
     pub fn might_be_dyn_thread_safe() -> bool {
         DYN_THREAD_SAFE_MODE.load(Ordering::Relaxed) != DYN_NOT_THREAD_SAFE
     }