Skip to main content

rkyv/util/
inline_vec.rs

1use core::{
2    borrow::{Borrow, BorrowMut},
3    fmt,
4    marker::PhantomData,
5    mem::MaybeUninit,
6    ops,
7    ptr::{self, NonNull},
8    slice::{self, from_raw_parts_mut},
9};
10
11/// A vector that uses inline-allocated memory.
12pub struct InlineVec<T, const N: usize> {
13    elements: [MaybeUninit<T>; N],
14    len: usize,
15}
16
17impl<T, const N: usize> Drop for InlineVec<T, N> {
18    fn drop(&mut self) {
19        self.clear()
20    }
21}
22
23// SAFETY: InlineVec is safe to send to another thread is T is safe to send to
24// another thread
25unsafe impl<T: Send, const N: usize> Send for InlineVec<T, N> {}
26
27// SAFETY: InlineVec is safe to share between threads if T is safe to share
28// between threads
29unsafe impl<T: Sync, const N: usize> Sync for InlineVec<T, N> {}
30
31impl<T, const N: usize> InlineVec<T, N> {
32    /// Constructs a new, empty `InlineVec`.
33    ///
34    /// The vector will be able to hold exactly `N` elements.
35    pub fn new() -> Self {
36        Self {
37            elements: unsafe { MaybeUninit::uninit().assume_init() },
38            len: 0,
39        }
40    }
41
42    /// Clears the vector, removing all values.
43    pub fn clear(&mut self) {
44        let len = self.len;
45        self.len = 0;
46
47        for i in 0..len {
48            unsafe {
49                self.elements[i].as_mut_ptr().drop_in_place();
50            }
51        }
52    }
53
54    /// Returns an unsafe mutable pointer to the vector's buffer.
55    ///
56    /// The caller must ensure that the vector outlives the pointer this
57    /// function returns, or else it will end up pointing to garbage.
58    pub fn as_mut_ptr(&mut self) -> *mut T {
59        self.elements.as_mut_ptr().cast()
60    }
61
62    /// Extracts a mutable slice of the entire vector.
63    ///
64    /// Equivalent to `&mut s[..]`.
65    pub fn as_mut_slice(&mut self) -> &mut [T] {
66        unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
67    }
68
69    /// Returns a raw pointer to the vector's buffer.
70    ///
71    /// The caller must ensure that the vector outlives the pointer this
72    /// functions returns, or else it will end up pointing to garbage.
73    ///
74    /// The caller must also ensure that the memory the pointer
75    /// (non-transitively) points to is never written to (except inside an
76    /// `UnsafeCell`) using this pointer or any pointer derived from it. If
77    /// you need to mutate the contents of the slice, use
78    /// [`as_mut_ptr`](Self::as_mut_ptr).
79    pub fn as_ptr(&self) -> *const T {
80        self.elements.as_ptr().cast()
81    }
82
83    /// Extracts a slice containing the entire vector.
84    ///
85    /// Equivalent to `&s[..]`.
86    pub fn as_slice(&self) -> &[T] {
87        unsafe { slice::from_raw_parts(self.as_ptr(), self.len) }
88    }
89
90    /// Returns the number of elements the vector can hole without reallocating.
91    pub const fn capacity(&self) -> usize {
92        N
93    }
94
95    /// Ensures that there is capacity for at least `additional` more elements
96    /// to be inserted into the `ScratchVec`.
97    ///
98    /// # Panics
99    ///
100    /// Panics if the required capacity exceeds the available capacity.
101    pub fn reserve(&mut self, additional: usize) {
102        if N - self.len < additional {
103            Self::out_of_space();
104        }
105    }
106
107    #[cold]
108    fn out_of_space() -> ! {
109        panic!(
110            "reserve requested more capacity than the InlineVec has available"
111        );
112    }
113
114    /// Returns `true` if the vector contains no elements.
115    pub fn is_empty(&self) -> bool {
116        self.len == 0
117    }
118
119    /// Returns the number of elements in the vector, also referred to as its
120    /// `length`.
121    pub fn len(&self) -> usize {
122        self.len
123    }
124
125    /// Copies and appends all elements in a slice to the `ScratchVec`.
126    ///
127    /// The elements of the slice are appended in-order.
128    pub fn extend_from_slice(&mut self, other: &[T])
129    where
130        T: Copy,
131    {
132        if !other.is_empty() {
133            self.reserve(other.len());
134            unsafe {
135                core::ptr::copy_nonoverlapping(
136                    other.as_ptr(),
137                    self.as_mut_ptr().add(self.len()),
138                    other.len(),
139                );
140            }
141            self.len += other.len();
142        }
143    }
144
145    /// Removes the last element from a vector and returns it, or `None` if it
146    /// is empty.
147    pub fn pop(&mut self) -> Option<T> {
148        if self.len == 0 {
149            None
150        } else {
151            unsafe {
152                self.len -= 1;
153                Some(self.as_ptr().add(self.len()).read())
154            }
155        }
156    }
157
158    /// Appends an element to the back of a collection without performing bounds
159    /// checking.
160    ///
161    /// # Safety
162    ///
163    /// The vector must have enough space reserved for the pushed element.
164    pub unsafe fn push_unchecked(&mut self, value: T) {
165        unsafe {
166            self.as_mut_ptr().add(self.len).write(value);
167            self.len += 1;
168        }
169    }
170
171    /// Appends an element to the back of a collection.
172    pub fn push(&mut self, value: T) {
173        if self.len == N {
174            Self::out_of_space()
175        } else {
176            unsafe {
177                self.push_unchecked(value);
178            }
179        }
180    }
181
182    /// Reserves the minimum capacity for exactly `additional` more elements to
183    /// be inserted in the given `AlignedVec`. After calling
184    /// `reserve_exact`, capacity will be greater than or equal
185    /// to `self.len() + additional`. Does nothing if the capacity is already
186    /// sufficient.
187    ///
188    /// # Panics
189    ///
190    /// Panics if the required capacity exceeds the available capacity.
191    pub fn reserve_exact(&mut self, additional: usize) {
192        self.reserve(additional);
193    }
194
195    /// Forces the length of the vector to `new_len`.
196    ///
197    /// This is a low-level operation that maintains none of the normal
198    /// invariants of the type.
199    ///
200    /// # Safety
201    ///
202    /// - `new_len` must be less than or equal to [`capacity()`](Self::capacity)
203    /// - The elements at `old_len..new_len` must be initialized
204    pub unsafe fn set_len(&mut self, new_len: usize) {
205        debug_assert!(new_len <= self.capacity());
206
207        self.len = new_len;
208    }
209
210    /// Creates a draining iterator that removes all of the elements from the
211    /// vector.
212    pub fn drain(&mut self) -> Drain<'_, T, N> {
213        let remaining = self.len();
214        unsafe {
215            self.set_len(0);
216        }
217
218        Drain {
219            current: unsafe { NonNull::new_unchecked(self.as_mut_ptr()) },
220            remaining,
221            _phantom: PhantomData,
222        }
223    }
224}
225
226impl<T, const N: usize> InlineVec<MaybeUninit<T>, N> {
227    /// Assuming that all the elements are initialized, removes the
228    /// `MaybeUninit` wrapper from the vector.
229    ///
230    /// # Safety
231    ///
232    /// It is up to the caller to guarantee that the `MaybeUninit<T>` elements
233    /// really are in an initialized state. Calling this when the content is
234    /// not yet fully initialized causes undefined behavior.
235    pub unsafe fn assume_init(self) -> InlineVec<T, N> {
236        let mut elements = unsafe {
237            MaybeUninit::<[MaybeUninit<T>; N]>::uninit().assume_init()
238        };
239        unsafe {
240            ptr::copy_nonoverlapping(
241                self.elements.as_ptr().cast(),
242                elements.as_mut_ptr(),
243                N,
244            );
245        }
246        InlineVec {
247            elements,
248            len: self.len,
249        }
250    }
251}
252
253impl<T, const N: usize> AsMut<[T]> for InlineVec<T, N> {
254    fn as_mut(&mut self) -> &mut [T] {
255        self.as_mut_slice()
256    }
257}
258
259impl<T, const N: usize> AsRef<[T]> for InlineVec<T, N> {
260    fn as_ref(&self) -> &[T] {
261        self.as_slice()
262    }
263}
264
265impl<T, const N: usize> Borrow<[T]> for InlineVec<T, N> {
266    fn borrow(&self) -> &[T] {
267        self.as_slice()
268    }
269}
270
271impl<T, const N: usize> BorrowMut<[T]> for InlineVec<T, N> {
272    fn borrow_mut(&mut self) -> &mut [T] {
273        self.as_mut_slice()
274    }
275}
276
277impl<T: fmt::Debug, const N: usize> fmt::Debug for InlineVec<T, N> {
278    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
279        self.as_slice().fmt(f)
280    }
281}
282
283impl<T, const N: usize> Default for InlineVec<T, N> {
284    fn default() -> Self {
285        Self::new()
286    }
287}
288
289impl<T, const N: usize> ops::Deref for InlineVec<T, N> {
290    type Target = [T];
291
292    fn deref(&self) -> &Self::Target {
293        self.as_slice()
294    }
295}
296
297impl<T, const N: usize> ops::DerefMut for InlineVec<T, N> {
298    fn deref_mut(&mut self) -> &mut Self::Target {
299        self.as_mut_slice()
300    }
301}
302
303impl<T, I: slice::SliceIndex<[T]>, const N: usize> ops::Index<I>
304    for InlineVec<T, N>
305{
306    type Output = <I as slice::SliceIndex<[T]>>::Output;
307
308    fn index(&self, index: I) -> &Self::Output {
309        &self.as_slice()[index]
310    }
311}
312
313impl<T, I: slice::SliceIndex<[T]>, const N: usize> ops::IndexMut<I>
314    for InlineVec<T, N>
315{
316    fn index_mut(&mut self, index: I) -> &mut Self::Output {
317        &mut self.as_mut_slice()[index]
318    }
319}
320
321/// A draining iterator for `InlineVec<T>`.
322///
323/// This `struct` is created by [`InlineVec::drain`]. See its documentation for
324/// more.
325pub struct Drain<'a, T: 'a, const N: usize> {
326    current: NonNull<T>,
327    remaining: usize,
328    _phantom: PhantomData<&'a mut InlineVec<T, N>>,
329}
330
331impl<T: fmt::Debug, const N: usize> fmt::Debug for Drain<'_, T, N> {
332    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
333        f.debug_tuple("Drain").field(&self.as_slice()).finish()
334    }
335}
336
337impl<T, const N: usize> Drain<'_, T, N> {
338    /// Returns the remaining items of this iterator as a slice.
339    pub fn as_slice(&self) -> &[T] {
340        unsafe { from_raw_parts_mut(self.current.as_ptr(), self.remaining) }
341    }
342}
343
344impl<T, const N: usize> AsRef<[T]> for Drain<'_, T, N> {
345    fn as_ref(&self) -> &[T] {
346        self.as_slice()
347    }
348}
349
350impl<T, const N: usize> Iterator for Drain<'_, T, N> {
351    type Item = T;
352
353    fn next(&mut self) -> Option<T> {
354        if self.remaining > 0 {
355            self.remaining -= 1;
356            let result = unsafe { self.current.as_ptr().read() };
357            self.current =
358                unsafe { NonNull::new_unchecked(self.current.as_ptr().add(1)) };
359            Some(result)
360        } else {
361            None
362        }
363    }
364
365    fn size_hint(&self) -> (usize, Option<usize>) {
366        (self.remaining, Some(self.remaining))
367    }
368}
369
370impl<T, const N: usize> DoubleEndedIterator for Drain<'_, T, N> {
371    fn next_back(&mut self) -> Option<T> {
372        if self.remaining > 0 {
373            self.remaining -= 1;
374            unsafe { Some(self.current.as_ptr().add(self.remaining).read()) }
375        } else {
376            None
377        }
378    }
379}
380
381impl<T, const N: usize> Drop for Drain<'_, T, N> {
382    fn drop(&mut self) {
383        for i in 0..self.remaining {
384            unsafe {
385                self.current.as_ptr().add(i).drop_in_place();
386            }
387        }
388    }
389}
390
391impl<T, const N: usize> ExactSizeIterator for Drain<'_, T, N> {}
392
393impl<T, const N: usize> core::iter::FusedIterator for Drain<'_, T, N> {}
394
395#[cfg(test)]
396mod tests {
397    use crate::util::InlineVec;
398
399    #[test]
400    fn drain() {
401        let mut vec = InlineVec::<_, 8>::new();
402
403        for i in 0..100 {
404            vec.push(i);
405            if vec.len() == vec.capacity() {
406                for j in vec.drain() {
407                    let _ = j;
408                }
409            }
410        }
411    }
412}