Add Freeze type and use it to store Definitions

Author: Zoxc

compiler/rustc_data_structures/src/sync.rs

@@ -61,6 +61,9 @@ pub use vec::{AppendOnlyIndexVec, AppendOnlyVec};
 
 mod vec;
 
+mod freeze;
+pub use freeze::{Freeze, FreezeReadGuard, FreezeWriteGuard};
+
 mod mode {
     use super::Ordering;
     use std::sync::atomic::AtomicU8;

compiler/rustc_data_structures/src/sync/freeze.rs

@@ -0,0 +1,104 @@
+use crate::sync::{AtomicBool, Lock, LockGuard};
+#[cfg(parallel_compiler)]
+use crate::sync::{DynSend, DynSync};
+use std::{
+    cell::UnsafeCell,
+    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.
+#[derive(Default)]
+pub struct Freeze<T> {
+    data: UnsafeCell<T>,
+    frozen: AtomicBool,
+    lock: Lock<()>,
+}
+
+#[cfg(parallel_compiler)]
+unsafe impl<T: DynSync + DynSend> DynSync for Freeze<T> {}
+
+impl<T> Freeze<T> {
+    #[inline]
+    pub fn new(value: T) -> Self {
+        Self { data: UnsafeCell::new(value), frozen: AtomicBool::new(false), lock: Lock::new(()) }
+    }
+
+    #[inline]
+    pub fn read(&self) -> FreezeReadGuard<'_, T> {
+        FreezeReadGuard {
+            _lock_guard: if self.frozen.load(Ordering::Acquire) {
+                None
+            } else {
+                Some(self.lock.lock())
+            },
+            // SAFETY: If this is not frozen, `_lock_guard` holds the lock to the `UnsafeCell` so
+            // this has shared access until the `FreezeReadGuard` is dropped. If this is frozen,
+            // the data cannot be modified and shared access is sound.
+            data: unsafe { &*self.data.get() },
+        }
+    }
+
+    #[inline]
+    #[track_caller]
+    pub fn write(&self) -> FreezeWriteGuard<'_, T> {
+        let _lock_guard = self.lock.lock();
+        assert!(!self.frozen.load(Ordering::Relaxed), "still mutable");
+        FreezeWriteGuard {
+            _lock_guard,
+            // SAFETY: `_lock_guard` holds the lock to the `UnsafeCell` so this has mutable access
+            // until the `FreezeWriteGuard` is dropped.
+            data: unsafe { &mut *self.data.get() },
+        }
+    }
+
+    #[inline]
+    pub fn freeze(&self) -> &T {
+        if !self.frozen.load(Ordering::Acquire) {
+            // Get the lock to ensure no concurrent writes.
+            let _lock = self.lock.lock();
+            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 `Freeze` which is in a locked state or frozen.
+#[must_use = "if unused the Freeze may immediately unlock"]
+pub struct FreezeReadGuard<'a, T> {
+    _lock_guard: Option<LockGuard<'a, ()>>,
+    data: &'a T,
+}
+
+impl<'a, T: 'a> Deref for FreezeReadGuard<'a, T> {
+    type Target = T;
+    #[inline]
+    fn deref(&self) -> &T {
+        self.data
+    }
+}
+
+/// A guard holding mutable access to a `Freeze` which is in a locked state or frozen.
+#[must_use = "if unused the Freeze may immediately unlock"]
+pub struct FreezeWriteGuard<'a, T> {
+    _lock_guard: LockGuard<'a, ()>,
+    data: &'a mut T,
+}
+
+impl<'a, T: 'a> Deref for FreezeWriteGuard<'a, T> {
+    type Target = T;
+    #[inline]
+    fn deref(&self) -> &T {
+        self.data
+    }
+}
+
+impl<'a, T: 'a> DerefMut for FreezeWriteGuard<'a, T> {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut T {
+        self.data
+    }
+}

compiler/rustc_interface/src/queries.rs

@@ -7,7 +7,7 @@ use rustc_codegen_ssa::traits::CodegenBackend;
 use rustc_codegen_ssa::CodegenResults;
 use rustc_data_structures::steal::Steal;
 use rustc_data_structures::svh::Svh;
-use rustc_data_structures::sync::{AppendOnlyIndexVec, Lrc, OnceLock, RwLock, WorkerLocal};
+use rustc_data_structures::sync::{AppendOnlyIndexVec, Freeze, Lrc, OnceLock, RwLock, WorkerLocal};
 use rustc_hir::def_id::{StableCrateId, CRATE_DEF_ID, LOCAL_CRATE};
 use rustc_hir::definitions::Definitions;
 use rustc_incremental::DepGraphFuture;
@@ -197,7 +197,7 @@ impl<'tcx> Queries<'tcx> {
                 self.codegen_backend().metadata_loader(),
                 stable_crate_id,
             )) as _);
