skrifa/
collections.rs

Help
1//! Internal "small" style collection types.
2
3use alloc::vec::Vec;
4use core::hash::{Hash, Hasher};
5
6/// A growable vector type with inline storage optimization.
7///
8/// Note that unlike the real `SmallVec`, this only works with types that
9/// are `Copy + Default` to simplify our implementation.
10#[derive(Clone)]
11pub(crate) struct SmallVec<T, const N: usize>(Storage<T, N>);
12
13impl<T, const N: usize> SmallVec<T, N>
14where
15    T: Copy + Default,
16{
17    /// Creates a new, empty `SmallVec<T>`.
18    pub fn new() -> Self {
19        Self(Storage::Inline([T::default(); N], 0))
20    }
21
22    /// Creates a new `SmallVec<T>` of the given length with each element
23    /// containing a copy of `value`.
24    pub fn with_len(len: usize, value: T) -> Self {
25        if len <= N {
26            Self(Storage::Inline([value; N], len))
27        } else {
28            let mut vec = Vec::new();
29            vec.resize(len, value);
30            Self(Storage::Heap(vec))
31        }
32    }
33
34    /// Clears the vector, removing all values.
35    pub fn clear(&mut self) {
36        match &mut self.0 {
37            Storage::Inline(_buf, len) => *len = 0,
38            Storage::Heap(vec) => vec.clear(),
39        }
40    }
41
42    /// Tries to reserve capacity for at least `additional` more elements.
43    pub fn try_reserve(&mut self, additional: usize) -> bool {
44        match &mut self.0 {
45            Storage::Inline(buf, len) => {
46                let new_cap = *len + additional;
47                if new_cap > N {
48                    let mut vec = Vec::new();
49                    if vec.try_reserve(new_cap).is_err() {
50                        return false;
51                    }
52                    vec.extend_from_slice(&buf[..*len]);
53                    self.0 = Storage::Heap(vec);
54                }
55            }
56            Storage::Heap(vec) => {
57                if vec.try_reserve(additional).is_err() {
58                    return false;
59                }
60            }
61        }
62        true
63    }
64
65    /// Appends an element to the back of the collection.
66    pub fn push(&mut self, value: T) {
67        match &mut self.0 {
68            Storage::Inline(buf, len) => {
69                if *len + 1 > N {
70                    let mut vec = Vec::with_capacity(*len + 1);
71                    vec.extend_from_slice(&buf[..*len]);
72                    vec.push(value);
73                    self.0 = Storage::Heap(vec);
74                } else {
75                    buf[*len] = value;
76                    *len += 1;
77                }
78            }
79            Storage::Heap(vec) => vec.push(value),
80        }
81    }
82
83    /// Removes and returns the value at the back of the collection.
84    pub fn pop(&mut self) -> Option<T> {
85        match &mut self.0 {
86            Storage::Inline(buf, len) => {
87                if *len > 0 {
88                    *len -= 1;
89                    Some(buf[*len])
90                } else {
91                    None
92                }
93            }
94            Storage::Heap(vec) => vec.pop(),
95        }
96    }
97
98    /// Shortens the vector, keeping the first `len` elements.
99    pub fn truncate(&mut self, len: usize) {
100        match &mut self.0 {
101            Storage::Inline(_buf, inline_len) => {
102                *inline_len = len.min(*inline_len);
103            }
104            Storage::Heap(vec) => vec.truncate(len),
105        }
106    }
107}
108
109impl<T, const N: usize> SmallVec<T, N> {
110    /// Extracts a slice containing the entire vector.
111    pub fn as_slice(&self) -> &[T] {
112        match &self.0 {
113            Storage::Inline(buf, len) => &buf[..*len],
114            Storage::Heap(vec) => vec.as_slice(),
115        }
116    }
117
118    /// Extracts a mutable slice containing the entire vector.
119    pub fn as_mut_slice(&mut self) -> &mut [T] {
120        match &mut self.0 {
121            Storage::Inline(buf, len) => &mut buf[..*len],
122            Storage::Heap(vec) => vec.as_mut_slice(),
123        }
124    }
125}
126
127impl<T, const N: usize> Default for SmallVec<T, N>
128where
129    T: Copy + Default,
130{
131    fn default() -> Self {
132        Self::new()
133    }
134}
135
136impl<T, const N: usize> core::ops::Deref for SmallVec<T, N> {
137    type Target = [T];
138
139    fn deref(&self) -> &Self::Target {
140        self.as_slice()
141    }
142}
143
144impl<T, const N: usize> core::ops::DerefMut for SmallVec<T, N> {
145    fn deref_mut(&mut self) -> &mut Self::Target {
146        self.as_mut_slice()
147    }
148}
149
150impl<T, const N: usize> Hash for SmallVec<T, N>
151where
152    T: Hash,
153{
154    fn hash<H: Hasher>(&self, state: &mut H) {
155        self.as_slice().hash(state);
156    }
157}
158
159impl<T, const N: usize> PartialEq for SmallVec<T, N>
160where
161    T: PartialEq,
162{
163    fn eq(&self, other: &Self) -> bool {
164        self.as_slice() == other.as_slice()
165    }
166}
167
168impl<T, const N: usize> PartialEq<[T]> for SmallVec<T, N>
169where
170    T: PartialEq,
171{
172    fn eq(&self, other: &[T]) -> bool {
173        self.as_slice() == other
174    }
175}
176
177impl<T, const N: usize> Eq for SmallVec<T, N> where T: Eq {}
178
179impl<T, const N: usize> core::fmt::Debug for SmallVec<T, N>
180where
181    T: core::fmt::Debug,
182{
183    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
184        f.debug_list().entries(self.as_slice().iter()).finish()
185    }
186}
187
188impl<'a, T, const N: usize> IntoIterator for &'a SmallVec<T, N> {
189    type IntoIter = core::slice::Iter<'a, T>;
190    type Item = &'a T;
191
192    fn into_iter(self) -> Self::IntoIter {
193        self.as_slice().iter()
194    }
195}
196
197impl<'a, T, const N: usize> IntoIterator for &'a mut SmallVec<T, N> {
198    type IntoIter = core::slice::IterMut<'a, T>;
199    type Item = &'a mut T;
200
201    fn into_iter(self) -> Self::IntoIter {
202        self.as_mut_slice().iter_mut()
203    }
204}
205
206impl<T, const N: usize> IntoIterator for SmallVec<T, N>
207where
208    T: Copy,
209{
210    type IntoIter = IntoIter<T, N>;
211    type Item = T;
212
213    fn into_iter(self) -> Self::IntoIter {
214        IntoIter { vec: self, pos: 0 }
215    }
216}
217
218#[derive(Clone)]
219pub(crate) struct IntoIter<T, const N: usize> {
220    vec: SmallVec<T, N>,
221    pos: usize,
222}
223
224impl<T, const N: usize> Iterator for IntoIter<T, N>
225where
226    T: Copy,
227{
228    type Item = T;
229
230    fn next(&mut self) -> Option<Self::Item> {
231        let value = self.vec.get(self.pos)?;
232        self.pos += 1;
233        Some(*value)
234    }
235}
236
237#[derive(Clone)]
238enum Storage<T, const N: usize> {
239    Inline([T; N], usize),
240    Heap(Vec<T>),
241}
242
243#[cfg(test)]
244mod test {
245    use super::{SmallVec, Storage};
246
247    #[test]
248    fn choose_inline() {
249        let vec = SmallVec::<_, 4>::with_len(4, 0);
250        assert!(matches!(vec.0, Storage::Inline(..)));
251        assert_eq!(vec.len(), 4);
252    }
253
254    #[test]
255    fn choose_heap() {
256        let vec = SmallVec::<_, 4>::with_len(5, 0);
257        assert!(matches!(vec.0, Storage::Heap(..)));
258        assert_eq!(vec.len(), 5);
259    }
260
261    #[test]
262    fn store_and_read_inline() {
263        let mut vec = SmallVec::<_, 8>::with_len(8, 0);
264        for (i, value) in vec.iter_mut().enumerate() {
265            *value = i * 2;
266        }
267        let expected = [0, 2, 4, 6, 8, 10, 12, 14];
268        assert_eq!(vec.as_slice(), &expected);
269        assert_eq!(format!("{vec:?}"), format!("{expected:?}"));
270    }
271
272    #[test]
273    fn store_and_read_heap() {
274        let mut vec = SmallVec::<_, 4>::with_len(8, 0);
275        for (i, value) in vec.iter_mut().enumerate() {
276            *value = i * 2;
277        }
278        let expected = [0, 2, 4, 6, 8, 10, 12, 14];
279        assert_eq!(vec.as_slice(), &expected);
280        assert_eq!(format!("{vec:?}"), format!("{expected:?}"));
281    }
282
283    #[test]
284    fn spill_to_heap() {
285        let mut vec = SmallVec::<_, 4>::new();
286        for i in 0..4 {
287            vec.push(i);
288        }
289        assert!(matches!(vec.0, Storage::Inline(..)));
290        vec.push(4);
291        assert!(matches!(vec.0, Storage::Heap(..)));
292        let expected = [0, 1, 2, 3, 4];
293        assert_eq!(vec.as_slice(), &expected);
294    }
295
296    #[test]
297    fn clear_inline() {
298        let mut vec = SmallVec::<_, 4>::new();
299        for i in 0..4 {
300            vec.push(i);
301        }
302        assert!(matches!(vec.0, Storage::Inline(..)));
303        assert_eq!(vec.len(), 4);
304        vec.clear();
305        assert_eq!(vec.len(), 0);
306    }
307
308    #[test]
309    fn clear_heap() {
310        let mut vec = SmallVec::<_, 3>::new();
311        for i in 0..4 {
312            vec.push(i);
313        }
314        assert!(matches!(vec.0, Storage::Heap(..)));
315        assert_eq!(vec.len(), 4);
316        vec.clear();
317        assert_eq!(vec.len(), 0);
318    }
319
320    #[test]
321    fn reserve() {
322        let mut vec = SmallVec::<_, 3>::new();
323        for i in 0..2 {
324            vec.push(i);
325        }
326        assert!(matches!(vec.0, Storage::Inline(..)));
327        assert!(vec.try_reserve(1));
328        // still inline after reserving 1
329        assert!(matches!(vec.0, Storage::Inline(..)));
330        assert!(vec.try_reserve(2));
331        // reserving 2 spills to heap
332        assert!(matches!(vec.0, Storage::Heap(..)));
333    }
334
335    #[test]
336    fn iter() {
337        let mut vec = SmallVec::<_, 3>::new();
338        for i in 0..3 {
339            vec.push(i);
340        }
341        assert!(&[0, 1, 2].iter().eq(vec.iter()));
342    }
343
344    #[test]
345    fn into_iter() {
346        let mut vec = SmallVec::<_, 3>::new();
347        for i in 0..3 {
348            vec.push(i);
349        }
350        assert!([0, 1, 2].into_iter().eq(vec.into_iter()));
351    }
352}
skrifa/collections.rs

skrifa/
collections.rs
