Auto merge of #115401 - Zoxc:freeze, r=oli-obk

Add `FreezeLock` type and use it to store `Definitions`

This adds a `FreezeLock` type which allows mutation using a lock until the value is frozen where it can be accessed lock-free. It's used to store `Definitions` in `Untracked` instead of a `RwLock`. Unlike the current scheme of leaking read guards this doesn't deadlock if definitions is written to after no mutation are expected.

2 files changed, 111 insertions, 0 deletions

diff --git a/compiler/rustc_data_structures/src/sync.rs b/compiler/rustc_data_structures/src/sync.rs
index e82b0f6d496..8eda9043e35 100644
--- a/compiler/rustc_data_structures/src/sync.rs
+++ b/compiler/rustc_data_structures/src/sync.rs
@@ -61,6 +61,9 @@ pub use vec::{AppendOnlyIndexVec, AppendOnlyVec};
 
 mod vec;
 
+mod freeze;
+pub use freeze::{FreezeLock, FreezeReadGuard, FreezeWriteGuard};
+
 mod mode {
     use super::Ordering;
     use std::sync::atomic::AtomicU8;
diff --git a/compiler/rustc_data_structures/src/sync/freeze.rs b/compiler/rustc_data_structures/src/sync/freeze.rs
new file mode 100644
index 00000000000..d9f1d72d851
--- /dev/null
+++ b/compiler/rustc_data_structures/src/sync/freeze.rs
@@ -0,0 +1,108 @@
+use crate::sync::{AtomicBool, ReadGuard, RwLock, WriteGuard};
+#[cfg(parallel_compiler)]
+use crate::sync::{DynSend, DynSync};
+use std::{
+    cell::UnsafeCell,
+    marker::PhantomData,
+    ops::{Deref, DerefMut},
+    sync::atomic::Ordering,
+};
+
+/// A type which allows mutation using a lock until
+/// the value is frozen and can be accessed lock-free.
+///
+/// Unlike `RwLock`, it can be used to prevent mutation past a point.
+#[derive(Default)]
+pub struct FreezeLock<T> {
+    data: UnsafeCell<T>,
+    frozen: AtomicBool,
+
+    /// This lock protects writes to the `data` and `frozen` fields.
+    lock: RwLock<()>,
+}
+
+#[cfg(parallel_compiler)]
+unsafe impl<T: DynSync + DynSend> DynSync for FreezeLock<T> {}
+
+impl<T> FreezeLock<T> {
+    #[inline]
+    pub fn new(value: T) -> Self {
+        Self { data: UnsafeCell::new(value), frozen: AtomicBool::new(false), lock: RwLock::new(()) }
+    }
+
+    #[inline]
+    pub fn read(&self) -> FreezeReadGuard<'_, T> {
+        FreezeReadGuard {
+            _lock_guard: if self.frozen.load(Ordering::Acquire) {
+                None
+            } else {
+                Some(self.lock.read())
+            },
+            lock: self,
+        }
+    }
+
+    #[inline]
+    #[track_caller]
+    pub fn write(&self) -> FreezeWriteGuard<'_, T> {
+        let _lock_guard = self.lock.write();
+        // Use relaxed ordering since we're in the write lock.
+        assert!(!self.frozen.load(Ordering::Relaxed), "still mutable");
+        FreezeWriteGuard { _lock_guard, lock: self, marker: PhantomData }
+    }
+
+    #[inline]
+    pub fn freeze(&self) -> &T {
+        if !self.frozen.load(Ordering::Acquire) {
+            // Get the lock to ensure no concurrent writes and that we release the latest write.
+            let _lock = self.lock.write();
+            self.frozen.store(true, Ordering::Release);
+        }
+
+        // SAFETY: This is frozen so the data cannot be modified and shared access is sound.
+        unsafe { &*self.data.get() }
+    }
+}
+
+/// A guard holding shared access to a `FreezeLock` which is in a locked state or frozen.
+#[must_use = "if unused the FreezeLock may immediately unlock"]
+pub struct FreezeReadGuard<'a, T> {
+    _lock_guard: Option<ReadGuard<'a, ()>>,
+    lock: &'a FreezeLock<T>,
+}
+
+impl<'a, T: 'a> Deref for FreezeReadGuard<'a, T> {
+    type Target = T;
+    #[inline]
+    fn deref(&self) -> &T {
+        // SAFETY: If `lock` is not frozen, `_lock_guard` holds the lock to the `UnsafeCell` so
+        // this has shared access until the `FreezeReadGuard` is dropped. If `lock` is frozen,
+        // the data cannot be modified and shared access is sound.
+        unsafe { &*self.lock.data.get() }
+    }
+}
+
+/// A guard holding mutable access to a `FreezeLock` which is in a locked state or frozen.
+#[must_use = "if unused the FreezeLock may immediately unlock"]
+pub struct FreezeWriteGuard<'a, T> {
+    _lock_guard: WriteGuard<'a, ()>,
+    lock: &'a FreezeLock<T>,
+    marker: PhantomData<&'a mut T>,
+}
+
+impl<'a, T: 'a> Deref for FreezeWriteGuard<'a, T> {
+    type Target = T;
+    #[inline]
+    fn deref(&self) -> &T {
+        // SAFETY: `self._lock_guard` holds the lock to the `UnsafeCell` so this has shared access.
+        unsafe { &*self.lock.data.get() }
+    }
+}
+
+impl<'a, T: 'a> DerefMut for FreezeWriteGuard<'a, T> {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut T {
+        // SAFETY: `self._lock_guard` holds the lock to the `UnsafeCell` so this has mutable access.
+        unsafe { &mut *self.lock.data.get() }
+    }
+}