nblfq 0.1.0

#![cfg_attr(not(any(test, feature = "std")), no_std)]
#![doc = include_str!("../README.md")]

use core::cell::UnsafeCell;
use core::mem::MaybeUninit;
use core::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};

/// Represents the byte-size of the sequence word.
const W: usize = usize::MAX;

/// Represents the inner [Queue] entry slot.
///
/// Provides safe wrappers around the dangerous `unsafe` bits.
#[derive(Debug)]
pub struct Slot<T> {
    sequence: AtomicUsize,
    ptr: AtomicPtr<T>,
    elem: UnsafeCell<Option<MaybeUninit<T>>>,
}

impl<T> Slot<T> {
    /// Creates a new [Slot].
    pub const fn new() -> Self {
        Self {
            sequence: AtomicUsize::new(0),
            ptr: AtomicPtr::new(core::ptr::null_mut()),
            elem: UnsafeCell::new(None),
        }
    }

    /// Creates a new [Slot].
    pub const fn new_sequence(sequence: usize) -> Self {
        Self {
            sequence: AtomicUsize::new(sequence),
            ptr: AtomicPtr::new(core::ptr::null_mut()),
            elem: UnsafeCell::new(None),
        }
    }

    /// Creates a new [Slot] from the provided parameters.
    pub fn create(sequence: usize, elem: T) -> Self {
        let s = Self {
            sequence: AtomicUsize::new(sequence),
            ptr: AtomicPtr::new(core::ptr::null_mut()),
            elem: UnsafeCell::new(Some(MaybeUninit::new(elem))),
        };

        // SAFETY: we just initialized the element, the pointer references valid memory.
        s.ptr.store(unsafe { s.elem_mut_ptr() }, Ordering::Release);

        s
    }

    /// Gets a reference to the sequence.
    pub const fn sequence(&self) -> &AtomicUsize {
        &self.sequence
    }

    /// Gets the sequence.
    pub fn load_sequence(&self) -> usize {
        self.sequence.load(Ordering::Acquire)
    }

    /// Sets the sequence.
    pub fn store_sequence(&self, seq: usize) {
        self.sequence.store(seq, Ordering::Release);
    }

    /// Loads the [Slot] element.
    ///
    /// Returns:
    ///
    /// - `Some(T)`: if the [Slot] has been initialized
    /// - `None`: if the [Slot] is empty.
    pub fn load_element(&self) -> Option<&T> {
        let ptr = self.ptr.load(Ordering::Acquire);
        if ptr.is_null() {
            None
        } else {
            // SAFETY: if non-null, `ptr` references an owned element.
            unsafe { Some(&*ptr) }
        }
    }

    /// Stores the [Slot] sequence and element.
    pub fn store_element(&self, sequence: usize, elem: T) {
        self.sequence.store(sequence, Ordering::Release);
        // SAFETY: we run a slight risk of overwriting an entry at the same position.
        // However, we are not in danger of undefined behavior.
        unsafe { (&mut *self.elem.get()).replace(MaybeUninit::new(elem)) };
        // SAFETY: we just iniitialized the element, so the pointer references valid memory.
        self.ptr
            .store(unsafe { self.elem_mut_ptr() }, Ordering::Release);
    }

    /// Gets a mutable pointer to the element.
    ///
    /// # Safety
    ///
    /// Caller must ensure exclusive access to the element.
    #[inline]
    unsafe fn elem_mut_ptr(&self) -> *mut T {
        unsafe { (&mut *self.elem.get()).as_mut().unwrap().as_mut_ptr() }
    }

    /// Clears the [Slot].
    ///
    /// Calls the destructor of the stored element, if it exists.
    pub fn clear(&self) -> Option<T> {
        self.sequence.store(0, Ordering::Release);

        if self.ptr.load(Ordering::Acquire).is_null() {
            None
        } else {
            self.ptr.store(core::ptr::null_mut(), Ordering::Release);
            // SAFETY: we only call `assume_init` on a populated element.
            unsafe { (&mut *self.elem.get()).take().map(|t| t.assume_init()) }
        }
    }