-            let definitions = RwLock::new(Definitions::new(stable_crate_id));
+            let definitions = Freeze::new(Definitions::new(stable_crate_id));
             let source_span = AppendOnlyIndexVec::new();
             let _id = source_span.push(krate.spans.inner_span);
             debug_assert_eq!(_id, CRATE_DEF_ID);

compiler/rustc_middle/src/hir/map/mod.rs

@@ -1207,7 +1207,7 @@ pub(super) fn crate_hash(tcx: TyCtxt<'_>, _: LocalCrate) -> Svh {
         source_file_names.hash_stable(&mut hcx, &mut stable_hasher);
         debugger_visualizers.hash_stable(&mut hcx, &mut stable_hasher);
         if tcx.sess.opts.incremental_relative_spans() {
-            let definitions = tcx.definitions_untracked();
+            let definitions = tcx.untracked().definitions.freeze();
             let mut owner_spans: Vec<_> = krate
                 .owners
                 .iter_enumerated()

compiler/rustc_middle/src/ty/context.rs

@@ -39,7 +39,9 @@ use rustc_data_structures::profiling::SelfProfilerRef;
 use rustc_data_structures::sharded::{IntoPointer, ShardedHashMap};
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::steal::Steal;
-use rustc_data_structures::sync::{self, Lock, Lrc, MappedReadGuard, ReadGuard, WorkerLocal};
+use rustc_data_structures::sync::{
+    self, FreezeReadGuard, Lock, Lrc, MappedReadGuard, ReadGuard, WorkerLocal,
+};
 use rustc_data_structures::unord::UnordSet;
 use rustc_errors::{
     DecorateLint, DiagnosticBuilder, DiagnosticMessage, ErrorGuaranteed, MultiSpan,
@@ -964,8 +966,8 @@ impl<'tcx> TyCtxt<'tcx> {
                 i += 1;
             }
 
-            // Leak a read lock once we finish iterating on definitions, to prevent adding new ones.
-            definitions.leak();
+            // Freeze definitions once we finish iterating on them, to prevent adding new ones.
+            definitions.freeze();
         })
     }
 
@@ -974,10 +976,9 @@ impl<'tcx> TyCtxt<'tcx> {
         // definitions change.
         self.dep_graph.read_index(DepNodeIndex::FOREVER_RED_NODE);
 
-        // Leak a read lock once we start iterating on definitions, to prevent adding new ones
+        // Freeze definitions once we start iterating on them, to prevent adding new ones
         // while iterating. If some query needs to add definitions, it should be `ensure`d above.
-        let definitions = self.untracked.definitions.leak();
-        definitions.def_path_table()
+        self.untracked.definitions.freeze().def_path_table()
     }
 
     pub fn def_path_hash_to_def_index_map(
@@ -986,10 +987,9 @@ impl<'tcx> TyCtxt<'tcx> {
         // Create a dependency to the crate to be sure we re-execute this when the amount of
         // definitions change.
         self.ensure().hir_crate(());
-        // Leak a read lock once we start iterating on definitions, to prevent adding new ones
+        // Freeze definitions once we start iterating on them, to prevent adding new ones
         // while iterating. If some query needs to add definitions, it should be `ensure`d above.
-        let definitions = self.untracked.definitions.leak();
-        definitions.def_path_hash_to_def_index_map()
+        self.untracked.definitions.freeze().def_path_hash_to_def_index_map()
     }
 
     /// Note that this is *untracked* and should only be used within the query
@@ -1006,7 +1006,7 @@ impl<'tcx> TyCtxt<'tcx> {
     /// Note that this is *untracked* and should only be used within the query
     /// system if the result is otherwise tracked through queries
     #[inline]
-    pub fn definitions_untracked(self) -> ReadGuard<'tcx, Definitions> {
+    pub fn definitions_untracked(self) -> FreezeReadGuard<'tcx, Definitions> {
         self.untracked.definitions.read()
     }
 

compiler/rustc_session/src/cstore.rs

@@ -7,7 +7,7 @@ use crate::utils::NativeLibKind;
 use crate::Session;
 use rustc_ast as ast;
 use rustc_data_structures::owned_slice::OwnedSlice;
-use rustc_data_structures::sync::{self, AppendOnlyIndexVec, RwLock};
+use rustc_data_structures::sync::{self, AppendOnlyIndexVec, Freeze, RwLock};
 use rustc_hir::def_id::{CrateNum, DefId, LocalDefId, StableCrateId, LOCAL_CRATE};
 use rustc_hir::definitions::{DefKey, DefPath, DefPathHash, Definitions};
 use rustc_span::hygiene::{ExpnHash, ExpnId};
@@ -261,5 +261,5 @@ pub struct Untracked {
     pub cstore: RwLock<Box<CrateStoreDyn>>,
     /// Reference span for definitions.
     pub source_span: AppendOnlyIndexVec<LocalDefId, Span>,
-    pub definitions: RwLock<Definitions>,
+    pub definitions: Freeze<Definitions>,
 }