Refactor the implementation of `RcuOption` for clarity

ostd/src/sync/rcu/mod.rs

@@ -58,14 +58,10 @@ pub use owner_ptr::OwnerPtr;
 ///
 /// assert_eq!(*rcu_guard, Some(&43));
 /// ```
-#[repr(transparent)]
+pub struct Rcu<P: OwnerPtr>(RcuInner<P>);
-pub struct Rcu<P: OwnerPtr, const NULLABLE: bool = false> {
+
-    ptr: AtomicPtr<<P as OwnerPtr>::Target>,
+/// A guard that allows access to the pointed data protected by a [`Rcu`].
-    // We want to implement Send and Sync explicitly.
+pub struct RcuReadGuard<'a, P: OwnerPtr>(RcuReadGuardInner<'a, P>);
-    // Having a pointer field prevents them from being implemented
-    // automatically by the compiler.
-    _marker: PhantomData<*const P::Target>,
-}
 
 /// A Read-Copy Update (RCU) cell for sharing a _nullable_ pointer.  
 ///
@@ -78,13 +74,10 @@ pub struct Rcu<P: OwnerPtr, const NULLABLE: bool = false> {
 /// ```
 /// use ostd::sync::RcuOption;
 ///
-/// // Also allows lazy initialization.
 /// static RCU: RcuOption<Box<usize>> = RcuOption::new_none();
 ///
-/// // Not initialized yet.
+/// assert!(RCU.read().is_none());
-/// assert!(RCU.read().try_get().is_none());
 ///
-/// // Initialize the data protected by RCU.
 /// RCU.update(Box::new(42));
 ///
 /// // Read the data protected by RCU
@@ -103,27 +96,43 @@ pub struct Rcu<P: OwnerPtr, const NULLABLE: bool = false> {
 ///     assert_eq!(*rcu_guard, 43);
 /// }
 /// ```
-pub type RcuOption<P> = Rcu<P, true>;
+pub struct RcuOption<P: OwnerPtr>(RcuInner<P>);
+
+/// A guard that allows access to the pointed data protected by a [`RcuOption`].
+pub struct RcuOptionReadGuard<'a, P: OwnerPtr>(RcuReadGuardInner<'a, P>);
+
+/// The inner implementation of both [`Rcu`] and [`RcuOption`].
+struct RcuInner<P: OwnerPtr> {
+    ptr: AtomicPtr<<P as OwnerPtr>::Target>,
+    // We want to implement Send and Sync explicitly.
+    // Having a pointer field prevents them from being implemented
+    // automatically by the compiler.
+    _marker: PhantomData<*const P::Target>,
+}
 
 // SAFETY: It is apparent that if `P::Target` is `Send`, then `Rcu<P>` is `Send`.
-unsafe impl<P: OwnerPtr, const NULLABLE: bool> Send for Rcu<P, NULLABLE> where
+unsafe impl<P: OwnerPtr> Send for RcuInner<P> where <P as OwnerPtr>::Target: Send {}
-    <P as OwnerPtr>::Target: Send
-{
-}
 
 // SAFETY: To implement `Sync` for `Rcu<P>`, we need to meet two conditions:
 //  1. `P::Target` must be `Sync` because `Rcu::get` allows concurrent access.
 //  2. `P::Target` must be `Send` because `Rcu::update` may obtain an object
 //     of `P` created on another thread.
-unsafe impl<P: OwnerPtr, const NULLABLE: bool> Sync for Rcu<P, NULLABLE> where
+unsafe impl<P: OwnerPtr> Sync for RcuInner<P>
-    <P as OwnerPtr>::Target: Send + Sync
+where
+    <P as OwnerPtr>::Target: Send + Sync,
+    P: Send,
 {
 }
 
-// Non-nullable RCU cell.
+impl<P: OwnerPtr> RcuInner<P> {
-impl<P: OwnerPtr> Rcu<P, false> {
+    const fn new_none() -> Self {
-    /// Creates a new RCU cell with the given pointer.
+        Self {
-    pub fn new(pointer: P) -> Self {
+            ptr: AtomicPtr::new(core::ptr::null_mut()),
+            _marker: PhantomData,
+        }
+    }
+
+    fn new(pointer: P) -> Self {
         let ptr = <P as OwnerPtr>::into_raw(pointer).as_ptr();
         let ptr = AtomicPtr::new(ptr);
         Self {
@@ -131,38 +140,14 @@ impl<P: OwnerPtr> Rcu<P, false> {
             _marker: PhantomData,
         }
     }
-}
 
-// Nullable RCU cell.
+    fn update(&self, new_ptr: Option<P>) {
-impl<P: OwnerPtr> Rcu<P, true> {
+        let new_ptr = if let Some(new_ptr) = new_ptr {
-    /// Creates a new uninitialized RCU cell.
+            <P as OwnerPtr>::into_raw(new_ptr).as_ptr()
-    ///
+        } else {
-    /// Initialization can be done by calling
+            core::ptr::null_mut()
-    /// [`RcuReadGuard::compare_exchange`] after getting a read
+        };
-    /// guard using [`Rcu::read`]. Then only the first initialization will be
-    /// successful. If initialization can be done multiple times, using
-    /// [`Rcu::update`] is fine.
-    pub const fn new_none() -> Self {
-        let ptr = AtomicPtr::new(core::ptr::null_mut());
-        Self {
-            ptr,
-            _marker: PhantomData,
-        }
-    }
-}
 
-impl<P: OwnerPtr + Send, const NULLABLE: bool> Rcu<P, NULLABLE> {
-    /// Replaces the current pointer with a new pointer.
-    ///
-    /// This function updates the pointer to the new pointer regardless of the
-    /// original pointer. If the original pointer is not NULL, it will be
-    /// dropped after the grace period.
-    ///
-    /// Oftentimes this function is not recommended unless you have
-    /// synchronized writes with locks. Otherwise, you can use [`Self::read`]
-    /// and then [`RcuReadGuard::compare_exchange`] to update the pointer.
-    pub fn update(&self, new_ptr: P) {
-        let new_ptr = <P as OwnerPtr>::into_raw(new_ptr).as_ptr();
         let old_raw_ptr = self.ptr.swap(new_ptr, AcqRel);
 
         if let Some(p) = NonNull::new(old_raw_ptr) {
@@ -171,17 +156,9 @@ impl<P: OwnerPtr + Send, const NULLABLE: bool> Rcu<P, NULLABLE> {
         }
     }
 
-    /// Retrieves a read guard for the RCU cell.
+    fn read(&self) -> RcuReadGuardInner<'_, P> {
-    ///
-    /// The guard allows read-only access to the data protected by RCU.
-    ///
-    /// If the RCU cell is nullable, the guard will be nullable and you can
-    /// only dereference it after checking with [`RcuReadGuard::try_get`].
-    /// If the RCU cell is non-nullable, the guard will be non-nullable and
-    /// you can dereference it directly.
-    pub fn read(&self) -> RcuReadGuard<'_, P, NULLABLE> {
         let guard = disable_preempt();
-        RcuReadGuard {
+        RcuReadGuardInner {
             obj_ptr: self.ptr.load(Acquire),
             rcu: self,
             _inner_guard: guard,
@@ -189,75 +166,49 @@ impl<P: OwnerPtr + Send, const NULLABLE: bool> Rcu<P, NULLABLE> {
     }
 }
 
-/// A guard that allows read-only access to the initialized data protected
+impl<P: OwnerPtr> Drop for RcuInner<P> {
-/// by the RCU mechanism.
+    fn drop(&mut self) {
-pub struct RcuReadGuard<'a, P: OwnerPtr, const NULLABLE: bool> {
+        let ptr = self.ptr.load(Acquire);
-    /// If maybe uninitialized, the pointer can be NULL.
+        if let Some(p) = NonNull::new(ptr) {
+            // SAFETY: It was previously returned by `into_raw` when creating
+            // the RCU primitive.
+            let pointer = unsafe { <P as OwnerPtr>::from_raw(p) };
+            // It is OK not to delay the drop because the RCU primitive is
+            // owned by nobody else.
+            drop(pointer);
+        }
+    }
+}
+
+/// The inner implementation of both [`RcuReadGuard`] and [`RcuOptionReadGuard`].
+struct RcuReadGuardInner<'a, P: OwnerPtr> {
     obj_ptr: *mut <P as OwnerPtr>::Target,
-    rcu: &'a Rcu<P, NULLABLE>,
+    rcu: &'a RcuInner<P>,
     _inner_guard: DisabledPreemptGuard,
 }
 
-// Non-nullable RCU guard can be directly dereferenced.
+impl<P: OwnerPtr> RcuReadGuardInner<'_, P> {
-impl<P: OwnerPtr> Deref for RcuReadGuard<'_, P, false> {
+    fn compare_exchange(self, new_ptr: Option<P>) -> Result<(), Option<P>> {
-    type Target = <P as OwnerPtr>::Target;
+        let new_ptr = if let Some(new_ptr) = new_ptr {
-
+            <P as OwnerPtr>::into_raw(new_ptr).as_ptr()
-    fn deref(&self) -> &Self::Target {
+        } else {
-        // SAFETY: Since the preemption is disabled, the pointer is valid
+            core::ptr::null_mut()
-        // because other writers won't release the allocation until this task
+        };
-        // passes the quiescent state.
-        // And this pointer is not NULL.
-        unsafe { &*self.obj_ptr }
-    }
-}
-
-// Nullable RCU guard can be dereferenced after checking.
-impl<'a, P: OwnerPtr> RcuReadGuard<'a, P, true> {
-    /// Tries to get the initialized read guard.
-    ///
-    /// If the RCU cell is not initialized, this function will return
-    /// [`Err`] with the guard itself unchanged. Otherwise a dereferenceable
-    /// read guard will be returned.
-    pub fn try_get(self) -> Result<RcuReadGuard<'a, P, false>, Self> {
-        if self.obj_ptr.is_null() {
-            return Err(self);
-        }
-        Ok(RcuReadGuard {
-            obj_ptr: self.obj_ptr,
-            // SAFETY: It is initialized. The layout is the same.
-            rcu: unsafe { core::mem::transmute::<&Rcu<P, true>, &Rcu<P, false>>(self.rcu) },
-            _inner_guard: self._inner_guard,
-        })
-    }
-}
-
-impl<P: OwnerPtr + Send, const NULLABLE: bool> RcuReadGuard<'_, P, NULLABLE> {
-    /// Tries to replace the already read pointer with a new pointer.
-    ///
-    /// If another thread has updated the pointer after the read, this
-    /// function will fail and return the new pointer. Otherwise, it will
-    /// replace the pointer with the new one and drop the old pointer after
-    /// the grace period.
-    ///
-    /// If spinning on this function, it is recommended to relax the CPU
-    /// or yield the task on failure. Otherwise contention will occur.
-    ///
-    /// This API does not help to avoid
-    /// [the ABA problem](https://en.wikipedia.org/wiki/ABA_problem).
-    pub fn compare_exchange(self, new_ptr: P) -> Result<(), P> {
-        let new_ptr = <P as OwnerPtr>::into_raw(new_ptr);
 
         if self
             .rcu
             .ptr
-            .compare_exchange(self.obj_ptr, new_ptr.as_ptr(), AcqRel, Acquire)
+            .compare_exchange(self.obj_ptr, new_ptr, AcqRel, Acquire)
             .is_err()
         {
+            let Some(new_ptr) = NonNull::new(new_ptr) else {
+                return Err(None);
+            };
             // SAFETY:
             // 1. It was previously returned by `into_raw`.
             // 2. The `compare_exchange` fails so the pointer will not
-            //    be used anymore.
+            //    be used by other threads via reading the RCU primitive.
-            return Err(unsafe { <P as OwnerPtr>::from_raw(new_ptr) });
+            return Err(Some(unsafe { <P as OwnerPtr>::from_raw(new_ptr) }));
         }
 
         if let Some(p) = NonNull::new(self.obj_ptr) {
@@ -269,16 +220,162 @@ impl<P: OwnerPtr + Send, const NULLABLE: bool> RcuReadGuard<'_, P, NULLABLE> {
     }
 }
 
+impl<P: OwnerPtr> Rcu<P> {
+    /// Creates a new RCU primitive with the given pointer.
+    pub fn new(pointer: P) -> Self {
+        Self(RcuInner::new(pointer))
+    }
+
+    /// Replaces the current pointer with a null pointer.
+    ///
+    /// This function updates the pointer to the new pointer regardless of the
+    /// original pointer. The original pointer will be dropped after the grace
+    /// period.
+    ///
+    /// Oftentimes this function is not recommended unless you have serialized
+    /// writes with locks. Otherwise, you can use [`Self::read`] and then
+    /// [`RcuReadGuard::compare_exchange`] to update the pointer.
+    pub fn update(&self, new_ptr: P) {
+        self.0.update(Some(new_ptr));
+    }
+
+    /// Retrieves a read guard for the RCU primitive.
+    ///
+    /// The guard allows read access to the data protected by RCU, as well
+    /// as the ability to do compare-and-exchange.
+    pub fn read(&self) -> RcuReadGuard<'_, P> {
+        RcuReadGuard(self.0.read())
+    }
+}
+
+impl<P: OwnerPtr> RcuOption<P> {
+    /// Creates a new RCU primitive with the given pointer.
+    pub fn new(pointer: Option<P>) -> Self {
+        if let Some(pointer) = pointer {
+            Self(RcuInner::new(pointer))
+        } else {
+            Self(RcuInner::new_none())
+        }
+    }
+
+    /// Creates a new RCU primitive that contains nothing.
+    ///
+    /// This is a constant equivalence to [`RcuOption::new(None)`].
+    pub const fn new_none() -> Self {
+        Self(RcuInner::new_none())
+    }
+
+    /// Replaces the current pointer with a null pointer.
+    ///
+    /// This function updates the pointer to the new pointer regardless of the
+    /// original pointer. If the original pointer is not NULL, it will be
+    /// dropped after the grace period.
+    ///
+    /// Oftentimes this function is not recommended unless you have
+    /// synchronized writes with locks. Otherwise, you can use [`Self::read`]
+    /// and then [`RcuOptionReadGuard::compare_exchange`] to update the pointer.
+    pub fn update(&self, new_ptr: Option<P>) {
+        self.0.update(new_ptr);
+    }
+
+    /// Retrieves a read guard for the RCU primitive.
+    ///
+    /// The guard allows read access to the data protected by RCU, as well
+    /// as the ability to do compare-and-exchange.
+    ///
+    /// The contained pointer can be NULL and you can only get a reference
+    /// (if checked non-NULL) via [`RcuOptionReadGuard::get`].
+    pub fn read(&self) -> RcuOptionReadGuard<'_, P> {
+        RcuOptionReadGuard(self.0.read())
+    }
+}
+
+// RCU guards that have a non-null pointer can be directly dereferenced.
+impl<P: OwnerPtr> Deref for RcuReadGuard<'_, P> {
+    type Target = <P as OwnerPtr>::Target;
+
+    fn deref(&self) -> &Self::Target {
+        // SAFETY:
+        // 1. This pointer is not NULL because the type is `RcuReadGuard`.
+        // 2. Since the preemption is disabled, the pointer is valid because
+        //    other writers won't release the allocation until this task passes
+        //    the quiescent state.
+        unsafe { &*self.0.obj_ptr }
+    }
+}
+
+impl<P: OwnerPtr> RcuReadGuard<'_, P> {
+    /// Tries to replace the already read pointer with a new pointer.
+    ///
+    /// If another thread has updated the pointer after the read, this
+    /// function will fail, and returns the given pointer back. Otherwise,
+    /// it will replace the pointer with the new one and drop the old pointer
+    /// after the grace period.
+    ///
+    /// If spinning on [`Rcu::read`] and this function, it is recommended
+    /// to relax the CPU or yield the task on failure. Otherwise contention
+    /// will occur.
+    ///
+    /// This API does not help to avoid
+    /// [the ABA problem](https://en.wikipedia.org/wiki/ABA_problem).
+    pub fn compare_exchange(self, new_ptr: P) -> Result<(), P> {
+        self.0
+            .compare_exchange(Some(new_ptr))
+            .map_err(|err| err.unwrap())
+    }
+}
+
+// RCU guards that may have a null pointer can be dereferenced after checking.
+impl<P: OwnerPtr> RcuOptionReadGuard<'_, P> {
+    /// Gets the reference of the protected data.
+    ///
+    /// If the RCU primitive protects nothing, this function returns `None`.
+    pub fn get(&self) -> Option<&<P as OwnerPtr>::Target> {
+        if self.0.obj_ptr.is_null() {
+            return None;
+        }
+        // SAFETY:
+        // 1. This pointer is not NULL.
+        // 2. Since the preemption is disabled, the pointer is valid because
+        //    other writers won't release the allocation until this task passes
+        //    the quiescent state.
+        Some(unsafe { &*self.0.obj_ptr })
+    }
+
+    /// Returns if the RCU primitive protects nothing when [`Rcu::read`] happens.
+    pub fn is_none(&self) -> bool {
+        self.0.obj_ptr.is_null()
+    }
+
+    /// Tries to replace the already read pointer with a new pointer
+    /// (or none).
+    ///
+    /// If another thread has updated the pointer after the read, this
+    /// function will fail, and returns the given pointer back. Otherwise,
+    /// it will replace the pointer with the new one and drop the old pointer
+    /// after the grace period.
+    ///
+    /// If spinning on [`RcuOption::read`] and this function, it is recommended
+    /// to relax the CPU or yield the task on failure. Otherwise contention
+    /// will occur.
+    ///
+    /// This API does not help to avoid
+    /// [the ABA problem](https://en.wikipedia.org/wiki/ABA_problem).
+    pub fn compare_exchange(self, new_ptr: Option<P>) -> Result<(), Option<P>> {
+        self.0.compare_exchange(new_ptr)
+    }
+}
+
 /// # Safety
 ///
 /// The pointer must be previously returned by `into_raw` and the pointer
 /// must be only be dropped once.
-unsafe fn delay_drop<P: OwnerPtr + Send>(pointer: NonNull<<P as OwnerPtr>::Target>) {
+unsafe fn delay_drop<P: OwnerPtr>(pointer: NonNull<<P as OwnerPtr>::Target>) {
     struct ForceSend<P: OwnerPtr>(NonNull<<P as OwnerPtr>::Target>);
     // SAFETY: Sending a raw pointer to another task is safe as long as
     // the pointer access in another task is safe (guaranteed by the trait
     // bound `P: Send`).
-    unsafe impl<P: OwnerPtr + Send> Send for ForceSend<P> {}
+    unsafe impl<P: OwnerPtr> Send for ForceSend<P> {}
 
     let pointer: ForceSend<P> = ForceSend(pointer);
 

ostd/src/sync/rcu/owner_ptr.rs

@@ -17,7 +17,7 @@ use crate::prelude::*;
 /// raw pointers.
 ///
 /// [`Rcu`]: super::Rcu
-pub unsafe trait OwnerPtr: 'static {
+pub unsafe trait OwnerPtr: Send + 'static {
     /// The target type that this pointer refers to.
     // TODO: allow ?Sized
     type Target;
@@ -49,7 +49,7 @@ pub unsafe trait OwnerPtr: 'static {
     unsafe fn from_raw(ptr: NonNull<Self::Target>) -> Self;
 }
 
-unsafe impl<T: 'static> OwnerPtr for Box<T> {
+unsafe impl<T: Send + 'static> OwnerPtr for Box<T> {
     type Target = T;
 
     fn new(value: Self::Target) -> Self {
@@ -71,7 +71,7 @@ unsafe impl<T: 'static> OwnerPtr for Box<T> {
     }
 }
 
-unsafe impl<T: 'static> OwnerPtr for Arc<T> {
+unsafe impl<T: Send + Sync + 'static> OwnerPtr for Arc<T> {
     type Target = T;
 
     fn new(value: Self::Target) -> Self {