Use `FreezeLock` for `CStore`

8 files changed, 63 insertions, 38 deletions

diff --git a/compiler/rustc_data_structures/src/sync/freeze.rs b/compiler/rustc_data_structures/src/sync/freeze.rs
index 58ab91237f4..466c44f59bb 100644
--- a/compiler/rustc_data_structures/src/sync/freeze.rs
+++ b/compiler/rustc_data_structures/src/sync/freeze.rs
@@ -6,6 +6,7 @@ use std::{
     intrinsics::likely,
     marker::PhantomData,
     ops::{Deref, DerefMut},
+    ptr::NonNull,
     sync::atomic::Ordering,
 };
 
@@ -79,7 +80,7 @@ impl<T> FreezeLock<T> {
             } else {
                 Some(self.lock.read())
             },
-            lock: self,
+            data: unsafe { NonNull::new_unchecked(self.data.get()) },
         }
     }
 
@@ -101,7 +102,12 @@ impl<T> FreezeLock<T> {
         if self.frozen.load(Ordering::Relaxed) {
             None
         } else {
-            Some(FreezeWriteGuard { _lock_guard, lock: self, marker: PhantomData })
+            Some(FreezeWriteGuard {
+                _lock_guard,
+                data: unsafe { NonNull::new_unchecked(self.data.get()) },
+                frozen: &self.frozen,
+                marker: PhantomData,
+            })
         }
     }
 
@@ -120,52 +126,75 @@ impl<T> FreezeLock<T> {
 
 /// 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> {
+pub struct FreezeReadGuard<'a, T: ?Sized> {
     _lock_guard: Option<ReadGuard<'a, ()>>,
-    lock: &'a FreezeLock<T>,
+    data: NonNull<T>,
 }
 
-impl<'a, T: 'a> Deref for FreezeReadGuard<'a, T> {
+impl<'a, T: ?Sized + '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,
+        // 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.lock.data.get() }
+        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> {
+pub struct FreezeWriteGuard<'a, T: ?Sized> {
     _lock_guard: WriteGuard<'a, ()>,
-    lock: &'a FreezeLock<T>,
+    frozen: &'a AtomicBool,
+    data: NonNull<T>,
     marker: PhantomData<&'a mut T>,
 }
 
 impl<'a, T> FreezeWriteGuard<'a, T> {
     pub fn freeze(self) -> &'a T {
-        self.lock.frozen.store(true, Ordering::Release);
+        self.frozen.store(true, Ordering::Release);
 
         // SAFETY: This is frozen so the data cannot be modified and shared access is sound.
-        unsafe { &*self.lock.data.get() }
+        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: 'a> Deref for FreezeWriteGuard<'a, T> {
+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.lock.data.get() }
+        unsafe { &*self.data.as_ptr() }
     }
 }
 
-impl<'a, T: 'a> DerefMut for FreezeWriteGuard<'a, T> {
+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.lock.data.get() }
+        unsafe { &mut *self.data.as_ptr() }
     }
 }
diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs
index 9fcaf643179..e5ae6d5b5d6 100644
--- a/compiler/rustc_interface/src/passes.rs
+++ b/compiler/rustc_interface/src/passes.rs
@@ -584,7 +584,7 @@ fn resolver_for_lowering<'tcx>(
     let krate = configure_and_expand(krate, &pre_configured_attrs, &mut resolver);
 
     // Make sure we don't mutate the cstore from here on.
