Skip to main content

wg_utils/lfs/
stack.rs

1use std::ptr;
2use std::sync::atomic::{AtomicPtr, Ordering};
3
4pub struct LockFreeStack<T> {
5    top: AtomicPtr<Node<T>>,
6}
7
8struct Node<T> {
9    value: T,
10    next: *mut Node<T>,
11}
12
13impl<T> LockFreeStack<T> {
14    pub fn new() -> Self {
15        LockFreeStack {
16            top: AtomicPtr::new(ptr::null_mut()),
17        }
18    }
19
20    #[inline(always)]
21    pub fn push(&self, value: T) {
22        let new_node = Box::into_raw(Box::new(Node {
23            value,
24            next: ptr::null_mut(),
25        }));
26
27        loop {
28            let top = self.top.load(Ordering::Acquire);
29            unsafe { (*new_node).next = top };
30
31            if self
32                .top
33                .compare_exchange_weak(top, new_node, Ordering::Release, Ordering::Relaxed)
34                .is_ok()
35            {
36                break;
37            }
38
39            core::hint::spin_loop();
40        }
41    }
42
43    #[inline(always)]
44    pub fn pop(&self) -> Option<T> {
45        loop {
46            let top = self.top.load(Ordering::Acquire);
47            if top.is_null() {
48                return None;
49            }
50
51            let next = unsafe { (*top).next };
52
53            if self
54                .top
55                .compare_exchange_weak(top, next, Ordering::Release, Ordering::Relaxed)
56                .is_ok()
57            {
58                let value = unsafe { ptr::read(&(*top).value) };
59                unsafe { drop(Box::from_raw(top)) };
60                return Some(value);
61            }
62
63            core::hint::spin_loop();
64        }
65    }
66}
67
68impl<T> Drop for LockFreeStack<T> {
69    fn drop(&mut self) {
70        while self.pop().is_some() {}
71    }
72}