tokio-rs · datdenkikniet · Mar 23, 2025 · datdenkikniet · Mar 23, 2025
diff --git a/src/cell/mod.rs b/src/cell/mod.rs
@@ -5,4 +5,4 @@ mod cell;
 mod unsafe_cell;
 
 pub use self::cell::Cell;
-pub use self::unsafe_cell::{ConstPtr, MutPtr, UnsafeCell};
+pub use self::unsafe_cell::{ConstPtr, MutPtr, MutRef, UnsafeCell};
diff --git a/src/cell/unsafe_cell.rs b/src/cell/unsafe_cell.rs
@@ -96,11 +96,26 @@ pub struct ConstPtr<T: ?Sized> {
 /// [here]: #correct-usage
 #[derive(Debug)]
 pub struct MutPtr<T: ?Sized> {
-    /// Drop guard representing the lifetime of the `ConstPtr`'s access.
+    /// Drop guard representing the lifetime of the `MutPtr`'s access.
     _guard: rt::cell::Writing,
     ptr: *mut T,
 }
 
+/// A checked mutable reference to an [`UnsafeCell`].
+///
+/// This type is essentially a [`&mut T`], but with the added ability to
+/// participate in Loom's [`UnsafeCell`] access tracking. While a `MutRef` to a
+/// given [`UnsafeCell`] exists, Loom will track that the [`UnsafeCell`] is
+/// being accessed mutably.
+///
+/// [`MutRef`]s are produced by [`UnsafeCell::as_mut`].
+#[derive(Debug)]
+pub struct MutRef<'a, T: ?Sized> {
+    /// Drop guard representing the lifetime of the `MutRef`'s access.
+    _guard: rt::cell::Writing,
+    ptr: &'a mut T,
+}
+
 impl<T> UnsafeCell<T> {
     /// Constructs a new instance of `UnsafeCell` which will wrap the specified value.
     #[track_caller]
@@ -200,6 +215,21 @@ impl<T: ?Sized> UnsafeCell<T> {
             ptr: self.data.get(),
         }
     }
+
+    /// Get a mutable reference to the wrapped value.
+    ///
+    /// This function returns a [`MutRef`] guard, which is analogous to a
+    /// `&mut T`, but tracked by Loom. As long as the returned `MutRef`
+    /// exists, Loom will consider the cell to be accessed mutably.
+    ///
+    /// This is analogous to [`core::cell::UnsafeCell::get_mut`].
+    #[track_caller]
+    pub fn as_mut(&mut self) -> MutRef<'_, T> {
+        MutRef {
+            _guard: self.state.start_write(location!()),
+            ptr: self.data.get_mut(),
+        }
+    }
 }
 
 impl<T: Default> Default for UnsafeCell<T> {
@@ -382,6 +412,20 @@ impl<T: ?Sized> ConstPtr<T> {
     }
 }
 
+impl<T: ?Sized> core::ops::Deref for MutRef<'_, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        self.ptr
+    }
+}
+
+impl<T: ?Sized> core::ops::DerefMut for MutRef<'_, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.ptr
+    }
+}
+
 impl<T: ?Sized> MutPtr<T> {
     /// Dereference the raw pointer.
     ///

diff --git a/tests/unsafe_cell.rs b/tests/unsafe_cell.rs
@@ -333,3 +333,12 @@ fn unsafe_cell_access_after_sync() {
         }
     });
 }
+
+#[test]
+fn as_mut_is_ergonomical_and_safe() {
+    loom::model(|| {
+        let mut value = UnsafeCell::new(-3i32);
+        let mut mut_ref = value.as_mut();
+        *mut_ref = mut_ref.abs();
+    });
+}