Skip to main content

rkyv/util/
ser_vec.rs

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