1 files changed, 200 insertions, 0 deletions
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 000000000..466c44f59
--- /dev/null
+++ b/compiler/rustc_data_structures/src/sync/freeze.rs
@@ -0,0 +1,200 @@
+use crate::sync::{AtomicBool, ReadGuard, RwLock, WriteGuard};
+#[cfg(parallel_compiler)]
+use crate::sync::{DynSend, DynSync};
+use std::{
+    cell::UnsafeCell,
+    intrinsics::likely,
+    marker::PhantomData,
+    ops::{Deref, DerefMut},
+    ptr::NonNull,
+    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::with(value, false)
+    }
+
+    #[inline]
+    pub fn frozen(value: T) -> Self {
+        Self::with(value, true)
+    }
+
+    #[inline]
+    pub fn with(value: T, frozen: bool) -> Self {
+        Self {
+            data: UnsafeCell::new(value),
+            frozen: AtomicBool::new(frozen),
+            lock: RwLock::new(()),
+        }
+    }
+
+    /// Clones the inner value along with the frozen state.
+    #[inline]
+    pub fn clone(&self) -> Self
+    where
+        T: Clone,
+    {
+        let lock = self.read();
+        Self::with(lock.clone(), self.is_frozen())
+    }
+
+    #[inline]
+    pub fn is_frozen(&self) -> bool {
+        self.frozen.load(Ordering::Acquire)
+    }
+
+    /// Get the inner value if frozen.
+    #[inline]
+    pub fn get(&self) -> Option<&T> {
+        if likely(self.frozen.load(Ordering::Acquire)) {
+            // SAFETY: This is frozen so the data cannot be modified.
+            unsafe { Some(&*self.data.get()) }
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    pub fn read(&self) -> FreezeReadGuard<'_, T> {
+        FreezeReadGuard {
+            _lock_guard: if self.frozen.load(Ordering::Acquire) {
+                None
+            } else {
+                Some(self.lock.read())
+            },
+            data: unsafe { NonNull::new_unchecked(self.data.get()) },
+        }
+    }
+
+    #[inline]
+    pub fn borrow(&self) -> FreezeReadGuard<'_, T> {
+        self.read()
+    }
+
+    #[inline]
+    #[track_caller]
+    pub fn write(&self) -> FreezeWriteGuard<'_, T> {
+        self.try_write().expect("still mutable")
+    }
+
+    #[inline]
+    pub fn try_write(&self) -> Option<FreezeWriteGuard<'_, T>> {
+        let _lock_guard = self.lock.write();
+        // Use relaxed ordering since we're in the write lock.
+        if self.frozen.load(Ordering::Relaxed) {
+            None
+        } else {
+            Some(FreezeWriteGuard {
+                _lock_guard,
+                data: unsafe { NonNull::new_unchecked(self.data.get()) },
+                frozen: &self.frozen,
+                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: ?Sized> {
+    _lock_guard: Option<ReadGuard<'a, ()>>,
+    data: NonNull<T>,
+}
+
+impl<'a, T: ?Sized + 'a> Deref for FreezeReadGuard<'a, T> {
+    type Target = T;
+    #[inline]
+    fn deref(&self) -> &T {
+        // SAFETY: If the lock is not frozen, `_lock_guard` holds the lock to the `UnsafeCell` so
+        // this has shared access until the `FreezeReadGuard` is dropped. If the lock is frozen,
+        // the data cannot be modified and shared access is sound.
+        unsafe { &*self.data.as_ptr() }
+    }
+}
+
+impl<'a, T: ?Sized> FreezeReadGuard<'a, T> {
+    #[inline]
+    pub fn map<U: ?Sized>(this: Self, f: impl FnOnce(&T) -> &U) -> FreezeReadGuard<'a, U> {
+        FreezeReadGuard { data: NonNull::from(f(&*this)), _lock_guard: this._lock_guard }
+    }
+}
+
+/// 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: ?Sized> {
+    _lock_guard: WriteGuard<'a, ()>,
+    frozen: &'a AtomicBool,
+    data: NonNull<T>,
+    marker: PhantomData<&'a mut T>,
+}
+
+impl<'a, T> FreezeWriteGuard<'a, T> {
+    pub fn freeze(self) -> &'a T {
+        self.frozen.store(true, Ordering::Release);
+
+        // SAFETY: This is frozen so the data cannot be modified and shared access is sound.
+        unsafe { &*self.data.as_ptr() }
+    }
+}
+
+impl<'a, T: ?Sized> FreezeWriteGuard<'a, T> {
+    #[inline]
+    pub fn map<U: ?Sized>(
+        mut this: Self,
+        f: impl FnOnce(&mut T) -> &mut U,
+    ) -> FreezeWriteGuard<'a, U> {
+        FreezeWriteGuard {
+            data: NonNull::from(f(&mut *this)),
+            _lock_guard: this._lock_guard,
+            frozen: this.frozen,
+            marker: PhantomData,
+        }
+    }
+}
+
+impl<'a, T: ?Sized + '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.data.as_ptr() }
+    }
+}
+
+impl<'a, T: ?Sized + '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.data.as_ptr() }
+    }
+}