-    tcx.untracked().cstore.leak();
+    tcx.untracked().cstore.freeze();
 
     let ty::ResolverOutputs {
         global_ctxt: untracked_resolutions,
diff --git a/compiler/rustc_interface/src/queries.rs b/compiler/rustc_interface/src/queries.rs
index e0d9998d919..29a5837eb7c 100644
--- a/compiler/rustc_interface/src/queries.rs
+++ b/compiler/rustc_interface/src/queries.rs
@@ -7,9 +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, FreezeLock, Lrc, OnceLock, RwLock, WorkerLocal,
-};
+use rustc_data_structures::sync::{AppendOnlyIndexVec, FreezeLock, Lrc, OnceLock, WorkerLocal};
 use rustc_hir::def_id::{StableCrateId, CRATE_DEF_ID, LOCAL_CRATE};
 use rustc_hir::definitions::Definitions;
 use rustc_incremental::DepGraphFuture;
@@ -195,7 +193,7 @@ impl<'tcx> Queries<'tcx> {
                 self.compiler.register_lints.as_deref(),
                 &pre_configured_attrs,
             ));
-            let cstore = RwLock::new(Box::new(CStore::new(
+            let cstore = FreezeLock::new(Box::new(CStore::new(
                 self.codegen_backend().metadata_loader(),
                 stable_crate_id,
             )) as _);
diff --git a/compiler/rustc_metadata/src/creader.rs b/compiler/rustc_metadata/src/creader.rs
index fce80ab37dd..69221475356 100644
--- a/compiler/rustc_metadata/src/creader.rs
+++ b/compiler/rustc_metadata/src/creader.rs
@@ -8,7 +8,7 @@ use rustc_ast::expand::allocator::{alloc_error_handler_name, global_fn_name, All
 use rustc_ast::{self as ast, *};
 use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::svh::Svh;
-use rustc_data_structures::sync::{MappedReadGuard, MappedWriteGuard, ReadGuard, WriteGuard};
+use rustc_data_structures::sync::{FreezeReadGuard, FreezeWriteGuard};
 use rustc_expand::base::SyntaxExtension;
 use rustc_hir::def_id::{CrateNum, LocalDefId, StableCrateId, StableCrateIdMap, LOCAL_CRATE};
 use rustc_hir::definitions::Definitions;
@@ -134,14 +134,14 @@ impl<'a> std::fmt::Debug for CrateDump<'a> {
 }
 
 impl CStore {
-    pub fn from_tcx(tcx: TyCtxt<'_>) -> MappedReadGuard<'_, CStore> {
-        ReadGuard::map(tcx.untracked().cstore.read(), |cstore| {
+    pub fn from_tcx(tcx: TyCtxt<'_>) -> FreezeReadGuard<'_, CStore> {
+        FreezeReadGuard::map(tcx.untracked().cstore.read(), |cstore| {
             cstore.as_any().downcast_ref::<CStore>().expect("`tcx.cstore` is not a `CStore`")
         })
     }
 
-    pub fn from_tcx_mut(tcx: TyCtxt<'_>) -> MappedWriteGuard<'_, CStore> {
-        WriteGuard::map(tcx.untracked().cstore.write(), |cstore| {
+    pub fn from_tcx_mut(tcx: TyCtxt<'_>) -> FreezeWriteGuard<'_, CStore> {
+        FreezeWriteGuard::map(tcx.untracked().cstore.write(), |cstore| {
             cstore.untracked_as_any().downcast_mut().expect("`tcx.cstore` is not a `CStore`")
         })
     }
diff --git a/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs b/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
index aeda8af6d2c..f964a8c76c0 100644
--- a/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
+++ b/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
@@ -131,7 +131,7 @@ macro_rules! provide_one {
                 $tcx.ensure().crate_hash($def_id.krate);
             }
 
-            let cdata = rustc_data_structures::sync::MappedReadGuard::map(CStore::from_tcx($tcx), |c| {
+            let cdata = rustc_data_structures::sync::FreezeReadGuard::map(CStore::from_tcx($tcx), |c| {
                 c.get_crate_data($def_id.krate).cdata
             });
             let $cdata = crate::creader::CrateMetadataRef {
@@ -510,7 +510,7 @@ pub(in crate::rmeta) fn provide(providers: &mut Providers) {
         crates: |tcx, ()| {
             // The list of loaded crates is now frozen in query cache,
             // so make sure cstore is not mutably accessed from here on.
-            tcx.untracked().cstore.leak();
+            tcx.untracked().cstore.freeze();
             tcx.arena.alloc_from_iter(CStore::from_tcx(tcx).iter_crate_data().map(|(cnum, _)| cnum))
         },
         ..*providers
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index 90d847804f0..eeb87d5561d 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/compiler/rustc_middle/src/ty/context.rs
@@ -39,9 +39,7 @@ 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, FreezeReadGuard, Lock, Lrc, MappedReadGuard, ReadGuard, WorkerLocal,
-};
+use rustc_data_structures::sync::{self, FreezeReadGuard, Lock, Lrc, WorkerLocal};
 use rustc_data_structures::unord::UnordSet;
 use rustc_errors::{
     DecorateLint, DiagnosticBuilder, DiagnosticMessage, ErrorGuaranteed, MultiSpan,
@@ -995,8 +993,8 @@ 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 cstore_untracked(self) -> MappedReadGuard<'tcx, CrateStoreDyn> {
-        ReadGuard::map(self.untracked.cstore.read(), |c| &**c)
+    pub fn cstore_untracked(self) -> FreezeReadGuard<'tcx, CrateStoreDyn> {
+        FreezeReadGuard::map(self.untracked.cstore.read(), |c| &**c)
     }
 
     /// Give out access to the untracked data without any sanity checks.
diff --git a/compiler/rustc_resolve/src/lib.rs b/compiler/rustc_resolve/src/lib.rs
index b757f42eaa0..a4175de464c 100644
--- a/compiler/rustc_resolve/src/lib.rs
+++ b/compiler/rustc_resolve/src/lib.rs
@@ -34,7 +34,7 @@ use rustc_ast::{AngleBracketedArg, Crate, Expr, ExprKind, GenericArg, GenericArg
 use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexMap, FxIndexSet};
 use rustc_data_structures::intern::Interned;
 use rustc_data_structures::steal::Steal;
-use rustc_data_structures::sync::{Lrc, MappedReadGuard};
+use rustc_data_structures::sync::{FreezeReadGuard, Lrc};
 use rustc_errors::{
     Applicability, DiagnosticBuilder, DiagnosticMessage, ErrorGuaranteed, SubdiagnosticMessage,
 };
@@ -1544,7 +1544,7 @@ impl<'a, 'tcx> Resolver<'a, 'tcx> {
         ))
     }
 
-    fn cstore(&self) -> MappedReadGuard<'_, CStore> {
+    fn cstore(&self) -> FreezeReadGuard<'_, CStore> {
         CStore::from_tcx(self.tcx)
     }
 
@@ -1599,7 +1599,7 @@ impl<'a, 'tcx> Resolver<'a, 'tcx> {
         });
 
         // Make sure we don't mutate the cstore from here on.
-        self.tcx.untracked().cstore.leak();
+        self.tcx.untracked().cstore.freeze();
     }
 
     fn traits_in_scope(
diff --git a/compiler/rustc_session/src/cstore.rs b/compiler/rustc_session/src/cstore.rs
index 90f57be9324..d816842b02b 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, FreezeLock, RwLock};
+use rustc_data_structures::sync::{self, AppendOnlyIndexVec, FreezeLock};
 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};
@@ -258,7 +258,7 @@ pub trait CrateStore: std::fmt::Debug {
 pub type CrateStoreDyn = dyn CrateStore + sync::DynSync + sync::DynSend;
 
 pub struct Untracked {
-    pub cstore: RwLock<Box<CrateStoreDyn>>,
+    pub cstore: FreezeLock<Box<CrateStoreDyn>>,
     /// Reference span for definitions.
     pub source_span: AppendOnlyIndexVec<LocalDefId, Span>,
     pub definitions: FreezeLock<Definitions>,