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        if !other.is_empty() {
130            self.reserve(other.len());
131            unsafe {
132                core::ptr::copy_nonoverlapping(
133                    other.as_ptr(),
134                    self.as_mut_ptr().add(self.len()),
135                    other.len(),
136                );
137            }
138            self.len += other.len();
139        }
140    }
141
142    /// Removes the last element from a vector and returns it, or `None` if it
143    /// is empty.
144    pub fn pop(&mut self) -> Option<T> {
145        if self.len == 0 {
146            None
147        } else {
148            unsafe {
149                self.len -= 1;
150                Some(self.as_ptr().add(self.len()).read())
151            }
152        }
153    }
154
155    /// Appends an element to the back of a collection without performing bounds
156    /// checking.
157    ///
158    /// # Safety
159    ///
160    /// The vector must have enough space reserved for the pushed element.
161    pub unsafe fn push_unchecked(&mut self, value: T) {
162        unsafe {
163            self.as_mut_ptr().add(self.len).write(value);
164            self.len += 1;
165        }
166    }
167
168    /// Appends an element to the back of a collection.
169    pub fn push(&mut self, value: T) {
170        if self.len == N {
171            Self::out_of_space()
172        } else {
173            unsafe {
174                self.push_unchecked(value);
175            }
176        }
177    }
178
179    /// Reserves the minimum capacity for exactly `additional` more elements to
180    /// be inserted in the given `AlignedVec`. After calling
181    /// `reserve_exact`, capacity will be greater than or equal
182    /// to `self.len() + additional`. Does nothing if the capacity is already
183    /// sufficient.
184    ///
185    /// # Panics
186    ///
187    /// Panics if the required capacity exceeds the available capacity.
188    pub fn reserve_exact(&mut self, additional: usize) {
189        self.reserve(additional);
190    }
191
192    /// Forces the length of the vector to `new_len`.
193    ///
194    /// This is a low-level operation that maintains none of the normal
195    /// invariants of the type.
196    ///
197    /// # Safety
198    ///
199    /// - `new_len` must be less than or equal to [`capacity()`](Self::capacity)
200    /// - The elements at `old_len..new_len` must be initialized
201    pub unsafe fn set_len(&mut self, new_len: usize) {
202        debug_assert!(new_len <= self.capacity());
203
204        self.len = new_len;
205    }
206
207    /// Creates a draining iterator that removes all of the elements from the
208    /// vector.
209    pub fn drain(&mut self) -> Drain<'_, T, N> {
210        let remaining = self.len();
211        unsafe {
212            self.set_len(0);
213        }
214
215        Drain {
216            current: unsafe { NonNull::new_unchecked(self.as_mut_ptr()) },
217            remaining,
218            _phantom: PhantomData,
219        }
220    }
221}
222
223impl<T, const N: usize> InlineVec<MaybeUninit<T>, N> {
224    /// Assuming that all the elements are initialized, removes the
225    /// `MaybeUninit` wrapper from the vector.
226    ///
227    /// # Safety
228    ///
229    /// It is up to the caller to guarantee that the `MaybeUninit<T>` elements
230    /// really are in an initialized state. Calling this when the content is
231    /// not yet fully initialized causes undefined behavior.
232    pub fn assume_init(self) -> InlineVec<T, N> {
233        let mut elements = unsafe {
234            MaybeUninit::<[MaybeUninit<T>; N]>::uninit().assume_init()
235        };
236        unsafe {
237            ptr::copy_nonoverlapping(
238                self.elements.as_ptr().cast(),
239                elements.as_mut_ptr(),
240                N,
241            );
242        }
243        InlineVec {
244            elements,
245            len: self.len,
246        }
247    }
248}
249
250impl<T, const N: usize> AsMut<[T]> for InlineVec<T, N> {
251    fn as_mut(&mut self) -> &mut [T] {
252        self.as_mut_slice()
253    }
254}
255
256impl<T, const N: usize> AsRef<[T]> for InlineVec<T, N> {
257    fn as_ref(&self) -> &[T] {
258        self.as_slice()
259    }
260}
261
262impl<T, const N: usize> Borrow<[T]> for InlineVec<T, N> {
263    fn borrow(&self) -> &[T] {
264        self.as_slice()
265    }
266}
267
268impl<T, const N: usize> BorrowMut<[T]> for InlineVec<T, N> {
269    fn borrow_mut(&mut self) -> &mut [T] {
270        self.as_mut_slice()
271    }
272}
273
274impl<T: fmt::Debug, const N: usize> fmt::Debug for InlineVec<T, N> {
275    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
276        self.as_slice().fmt(f)
277    }
278}
279
280impl<T, const N: usize> Default for InlineVec<T, N> {
281    fn default() -> Self {
282        Self::new()
283    }
284}
285
286impl<T, const N: usize> ops::Deref for InlineVec<T, N> {
287    type Target = [T];
288
289    fn deref(&self) -> &Self::Target {
290        self.as_slice()
291    }
292}
293
294impl<T, const N: usize> ops::DerefMut for InlineVec<T, N> {
295    fn deref_mut(&mut self) -> &mut Self::Target {
296        self.as_mut_slice()
297    }
298}
299
300impl<T, I: slice::SliceIndex<[T]>, const N: usize> ops::Index<I>
301    for InlineVec<T, N>
302{
303    type Output = <I as slice::SliceIndex<[T]>>::Output;
304
305    fn index(&self, index: I) -> &Self::Output {
306        &self.as_slice()[index]
307    }
308}
309
310impl<T, I: slice::SliceIndex<[T]>, const N: usize> ops::IndexMut<I>
311    for InlineVec<T, N>
312{
313    fn index_mut(&mut self, index: I) -> &mut Self::Output {
314        &mut self.as_mut_slice()[index]
315    }
316}
317
318/// A draining iterator for `InlineVec<T>`.
319///
320/// This `struct` is created by [`InlineVec::drain`]. See its documentation for
321/// more.
322pub struct Drain<'a, T: 'a, const N: usize> {
323    current: NonNull<T>,
324    remaining: usize,
325    _phantom: PhantomData<&'a mut InlineVec<T, N>>,
326}
327
328impl<T: fmt::Debug, const N: usize> fmt::Debug for Drain<'_, T, N> {
329    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
330        f.debug_tuple("Drain").field(&self.as_slice()).finish()
331    }
332}
333
334impl<T, const N: usize> Drain<'_, T, N> {
335    /// Returns the remaining items of this iterator as a slice.
336    pub fn as_slice(&self) -> &[T] {
337        unsafe { from_raw_parts_mut(self.current.as_ptr(), self.remaining) }
338    }
339}
340
341impl<T, const N: usize> AsRef<[T]> for Drain<'_, T, N> {
342    fn as_ref(&self) -> &[T] {
343        self.as_slice()
344    }
345}
346
347impl<T, const N: usize> Iterator for Drain<'_, T, N> {
348    type Item = T;
349
350    fn next(&mut self) -> Option<T> {
351        if self.remaining > 0 {
352            self.remaining -= 1;
353            let result = unsafe { self.current.as_ptr().read() };
354            self.current =
355                unsafe { NonNull::new_unchecked(self.current.as_ptr().add(1)) };
356            Some(result)
357        } else {
358            None
359        }
360    }
361
362    fn size_hint(&self) -> (usize, Option<usize>) {
363        (self.remaining, Some(self.remaining))
364    }
365}
366
367impl<T, const N: usize> DoubleEndedIterator for Drain<'_, T, N> {
368    fn next_back(&mut self) -> Option<T> {
369        if self.remaining > 0 {
370            self.remaining -= 1;
371            unsafe { Some(self.current.as_ptr().add(self.remaining).read()) }
372        } else {
373            None
374        }
375    }
376}
377
378impl<T, const N: usize> Drop for Drain<'_, T, N> {
379    fn drop(&mut self) {
380        for i in 0..self.remaining {
381            unsafe {
382                self.current.as_ptr().add(i).drop_in_place();
383            }
384        }
385    }
386}
387
388impl<T, const N: usize> ExactSizeIterator for Drain<'_, T, N> {}
389
390impl<T, const N: usize> core::iter::FusedIterator for Drain<'_, T, N> {}
391
392#[cfg(test)]
393mod tests {
394    use crate::util::InlineVec;
395
396    #[test]
397    fn drain() {
398        let mut vec = InlineVec::<_, 8>::new();
399
400        for i in 0..100 {
401            vec.push(i);
402            if vec.len() == vec.capacity() {
403                for j in vec.drain() {
404                    let _ = j;
405                }
406            }
407        }
408    }
409}
rkyv/util/inline_vec.rs

rkyv/util/
inline_vec.rs
