peg_pack/runtime/
small_vec.rs

1use std::hint::unreachable_unchecked;
2use std::mem;
3
4use super::array_vec::ArrayVec;
5
6pub struct SmallVec<T, const N: usize> {
7    data: Data<T, N>,
8}
9
10enum Data<T, const N: usize> {
11    Stack(ArrayVec<T, N>),
12    Heap(Vec<T>),
13}
14
15impl<T, const N: usize> SmallVec<T, N> {
16    pub fn new() -> Self {
17        Self {
18            data: Data::Stack(ArrayVec::new()),
19        }
20    }
21
22    pub fn push(&mut self, value: T) {
23        match &mut self.data {
24            Data::Stack(stack) => unsafe {
25                if stack.len() < N {
26                    stack.push_unchecked(value);
27                } else {
28                    self.promote();
29                    self.push_heap(value);
30                }
31            },
32            Data::Heap(_) => unsafe {
33                self.push_heap(value);
34            },
35        }
36    }
37
38    pub fn pop(&mut self) -> Option<T> {
39        match &mut self.data {
40            Data::Stack(stack) => stack.pop(),
41            Data::Heap(heap) => heap.pop(),
42        }
43    }
44
45    pub fn last(&self) -> Option<&T> {
46        match &self.data {
47            Data::Stack(stack) => stack.last(),
48            Data::Heap(heap) => heap.last(),
49        }
50    }
51
52    pub fn last_mut(&mut self) -> Option<&mut T> {
53        match &mut self.data {
54            Data::Stack(stack) => stack.last_mut(),
55            Data::Heap(heap) => heap.last_mut(),
56        }
57    }
58
59    unsafe fn push_heap(&mut self, value: T) {
60        match &mut self.data {
61            Data::Heap(heap) => heap.push(value),
62            Data::Stack(_) => unreachable_unchecked(),
63        }
64    }
65
66    unsafe fn promote(&mut self) {
67        match &mut self.data {
68            Data::Stack(stack) => {
69                let mut vector = Vec::with_capacity(N * 2);
70
71                for value in mem::replace(stack, ArrayVec::new()) {
72                    vector.push(value);
73                }
74
75                self.data = Data::Heap(vector);
76            }
77            Data::Heap(_) => unreachable_unchecked(),
78        }
79    }
80}