Refactor Lock implementation

rust-lang · Sep 3, 2023 · d500310 · d500310
1 parent 7281a8a
commit d500310
Show file tree

Hide file tree

Showing 3 changed files with 216 additions and 246 deletions.
diff --git a/compiler/rustc_data_structures/src/sharded.rs b/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::{Assume, Lock, LockGuard};
 #[cfg(parallel_compiler)]
 use itertools::Either;
 use std::borrow::Borrow;
@@ -75,6 +75,7 @@ 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) => {
@@ -83,19 +84,21 @@ impl<T> Sharded<T> {
 
                 // 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_no_sync() }
+                unsafe { single.lock_assume(Assume::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) => {
@@ -104,7 +107,7 @@ impl<T> Sharded<T> {
 
                 // 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_no_sync() }
+                unsafe { single.lock_assume(Assume::NoSync) }
             }
             #[cfg(parallel_compiler)]
             Self::Shards(shards) => {
@@ -115,7 +118,7 @@ impl<T> Sharded<T> {
                 // 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_sync() }
+                unsafe { shards.get_unchecked(_i & (SHARDS - 1)).0.lock_assume(Assume::Sync) }
             }
         }
     }

diff --git a/compiler/rustc_data_structures/src/sync.rs b/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::{Assume, Lock, LockGuard};
 
 mod worker_local;
 pub use worker_local::{Registry, WorkerLocal};
@@ -83,7 +83,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
     }