housekeeping 0.0.3

//! A simple concurrent stack.

use core::{
    iter::FusedIterator,
    mem,
    ptr::{self, NonNull},
};

use alloc::boxed::Box;

use crate::loomish::sync::atomic::{AtomicPtr, Ordering::*};

//----------- Stack ------------------------------------------------------------

/// A simple concurrent stack.
///
/// This is a concurrent first-in-first-out (FIFO) stack of arbitrary data. It
/// allows items to be pushed and popped across multiple threads. It requires
/// special care: when a node is popped, it cannot be deallocated until every
/// pop operation that may have occurred concurrently is finished.
pub struct Stack<T> {
    /// The head of the stack.
    head: AtomicPtr<Node<T>>,
}

impl<T> Stack<T> {
    /// Construct a new, empty [`Stack`].
    #[must_use]
    pub fn new() -> Self {
        Self {
            head: AtomicPtr::new(ptr::null_mut()),
        }
    }

    /// Push a node into the stack.
    ///
    /// ## Safety
    ///
    /// `self.push(node)` is sound if and only if:
    ///
    /// - `node` is a unique pointer to a heap-allocated [`Node`].
    pub unsafe fn push(&self, node: NonNull<Node<T>>) {
        let node = node.as_ptr();
        self.head
            .fetch_update(Release, Relaxed, |head| {
                // Link `node` to `head`.
                unsafe { (*node).next.store(head, Relaxed) };

                Some(node)
            })
            .unwrap();
    }

    /// Pop a node from the stack.
    ///
    /// If the returned pointer is null, the stack was empty.
    ///
    /// ## Safety
    ///
    /// `node = self.pop()` is sound if and only if:
    ///
    /// - The caller does not deallocate `node` until all possible concurrently
    ///   executing pop/swap operations are finished.
    #[must_use]
    pub unsafe fn pop(&self) -> Option<NonNull<Node<T>>> {
        self.head
            .fetch_update(Relaxed, Acquire, |head| {
                // If the node exists, find whatever comes after it.
                (!head.is_null()).then(|| unsafe { (*head).next.load(Relaxed) })
            })
            .ok()
            .map(|node| unsafe { NonNull::new_unchecked(node) })
    }

    /// Swap this stack for another, atomically.
    ///
    /// ## Safety
    ///
    /// `stack = self.swap(other)` is sound if and only if:
    ///
    /// - The caller does not deallocate the nodes in `stack` until all possible
    ///   concurrently executing pop/swap operations are finished.
    pub unsafe fn swap(&self, mut other: Self) -> Self {
        let other = mem::take(&mut other.head).into_inner();
        let head = self.head.swap(other, AcqRel);
        Self {
            head: AtomicPtr::new(head),
        }
    }
}

impl<T> Default for Stack<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Drop for Stack<T> {
    fn drop(&mut self) {
        let mut head = mem::take(&mut self.head).into_inner();
        while !head.is_null() {
            head = unsafe { *Box::from_raw(head) }.next.into_inner();
        }
    }
}

unsafe impl<T: Send> Send for Stack<T> {}
unsafe impl<T: Send> Sync for Stack<T> {}

impl<T> IntoIterator for Stack<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;

    fn into_iter(mut self) -> Self::IntoIter {
        let head = mem::take(&mut self.head).into_inner();
        IntoIter { head }
    }
}

//----------- Node -------------------------------------------------------------

/// A node in a [`Stack`].
#[repr(C)] // So 'Node<MaybeUninit<T>>' and 'Node<T>' are compatible
pub struct Node<T> {
    /// The next item in the stack, if any.
    ///
    /// This field is either null (indicating that `self` is the last item in
    /// the set) or non-null (pointing to the next item). It is effectively
    /// `Option<Box<Self>>`, but it cannot be represented as such due to
    /// potential non-uniqueness during insertion.
    ///
    /// After a node is taken from [`Stack::head`], its `next` field may be
    /// modified. But other threads that were racing to pop the same node may be
    /// accessing `next` concurrently. During this time, `next` requires atomic
    /// operation to prevent data races.
    pub next: AtomicPtr<Self>,

    /// The item in this node.
    pub item: T,
}

impl<T> Node<T> {
    /// Construct a new [`Node`] on the heap.
    ///
    /// The `next` field will be null.
    #[must_use]
    pub fn new(item: T) -> NonNull<Self> {
        let next = AtomicPtr::new(ptr::null_mut());
        let this = Box::new(Self { next, item });
        unsafe { NonNull::new_unchecked(Box::into_raw(this)) }
    }
}

//----------- IntoIter ---------------------------------------------------------

/// A destructive iterator of the items in a [`Stack`].
pub struct IntoIter<T> {
    /// The head of the stack.
    head: *mut Node<T>,
}

impl<T> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        if self.head.is_null() {
            return None;
        }

        let Node { next, item } = unsafe { *Box::from_raw(self.head) };
        self.head = next.into_inner();
        Some(item)
    }
}

impl<T> FusedIterator for IntoIter<T> {}

//============ Tests ===========================================================

#[cfg(test)]
mod tests {
    use alloc::boxed::Box;

    use super::{Node, Stack};

    #[test]
    fn push_pop() {
        let stack = Stack::new();
        unsafe {
            stack.push(Node::new(47));
            stack.push(Node::new(752));
            stack.push(Node::new(1367));
            assert_eq!(Box::from_raw(stack.pop().unwrap().as_ptr()).item, 1367);
            assert_eq!(Box::from_raw(stack.pop().unwrap().as_ptr()).item, 752);
            assert_eq!(Box::from_raw(stack.pop().unwrap().as_ptr()).item, 47);
            assert!(stack.pop().is_none());
        }
    }
}