Add UnsafeCell guarded ref and mut access

src/cell/mod.rs

@@ -2,4 +2,4 @@
 
 mod unsafe_cell;
 
-pub use self::unsafe_cell::UnsafeCell;
+pub use self::unsafe_cell::{UnsafeCell, UnsafeCellMutToken, UnsafeCellRefToken};

src/cell/unsafe_cell.rs

@@ -12,6 +12,24 @@ pub struct UnsafeCell<T: ?Sized> {
     data: std::cell::UnsafeCell<T>,
 }
 
+/// Token for tracking immutable access to the wrapped value.
+///
+/// This token must be held for as long as the access lasts.
+#[derive(Debug)]
+pub struct UnsafeCellRefToken<T: ?Sized> {
+    data: *const T,
+    token: rt::CellRefToken,
+}
+
+/// Token for tracking mutable access to the wrapped value.
+///
+/// This token must be held for as long as the access lasts.
+#[derive(Debug)]
+pub struct UnsafeCellMutToken<T: ?Sized> {
+    data: *mut T,
+    token: rt::CellMutToken,
+}
+
 impl<T> UnsafeCell<T> {
     /// Constructs a new instance of `UnsafeCell` which will wrap the specified value.
     #[cfg_attr(loom_nightly, track_caller)]
@@ -54,6 +72,36 @@ impl<T: ?Sized> UnsafeCell<T> {
     {
         self.state.with_mut(location!(), || f(self.data.get()))
     }
+
+    /// Get an `UnsafeCellRefToken` that provides scoped immutable access to the
+    /// wrapped value.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the access is not valid under the Rust memory
+    /// model.
+    pub unsafe fn guard_ref_access(&self) -> UnsafeCellRefToken<T> {
+        let token = self.state.guard_ref_access(location!());
+        UnsafeCellRefToken {
+            data: self.data.get() as *const T,
+            token,
+        }
+    }
+
+    /// Get an `UnsafeCellMutToken` that provides scoped mutable access to the
+    /// wrapped value.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the access is not valid under the Rust memory
+    /// model.
+    pub unsafe fn guard_mut_access(&self) -> UnsafeCellMutToken<T> {
+        let token = self.state.guard_mut_access(location!());
+        UnsafeCellMutToken {
+            data: self.data.get(),
+            token,
+        }
+    }
 }
 
 impl<T: Default> Default for UnsafeCell<T> {
@@ -67,3 +115,17 @@ impl<T> From<T> for UnsafeCell<T> {
         UnsafeCell::new(src)
     }
 }
+
+impl<T: ?Sized> UnsafeCellRefToken<T> {
+    /// Get an immutable pointer to the wrapped value.
+    pub const fn get(&self) -> *const T {
+        self.data
+    }
+}
+
+impl<T: ?Sized> UnsafeCellMutToken<T> {
+    /// Get a mutable pointer to the wrapped value.
+    pub const fn get(&self) -> *mut T {
+        self.data
+    }
+}

src/rt/cell.rs

@@ -7,6 +7,18 @@ pub(crate) struct Cell {
     state: object::Ref<State>,
 }
 
+/// Tracks immutable access to `Cell`
+#[derive(Debug)]
+pub(crate) struct CellRefToken {
+    state: object::Ref<State>,
+}
+
+/// Tracks mutable access to `Cell`
+#[derive(Debug)]
+pub(crate) struct CellMutToken {
+    state: object::Ref<State>,
+}
+
 #[derive(Debug)]
 pub(super) struct State {
     /// Where the cell was created
@@ -43,78 +55,94 @@ impl Cell {
     }
 
     pub(crate) fn with<R>(&self, location: Location, f: impl FnOnce() -> R) -> R {
-        struct Reset {
-            state: object::Ref<State>,
-        }
+        // Enter the read closure
+        let _token = CellRefToken::new(self.state, location);
 
-        impl Drop for Reset {
-            fn drop(&mut self) {
-                rt::execution(|execution| {
-                    let state = self.state.get_mut(&mut execution.objects);
+        f()
+    }
 
-                    assert!(state.is_reading > 0);
-                    assert!(!state.is_writing);
+    pub(crate) fn with_mut<R>(&self, location: Location, f: impl FnOnce() -> R) -> R {
+        // Enter the write closure
+        let _token = CellMutToken::new(self.state, location);
 
-                    state.is_reading -= 1;
+        f()
+    }
 
-                    if !std::thread::panicking() {
-                        state.track_read(&execution.threads);
-                    }
-                })
-            }
-        }
+    pub(crate) fn guard_ref_access(&self, location: Location) -> CellRefToken {
+        CellRefToken::new(self.state, location)
+    }
 
-        // Enter the read closure
-        let _reset = rt::synchronize(|execution| {
-            let state = self.state.get_mut(&mut execution.objects);
+    pub(crate) fn guard_mut_access(&self, location: Location) -> CellMutToken {
+        CellMutToken::new(self.state, location)
+    }
+}
 
-            assert!(!state.is_writing, "currently writing to cell");
+impl CellRefToken {
+    fn new(state: object::Ref<State>, location: Location) -> Self {
+        rt::synchronize(|execution| {
+            let state_obj = state.get_mut(&mut execution.objects);
 
-            state.is_reading += 1;
-            state.read_locations.track(location, &execution.threads);
-            state.track_read(&execution.threads);
+            assert!(!state_obj.is_writing, "currently writing to cell");
 
-            Reset { state: self.state }
-        });
+            state_obj.is_reading += 1;
+            state_obj.read_locations.track(location, &execution.threads);
+            state_obj.track_read(&execution.threads);
 
-        f()
+            Self { state }
+        })
     }
+}
 
-    pub(crate) fn with_mut<R>(&self, location: Location, f: impl FnOnce() -> R) -> R {
-        struct Reset(object::Ref<State>);
-
-        impl Drop for Reset {
-            fn drop(&mut self) {
-                rt::execution(|execution| {
-                    let state = self.0.get_mut(&mut execution.objects);
+impl Drop for CellRefToken {
+    fn drop(&mut self) {
+        rt::synchronize(|execution| {
+            let state_obj = self.state.get_mut(&mut execution.objects);
 
-                    assert!(state.is_writing);
-                    assert!(state.is_reading == 0);
+            assert!(state_obj.is_reading > 0);
+            assert!(!state_obj.is_writing);
 
-                    state.is_writing = false;
+            state_obj.is_reading -= 1;
 
-                    if !std::thread::panicking() {
-                        state.track_write(&execution.threads);
-                    }
-                })
+            if !std::thread::panicking() {
+                state_obj.track_read(&execution.threads);
             }
-        }
+        })
+    }
+}
 
-        // Enter the read closure
-        let _reset = rt::synchronize(|execution| {
-            let state = self.state.get_mut(&mut execution.objects);
+impl CellMutToken {
+    fn new(state: object::Ref<State>, location: Location) -> Self {
+        rt::execution(|execution| {
+            let state_obj = state.get_mut(&mut execution.objects);
 
-            assert!(state.is_reading == 0, "currently reading from cell");
-            assert!(!state.is_writing, "currently writing to cell");
+            assert!(state_obj.is_reading == 0, "currently reading from cell");
+            assert!(!state_obj.is_writing, "currently writing to cell");
 
-            state.is_writing = true;
-            state.write_locations.track(location, &execution.threads);
-            state.track_write(&execution.threads);
+            state_obj.is_writing = true;
+            state_obj
+                .write_locations
+                .track(location, &execution.threads);
+            state_obj.track_write(&execution.threads);
 
-            Reset(self.state)
-        });
+            Self { state }
+        })
+    }
+}
 
-        f()
+impl Drop for CellMutToken {
+    fn drop(&mut self) {
+        rt::execution(|execution| {
+            let state_obj = self.state.get_mut(&mut execution.objects);
+
+            assert!(state_obj.is_writing);
+            assert!(state_obj.is_reading == 0);
+
+            state_obj.is_writing = false;
+
+            if !std::thread::panicking() {
+                state_obj.track_write(&execution.threads);
+            }
+        })
     }
 }
 

src/rt/mod.rs

@@ -15,7 +15,7 @@ mod location;
 pub(crate) use self::location::Location;
 
 mod cell;
-pub(crate) use self::cell::Cell;
+pub(crate) use self::cell::{Cell, CellMutToken, CellRefToken};
 
 mod condvar;
 pub(crate) use self::condvar::Condvar;