lfstack.rs

//! This is a (hopefully) correct implementation of a lockfree stack.
//!
//! Doing this without a memory reclamation implementation requires pointer
//! tagging, which is slightly architecture specific. This supports aarch64 and
//! x86_64 so long as 5-level page tables aren't in use (rare to the point of
//! nonexistence), and arm64e (or other hardware pointer authentication) is not
//! enabled, which is not currently supported by Rust. On x86_64/aarch64 where
//! this doesn't hold, it will panic.
//!
//! That said, while this implementation is believed to be correct, it could
//! have bugs so general use externally is somewhat discouraged — (it certainly
//! has inefficiencies, for example).

use core::marker::PhantomData;
use core::sync::atomic::{AtomicPtr, AtomicU64, Ordering::*};

type AtomicTaggedPtr = AtomicU64;

#[cfg(target_pointer_width = "64")]
const COUNTER_BITS: usize = 16;

#[cfg(target_pointer_width = "32")]
const COUNTER_BITS: usize = 32;

const PTR_MASK: u64 = (!0u64) >> COUNTER_BITS;
const COUNTER_MASK: u64 = !PTR_MASK;
const COUNTER_INC: u64 = PTR_MASK + 1;

pub struct Stack<T> {
    head: AtomicTaggedPtr,
    _boo: PhantomData<T>,
}

struct Node<T> {
    data: T,
    next: AtomicPtr<Node<T>>,
}

impl<T> Drop for Stack<T> {
    fn drop(&mut self) {
        while self.pop().is_some() {}
    }
}

impl<T> Stack<T> {
    pub fn new() -> Self {
        Self {
            head: AtomicTaggedPtr::new(0),
            _boo: PhantomData,
        }
    }

    pub fn push(&self, data: T) {
        let n = Box::into_raw(Box::new(Node {
            next: AtomicPtr::new(core::ptr::null_mut()),
            data,
        }));
        assert_eq!((n as u64) & COUNTER_MASK, 0, "ptr: {:p}", n);
        let mut cmp = self.head.load(Relaxed);
        loop {
            let counter = (cmp & COUNTER_MASK).wrapping_add(COUNTER_INC);
            assert_eq!(counter & PTR_MASK, 0);
            let swap = (n as u64) | counter;
            // Safety: We fully own `n` still, so this is fine
            unsafe {
                (*n).next = AtomicPtr::new((cmp & PTR_MASK) as usize as *mut Nod
e<T>);
            }
            match self.head.compare_exchange_weak(cmp, swap, Release, Relaxed) {
                Ok(_) => break,
                Err(new) => cmp = new,
            }
        }
    }

    pub fn pop(&self) -> Option<T> {
        let mut cmp = self.head.load(Acquire);
        let cur = loop {
            let cur = (cmp & PTR_MASK) as usize as *mut Node<T>;
            if cur.is_null() {
                return None;
            }
            let next = unsafe { (*cur).next.load(Relaxed) };
            let counter = cmp & COUNTER_MASK;
            let swap = (next as u64) | counter;
            match self.head.compare_exchange_weak(cmp, swap, Acquire, Relaxed) {
                Ok(_) => break cur,
                Err(h) => cmp = h,
            }
        };
        assert!(!cur.is_null());
        unsafe {
            let data = core::ptr::read(&(*cur).data);
            std::alloc::dealloc(cur.cast::<u8>(), std::alloc::Layout::new::<Node
<T>>());
            Some(data)
        }
    }
}

// send+sync for sendable data.
unsafe impl<T> Send for Stack<T> where T: Send + 'static {}
unsafe impl<T> Sync for Stack<T> where T: Send + 'static {}

#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
fn main() {
    cobb::run_test(cobb::TestCfg::<Stack<usize>> {
        threads: if cfg!(miri) { 8 } else { 16 },
        iterations: if cfg!(miri) { 100 } else { 1000 },
        sub_iterations: if cfg!(miri) { 10 } else { 20 },
        setup: || Stack::new(),
        test: |stk, tctx| {
            stk.push(tctx.thread_index());
            let _ = stk.pop();
        },
        ..Default::default()
    });
}
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
fn main() {
    panic!("Sorry, this architecture is unsupported");
}