skrifa/
collections.rs

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    /// Resizes the vector to `len` elements, filling with `value`.
109    /// Reuses existing heap allocation when possible.
110    pub fn resize_and_fill(&mut self, len: usize, value: T) {
111        match &mut self.0 {
112            Storage::Inline(buf, inline_len) => {
113                if len <= N {
114                    buf[..len].fill(value);
115                    *inline_len = len;
116                } else {
117                    // Need to spill to heap
118                    let mut vec = Vec::with_capacity(len);
119                    vec.resize(len, value);
120                    self.0 = Storage::Heap(vec);
121                }
122            }
123            Storage::Heap(vec) => {
124                vec.clear();
125                vec.resize(len, value);
126            }
127        }
128    }
129}
130
131impl<T, const N: usize> SmallVec<T, N> {
132    /// Extracts a slice containing the entire vector.
133    pub fn as_slice(&self) -> &[T] {
134        match &self.0 {
135            Storage::Inline(buf, len) => &buf[..*len],
136            Storage::Heap(vec) => vec.as_slice(),
137        }
138    }
139
140    /// Extracts a mutable slice containing the entire vector.
141    pub fn as_mut_slice(&mut self) -> &mut [T] {
142        match &mut self.0 {
143            Storage::Inline(buf, len) => &mut buf[..*len],
144            Storage::Heap(vec) => vec.as_mut_slice(),
145        }
146    }
147}
148
149impl<T, const N: usize> Default for SmallVec<T, N>
150where
151    T: Copy + Default,
152{
153    fn default() -> Self {
154        Self::new()
155    }
156}
157
158impl<T, const N: usize> core::ops::Deref for SmallVec<T, N> {
159    type Target = [T];
160
161    fn deref(&self) -> &Self::Target {
162        self.as_slice()
163    }
164}
165
166impl<T, const N: usize> core::ops::DerefMut for SmallVec<T, N> {
167    fn deref_mut(&mut self) -> &mut Self::Target {
168        self.as_mut_slice()
169    }
170}
171
172impl<T, const N: usize> Hash for SmallVec<T, N>
173where
174    T: Hash,
175{
176    fn hash<H: Hasher>(&self, state: &mut H) {
177        self.as_slice().hash(state);
178    }
179}
180
181impl<T, const N: usize> PartialEq for SmallVec<T, N>
182where
183    T: PartialEq,
184{
185    fn eq(&self, other: &Self) -> bool {
186        self.as_slice() == other.as_slice()
187    }
188}
189
190impl<T, const N: usize> PartialEq<[T]> for SmallVec<T, N>
191where
192    T: PartialEq,
193{
194    fn eq(&self, other: &[T]) -> bool {
195        self.as_slice() == other
196    }
197}
198
199impl<T, const N: usize> Eq for SmallVec<T, N> where T: Eq {}
200
201impl<T, const N: usize> core::fmt::Debug for SmallVec<T, N>
202where
203    T: core::fmt::Debug,
204{
205    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
206        f.debug_list().entries(self.as_slice().iter()).finish()
207    }
208}
209
210impl<'a, T, const N: usize> IntoIterator for &'a SmallVec<T, N> {
211    type IntoIter = core::slice::Iter<'a, T>;
212    type Item = &'a T;
213
214    fn into_iter(self) -> Self::IntoIter {
215        self.as_slice().iter()
216    }
217}
218
219impl<'a, T, const N: usize> IntoIterator for &'a mut SmallVec<T, N> {
220    type IntoIter = core::slice::IterMut<'a, T>;
221    type Item = &'a mut T;
222
223    fn into_iter(self) -> Self::IntoIter {
224        self.as_mut_slice().iter_mut()
225    }
226}
227
228impl<T, const N: usize> IntoIterator for SmallVec<T, N>
229where
230    T: Copy,
231{
232    type IntoIter = IntoIter<T, N>;
233    type Item = T;
234
235    fn into_iter(self) -> Self::IntoIter {
236        IntoIter { vec: self, pos: 0 }
237    }
238}
239
240#[derive(Clone)]
241pub(crate) struct IntoIter<T, const N: usize> {
242    vec: SmallVec<T, N>,
243    pos: usize,
244}
245
246impl<T, const N: usize> Iterator for IntoIter<T, N>
247where
248    T: Copy,
249{
250    type Item = T;
251
252    fn next(&mut self) -> Option<Self::Item> {
253        let value = self.vec.get(self.pos)?;
254        self.pos += 1;
255        Some(*value)
256    }
257}
258
259#[derive(Clone)]
260enum Storage<T, const N: usize> {
261    Inline([T; N], usize),
262    Heap(Vec<T>),
263}
264
265#[cfg(test)]
266mod test {
267    use super::{SmallVec, Storage};
268
269    #[test]
270    fn choose_inline() {
271        let vec = SmallVec::<_, 4>::with_len(4, 0);
272        assert!(matches!(vec.0, Storage::Inline(..)));
273        assert_eq!(vec.len(), 4);
274    }
275
276    #[test]
277    fn choose_heap() {
278        let vec = SmallVec::<_, 4>::with_len(5, 0);
279        assert!(matches!(vec.0, Storage::Heap(..)));
280        assert_eq!(vec.len(), 5);
281    }
282
283    #[test]
284    fn store_and_read_inline() {
285        let mut vec = SmallVec::<_, 8>::with_len(8, 0);
286        for (i, value) in vec.iter_mut().enumerate() {
287            *value = i * 2;
288        }
289        let expected = [0, 2, 4, 6, 8, 10, 12, 14];
290        assert_eq!(vec.as_slice(), &expected);
291        assert_eq!(format!("{vec:?}"), format!("{expected:?}"));
292    }
293
294    #[test]
295    fn store_and_read_heap() {
296        let mut vec = SmallVec::<_, 4>::with_len(8, 0);
297        for (i, value) in vec.iter_mut().enumerate() {
298            *value = i * 2;
299        }
300        let expected = [0, 2, 4, 6, 8, 10, 12, 14];
301        assert_eq!(vec.as_slice(), &expected);
302        assert_eq!(format!("{vec:?}"), format!("{expected:?}"));
303    }
304
305    #[test]
306    fn spill_to_heap() {
307        let mut vec = SmallVec::<_, 4>::new();
308        for i in 0..4 {
309            vec.push(i);
310        }
311        assert!(matches!(vec.0, Storage::Inline(..)));
312        vec.push(4);
313        assert!(matches!(vec.0, Storage::Heap(..)));
314        let expected = [0, 1, 2, 3, 4];
315        assert_eq!(vec.as_slice(), &expected);
316    }
317
318    #[test]
319    fn clear_inline() {
320        let mut vec = SmallVec::<_, 4>::new();
321        for i in 0..4 {
322            vec.push(i);
323        }
324        assert!(matches!(vec.0, Storage::Inline(..)));
325        assert_eq!(vec.len(), 4);
326        vec.clear();
327        assert_eq!(vec.len(), 0);
328    }
329
330    #[test]
331    fn clear_heap() {
332        let mut vec = SmallVec::<_, 3>::new();
333        for i in 0..4 {
334            vec.push(i);
335        }
336        assert!(matches!(vec.0, Storage::Heap(..)));
337        assert_eq!(vec.len(), 4);
338        vec.clear();
339        assert_eq!(vec.len(), 0);
340    }
341
342    #[test]
343    fn reserve() {
344        let mut vec = SmallVec::<_, 3>::new();
345        for i in 0..2 {
346            vec.push(i);
347        }
348        assert!(matches!(vec.0, Storage::Inline(..)));
349        assert!(vec.try_reserve(1));
350        // still inline after reserving 1
351        assert!(matches!(vec.0, Storage::Inline(..)));
352        assert!(vec.try_reserve(2));
353        // reserving 2 spills to heap
354        assert!(matches!(vec.0, Storage::Heap(..)));
355    }
356
357    #[test]
358    fn iter() {
359        let mut vec = SmallVec::<_, 3>::new();
360        for i in 0..3 {
361            vec.push(i);
362        }
363        assert!(&[0, 1, 2].iter().eq(vec.iter()));
364    }
365
366    #[test]
367    fn into_iter() {
368        let mut vec = SmallVec::<_, 3>::new();
369        for i in 0..3 {
370            vec.push(i);
371        }
372        assert!([0, 1, 2].into_iter().eq(vec.into_iter()));
373    }
374}
skrifa/collections.rs

skrifa/
collections.rs
