tl/inline/
vec.rs

1use std::fmt::{Debug, Formatter};
2use std::mem::MaybeUninit;
3use std::ops::Index;
4use std::ptr;
5
6/// A wrapper around a `Vec<T>` that lives on the stack if it is small enough.
7#[derive(Debug, Clone)]
8pub struct InlineVec<T, const N: usize>(InlineVecInner<T, N>);
9
10impl<T, const N: usize> InlineVec<T, N> {
11    /// Creates a new InlineVec
12    #[inline]
13    pub(crate) fn new() -> Self {
14        Self(InlineVecInner::new())
15    }
16
17    /// Returns the number of elements in the vector
18    #[inline]
19    pub fn len(&self) -> usize {
20        self.0.len()
21    }
22
23    /// Returns true if the vector contains no elements
24    #[inline]
25    pub fn is_empty(&self) -> bool {
26        self.len() == 0
27    }
28
29    /// Checks whether this vector is allocated on the heap
30    #[inline]
31    pub fn is_heap_allocated(&self) -> bool {
32        self.0.is_heap_allocated()
33    }
34
35    /// Copies `self` into a new `Vec<T>`
36    #[inline]
37    pub fn to_vec(&self) -> Vec<T>
38    where
39        T: Clone,
40    {
41        self.0.to_vec()
42    }
43
44    /// Inserts a new element into the vector
45    #[inline]
46    pub fn push(&mut self, value: T) {
47        self.0.push(value)
48    }
49
50    /// Returns a reference to the value at the given index
51    #[inline]
52    pub fn get(&self, index: usize) -> Option<&T> {
53        self.0.get(index)
54    }
55
56    /// Returns a mutable reference to the value at the given index
57    #[inline]
58    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
59        self.0.get_mut(index)
60    }
61
62    /// Removes an element at a given index
63    ///
64    /// # Panics
65    /// Just like `Vec::remove`, this method will panic if the index is out of bounds.
66    #[inline]
67    pub fn remove(&mut self, index: usize) -> T {
68        self.0.remove(index)
69    }
70
71    /// Returns an iterator over the elements of this vector
72    #[inline]
73    pub fn iter(&self) -> InlineVecIter<'_, T, N> {
74        self.0.iter()
75    }
76
77    /// Returns a slice to the contents of this vector
78    #[inline]
79    pub fn as_slice(&self) -> &[T] {
80        self.0.as_slice()
81    }
82}
83
84enum InlineVecInner<T, const N: usize> {
85    Inline {
86        len: usize,
87        data: [MaybeUninit<T>; N],
88    },
89    Heap(Vec<T>),
90}
91
92impl<T, const N: usize> Debug for InlineVecInner<T, N>
93where
94    T: Debug,
95{
96    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
97        write!(f, "InlineVec<{} items>", self.len())
98    }
99}
100
101impl<T, const N: usize> Clone for InlineVecInner<T, N>
102where
103    T: Clone,
104{
105    fn clone(&self) -> Self {
106        match self {
107            Self::Heap(m) => Self::Heap(m.clone()),
108            Self::Inline { len, data } => {
109                let mut new_data = super::uninit_array();
110
111                let iter = data.iter().take(*len).enumerate();
112
113                for (idx, element) in iter {
114                    let element = unsafe { &*element.as_ptr() };
115                    new_data[idx] = MaybeUninit::new(T::clone(element));
116                }
117
118                Self::Inline {
119                    len: *len,
120                    data: new_data,
121                }
122            }
123        }
124    }
125}
126
127impl<T, const N: usize> InlineVecInner<T, N> {
128    #[inline]
129    pub(crate) fn new() -> Self {
130        Self::Inline {
131            len: 0,
132            data: super::uninit_array(),
133        }
134    }
135
136    pub fn as_slice(&self) -> &[T] {
137        match self {
138            Self::Heap(v) => v.as_slice(),
139            Self::Inline { len, data } => unsafe {
140                std::slice::from_raw_parts(data.as_ptr() as *const T, *len)
141            },
142        }
143    }
144
145    pub fn to_vec(&self) -> Vec<T>
146    where
147        T: Clone,
148    {
149        match &self {
150            InlineVecInner::Heap(m) => m.to_vec(),
151            InlineVecInner::Inline { len, data } => {
152                let mut new_data = Vec::with_capacity(*len);
153
154                let iter = data.iter().take(*len);
155
156                for element in iter {
157                    new_data.push(unsafe { T::clone(&*element.as_ptr()) });
158                }
159
160                new_data
161            }
162        }
163    }
164
165    #[inline]
166    pub fn iter(&self) -> InlineVecIter<'_, T, N> {
167        InlineVecIter { idx: 0, vec: self }
168    }
169
170    #[inline]
171    pub fn len(&self) -> usize {
172        match self {
173            Self::Inline { len, .. } => *len,
174            Self::Heap(vec) => vec.len(),
175        }
176    }
177
178    pub fn get(&self, idx: usize) -> Option<&T> {
179        match self {
180            Self::Inline { data, len } => {
181                if idx < *len {
182                    Some(unsafe { &*data.get_unchecked(idx).as_ptr() })
183                } else {
184                    None
185                }
186            }
187            Self::Heap(vec) => vec.get(idx),
188        }
189    }
190
191    pub fn get_mut(&mut self, idx: usize) -> Option<&mut T> {
192        match self {
193            Self::Inline { data, len } => {
194                if idx < *len {
195                    Some(unsafe { &mut *data.get_unchecked_mut(idx).as_mut_ptr() })
196                } else {
197                    None
198                }
199            }
200            Self::Heap(vec) => vec.get_mut(idx),
201        }
202    }
203
204    pub fn remove(&mut self, idx: usize) -> T {
205        match self {
206            Self::Inline { data, len } => {
207                assert!(idx < *len);
208
209                // at this point we know idx is in bounds
210                // carefully replace the value with MaybeUninit::uninit(), so it can be returned
211                let element = unsafe {
212                    std::mem::replace(data.get_unchecked_mut(idx), MaybeUninit::uninit())
213                };
214
215                for i in idx + 1..*len {
216                    // TODO(y21): data.swap_unchecked() worth it?
217                    data.swap(i, i - 1);
218                }
219
220                *len -= 1;
221
222                // we've made sure that idx is in bounds and if idx is in bounds, then `T` must be initialized
223                unsafe { element.assume_init() }
224            }
225            Self::Heap(h) => h.remove(idx),
226        }
227    }
228
229    pub fn push(&mut self, value: T) {
230        let (array, len) = match self {
231            Self::Inline { data, len } => (data, len),
232            Self::Heap(vec) => {
233                vec.push(value);
234                return;
235            }
236        };
237
238        if *len >= N {
239            let mut vec = Vec::with_capacity(*len + 1);
240
241            // move old elements to heap
242            for element in array.iter_mut().take(*len) {
243                let element = std::mem::replace(element, MaybeUninit::uninit());
244
245                vec.push(unsafe { element.assume_init() });
246            }
247
248            // push the new element
249            vec.push(value);
250            let new_heap = InlineVecInner::Heap(vec);
251
252            // do not call the destructor!
253            unsafe { ptr::write(self, new_heap) };
254        } else {
255            array[*len].write(value);
256            *len += 1;
257        }
258    }
259
260    #[inline]
261    pub fn is_heap_allocated(&self) -> bool {
262        matches!(self, Self::Heap(_))
263    }
264}
265
266impl<T, const N: usize> Index<usize> for InlineVec<T, N> {
267    type Output = T;
268
269    fn index(&self, idx: usize) -> &Self::Output {
270        self.0.get(idx).expect("index out of bounds")
271    }
272}
273
274/// An iterator over the elements stored in an [`InlineVec`]
275pub struct InlineVecIter<'a, T, const N: usize> {
276    vec: &'a InlineVecInner<T, N>,
277    idx: usize,
278}
279
280impl<'a, T, const N: usize> Iterator for InlineVecIter<'a, T, N> {
281    type Item = &'a T;
282
283    fn next(&mut self) -> Option<Self::Item> {
284        self.idx += 1;
285        self.vec.get(self.idx - 1)
286    }
287}
288
289impl<T, const N: usize> Drop for InlineVecInner<T, N> {
290    fn drop(&mut self) {
291        if let Self::Inline { len, data } = self {
292            for element in data.iter_mut().take(*len) {
293                unsafe { ptr::drop_in_place(element.as_mut_ptr()) };
294            }
295        }
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    #[test]
304    fn inlinevec_to_vec_stack() {
305        let mut x = InlineVec::<usize, 4>::new();
306
307        for i in 0..4 {
308            x.push(i * 2);
309        }
310
311        assert!(!x.is_heap_allocated());
312        assert_eq!(x.len(), 4);
313
314        let xx = x.to_vec();
315        assert_eq!(xx.as_slice(), &[0, 2, 4, 6]);
316
317        x.push(42);
318        assert!(x.is_heap_allocated());
319        assert_eq!(x.as_slice(), &[0, 2, 4, 6, 42]);
320        assert_eq!(x.get(4), Some(&42));
321
322        let xx = x.to_vec();
323        assert_eq!(xx.as_slice(), &[0, 2, 4, 6, 42]);
324    }
325
326    #[test]
327    fn inlinevec_to_vec_heap() {
328        let mut x = InlineVec::<String, 4>::new();
329
330        for i in 0..4u8 {
331            x.push(i.to_string());
332        }
333
334        assert!(!x.is_heap_allocated());
335        assert_eq!(x.len(), 4);
336
337        let xx = x.to_vec();
338        assert_eq!(xx.as_slice(), &["0", "1", "2", "3"]);
339
340        x.push("1337".into());
341        assert!(x.is_heap_allocated());
342        assert_eq!(x.as_slice(), &["0", "1", "2", "3", "1337"]);
343        assert_eq!(x.get(4).map(|x| &**x), Some("1337"));
344
345        let xx = x.to_vec();
346        assert_eq!(xx.as_slice(), &["0", "1", "2", "3", "1337"]);
347    }
348
349    #[test]
350    fn inlinevec_drop_stack() {
351        let mut x = InlineVec::<String, 4>::new();
352
353        for i in 0..3u8 {
354            x.push(i.to_string());
355        }
356
357        assert_eq!(x.as_slice(), &["0", "1", "2"]);
358        assert!(!x.is_heap_allocated());
359    }
360
361    #[test]
362    fn inlinehashmap_drop_heap() {
363        let mut x = InlineVec::<String, 4>::new();
364
365        for i in 0..8u8 {
366            x.push(i.to_string());
367        }
368
369        assert_eq!(x.as_slice(), &["0", "1", "2", "3", "4", "5", "6", "7"]);
370        assert!(x.is_heap_allocated());
371    }
372
373    #[test]
374    fn inlinevec_iter() {
375        let mut x = InlineVecInner::<usize, 2>::new();
376        x.push(13);
377        x.push(42);
378        x.push(17);
379        x.push(19);
380        let mut iter = x.iter();
381        assert_eq!(iter.next(), Some(&13));
382        assert_eq!(iter.next(), Some(&42));
383        assert_eq!(iter.next(), Some(&17));
384        assert_eq!(iter.next(), Some(&19));
385        assert_eq!(iter.next(), None);
386    }
387
388    #[test]
389    fn inlinevec_remove() {
390        let mut x = InlineVecInner::<usize, 4>::new();
391        x.push(789);
392        assert_eq!(x.len(), 1);
393        assert_eq!(x.get(0), Some(&789));
394        assert_eq!(x.remove(0), 789);
395        assert_eq!(x.len(), 0);
396
397        {
398            let mut xc = x.clone();
399            // out of bounds index must panic
400            assert!(std::panic::catch_unwind(move || xc.remove(0)).is_err());
401        }
402
403        for i in 0..4 {
404            x.push(i * 2);
405        }
406
407        assert!(!x.is_heap_allocated());
408        assert_eq!(x.as_slice(), &[0, 2, 4, 6]);
409
410        assert_eq!(x.remove(2), 4);
411        assert_eq!(x.as_slice(), &[0, 2, 6]);
412
413        assert_eq!(x.remove(2), 6);
414        assert_eq!(x.as_slice(), &[0, 2]);
415
416        assert_eq!(x.remove(1), 2);
417        assert_eq!(x.as_slice(), &[0]);
418
419        assert_eq!(x.remove(0), 0);
420        assert_eq!(x.as_slice(), &[]);
421        assert!(!x.is_heap_allocated());
422
423        // trigger heap allocation
424        for i in 0..8 {
425            x.push(i * 2);
426        }
427        assert!(x.is_heap_allocated());
428        assert_eq!(x.as_slice(), &[0, 2, 4, 6, 8, 10, 12, 14]);
429
430        assert_eq!(x.remove(7), 14);
431        assert_eq!(x.remove(0), 0);
432    }
433
434    #[test]
435    fn inlinevec_remove_heap() {
436        let mut x = InlineVecInner::<String, 4>::new();
437        x.push("test".into());
438        assert_eq!(x.len(), 1);
439        assert_eq!(x.remove(0), "test");
440        assert_eq!(x.len(), 0);
441    }
442
443    #[test]
444    fn inlinevec() {
445        let mut x = InlineVecInner::<usize, 4>::new();
446        assert_eq!(x.len(), 0);
447        assert_eq!(x.get(0), None);
448        assert!(!x.is_heap_allocated());
449
450        x.push(1337);
451        assert_eq!(x.len(), 1);
452        assert_eq!(x.get(0), Some(&1337));
453        assert!(!x.is_heap_allocated());
454
455        for v in 0..3 {
456            x.push(v);
457        }
458
459        assert_eq!(x.len(), 4);
460
461        // this call should move the vector to the heap
462        x.push(42);
463        assert_eq!(x.len(), 5);
464        assert!(x.is_heap_allocated());
465
466        // check that the old vector is still valid
467        assert_eq!(x.get(0), Some(&1337));
468
469        for v in 0..500 {
470            x.push(v);
471        }
472
473        assert_eq!(x.len(), 505);
474        assert!(x.is_heap_allocated());
475
476        assert_eq!(x.get(1337), None);
477
478        *x.get_mut(0).unwrap() = 444;
479        assert_eq!(x.get(0), Some(&444));
480        assert_eq!(x.get_mut(99999 /* out of bounds */), None);
481    }
482
483    #[test]
484    fn inlinevec_as_slice() {
485        let mut x = InlineVecInner::<usize, 4>::new();
486        x.push(1337);
487        x.push(42);
488        x.push(17);
489        assert_eq!(x.as_slice(), &[1337, 42, 17]);
490        x.push(19);
491        x.push(34);
492        assert_eq!(x.as_slice(), &[1337, 42, 17, 19, 34]);
493    }
494}
tl/inline/vec.rs

tl/inline/
vec.rs
