Made Node trait functions unsafe and guaranteed correct initialization in ArcPool and BoxPool

Author: pc-for-science

src/pool/arc.rs

@@ -72,7 +72,7 @@
 use core::{
     fmt,
     hash::{Hash, Hasher},
-    mem::{ManuallyDrop, MaybeUninit},
+    mem::MaybeUninit,
     ops, ptr,
 };
 
@@ -188,6 +188,10 @@ impl<T> ArcPoolImpl<T> {
     fn manage(&self, block: &'static mut ArcBlock<T>) {
         let node: &'static mut _ = &mut block.node;
 
+        // SAFETY: The node within an `ArcBlock` is always properly initialized for linking because the only way for
+        //         client code to construct an `ArcBlock` is through `ArcBlock::new`. The `NonNullPtr` comes from a
+        //         reference, so it is guaranteed to be dereferencable. It is also unique because the `ArcBlock` itself
+        //         is passed as a `&mut`
         unsafe { self.stack.push(NonNullPtr::from_static_mut_ref(node)) }
     }
 }
@@ -383,9 +387,7 @@ impl<T> ArcBlock<T> {
     /// Creates a new memory block
     pub const fn new() -> Self {
         Self {
-            node: UnionNode {
-                data: ManuallyDrop::new(MaybeUninit::uninit()),
-            },
+            node: UnionNode::unlinked(),
         }
     }
 }

src/pool/boxed.rs

@@ -376,6 +376,10 @@ impl<T> BoxPoolImpl<T> {
     fn manage(&self, block: &'static mut BoxBlock<T>) {
         let node: &'static mut _ = &mut block.node;
 
+        // SAFETY: The node within a `BoxBlock` is always properly initialized for linking because the only way for
+        //         client code to construct a `BoxBlock` is through `BoxBlock::new`. The `NonNullPtr` comes from a
+        //         reference, so it is guaranteed to be dereferencable. It is also unique because the `BoxBlock` itself
+        //         is passed as a `&mut`
         unsafe { self.stack.push(NonNullPtr::from_static_mut_ref(node)) }
     }
 }
@@ -391,9 +395,7 @@ impl<T> BoxBlock<T> {
     /// Creates a new memory block
     pub const fn new() -> Self {
         Self {
-            node: UnionNode {
-                data: ManuallyDrop::new(MaybeUninit::uninit()),
-            },
+            node: UnionNode::unlinked(),
         }
     }
 }

src/pool/treiber.rs

@@ -26,6 +26,8 @@ where
     /// # Safety
     /// - `node` must be a valid pointer
     /// - aliasing rules must be enforced by the caller. e.g, the same `node` may not be pushed more than once
+    /// - It must be valid to call `next()` on the `node`, meaning it must be properly initialized for insertion
+    ///   into a stack/linked list
     pub unsafe fn push(&self, node: NonNullPtr<N>) {
         impl_::push(self, node);
     }
@@ -38,10 +40,22 @@ where
 pub trait Node: Sized {
     type Data;
 
-    fn next(&self) -> &AtomicPtr<Self>;
+    /// Returns a reference to the atomic pointer that stores the link to the next `Node`
+    ///
+    /// # Safety
+    ///
+    /// It must be valid to obtain a reference to the next link pointer, e.g. in the case of
+    /// `UnionNode`, the `next` field must be properly initialized when calling this function
+    unsafe fn next(&self) -> &AtomicPtr<Self>;
 
+    /// Returns a mutable reference to the atomic pointer that stores the link to the next `Node`
+    ///
+    /// # Safety
+    ///
+    /// It must be valid to obtain a reference to the next link pointer, e.g. in the case of
+    /// `UnionNode`, the `next` field must be properly initialized when calling this function
     #[allow(dead_code)] // used conditionally
-    fn next_mut(&mut self) -> &mut AtomicPtr<Self>;
+    unsafe fn next_mut(&mut self) -> &mut AtomicPtr<Self>;
 }
 
 #[repr(C)]
@@ -50,14 +64,27 @@ pub union UnionNode<T> {
     pub data: ManuallyDrop<T>,
 }
 
+impl<T> UnionNode<T> {
+    /// Returns a new `UnionNode` that does not contain data and is not linked to any other nodes.
+    /// The return value of this function is guaranteed to have the `next` field properly initialized.
+    /// Use this function if you want to insert a new `UnionNode` into a linked list
+    pub const fn unlinked() -> Self {
+        Self {
+            next: ManuallyDrop::new(AtomicPtr::null()),
+        }
+    }
+}
+
 impl<T> Node for UnionNode<T> {
     type Data = T;
 
-    fn next(&self) -> &AtomicPtr<Self> {
+    unsafe fn next(&self) -> &AtomicPtr<Self> {
+        // SAFETY: Caller ensures that `self.next` is properly initialized
         unsafe { &self.next }
     }
 
-    fn next_mut(&mut self) -> &mut AtomicPtr<Self> {
+    unsafe fn next_mut(&mut self) -> &mut AtomicPtr<Self> {
+        // SAFETY: Caller ensures that `self.next` is properly initialized
         unsafe { &mut self.next }
     }
 }
@@ -70,11 +97,11 @@ pub struct StructNode<T> {
 impl<T> Node for StructNode<T> {
     type Data = T;
 
-    fn next(&self) -> &AtomicPtr<Self> {
+    unsafe fn next(&self) -> &AtomicPtr<Self> {
         &self.next
     }
 
-    fn next_mut(&mut self) -> &mut AtomicPtr<Self> {
+    unsafe fn next_mut(&mut self) -> &mut AtomicPtr<Self> {
         &mut self.next
     }
 }

src/pool/treiber/cas.rs

@@ -188,6 +188,13 @@ const fn initial_tag() -> Tag {
     Tag::MIN
 }
 
+/// Pushes the given node on top of the stack
+///
+/// # Safety
+///
+/// - `new_top` must point to a node that is properly initialized for linking, i.e. `new_top.as_ref().next()`
+///   must be valid to call (see [`Node::next`])
+/// - `new_top` must be convertible to a reference (see [`NonNull::as_ref`])
 pub unsafe fn push<N>(stack: &Stack<N>, new_top: NonNullPtr<N>)
 where
     N: Node,
@@ -198,6 +205,7 @@ where
         new_top
             .non_null()
             .as_ref()
+            // SAFETY: Caller ensures that it is safe to call `next`
             .next()
             .store(top, Ordering::Relaxed);
 

src/pool/treiber/llsc.rs

@@ -67,6 +67,13 @@ where
 
 impl<N> Copy for NonNullPtr<N> where N: Node {}
 
+/// Pushes the given node on top of the stack
+///
+/// # Safety
+///
+/// - `node` must point to a node that is properly initialized for linking, i.e. `node.as_mut().next_mut()`
+///   must be valid to call (see [`Node::next_mut`])
+/// - `node` must be convertible to a reference (see [`NonNull::as_mut`])
 pub unsafe fn push<N>(stack: &Stack<N>, mut node: NonNullPtr<N>)
 where
     N: Node,
@@ -78,9 +85,11 @@ where
 
         node.inner
             .as_mut()
+            // SAFETY: Caller guarantees that it is valid to call `next_mut`
             .next_mut()
             .inner
             .get()
+            // SAFETY: The pointer comes from `AtomicPtr::inner`, which is valid for writes
             .write(NonNull::new(top as *mut _));
 
         if arch::store_conditional(node.inner.as_ptr() as usize, top_addr).is_ok() {