Improve documentation for Treiber stack.

Author: reitermarkus

src/pool/treiber/cas.rs

@@ -68,13 +68,16 @@ where
                 failure,
             )
             .map(drop)
-            .map_err(NonNullPtr::from_inner)
+            .map_err(|value| {
+                // SAFETY: `value` cam from a `NonNullPtr::into_inner` call.
+                unsafe { NonNullPtr::from_inner(value) }
+            })
     }
 
     #[inline]
     fn load(&self, order: Ordering) -> Option<NonNullPtr<N>> {
-        InnerNonZero::new(self.inner.load(order)).map(|inner| NonNullPtr {
-            inner,
+        Some(NonNullPtr {
+            inner: InnerNonZero::new(self.inner.load(order))?,
             _marker: PhantomData,
         })
     }
@@ -115,33 +118,41 @@ where
     }
 
     #[inline]
-    pub fn from_static_mut_ref(ref_: &'static mut N) -> NonNullPtr<N> {
-        let non_null = NonNull::from(ref_);
-        Self::from_non_null(non_null)
+    pub fn from_static_mut_ref(reference: &'static mut N) -> NonNullPtr<N> {
+        // SAFETY: `reference` is a static mutable reference, i.e. a valid pointer.
+        unsafe { Self::new_unchecked(initial_tag(), NonNull::from(reference)) }
     }
 
-    fn from_non_null(ptr: NonNull<N>) -> Self {
-        let address = ptr.as_ptr() as Address;
-        let tag = initial_tag().get();
-
-        let value = (Inner::from(tag) << Address::BITS) | Inner::from(address);
+    /// # Safety
+    ///
+    /// - `ptr` must be a valid pointer.
+    #[inline]
+    unsafe fn new_unchecked(tag: Tag, ptr: NonNull<N>) -> Self {
+        let value =
+            (Inner::from(tag.get()) << Address::BITS) | Inner::from(ptr.as_ptr() as Address);
 
         Self {
+            // SAFETY: `value` is constructed from a `Tag` which is non-zero and half the
+            //         size of the `InnerNonZero` type, and a `NonNull<N>` pointer.
             inner: unsafe { InnerNonZero::new_unchecked(value) },
             _marker: PhantomData,
         }
     }
 
+    /// # Safety
+    ///
+    /// - `value` must come from a `Self::into_inner` call.
     #[inline]
-    fn from_inner(value: Inner) -> Option<Self> {
-        InnerNonZero::new(value).map(|inner| Self {
-            inner,
+    unsafe fn from_inner(value: Inner) -> Option<Self> {
+        Some(Self {
+            inner: InnerNonZero::new(value)?,
             _marker: PhantomData,
         })
     }
 
     #[inline]
     fn non_null(&self) -> NonNull<N> {
+        // SAFETY: `Self` can only be constructed using a `NonNull<N>`.
         unsafe { NonNull::new_unchecked(self.as_ptr()) }
     }
 
@@ -152,17 +163,15 @@ where
 
     #[inline]
     fn tag(&self) -> Tag {
+        // SAFETY: `self.inner` was constructed from a non-zero `Tag`.
         unsafe { Tag::new_unchecked((self.inner.get() >> Address::BITS) as Address) }
     }
 
-    fn increase_tag(&mut self) {
-        let address = self.as_ptr() as Address;
-
-        let new_tag = self.tag().checked_add(1).unwrap_or_else(initial_tag).get();
-
-        let value = (Inner::from(new_tag) << Address::BITS) | Inner::from(address);
+    fn increment_tag(&mut self) {
+        let new_tag = self.tag().checked_add(1).unwrap_or_else(initial_tag);
 
-        self.inner = unsafe { InnerNonZero::new_unchecked(value) };
+        // SAFETY: `self.non_null()` is a valid pointer.
+        *self = unsafe { Self::new_unchecked(new_tag, self.non_null()) };
     }
 }
 
@@ -210,7 +219,40 @@ where
                 .compare_and_exchange_weak(Some(top), next, Ordering::Release, Ordering::Relaxed)
                 .is_ok()
             {
-                top.increase_tag();
+                // Prevent the ABA problem (https://en.wikipedia.org/wiki/Treiber_stack#Correctness).
+                //
+                // Without this, the following would be possible:
+                //
+                // | Thread 1                      | Thread 2                | Stack                        |
+                // |-------------------------------|-------------------------|------------------------------|
+                // | push((1, 1))                  |                         | (1, 1)                       |
+                // | push((1, 2))                  |                         | (1, 2) -> (1, 1)             |
+                // | p = try_pop()::load // (1, 2) |                         | (1, 2) -> (1, 1)             |
+                // |                               | p = try_pop() // (1, 2) | (1, 1)                       |
+                // |                               | push((1, 3))            | (1, 3) -> (1, 1)             |
+                // |                               | push(p)                 | (1, 2) -> (1, 3) -> (1, 1)   |
+                // | try_pop()::cas(p, p.next)     |                         | (1, 1)                       |
+                //
+                // As can be seen, the `cas` operation succeeds, wrongly removing pointer `3` from the stack.
+                //
+                // By incrementing the tag before returning the pointer, it cannot be pushed again with the,
+                // same tag, preventing the `try_pop()::cas(p, p.next)` operation from succeeding.
+                //
+                // With this fix, `try_pop()` in thread 2 returns `(2, 2)` and the comparison between
+                // `(1, 2)` and `(2, 2)` fails, restarting the loop and correctly removing the new top:
+                //
+                // | Thread 1                      | Thread 2                | Stack                        |
+                // |-------------------------------|-------------------------|------------------------------|
+                // | push((1, 1))                  |                         | (1, 1)                       |
+                // | push((1, 2))                  |                         | (1, 2) -> (1, 1)             |
+                // | p = try_pop()::load // (1, 2) |                         | (1, 2) -> (1, 1)             |
+                // |                               | p = try_pop() // (2, 2) | (1, 1)                       |
+                // |                               | push((1, 3))            | (1, 3) -> (1, 1)             |
+                // |                               | push(p)                 | (2, 2) -> (1, 3) -> (1, 1)   |
+                // | try_pop()::cas(p, p.next)     |                         | (2, 2) -> (1, 3) -> (1, 1)   |
+                // | p = try_pop()::load // (2, 2) |                         | (2, 2) -> (1, 3) -> (1, 1)   |
+                // | try_pop()::cas(p, p.next)     |                         | (1, 3) -> (1, 1)             |
+                top.increment_tag();
 
                 return Some(top);
             }

src/pool/treiber/llsc.rs

@@ -125,7 +125,8 @@ mod arch {
     }
 
     /// # Safety
-    /// - `addr` must be a valid pointer
+    ///
+    /// - `addr` must be a valid pointer.
     #[inline(always)]
     pub unsafe fn load_link(addr: *const usize) -> usize {
         let value;
@@ -134,7 +135,8 @@ mod arch {
     }
 
     /// # Safety
-    /// - `addr` must be a valid pointer
+    ///
+    /// - `addr` must be a valid pointer.
     #[inline(always)]
     pub unsafe fn store_conditional(value: usize, addr: *mut usize) -> Result<(), ()> {
         let outcome: usize;