polars_utils/
enum_unit_vec.rs

1use either::Either;
2
3/// List of items where a single item is stored on the stack.
4/// Similar to UnitVec, but without size / alignment limitations.
5#[derive(Debug, Clone, PartialEq, Hash)]
6pub struct EUnitVec<T>(Either<[T; 1], Vec<T>>);
7
8impl<T> EUnitVec<T> {
9    pub const fn new() -> Self {
10        Self(Either::Right(Vec::new()))
11    }
12
13    pub const fn new_single(value: T) -> Self {
14        Self(Either::Left([value]))
15    }
16
17    #[inline]
18    pub fn capacity(&self) -> usize {
19        match &self.0 {
20            Either::Left(_) => 1,
21            Either::Right(vec) => vec.capacity().max(1),
22        }
23    }
24
25    #[inline]
26    pub fn clear(&mut self) {
27        match &mut self.0 {
28            Either::Left(_) => *self = Self::new(),
29            Either::Right(vec) => vec.clear(),
30        }
31    }
32
33    #[inline]
34    pub fn push(&mut self, value: T) {
35        match &mut self.0 {
36            Either::Left(_) => {
37                let first = self.single_to_vec(2);
38
39                let vec = self.as_mut_vec().unwrap();
40                vec.push(first);
41                vec.push(value);
42            },
43            Either::Right(vec) => {
44                if vec.capacity() == 0 {
45                    self.0 = Either::Left([value])
46                } else {
47                    vec.push(value)
48                }
49            },
50        }
51    }
52
53    pub fn reserve(&mut self, additional: usize) {
54        match &mut self.0 {
55            Either::Left(_) => {
56                let new_len = self.len().checked_add(additional).unwrap();
57                if new_len > 1 {
58                    let first = self.single_to_vec(new_len);
59                    self.as_mut_vec().unwrap().push(first);
60                }
61            },
62            Either::Right(vec) => vec.reserve(additional),
63        }
64    }
65
66    pub fn with_capacity(capacity: usize) -> Self {
67        if capacity <= 1 {
68            Self::new()
69        } else {
70            Self(Either::Right(Vec::with_capacity(capacity)))
71        }
72    }
73
74    /// Replaces an inner repr of `Left([T])` with `Right(Vec::with_capacity(capacity))`, and returns
75    /// `T`.
76    ///
77    /// # Panics
78    /// Panics if `self` does not contain `Left([T])`.
79    #[inline]
80    fn single_to_vec(&mut self, capacity: usize) -> T {
81        let Either::Left([first]) =
82            std::mem::replace(&mut self.0, Either::Right(Vec::with_capacity(capacity)))
83        else {
84            panic!()
85        };
86
87        first
88    }
89
90    #[inline]
91    fn as_mut_vec(&mut self) -> Option<&mut Vec<T>> {
92        if let Either::Right(vec) = &mut self.0 {
93            Some(vec)
94        } else {
95            None
96        }
97    }
98
99    #[inline]
100    pub fn as_slice(&self) -> &[T] {
101        self.as_ref()
102    }
103
104    #[inline]
105    pub fn as_mut_slice(&mut self) -> &mut [T] {
106        self.as_mut()
107    }
108
109    #[inline]
110    pub fn pop(&mut self) -> Option<T> {
111        match &mut self.0 {
112            Either::Left(_) => Some(self.single_to_vec(0)),
113            Either::Right(vec) => vec.pop(),
114        }
115    }
116}
117
118impl<T> Default for EUnitVec<T> {
119    fn default() -> Self {
120        Self::new()
121    }
122}
123
124impl<T> std::ops::Deref for EUnitVec<T> {
125    type Target = [T];
126
127    fn deref(&self) -> &Self::Target {
128        AsRef::as_ref(&self.0)
129    }
130}
131
132impl<T> std::ops::DerefMut for EUnitVec<T> {
133    fn deref_mut(&mut self) -> &mut Self::Target {
134        AsMut::as_mut(&mut self.0)
135    }
136}
137
138impl<T> From<Vec<T>> for EUnitVec<T> {
139    fn from(value: Vec<T>) -> Self {
140        Self(Either::Right(value))
141    }
142}
143
144impl<T, const N: usize> From<[T; N]> for EUnitVec<T> {
145    fn from(value: [T; N]) -> Self {
146        value.into_iter().collect()
147    }
148}
149
150impl<T> IntoIterator for EUnitVec<T> {
151    type IntoIter = Either<<[T; 1] as IntoIterator>::IntoIter, <Vec<T> as IntoIterator>::IntoIter>;
152    type Item = T;
153
154    fn into_iter(self) -> Self::IntoIter {
155        match self.0 {
156            Either::Left(v) => Either::Left(v.into_iter()),
157            Either::Right(v) => Either::Right(v.into_iter()),
158        }
159    }
160}
161
162impl<T> FromIterator<T> for EUnitVec<T> {
163    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
164        let mut iter = iter.into_iter();
165
166        let Some(first) = iter.next() else {
167            return Self::new();
168        };
169
170        let Some(second) = iter.next() else {
171            return Self::new_single(first);
172        };
173
174        let mut vec = Vec::with_capacity(iter.size_hint().0 + 2);
175        vec.push(first);
176        vec.push(second);
177        vec.extend(iter);
178        Self::from(vec)
179    }
180}
181
182impl<T> Extend<T> for EUnitVec<T> {
183    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
184        let mut iter = iter.into_iter();
185        self.reserve(iter.size_hint().0);
186
187        loop {
188            if self.as_mut_vec().is_some() {
189                break;
190            }
191
192            let Some(item) = iter.next() else {
193                return;
194            };
195
196            self.push(item);
197        }
198
199        self.as_mut_vec().unwrap().extend(iter)
200    }
201}
202
203#[macro_export]
204macro_rules! eunitvec {
205    () => {{
206        $crate::enum_unit_vec::EUnitVec::new()
207    }};
208    ($elem:expr; $n:expr) => {{
209        let mut new = $crate::enum_unit_vec::EUnitVec::new();
210        for _ in 0..$n {
211            new.push($elem)
212        }
213        new
214    }};
215    ($elem:expr) => {{
216        $crate::enum_unit_vec::EUnitVec::new_single($elem)
217    }};
218    ($($x:expr),+ $(,)?) => {{
219        vec![$($x),+].into()
220    }};
221}
222
223mod tests {
224
225    #[test]
226    fn test_enum_unitvec_clone() {
227        {
228            let v = eunitvec![1usize];
229            assert_eq!(v, v.clone());
230        }
231
232        for n in [
233            26903816120209729usize,
234            42566276440897687,
235            44435161834424652,
236            49390731489933083,
237            51201454727649242,
238            83861672190814841,
239            92169290527847622,
240            92476373900398436,
241            95488551309275459,
242            97499984126814549,
243        ] {
244            let v = eunitvec![n];
245            assert_eq!(v, v.clone());
246        }
247    }
248
249    #[test]
250    fn test_enum_unitvec_repeat_n() {
251        assert_eq!(eunitvec![5; 3].as_slice(), &[5, 5, 5])
252    }
253}
polars_utils/enum_unit_vec.rs

polars_utils/
enum_unit_vec.rs
