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

skrifa/
collections.rs
