manual borrow counting in CmRefCell so we don't need to match in CmRef::deref

Author: LorrensP-2158466

compiler/rustc_resolve/src/lib.rs

@@ -18,10 +18,10 @@
 #![feature(option_into_flat_iter)]
 #![feature(rustc_attrs)]
 #![feature(trim_prefix_suffix)]
+#![feature(unsafe_cell_access)]
 #![recursion_limit = "256"]
 // tidy-alphabetical-end
 
-use std::cell::RefMut;
 use std::collections::BTreeSet;
 use std::ops::ControlFlow;
 use std::sync::{Arc, OnceLock};
@@ -81,7 +81,7 @@ use crate::diagnostics::impls::{
     ImportSuggestion, LabelSuggestion, OnUnknownData, StructCtor, Suggestion,
 };
 use crate::imports::{ImportResolution, NameResolutionRef};
-use crate::ref_mut::{CmCell, CmRef, CmRefCell};
+use crate::ref_mut::{CmCell, CmRef, CmRefCell, RefMut};
 
 mod build_reduced_graph;
 mod check_unused;
@@ -2182,9 +2182,9 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
         match &module.0.0.lazy_resolutions {
             Resolutions::Local(local_res) => local_res.borrow(self),
             Resolutions::Extern(extern_res) => {
-                // It is fine to return a `CmRef::Untracked`, we never give out a `&mut`
+                // It is fine to return a `CmRef::untracked`, we never give out a `&mut`
                 // to an external table.
-                CmRef::Untracked(
+                CmRef::untracked(
                     // As long as 1 thread is building this external table, all other threads will wait.
                     extern_res
                         .get_or_init(|| self.build_reduced_graph_external(module.expect_extern())),
@@ -2831,9 +2831,11 @@ type CmResolver<'r, 'ra, 'tcx> = ref_mut::RefOrMut<'r, Resolver<'ra, 'tcx>>;
 use std::cell::{Cell as CacheCell, RefCell as CacheRefCell};
 
 mod ref_mut {
-    use std::cell::{BorrowMutError, Cell, Ref, RefCell, RefMut};
+    use std::cell::{Cell, UnsafeCell};
     use std::fmt;
-    use std::ops::Deref;
+    use std::marker::PhantomData;
+    use std::ops::{Deref, DerefMut};
+    use std::ptr::NonNull;
 
     use crate::Resolver;
 
@@ -2946,31 +2948,82 @@ mod ref_mut {
         }
     }
 
-    pub(crate) enum CmRef<'b, T> {
-        /// A tracked borrow of a [`CmRefCell`]
-        Tracked(Ref<'b, T>),
-        /// An untracked or normal reference (not dynamically borrow-checked by `RefCell`)
-        Untracked(&'b T),
+    const UNUSED: isize = 0;
+
+    pub(crate) struct CmRef<'b, T> {
+        value: NonNull<T>,
+        borrow: Option<&'b Cell<isize>>,
+        marker: PhantomData<&'b T>,
+    }
+
+    impl<'b, T> CmRef<'b, T> {
+        pub(crate) fn tracked(value: &'b T, borrow: &'b Cell<isize>) -> Self {
+            Self { value: NonNull::from_ref(value), borrow: Some(borrow), marker: PhantomData }
+        }
+
+        pub(crate) fn untracked(value: &'b T) -> Self {
+            Self { value: NonNull::from_ref(value), borrow: None, marker: PhantomData }
+        }
+    }
+
+    impl<'b, T> Drop for CmRef<'b, T> {
+        fn drop(&mut self) {
+            if let Some(borrow) = self.borrow {
+                borrow.update(|b| b - 1);
+            }
+        }
     }
 
     impl<'b, T> Deref for CmRef<'b, T> {
         type Target = T;
 
         fn deref(&self) -> &Self::Target {
-            match self {
-                CmRef::Tracked(r) => r,
-                CmRef::Untracked(r) => r,
-            }
+            // SAFETY: This type is proof we have shared access
+            unsafe { self.value.as_ref() }
+        }
+    }
+
+    pub(crate) struct RefMut<'b, T> {
+        value: NonNull<T>,
+        borrow: &'b Cell<isize>,
+        marker: PhantomData<&'b mut T>,
+    }
+
+    impl<'b, T> Drop for RefMut<'b, T> {
+        fn drop(&mut self) {
+            self.borrow.update(|b| {
+                debug_assert!(b < UNUSED, "`RefMut` exists but no borrow is counted");
+                UNUSED // `RefMut` does not impl clone, so UNUSED is enough here
+            });
+        }
+    }
+
+    impl<'b, T> Deref for RefMut<'b, T> {
+        type Target = T;
+
+        fn deref(&self) -> &Self::Target {
+            // SAFETY: This type is proof we have exclusive access
+            unsafe { self.value.as_ref() }
+        }
+    }
+
+    impl<'b, T> DerefMut for RefMut<'b, T> {
+        fn deref_mut(&mut self) -> &mut Self::Target {
+            // SAFETY: This type is proof we have exclusive access
+            unsafe { self.value.as_mut() }
         }
     }
 
     /// A wrapper around a [`RefCell`] that only allows writes (mutable borrows) based on a condition in the resolver.
     #[derive(Default)]
-    pub(crate) struct CmRefCell<T>(RefCell<T>);
+    pub(crate) struct CmRefCell<T> {
+        value: UnsafeCell<T>,
+        borrow: Cell<isize>,
+    }
 
     impl<T> CmRefCell<T> {
         pub(crate) fn new(value: T) -> CmRefCell<T> {
-            CmRefCell(RefCell::new(value))
+            CmRefCell { value: UnsafeCell::new(value), borrow: Cell::new(UNUSED) }
         }
 
         #[track_caller]
@@ -2982,37 +3035,76 @@ mod ref_mut {
         pub(crate) fn try_borrow_mut<'ra, 'tcx>(
             &self,
             r: &Resolver<'ra, 'tcx>,
-        ) -> Result<RefMut<'_, T>, BorrowMutError> {
+        ) -> Result<RefMut<'_, T>, ()> {
             if r.assert_speculative {
                 panic!("not allowed to mutably borrow a `CmRefCell` during speculative resolution");
             }
-            self.0.try_borrow_mut()
+            // SAFETY: we just checked that we are not in speculative resolution
+            unsafe { self.try_borrow_mut_inner() }
+        }
+
+        /// # SAFETY
+        ///
+        /// Caller must guarentee that we are not in speculative resolution
+        unsafe fn try_borrow_mut_inner(&self) -> Result<RefMut<'_, T>, ()> {
+            let borrow = self.borrow.get();
+            if borrow != UNUSED {
+                return Err(());
+            }
+            self.borrow.update(|_| UNUSED - 1);
+            // SAFETY: We can take exclusive access since this is unused.
+            let value = unsafe { self.value.as_mut_unchecked() };
+            Ok(RefMut {
+                value: NonNull::from_mut(value),
+                borrow: &self.borrow,
+                marker: PhantomData,
+            })
         }
 
         #[track_caller]
-        pub(crate) fn borrow_mut_with_token(&self, _: token::CmToken<'_>) -> RefMut<'_, T> {
-            self.0.borrow_mut()
+        pub(crate) fn borrow_mut_with_token(&self, t: token::CmToken<'_>) -> RefMut<'_, T> {
+            self.try_borrow_mut_with_token(t).unwrap()
         }
 
         #[track_caller]
         pub(crate) fn try_borrow_mut_with_token(
             &self,
             _: token::CmToken<'_>,
-        ) -> Result<RefMut<'_, T>, BorrowMutError> {
-            self.0.try_borrow_mut()
+        ) -> Result<RefMut<'_, T>, ()> {
+            // SAFETY: the given token guarentees us we are not in speculative resolution.
+            unsafe { self.try_borrow_mut_inner() }
+        }
+
+        /// # SAFETY
+        ///
+        /// Caller must guarentee that we are not in speculative resolution
+        pub(crate) unsafe fn borrow_inner(&self) -> CmRef<'_, T> {
+            let borrow = self.borrow.get();
+            if borrow < UNUSED {
+                panic!("`CmRefCell` already mutably borrowed")
+            }
+            self.borrow.update(|b| b + 1);
+            // SAFETY: We can take shared access since this is not mutably borrowed.
+            let value = unsafe { self.value.as_ref_unchecked() };
+
+            CmRef::tracked(value, &self.borrow)
         }
 
         #[track_caller]
-        pub(crate) fn borrow_with_token(&self, _: token::CmToken<'_>) -> Ref<'_, T> {
-            self.0.borrow()
+        pub(crate) fn borrow_with_token(&self, _: token::CmToken<'_>) -> CmRef<'_, T> {
+            // SAFETY: token guarentees us that we are not in speculative resolution.
+            unsafe { self.borrow_inner() }
         }
 
         #[track_caller]
         pub(crate) fn borrow<'ra, 'tcx>(&self, r: &Resolver<'ra, 'tcx>) -> CmRef<'_, T> {
             if r.assert_speculative {
-                // `try_borrow_unguarded` is unsafe because it returns a `&T` instead
-                // of `Ref<'_, T>`. It does provides an extra check to make sure no live
-                // `RefMut`s are still alive, but the other way can not be checked, so:
+                if self.borrow.get() < UNUSED {
+                    panic!("`CmRefCell` already mutably borrowed");
+                }
+                // mimics `RefCell::try_borrow_unguarded`, which is unsafe because it returns a
+                // `&T` instead of `Ref<'_, T>`. It does provides an extra check to make sure no
+                // live `RefMut`s are still alive, but the other way can not be checked, so:
                 //
                 // SAFETY: This is only safe because we know that every `Untracked` borrow
                 // is only created during the import resolutions phase:
@@ -3027,10 +3119,11 @@ mod ref_mut {
                 //
                 // We know that none of these `Untracked` borrows are alive after the import
                 // resolution phase is done (`assert_speculative = false`); so we deem this "safe".
-                let unguarded_borrow = unsafe { self.0.try_borrow_unguarded().unwrap() };
-                CmRef::Untracked(unguarded_borrow)
+                let unguarded_borrow = unsafe { self.value.as_ref_unchecked() };
+                CmRef::untracked(unguarded_borrow)
             } else {
-                CmRef::Tracked(self.0.borrow())
+                // SAFETY: we just checked that we are not in speculative resolution.
+                unsafe { self.borrow_inner() }
             }
         }
     }
@@ -3040,7 +3133,7 @@ mod ref_mut {
             if r.assert_speculative {
                 panic!("not allowed to mutate a CmRefCell during speculative resolution");
             }
-            self.0.take()
+            std::mem::take(&mut *self.borrow_mut(r))
         }
     }
 }