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