obvhs/
faststack.rs

1//! Stack data structures implemented on the stack or heap
2use core::ops::{Deref, DerefMut};
3
4// TODO could allow Clone + Zeroable instead of Copy
5pub trait FastStack<T: Copy + Default> {
6    fn push(&mut self, v: T);
7    fn pop_fast(&mut self) -> T;
8    fn pop(&mut self) -> Option<T>;
9    fn len(&self) -> usize;
10    fn is_empty(&self) -> bool;
11    fn clear(&mut self);
12}
13
14/// Creates a stack- or heap-allocated stack based on runtime size.
15///
16/// This macro chooses between `StackStack<T, N>` (stack-allocated) and `HeapStack<T>`
17/// (heap-allocated) depending on the provided size and threshold values.
18///
19/// # Usage
20/// fast_stack!(
21///     u32,             // type
22///     (64, 128, 256),  // thresholds
23///     size,            // runtime size
24///     stack,           // variable name
25///     {
26///         stack.push(42);
27///     }
28/// );
29///
30/// - Falls back to `HeapStack<T>::new_with_capacity(size)` if `size` exceeds thresholds.
31#[macro_export]
32macro_rules! fast_stack {
33    ( $ty:ty,
34      ($first:expr $(, $rest:expr)* $(,)?),
35      $size:expr,
36      $stack_ident:ident,
37      $body:block
38    ) => {{
39        match $size {
40            s if s <= $first => {
41                let mut $stack_ident = $crate::faststack::StackStack::<$ty, $first>::default();
42                $body
43            }
44            $(
45                s if s <= $rest => {
46                    let mut $stack_ident = $crate::faststack::StackStack::<$ty, $rest>::default();
47                    $body
48                }
49            )*
50            _ => {
51                let mut $stack_ident = $crate::faststack::HeapStack::<$ty>::new_with_capacity($size);
52                $body
53            }
54        }
55    }};
56}
57
58/// A stack data structure implemented on the heap with adjustable capacity.
59///
60/// This structure allows pushing and popping elements and will never automatically
61/// allocate or deallocate. The only functions that will result in allocation are
62/// `HeapStack::new_with_capacity` and `HeapStack::reserve`.
63///
64/// The elements must implement the `Clone` and `Default` traits.
65#[derive(Clone)]
66pub struct HeapStack<T: Copy + Default> {
67    data: Vec<T>,
68    index: usize,
69}
70
71impl<T: Copy + Default> Default for HeapStack<T> {
72    // For safety, HeapStack cannot have a capacity of zero.
73    fn default() -> Self {
74        Self::new_with_capacity(1)
75    }
76}
77
78impl<T: Copy + Default> HeapStack<T> {
79    /// Creates a new `HeapStack` with the specified initial capacity.
80    ///
81    /// # Arguments
82    /// * `cap` - The initial capacity of the stack. Must be greater than zero.
83    ///
84    /// # Returns
85    /// A `HeapStack` with pre-allocated space for `cap` elements.
86    ///
87    /// # Panics
88    /// This function will panic if `cap` is zero.
89    #[inline(always)]
90    pub fn new_with_capacity(cap: usize) -> Self {
91        assert!(cap > 0);
92        HeapStack {
93            data: vec![Default::default(); cap],
94            index: 0,
95        }
96    }
97
98    /// Returns the capacity of the stack.
99    ///
100    /// # Returns
101    /// The capacity of the stack.
102    #[inline(always)]
103    pub fn cap(&self) -> usize {
104        self.data.len()
105    }
106
107    /// Reserves capacity for at least `cap` elements.
108    ///
109    /// # Arguments
110    /// * `cap` - The desired capacity.
111    ///
112    /// If the new capacity is smaller than the current capacity, this function does nothing.
113    #[inline(always)]
114    pub fn reserve(&mut self, cap: usize) {
115        if cap < self.data.len() {
116            return;
117        }
118        self.data.resize(cap, Default::default());
119    }
120}
121
122impl<T: Copy + Default> Deref for HeapStack<T> {
123    type Target = [T];
124
125    fn deref(&self) -> &Self::Target {
126        &self.data[..self.index]
127    }
128}
129
130impl<T: Copy + Default> DerefMut for HeapStack<T> {
131    fn deref_mut(&mut self) -> &mut Self::Target {
132        &mut self.data[..self.index]
133    }
134}
135
136impl<T: Copy + Default> FastStack<T> for HeapStack<T> {
137    /// Pushes a value onto the stack.
138    ///
139    /// # Arguments
140    /// * `v` - The value to be pushed onto the stack.
141    ///
142    /// # Panics
143    /// This function will panic if the stack is full.
144    #[inline(always)]
145    fn push(&mut self, v: T) {
146        if self.index < self.data.len() {
147            *unsafe { self.data.get_unchecked_mut(self.index) } = v;
148            self.index += 1;
149        } else {
150            panic!(
151                "Index out of bounds: the HeapStack is full (length: {}) and cannot accommodate more elements",
152                self.data.len()
153            );
154        }
155    }
156
157    /// Pops a value from the stack.
158    ///
159    /// # Returns
160    /// `Some(T)` if the stack is not empty, otherwise `None`.
161    #[inline(always)]
162    fn pop(&mut self) -> Option<T> {
163        if self.index > 0 {
164            self.index = self.index.saturating_sub(1);
165            Some(self.data[self.index])
166        } else {
167            None
168        }
169    }
170
171    /// Pops a value from the stack without checking bounds.
172    ///
173    /// This function is safe to call because a `HeapStack` cannot have a capacity of zero.
174    /// However, if the stack is empty when this function is called, it will access what was previously
175    /// the first value in the stack, which may not be the expected behavior.
176    ///
177    /// # Returns
178    /// The value at the top of the stack.
179    #[inline(always)]
180    fn pop_fast(&mut self) -> T {
181        self.index = self.index.saturating_sub(1);
182        let v = unsafe { self.data.get_unchecked(self.index) };
183        *v
184    }
185
186    /// Returns the number of elements in the stack.
187    ///
188    /// # Returns
189    /// The length of the stack.
190    #[inline(always)]
191    fn len(&self) -> usize {
192        self.index
193    }
194
195    /// Returns true if the stack is empty.
196    #[inline(always)]
197    fn is_empty(&self) -> bool {
198        self.index == 0
199    }
200
201    /// Clears the stack, removing all elements.
202    ///
203    /// This operation does not deallocate the stack's capacity.
204    #[inline(always)]
205    fn clear(&mut self) {
206        self.index = 0;
207    }
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213
214    #[test]
215    fn test_new_with_capacity() {
216        let stack: HeapStack<i32> = HeapStack::new_with_capacity(10);
217        assert_eq!(stack.len(), 0);
218        assert_eq!(stack.data.len(), 10);
219    }
220
221    #[test]
222    fn test_push_and_pop() {
223        let mut stack: HeapStack<i32> = HeapStack::new_with_capacity(10);
224        stack.push(1);
225        stack.push(2);
226        stack.push(3);
227
228        assert_eq!(stack.len(), 3);
229        assert_eq!(stack.pop(), Some(3));
230        assert_eq!(stack.pop(), Some(2));
231        assert_eq!(stack.pop(), Some(1));
232        assert_eq!(stack.pop(), None);
233    }
234
235    #[test]
236    #[should_panic(expected = "Index out of bounds: the HeapStack is full")]
237    fn test_push_panic() {
238        let mut stack: HeapStack<i32> = HeapStack::new_with_capacity(2);
239        stack.push(1);
240        stack.push(2);
241        stack.push(3); // This should panic
242    }
243
244    #[test]
245    fn test_pop_fast() {
246        let mut stack: HeapStack<i32> = HeapStack::new_with_capacity(10);
247        stack.push(1);
248        stack.push(2);
249        stack.push(3);
250
251        assert_eq!(stack.pop_fast(), 3);
252        assert_eq!(stack.pop_fast(), 2);
253        assert_eq!(stack.pop_fast(), 1);
254    }
255
256    #[test]
257    fn test_clear() {
258        let mut stack: HeapStack<i32> = HeapStack::new_with_capacity(10);
259        stack.push(1);
260        stack.push(2);
261        stack.push(3);
262
263        stack.clear();
264        assert_eq!(stack.len(), 0);
265        assert_eq!(stack.pop(), None);
266    }
267
268    #[test]
269    fn test_reserve() {
270        let mut stack: HeapStack<i32> = HeapStack::new_with_capacity(5);
271        assert_eq!(stack.data.len(), 5);
272
273        stack.reserve(10);
274        assert_eq!(stack.data.len(), 10);
275
276        stack.reserve(3); // Should not shrink the capacity
277        assert_eq!(stack.data.len(), 10);
278    }
279}
280
281/// A stack data structure implemented on the stack with fixed capacity.
282pub struct StackStack<T: Copy + Default, const STACK_SIZE: usize> {
283    data: [T; STACK_SIZE],
284    index: usize,
285}
286
287impl<T: Copy + Default, const STACK_SIZE: usize> Default for StackStack<T, STACK_SIZE> {
288    fn default() -> Self {
289        Self {
290            data: [Default::default(); STACK_SIZE],
291            index: Default::default(),
292        }
293    }
294}
295
296impl<T: Copy + Default, const STACK_SIZE: usize> FastStack<T> for StackStack<T, STACK_SIZE> {
297    /// Pushes a value onto the stack. If the stack is full it will overwrite the value in the last position.
298    #[inline(always)]
299    fn push(&mut self, v: T) {
300        // TODO: possibly check bounds in debug or make a push_fast()
301        *unsafe { self.data.get_unchecked_mut(self.index) } = v;
302        self.index = (self.index + 1).min(STACK_SIZE - 1);
303    }
304    /// Pops a value from the stack without checking bounds. If the stack is empty it will return the value in the first position.
305    #[inline(always)]
306    fn pop_fast(&mut self) -> T {
307        self.index = self.index.saturating_sub(1);
308        let v = unsafe { self.data.get_unchecked(self.index) };
309        *v
310    }
311    /// Pops a value from the stack.
312    #[inline(always)]
313    fn pop(&mut self) -> Option<T> {
314        if self.index > 0 {
315            self.index = self.index.saturating_sub(1);
316            Some(self.data[self.index])
317        } else {
318            None
319        }
320    }
321    /// Returns the number of elements in the stack.
322    #[inline(always)]
323    fn len(&self) -> usize {
324        self.index
325    }
326    /// Returns true if the stack is empty.
327    #[inline(always)]
328    fn is_empty(&self) -> bool {
329        self.index == 0
330    }
331    /// Clears the stack, removing all elements.
332    #[inline(always)]
333    fn clear(&mut self) {
334        self.index = 0;
335    }
336}
337
338impl<T: Copy + Default> FastStack<T> for Vec<T> {
339    fn push(&mut self, v: T) {
340        self.push(v);
341    }
342    fn pop_fast(&mut self) -> T {
343        self.pop().unwrap()
344    }
345    fn pop(&mut self) -> Option<T> {
346        self.pop()
347    }
348    fn len(&self) -> usize {
349        self.len()
350    }
351    fn is_empty(&self) -> bool {
352        self.is_empty()
353    }
354    fn clear(&mut self) {
355        self.clear();
356    }
357}