    /// Performs a compare-and-exchange (CAS) operation on the [Slot].
    ///
    /// Returns:
    ///
    /// - `Ok(Slot)`: returns the previous held entry on success
    /// - `Err(Slot)`: returns the replacement entry on error.
    pub fn compare_exchange(&self, cmp: &Self, rep: Self) -> Result<Self, Self> {
        let Ok(s) = self.sequence.compare_exchange(
            cmp.load_sequence(),
            rep.load_sequence(),
            Ordering::AcqRel,
            Ordering::Acquire,
        ) else {
            return Err(rep);
        };

        let Ok(p) = self.ptr.compare_exchange(
            cmp.ptr.load(Ordering::Acquire),
            rep.ptr.load(Ordering::Acquire),
            Ordering::AcqRel,
            Ordering::Acquire,
        ) else {
            return Err(rep);
        };

        let t = if p.is_null() {
            None
        } else {
            unsafe { (&mut *self.elem.get()).take() }
        };

        if !rep.ptr.load(Ordering::Acquire).is_null() {
            // SAFETY: we only call `assume_init` on a populated element.
            unsafe {
                if let Some(elem) = rep.elem.into_inner() {
                    (&mut *self.elem.get()).replace(MaybeUninit::new(elem.assume_init()));
                }
            };
        }

        if let Some(elem) = t {
            // SAFETY: we only call `assume_init` on a populated element.
            Ok(Self::create(s, unsafe { elem.assume_init() }))
        } else {
            Ok(Self::new_sequence(s))
        }
    }

    /// Compares if [Slot] at position `i` is before [Slot] at position `j`.
    pub fn comp(&self, i: usize, oth: &Self, j: usize) -> bool {
        let l_seq = self.load_sequence();
        let r_seq = oth.load_sequence();

        if l_seq == r_seq {
            i < j
        } else {
            r_seq.wrapping_add(W).wrapping_sub(l_seq) < (1 << (usize::BITS - 1))
        }
    }

    /// Destructs the [Slot] into its inner element.
    ///
    /// Returns:
    ///
    /// - `Some(T)`: if the [Slot] was initialized.
    /// - `None`: if the [Slot] was uninitialized.
    pub fn into_inner(self) -> Option<T> {
        let ptr = self.ptr.load(Ordering::Acquire);
        if ptr.is_null() {
            None
        } else {
            // SAFETY: if non-null, `ptr` references an owned element.
            unsafe { self.elem.into_inner().map(|e| e.assume_init()) }
        }
    }
}

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

impl<T: PartialEq> PartialEq for Slot<T> {
    fn eq(&self, oth: &Self) -> bool {
        self.load_sequence() == oth.load_sequence() && self.load_element() == oth.load_element()
    }
}

/// Represents an MPMC queue.
pub struct Queue<T, const N: usize> {
    head: AtomicUsize,
    tail: AtomicUsize,
    clones: AtomicUsize,
    buffer: [Slot<T>; N],
}

impl<T, const N: usize> Queue<T, N> {
    /// Creates a new [Queue].
    #[cfg(not(loom))]
    pub const fn new() -> Self {
        Self {
            head: AtomicUsize::new(0),
            tail: AtomicUsize::new(0),
            clones: AtomicUsize::new(1),
            buffer: [const { Slot::new() }; N],
        }
    }

    /// Creates a new [Queue].
    #[cfg(loom)]
    pub fn new() -> Self {
        Self {
            head: AtomicUsize::new(0),
            tail: AtomicUsize::new(0),
            clones: AtomicUsize::new(1),
            buffer: core::array::from_fn(|_| Slot::new()),
        }
    }

    /// Helper function to get the previous queue position.
    pub const fn prev(i: usize) -> usize {
        (i + N - 1) % N
    }

    /// Gets a reference to the head position.
    pub const fn head(&self) -> &AtomicUsize {
        &self.head
    }

    /// Gets the head position.
    pub fn load_head(&self) -> usize {
        self.head.load(Ordering::Acquire)
    }

    /// Sets the head position.
    pub fn store_head(&self, head: usize) {
        self.head.store(head, Ordering::Release);
    }

    /// Gets a reference to the tail position.
    pub const fn tail(&self) -> &AtomicUsize {
        &self.tail
    }

    /// Gets the tail position.
    pub fn load_tail(&self) -> usize {
        self.tail.load(Ordering::Acquire)
    }

    /// Sets the tail position.
    pub fn store_tail(&self, tail: usize) {
        self.tail.store(tail, Ordering::Release);
    }

