Skip to main content

extended_collections/sync/
stack.rs

1use epoch::{self, Atomic, Owned};
2use std::ptr;
3use std::sync::atomic::{AtomicUsize, Ordering};
4
5struct Node<T> {
6    value: T,
7    next: Atomic<Node<T>>,
8}
9
10/// A concurrent and lock-free stack using Treiber's algorithm.
11///
12/// The Treiber Stack is a simple concurrent data structure that uses the fine-grained
13/// "compare-and-swap" concurrency primitive.
14///
15/// # Examples
16/// ```
17/// use extended_collections::sync::Stack;
18///
19/// let mut s = Stack::new();
20///
21/// s.push(0);
22/// s.push(1);
23/// assert_eq!(s.len(), 2);
24///
25/// assert_eq!(s.try_pop(), Some(1));
26/// assert_eq!(s.try_pop(), Some(0));
27/// assert_eq!(s.len(), 0);
28/// ```
29pub struct Stack<T> {
30    head: Atomic<Node<T>>,
31    len: AtomicUsize,
32}
33
34impl<T> Stack<T> {
35    /// Constructs a new, empty `Stack<T>`.
36    ///
37    /// # Examples
38    /// ```
39    /// use extended_collections::sync::Stack;
40    ///
41    /// let s: Stack<u32> = Stack::new();
42    /// ```
43    pub fn new() -> Self {
44        Stack {
45            head: Atomic::null(),
46            len: AtomicUsize::new(0),
47        }
48    }
49
50    /// Pushes an item onto the stack.
51    ///
52    /// # Examples
53    /// ```
54    /// use extended_collections::sync::Stack;
55    ///
56    /// let mut s = Stack::new();
57    /// s.push(0);
58    /// ```
59    pub fn push(&self, value: T) {
60        let mut new_node = Owned::new(Node {
61            value,
62            next: Atomic::null(),
63        });
64
65        let guard = &epoch::pin();
66        loop {
67            let head_shared = self.head.load(Ordering::Relaxed, guard);
68            new_node.next.store(head_shared, Ordering::Relaxed);
69            match self.head.compare_and_set(head_shared, new_node, Ordering::Release, guard) {
70                Ok(_) => {
71                    self.len.fetch_add(1, Ordering::Release);
72                    break;
73                },
74                Err(e) => new_node = e.new,
75            }
76        }
77    }
78
79    /// Attempts to pop the top element of the stack. Returns `None` if it was unable to pop the
80    /// top element.
81    ///
82    /// # Examples
83    /// ```
84    /// use extended_collections::sync::Stack;
85    ///
86    /// let mut s = Stack::new();
87    ///
88    /// s.push(0);
89    ///
90    /// assert_eq!(s.try_pop(), Some(0));
91    /// assert_eq!(s.try_pop(), None);
92    /// ```
93    pub fn try_pop(&self) -> Option<T> {
94        let guard = &epoch::pin();
95        loop {
96            let head_shared = self.head.load(Ordering::Acquire, guard);
97            match unsafe { head_shared.as_ref() } {
98                Some(head) => {
99                    let next = head.next.load(Ordering::Relaxed, guard);
100                    if self.head.compare_and_set(head_shared, next, Ordering::Release, guard).is_ok() {
101                        unsafe {
102                            self.len.fetch_sub(1, Ordering::Release);
103                            guard.defer(move || head_shared.into_owned());
104                            return Some(ptr::read(&(*head).value));
105                        }
106                    }
107                },
108                None => return None,
109            }
110        }
111    }
112
113    /// Returns the approximate number of elements in the stack.
114    ///
115    /// # Examples
116    /// ```
117    /// use extended_collections::sync::Stack;
118    ///
119    /// let mut s = Stack::new();
120    /// assert_eq!(s.len(), 0);
121    ///
122    /// s.push(0);
123    /// assert_eq!(s.len(), 1);
124    /// ```
125    pub fn len(&self) -> usize {
126        self.len.load(Ordering::Acquire)
127    }
128
129    /// Returns `true` if the approximate number of elements in the stack is zero.
130    ///
131    /// # Examples
132    /// ```
133    /// use extended_collections::sync::Stack;
134    ///
135    /// let mut s = Stack::new();
136    /// assert!(s.is_empty());
137    ///
138    /// s.push(0);
139    /// assert!(!s.is_empty());
140    /// ```
141    pub fn is_empty(&self) -> bool {
142        self.len() == 0
143    }
144}
145
146impl<T> Default for Stack<T> {
147    fn default() -> Self {
148        Self::new()
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::Stack;
155
156    #[test]
157    fn test_len_empty() {
158        let stack: Stack<u32> = Stack::new();
159        assert_eq!(stack.len(), 0);
160    }
161
162    #[test]
163    fn test_is_empty() {
164        let stack: Stack<u32> = Stack::new();
165        assert!(stack.is_empty());
166    }
167
168    #[test]
169    fn test_push_pop() {
170        let stack = Stack::new();
171        stack.push(0);
172        stack.push(1);
173
174        assert_eq!(stack.len(), 2);
175        assert_eq!(stack.try_pop(), Some(1));
176        assert_eq!(stack.try_pop(), Some(0));
177        assert_eq!(stack.len(), 0);
178    }
179}