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.

9 files changed, 148 insertions, 32 deletions

diff --git a/compiler/rustc_ast_lowering/src/index.rs b/compiler/rustc_ast_lowering/src/index.rs
index ce847906fb9..eff362f3ff0 100644
--- a/compiler/rustc_ast_lowering/src/index.rs
+++ b/compiler/rustc_ast_lowering/src/index.rs
@@ -2,19 +2,17 @@ use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sorted_map::SortedMap;
 use rustc_hir as hir;
 use rustc_hir::def_id::LocalDefId;
-use rustc_hir::definitions;
 use rustc_hir::intravisit::{self, Visitor};
 use rustc_hir::*;
 use rustc_index::{Idx, IndexVec};
 use rustc_middle::span_bug;
-use rustc_session::Session;
-use rustc_span::source_map::SourceMap;
+use rustc_middle::ty::TyCtxt;
 use rustc_span::{Span, DUMMY_SP};
 
 /// A visitor that walks over the HIR and collects `Node`s into a HIR map.
 pub(super) struct NodeCollector<'a, 'hir> {
-    /// Source map
-    source_map: &'a SourceMap,
+    tcx: TyCtxt<'hir>,
+
     bodies: &'a SortedMap<ItemLocalId, &'hir Body<'hir>>,
 
     /// Outputs
@@ -25,14 +23,11 @@ pub(super) struct NodeCollector<'a, 'hir> {
     parent_node: hir::ItemLocalId,
 
     owner: OwnerId,
-
-    definitions: &'a definitions::Definitions,
 }
 
-#[instrument(level = "debug", skip(sess, definitions, bodies))]
+#[instrument(level = "debug", skip(tcx, bodies))]
 pub(super) fn index_hir<'hir>(
-    sess: &Session,
-    definitions: &definitions::Definitions,
+    tcx: TyCtxt<'hir>,
     item: hir::OwnerNode<'hir>,
     bodies: &SortedMap<ItemLocalId, &'hir Body<'hir>>,
 ) -> (IndexVec<ItemLocalId, Option<ParentedNode<'hir>>>, FxHashMap<LocalDefId, ItemLocalId>) {
@@ -42,8 +37,7 @@ pub(super) fn index_hir<'hir>(
     // used.
     nodes.push(Some(ParentedNode { parent: ItemLocalId::INVALID, node: item.into() }));
     let mut collector = NodeCollector {
-        source_map: sess.source_map(),
-        definitions,
+        tcx,
         owner: item.def_id(),
         parent_node: ItemLocalId::new(0),
         nodes,
@@ -79,11 +73,17 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
                     span,
                     "inconsistent HirId at `{:?}` for `{:?}`: \
                      current_dep_node_owner={} ({:?}), hir_id.owner={} ({:?})",
-                    self.source_map.span_to_diagnostic_string(span),
+                    self.tcx.sess.source_map().span_to_diagnostic_string(span),
                     node,
-                    self.definitions.def_path(self.owner.def_id).to_string_no_crate_verbose(),
+                    self.tcx
+                        .definitions_untracked()
+                        .def_path(self.owner.def_id)
+                        .to_string_no_crate_verbose(),
                     self.owner,
-                    self.definitions.def_path(hir_id.owner.def_id).to_string_no_crate_verbose(),
+                    self.tcx
+                        .definitions_untracked()
+                        .def_path(hir_id.owner.def_id)
+                        .to_string_no_crate_verbose(),
                     hir_id.owner,
                 )
             }
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
index a6d1ef33f40..41db61d391a 100644
--- a/compiler/rustc_ast_lowering/src/lib.rs
+++ b/compiler/rustc_ast_lowering/src/lib.rs
@@ -671,8 +671,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         } else {
             (None, None)
         };
-        let (nodes, parenting) =
-            index::index_hir(self.tcx.sess, &*self.tcx.definitions_untracked(), node, &bodies);
+        let (nodes, parenting) = index::index_hir(self.tcx, node, &bodies);
         let nodes = hir::OwnerNodes { opt_hash_including_bodies, nodes, bodies };
         let attrs = hir::AttributeMap { map: attrs, opt_hash: attrs_hash };
 
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() }
+    }
+}
diff --git a/compiler/rustc_hir_analysis/src/lib.rs b/compiler/rustc_hir_analysis/src/lib.rs
index 4f95174f869..b0f333b79ca 100644
--- a/compiler/rustc_hir_analysis/src/lib.rs
+++ b/compiler/rustc_hir_analysis/src/lib.rs
@@ -237,6 +237,10 @@ pub fn check_crate(tcx: TyCtxt<'_>) -> Result<(), ErrorGuaranteed> {
         tcx.hir().for_each_module(|module| tcx.ensure().check_mod_item_types(module))
     });
 
+    // Freeze definitions as we don't add new ones at this point. This improves performance by
+    // allowing lock-free access to them.
+    tcx.untracked().definitions.freeze();
+
     // FIXME: Remove this when we implement creating `DefId`s
     // for anon constants during their parents' typeck.
     // Typeck all body owners in parallel will produce queries
diff --git a/compiler/rustc_interface/src/queries.rs b/compiler/rustc_interface/src/queries.rs
index 2db7aa0e367..e0d9998d919 100644
--- a/compiler/rustc_interface/src/queries.rs
+++ b/compiler/rustc_interface/src/queries.rs
@@ -7,7 +7,9 @@ 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, FreezeLock, 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 +199,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 = FreezeLock::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);
diff --git a/compiler/rustc_middle/src/hir/map/mod.rs b/compiler/rustc_middle/src/hir/map/mod.rs
index 3e8f07ed233..e20a202561f 100644
--- a/compiler/rustc_middle/src/hir/map/mod.rs
+++ b/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()
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index e2f9e299566..90d847804f0 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/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()
     }
 
diff --git a/compiler/rustc_session/src/cstore.rs b/compiler/rustc_session/src/cstore.rs
index c53a355b533..90f57be9324 100644
--- a/compiler/rustc_session/src/cstore.rs
+++ b/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, FreezeLock, 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: FreezeLock<Definitions>,
 }