    /// Gets a reference to the clone count.
    pub const fn clones(&self) -> &AtomicUsize {
        &self.clones
    }

    /// Gets the clone count.
    pub fn load_clones(&self) -> usize {
        self.clones.load(Ordering::Acquire)
    }

    /// Increments the clone count.
    #[cfg(feature = "std")]
    pub(crate) fn inc_clones(&self) {
        self.clones.fetch_add(1, Ordering::AcqRel);
    }

    /// Attempts to enqueue an item into the [Queue].
    pub fn enqueue(&self, item: T) -> Result<(), T> {
        let res;
        let mut new_slot = Slot::create(0, item);

        'retry: loop {
            let mut pos = self.load_head();
            let mut slot;
            let mut prev_slot;

            'chase: loop {
                slot = &self.buffer[pos];
                let prev_pos = Self::prev(pos);
                prev_slot = &self.buffer[prev_pos];

                if prev_slot.load_element().is_some() && slot.load_element().is_none() {
                    /* null cell, non-empty predecessor */
                    break 'chase;
                }

                if !prev_slot.comp(prev_pos, slot, pos) {
                    /* found step */
                    if prev_slot.load_element().is_none() && slot.load_element().is_none() {
                        /* empty list */
                        break 'chase;
                    }

                    if prev_slot.load_element().is_some() && slot.load_element().is_some() {
                        /* full list */
                        res = Err(());
                        break 'retry;
                    }
                }

                pos = pos.wrapping_add(1) % N;
            }

            let mut seq = prev_slot.load_sequence();

            if prev_slot.load_element().is_none() {
                seq = seq.wrapping_add(W);
            }

            if pos == 0 {
                seq = seq.wrapping_add(1);
            }

            new_slot.store_sequence(seq);

            match self.buffer[pos].compare_exchange(&Slot::new_sequence(seq), new_slot) {
                Ok(_) => {
                    self.store_head((pos + 1) % N);
                    return Ok(());
                }
                Err(v) => new_slot = v,
            }
        }

        // SAFETY: we only call `unwrap` and `assume_init` on populated element.
        res.map_err(|_| new_slot.into_inner().unwrap())
    }

    /// Dequeues an item from the [Queue].
    pub fn dequeue(&self) -> Option<T> {
        /* retry loop */
        loop {
            let mut tail = self.load_tail();
            let prev_tail = Self::prev(tail);

            let mut prev_slot = &self.buffer[prev_tail];
            let mut slot = &self.buffer[tail];

            /* chase the tail */
            while prev_slot.comp(prev_tail, slot, tail) {
                tail = tail.wrapping_add(1) % N;
                prev_slot = slot;
                slot = &self.buffer[tail];
            }

            if prev_slot.load_element().is_none() && slot.load_element().is_none() {
                /* empty queue */
                return None;
            }

            let seq = slot.load_sequence().wrapping_add(1);

            if let Ok(slot) = self.buffer[tail].compare_exchange(slot, Slot::new_sequence(seq)) {
                self.store_tail(tail.wrapping_add(1) % N);
                return slot.clear();
            }
        }
    }

    #[cfg(feature = "std")]
    pub fn clone<'a>(self: &'a std::sync::Arc<Self>) -> std::sync::Arc<Self> {
        let c = std::sync::Arc::clone(self);
        c.inc_clones();
        c
    }
}

impl<T, const N: usize> Default for Queue<T, N> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const N: usize = 4;

    #[test]
    fn test_slot() {
        let _empty = Slot::<()>::new();
        let _full = Slot::<()>::create(0, ());
    }

    #[test]
    fn test_enqueue() {
        let queue = Queue::<(), N>::new();

        assert!(queue.enqueue(()).is_ok());

        let pos = queue.load_head() - 1;
        assert_eq!(&queue.buffer[pos], &Slot::create(pos, ()));
    }

    #[test]
    fn test_dequeue() {
        let queue = Queue::<(), N>::new();

        for _ in 0..N {
            assert!(queue.enqueue(()).is_ok());
        }

        assert!(queue.enqueue(()).is_err());

        for _i in 0..N {
            assert_eq!(queue.dequeue(), Some(()), "{_i}");
        }

        assert_eq!(queue.dequeue(), None);
    }
}