Skip to main content

rkyv/util/alloc/
aligned_vec.rs

1use core::{
2    alloc::Layout,
3    borrow::{Borrow, BorrowMut},
4    fmt,
5    mem::ManuallyDrop,
6    ops::{Deref, DerefMut, Index, IndexMut},
7    ptr::NonNull,
8    slice,
9};
10
11use rancor::Fallible;
12
13use crate::{
14    alloc::{
15        alloc::{alloc, dealloc, handle_alloc_error, realloc},
16        boxed::Box,
17        vec::Vec,
18    },
19    ser::{Allocator, Writer},
20    vec::{ArchivedVec, VecResolver},
21    with::{ArchiveWith, AsVec, DeserializeWith, SerializeWith},
22    Place,
23};
24
25/// A vector of bytes that aligns its memory to the specified alignment.
26///
27/// ```
28/// # use rkyv::util::AlignedVec;
29/// let bytes = AlignedVec::<4096>::with_capacity(1);
30/// assert_eq!(bytes.as_ptr() as usize % 4096, 0);
31/// ```
32pub struct AlignedVec<const ALIGNMENT: usize = 16> {
33    ptr: NonNull<u8>,
34    cap: usize,
35    len: usize,
36}
37
38impl<const A: usize> Drop for AlignedVec<A> {
39    fn drop(&mut self) {
40        if self.cap != 0 {
41            unsafe {
42                dealloc(self.ptr.as_ptr(), self.layout());
43            }
44        }
45    }
46}
47
48impl<const ALIGNMENT: usize> AlignedVec<ALIGNMENT> {
49    /// The alignment of the vector
50    pub const ALIGNMENT: usize = ALIGNMENT;
51
52    /// Maximum capacity of the vector.
53    ///
54    /// Dictated by the requirements of [`Layout`]. "`size`, when rounded up to
55    /// the nearest multiple of `align`, must not overflow `isize` (i.e. the
56    /// rounded value must be less than or equal to `isize::MAX`)".
57    pub const MAX_CAPACITY: usize = isize::MAX as usize - (Self::ALIGNMENT - 1);
58
59    /// Constructs a new, empty `AlignedVec`.
60    ///
61    /// The vector will not allocate until elements are pushed into it.
62    ///
63    /// # Examples
64    /// ```
65    /// # use rkyv::util::AlignedVec;
66    /// let mut vec = AlignedVec::<16>::new();
67    /// ```
68    pub fn new() -> Self {
69        Self::with_capacity(0)
70    }
71
72    /// Constructs a new, empty `AlignedVec` with the specified capacity.
73    ///
74    /// The vector will be able to hold exactly `capacity` bytes without
75    /// reallocating. If `capacity` is 0, the vector will not allocate.
76    ///
77    /// # Examples
78    /// ```
79    /// # use rkyv::util::AlignedVec;
80    /// let mut vec = AlignedVec::<16>::with_capacity(10);
81    ///
82    /// // The vector contains no items, even though it has capacity for more
83    /// assert_eq!(vec.len(), 0);
84    /// assert_eq!(vec.capacity(), 10);
85    ///
86    /// // These are all done without reallocating...
87    /// for i in 0..10 {
88    ///     vec.push(i);
89    /// }
90    /// assert_eq!(vec.len(), 10);
91    /// assert_eq!(vec.capacity(), 10);
92    ///
93    /// // ...but this may make the vector reallocate
94    /// vec.push(11);
95    /// assert_eq!(vec.len(), 11);
96    /// assert!(vec.capacity() >= 11);
97    /// ```
98    pub fn with_capacity(capacity: usize) -> Self {
99        assert!(ALIGNMENT > 0, "ALIGNMENT must be 1 or more");
100        assert!(
101            ALIGNMENT.is_power_of_two(),
102            "ALIGNMENT must be a power of 2"
103        );
104        // As `ALIGNMENT` has to be a power of 2, this caps `ALIGNMENT` at a max
105        // of `(isize::MAX + 1) / 2` (1 GiB on 32-bit systems).
106        assert!(
107            ALIGNMENT < isize::MAX as usize,
108            "ALIGNMENT must be less than isize::MAX"
109        );
110
111        if capacity == 0 {
112            Self {
113                ptr: NonNull::dangling(),
114                cap: 0,
115                len: 0,
116            }
117        } else {
118            assert!(
119                capacity <= Self::MAX_CAPACITY,
120                "`capacity` cannot exceed `Self::MAX_CAPACITY`"
121            );
122
123            let ptr = unsafe {
124                let layout = Layout::from_size_align_unchecked(
125                    capacity,
126                    Self::ALIGNMENT,
127                );
128                let ptr = alloc(layout);
129                if ptr.is_null() {
130                    handle_alloc_error(layout);
131                }
132                NonNull::new_unchecked(ptr)
133            };
134
135            Self {
136                ptr,
137                cap: capacity,
138                len: 0,
139            }
140        }
141    }
142
143    fn layout(&self) -> Layout {
144        unsafe { Layout::from_size_align_unchecked(self.cap, Self::ALIGNMENT) }
145    }
146
147    /// Clears the vector, removing all values.
148    ///
149    /// Note that this method has no effect on the allocated capacity of the
150    /// vector.
151    ///
152    /// # Examples
153    /// ```
154    /// # use rkyv::util::AlignedVec;
155    /// let mut v = AlignedVec::<16>::new();
156    /// v.extend_from_slice(&[1, 2, 3, 4]);
157    ///
158    /// v.clear();
159    ///
160    /// assert!(v.is_empty());
161    /// ```
162    pub fn clear(&mut self) {
163        self.len = 0;
164    }
165
166    /// Change capacity of vector.
167    ///
168    /// Will set capacity to exactly `new_cap`.
169    /// Can be used to either grow or shrink capacity.
170    /// Backing memory will be reallocated.
171    ///
172    /// Usually the safe methods `reserve` or `reserve_exact` are a better
173    /// choice. This method only exists as a micro-optimization for very
174    /// performance-sensitive code where where the calculation of capacity
175    /// required has already been performed, and you want to avoid doing it
176    /// again, or if you want to implement a different growth strategy.
177    ///
178    /// # Safety
179    ///
180    /// - `new_cap` must be less than or equal to
181    ///   [`MAX_CAPACITY`](AlignedVec::MAX_CAPACITY)
182    /// - `new_cap` must be greater than or equal to [`len()`](AlignedVec::len)
183    pub unsafe fn change_capacity(&mut self, new_cap: usize) {
184        debug_assert!(new_cap <= Self::MAX_CAPACITY);
185        debug_assert!(new_cap >= self.len);
186
187        if new_cap > 0 {
188            let new_ptr = if self.cap > 0 {
189                // SAFETY:
190                // - `self.ptr` is currently allocated because `self.cap` is
191                //   greater than zero.
192                // - `self.layout()` always matches the layout used to allocate
193                //   the current block of memory.
194                // - We checked that `new_cap` is greater than zero.
195                let new_ptr = unsafe {
196                    realloc(self.ptr.as_ptr(), self.layout(), new_cap)
197                };
198                if new_ptr.is_null() {
199                    // SAFETY:
200                    // - `ALIGNMENT` is always guaranteed to be a nonzero power
201                    //   of two.
202                    // - We checked that `new_cap` doesn't overflow `isize` when
203                    //   rounded up to the nearest power of two.
204                    let layout = unsafe {
205                        Layout::from_size_align_unchecked(
206                            new_cap,
207                            Self::ALIGNMENT,
208                        )
209                    };
210                    handle_alloc_error(layout);
211                }
212                new_ptr
213            } else {
214                // SAFETY:
215                // - `ALIGNMENT` is always guaranteed to be a nonzero power of
216                //   two.
217                // - We checked that `new_cap` doesn't overflow `isize` when
218                //   rounded up to the nearest power of two.
219                let layout = unsafe {
220                    Layout::from_size_align_unchecked(new_cap, Self::ALIGNMENT)
221                };
222                // SAFETY: We checked that `new_cap` has non-zero size.
223                let new_ptr = unsafe { alloc(layout) };
224                if new_ptr.is_null() {
225                    handle_alloc_error(layout);
226                }
227                new_ptr
228            };
229            // SAFETY: We checked that `new_ptr` is non-null in each of the
230            // branches.
231            self.ptr = unsafe { NonNull::new_unchecked(new_ptr) };
232            self.cap = new_cap;
233        } else if self.cap > 0 {
234            // SAFETY: Because the capacity is nonzero, `self.ptr` points to a
235            // currently-allocated memory block. All memory blocks are allocated
236            // with a layout of `self.layout()`.
237            unsafe {
238                dealloc(self.ptr.as_ptr(), self.layout());
239            }
240            self.ptr = NonNull::dangling();
241            self.cap = 0;
242        }
243    }
244
245    /// Shrinks the capacity of the vector as much as possible.
246    ///
247    /// It will drop down as close as possible to the length but the allocator
248    /// may still inform the vector that there is space for a few more
249    /// elements.
250    ///
251    /// # Examples
252    /// ```
253    /// # use rkyv::util::AlignedVec;
254    /// let mut vec = AlignedVec::<16>::with_capacity(10);
255    /// vec.extend_from_slice(&[1, 2, 3]);
256    /// assert_eq!(vec.capacity(), 10);
257    /// vec.shrink_to_fit();
258    /// assert!(vec.capacity() >= 3);
259    ///
260    /// vec.clear();
261    /// vec.shrink_to_fit();
262    /// assert!(vec.capacity() == 0);
263    /// ```
264    pub fn shrink_to_fit(&mut self) {
265        if self.cap != self.len {
266            // New capacity cannot exceed max as it's shrinking
267            unsafe { self.change_capacity(self.len) };
268        }
269    }
270
271    /// Returns an unsafe mutable pointer to the vector's buffer.
272    ///
273    /// The caller must ensure that the vector outlives the pointer this
274    /// function returns, or else it will end up pointing to garbage.
275    /// Modifying the vector may cause its buffer to be reallocated, which
276    /// would also make any pointers to it invalid.
277    ///
278    /// # Examples
279    /// ```
280    /// # use rkyv::util::AlignedVec;
281    /// // Allocate 1-aligned vector big enough for 4 bytes.
282    /// let size = 4;
283    /// let mut x = AlignedVec::<1>::with_capacity(size);
284    /// let x_ptr = x.as_mut_ptr();
285    ///
286    /// // Initialize elements via raw pointer writes, then set length.
287    /// unsafe {
288    ///     for i in 0..size {
289    ///         *x_ptr.add(i) = i as u8;
290    ///     }
291    ///     x.set_len(size);
292    /// }
293    /// assert_eq!(&*x, &[0, 1, 2, 3]);
294    /// ```
295    pub fn as_mut_ptr(&mut self) -> *mut u8 {
296        self.ptr.as_ptr()
297    }
298
299    /// Extracts a mutable slice of the entire vector.
300    ///
301    /// Equivalent to `&mut s[..]`.
302    ///
303    /// # Examples
304    /// ```
305    /// # use rkyv::util::AlignedVec;
306    /// let mut vec = AlignedVec::<16>::new();
307    /// vec.extend_from_slice(&[1, 2, 3, 4, 5]);
308    /// assert_eq!(vec.as_mut_slice().len(), 5);
309    /// for i in 0..5 {
310    ///     assert_eq!(vec.as_mut_slice()[i], i as u8 + 1);
311    ///     vec.as_mut_slice()[i] = i as u8;
312    ///     assert_eq!(vec.as_mut_slice()[i], i as u8);
313    /// }
314    /// ```
315    pub fn as_mut_slice(&mut self) -> &mut [u8] {
316        unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
317    }
318
319    /// Returns a raw pointer to the vector's buffer.
320    ///
321    /// The caller must ensure that the vector outlives the pointer this
322    /// function returns, or else it will end up pointing to garbage.
323    /// Modifying the vector may cause its buffer to be reallocated, which
324    /// would also make any pointers to it invalid.
325    ///
326    /// The caller must also ensure that the memory the pointer
327    /// (non-transitively) points to is never written to (except inside an
328    /// `UnsafeCell`) using this pointer or any pointer derived from it. If
329    /// you need to mutate the contents of the slice, use
330    /// [`as_mut_ptr`](AlignedVec::as_mut_ptr).
331    ///
332    /// # Examples
333    /// ```
334    /// # use rkyv::util::AlignedVec;
335    /// let mut x = AlignedVec::<16>::new();
336    /// x.extend_from_slice(&[1, 2, 4]);
337    /// let x_ptr = x.as_ptr();
338    ///
339    /// unsafe {
340    ///     for i in 0..x.len() {
341    ///         assert_eq!(*x_ptr.add(i), 1 << i);
342    ///     }
343    /// }
344    /// ```
345    pub fn as_ptr(&self) -> *const u8 {
346        self.ptr.as_ptr()
347    }
348
349    /// Extracts a slice containing the entire vector.
350    ///
351    /// Equivalent to `&s[..]`.
352    ///
353    /// # Examples
354    /// ```
355    /// # use rkyv::util::AlignedVec;
356    /// let mut vec = AlignedVec::<16>::new();
357    /// vec.extend_from_slice(&[1, 2, 3, 4, 5]);
358    /// assert_eq!(vec.as_slice().len(), 5);
359    /// for i in 0..5 {
360    ///     assert_eq!(vec.as_slice()[i], i as u8 + 1);
361    /// }
362    /// ```
363    pub fn as_slice(&self) -> &[u8] {
364        unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
365    }
366
367    /// Returns the number of elements the vector can hold without reallocating.
368    ///
369    /// # Examples
370    /// ```
371    /// # use rkyv::util::AlignedVec;
372    /// let vec = AlignedVec::<16>::with_capacity(10);
373    /// assert_eq!(vec.capacity(), 10);
374    /// ```
375    pub fn capacity(&self) -> usize {
376        self.cap
377    }
378
379    /// Reserves capacity for at least `additional` more bytes to be inserted
380    /// into the given `AlignedVec`. The collection may reserve more space
381    /// to avoid frequent reallocations. After calling `reserve`, capacity
382    /// will be greater than or equal to `self.len() + additional`. Does
383    /// nothing if capacity is already sufficient.
384    ///
385    /// # Panics
386    ///
387    /// Panics if the new capacity exceeds `Self::MAX_CAPACITY` bytes.
388    ///
389    /// # Examples
390    /// ```
391    /// # use rkyv::util::AlignedVec;
392    ///
393    /// let mut vec = AlignedVec::<16>::new();
394    /// vec.push(1);
395    /// vec.reserve(10);
396    /// assert!(vec.capacity() >= 11);
397    /// ```
398    pub fn reserve(&mut self, additional: usize) {
399        // Cannot wrap because capacity always exceeds len,
400        // but avoids having to handle potential overflow here
401        let remaining = self.cap.wrapping_sub(self.len);
402        if additional > remaining {
403            self.do_reserve(additional);
404        }
405    }
406
407    /// Extend capacity after `reserve` has found it's necessary.
408    ///
409    /// Actually performing the extension is in this separate function marked
410    /// `#[cold]` to hint to compiler that this branch is not often taken.
411    /// This keeps the path for common case where capacity is already sufficient
412    /// as fast as possible, and makes `reserve` more likely to be inlined.
413    /// This is the same trick that Rust's `Vec::reserve` uses.
414    #[cold]
415    fn do_reserve(&mut self, additional: usize) {
416        let new_cap = self
417            .len
418            .checked_add(additional)
419            .expect("cannot reserve a larger AlignedVec");
420        unsafe { self.grow_capacity_to(new_cap) };
421    }
422
423    /// Grows total capacity of vector to `new_cap` or more.
424    ///
425    /// Capacity after this call will be `new_cap` rounded up to next power of
426    /// 2, unless that would exceed maximum capacity, in which case capacity
427    /// is capped at the maximum.
428    ///
429    /// This is same growth strategy used by `reserve`, `push` and
430    /// `extend_from_slice`.
431    ///
432    /// Usually the safe methods `reserve` or `reserve_exact` are a better
433    /// choice. This method only exists as a micro-optimization for very
434    /// performance-sensitive code where where the calculation of capacity
435    /// required has already been performed, and you want to avoid doing it
436    /// again.
437    ///
438    /// Maximum capacity is `isize::MAX + 1 - Self::ALIGNMENT` bytes.
439    ///
440    /// # Panics
441    ///
442    /// Panics if `new_cap` exceeds `Self::MAX_CAPACITY` bytes.
443    ///
444    /// # Safety
445    ///
446    /// - `new_cap` must be greater than current
447    ///   [`capacity()`](AlignedVec::capacity)
448    ///
449    /// # Examples
450    /// ```
451    /// # use rkyv::util::AlignedVec;
452    ///
453    /// let mut vec = AlignedVec::<16>::new();
454    /// vec.push(1);
455    /// unsafe { vec.grow_capacity_to(50) };
456    /// assert_eq!(vec.len(), 1);
457    /// assert_eq!(vec.capacity(), 64);
458    /// ```
459    pub unsafe fn grow_capacity_to(&mut self, new_cap: usize) {
460        debug_assert!(new_cap > self.cap);
461
462        let new_cap = if new_cap > (isize::MAX as usize + 1) >> 1 {
463            // Rounding up to next power of 2 would result in `isize::MAX + 1`
464            // or higher, which exceeds max capacity. So cap at max
465            // instead.
466            assert!(
467                new_cap <= Self::MAX_CAPACITY,
468                "cannot reserve a larger AlignedVec"
469            );
470            Self::MAX_CAPACITY
471        } else {
472            // Cannot overflow due to check above
473            new_cap.next_power_of_two()
474        };
475        // SAFETY: We just checked that `new_cap` is greater than or equal to
476        // `len` and less than or equal to `MAX_CAPACITY`.
477        unsafe {
478            self.change_capacity(new_cap);
479        }
480    }
481
482    /// Resizes the Vec in-place so that len is equal to new_len.
483    ///
484    /// If new_len is greater than len, the Vec is extended by the difference,
485    /// with each additional slot filled with value. If new_len is less than
486    /// len, the Vec is simply truncated.
487    ///
488    /// # Panics
489    ///
490    /// Panics if the new length exceeds `Self::MAX_CAPACITY` bytes.
491    ///
492    /// # Examples
493    /// ```
494    /// # use rkyv::util::AlignedVec;
495    ///
496    /// let mut vec = AlignedVec::<16>::new();
497    /// vec.push(3);
498    /// vec.resize(3, 2);
499    /// assert_eq!(vec.as_slice(), &[3, 2, 2]);
500    ///
501    /// let mut vec = AlignedVec::<16>::new();
502    /// vec.extend_from_slice(&[1, 2, 3, 4]);
503    /// vec.resize(2, 0);
504    /// assert_eq!(vec.as_slice(), &[1, 2]);
505    /// ```
506    pub fn resize(&mut self, new_len: usize, value: u8) {
507        if new_len > self.len {
508            let additional = new_len - self.len;
509            self.reserve(additional);
510            unsafe {
511                core::ptr::write_bytes(
512                    self.ptr.as_ptr().add(self.len),
513                    value,
514                    additional,
515                );
516            }
517        }
518        unsafe {
519            self.set_len(new_len);
520        }
521    }
522
523    /// Returns `true` if the vector contains no elements.
524    ///
525    /// # Examples
526    /// ```
527    /// # use rkyv::util::AlignedVec;
528    ///
529    /// let mut v = Vec::new();
530    /// assert!(v.is_empty());
531    ///
532    /// v.push(1);
533    /// assert!(!v.is_empty());
534    /// ```
535    pub fn is_empty(&self) -> bool {
536        self.len == 0
537    }
538
539    /// Returns the number of elements in the vector, also referred to as its
540    /// 'length'.
541    ///
542    /// # Examples
543    /// ```
544    /// # use rkyv::util::AlignedVec;
545    ///
546    /// let mut a = AlignedVec::<16>::new();
547    /// a.extend_from_slice(&[1, 2, 3]);
548    /// assert_eq!(a.len(), 3);
549    /// ```
550    pub fn len(&self) -> usize {
551        self.len
552    }
553
554    /// Consumes and leaks the `AlignedVec`, returning a mutable reference to
555    /// the contents, `&'static mut [u8]`.
556    ///
557    /// This method does not reallocate or shrink the `AlignedVec`, so the
558    /// leaked allocation may include unused capacity that is not part of the
559    /// returned slice.
560    ///
561    /// This function is mainly useful for data that lives for the remainder of
562    /// the program's life. Dropping the returned reference will cause a memory
563    /// leak.
564    ///
565    /// # Examples
566    ///
567    /// Simple usage:
568    ///
569    /// ```
570    /// # use std::alloc::{Layout, dealloc};
571    /// # use rkyv::util::AlignedVec;
572    ///
573    /// let mut x = AlignedVec::<16>::new();
574    /// x.extend_from_slice(&[1, 2, 3]);
575    /// # let layout = Layout::from_size_align(x.capacity(), 16).unwrap();
576    /// let static_ref: &'static mut [u8] = x.leak();
577    /// static_ref[0] += 1;
578    /// assert_eq!(static_ref, &[2, 2, 3]);
579    /// # // Need to manually dealloc to avoid triggering Miri's leak check
580    /// # unsafe {
581    /// #     dealloc(static_ref.as_mut_ptr(), layout);
582    /// # }
583    /// ```
584    pub fn leak(self) -> &'static mut [u8] {
585        let mut me = ManuallyDrop::new(self);
586        unsafe { slice::from_raw_parts_mut(me.as_mut_ptr(), me.len) }
587    }
588
589    /// Copies and appends all bytes in a slice to the `AlignedVec`.
590    ///
591    /// The elements of the slice are appended in-order.
592    ///
593    /// # Examples
594    /// ```
595    /// # use rkyv::util::AlignedVec;
596    ///
597    /// let mut vec = AlignedVec::<16>::new();
598    /// vec.push(1);
599    /// vec.extend_from_slice(&[2, 3, 4]);
600    /// assert_eq!(vec.as_slice(), &[1, 2, 3, 4]);
601    /// ```
602    pub fn extend_from_slice(&mut self, other: &[u8]) {
603        self.reserve(other.len());
604        unsafe {
605            core::ptr::copy_nonoverlapping(
606                other.as_ptr(),
607                self.as_mut_ptr().add(self.len()),
608                other.len(),
609            );
610        }
611        self.len += other.len();
612    }
613
614    /// Removes the last element from a vector and returns it, or `None` if it
615    /// is empty.
616    ///
617    /// # Examples
618    /// ```
619    /// # use rkyv::util::AlignedVec;
620    ///
621    /// let mut vec = AlignedVec::<16>::new();
622    /// vec.extend_from_slice(&[1, 2, 3]);
623    /// assert_eq!(vec.pop(), Some(3));
624    /// assert_eq!(vec.as_slice(), &[1, 2]);
625    /// ```
626    pub fn pop(&mut self) -> Option<u8> {
627        if self.len == 0 {
628            None
629        } else {
630            let result = self[self.len - 1];
631            self.len -= 1;
632            Some(result)
633        }
634    }
635
636    /// Appends an element to the back of a collection.
637    ///
638    /// # Panics
639    ///
640    /// Panics if the new capacity exceeds `Self::MAX_CAPACITY` bytes.
641    ///
642    /// # Examples
643    /// ```
644    /// # use rkyv::util::AlignedVec;
645    ///
646    /// let mut vec = AlignedVec::<16>::new();
647    /// vec.extend_from_slice(&[1, 2]);
648    /// vec.push(3);
649    /// assert_eq!(vec.as_slice(), &[1, 2, 3]);
650    /// ```
651    pub fn push(&mut self, value: u8) {
652        if self.len == self.cap {
653            self.reserve_for_push();
654        }
655
656        unsafe {
657            self.as_mut_ptr().add(self.len).write(value);
658            self.len += 1;
659        }
660    }
661
662    /// Extend capacity by at least 1 byte after `push` has found it's
663    /// necessary.
664    ///
665    /// Actually performing the extension is in this separate function marked
666    /// `#[cold]` to hint to compiler that this branch is not often taken.
667    /// This keeps the path for common case where capacity is already sufficient
668    /// as fast as possible, and makes `push` more likely to be inlined.
669    /// This is the same trick that Rust's `Vec::push` uses.
670    #[cold]
671    fn reserve_for_push(&mut self) {
672        // `len` is always less than `isize::MAX`, so no possibility of overflow
673        // here
674        let new_cap = self.len + 1;
675        unsafe { self.grow_capacity_to(new_cap) };
676    }
677
678    /// Reserves the minimum capacity for exactly `additional` more elements to
679    /// be inserted in the given `AlignedVec`. After calling
680    /// `reserve_exact`, capacity will be greater than or equal
681    /// to `self.len() + additional`. Does nothing if the capacity is already
682    /// sufficient.
683    ///
684    /// Note that the allocator may give the collection more space than it
685    /// requests. Therefore, capacity can not be relied upon to be precisely
686    /// minimal. Prefer reserve if future insertions are expected.
687    ///
688    /// # Panics
689    ///
690    /// Panics if the new capacity exceeds `Self::MAX_CAPACITY`.
691    ///
692    /// # Examples
693    /// ```
694    /// # use rkyv::util::AlignedVec;
695    ///
696    /// let mut vec = AlignedVec::<16>::new();
697    /// vec.push(1);
698    /// vec.reserve_exact(10);
699    /// assert!(vec.capacity() >= 11);
700    /// ```
701    pub fn reserve_exact(&mut self, additional: usize) {
702        // This function does not use the hot/cold paths trick that `reserve`
703        // and `push` do, on assumption that user probably knows this will
704        // require an increase in capacity. Otherwise, they'd likely use
705        // `reserve`.
706        let new_cap = self
707            .len
708            .checked_add(additional)
709            .expect("cannot reserve a larger AlignedVec");
710        if new_cap > self.cap {
711            assert!(
712                new_cap <= Self::MAX_CAPACITY,
713                "cannot reserve a larger AlignedVec"
714            );
715            unsafe { self.change_capacity(new_cap) };
716        }
717    }
718
719    /// Forces the length of the vector to `new_len`.
720    ///
721    /// This is a low-level operation that maintains none of the normal
722    /// invariants of the type.
723    ///
724    /// # Safety
725    ///
726    /// - `new_len` must be less than or equal to
727    ///   [`capacity()`](AlignedVec::capacity)
728    /// - The elements at `old_len..new_len` must be initialized
729    ///
730    /// # Examples
731    /// ```
732    /// # use rkyv::util::AlignedVec;
733    /// let mut vec = AlignedVec::<16>::with_capacity(3);
734    /// vec.extend_from_slice(&[1, 2, 3]);
735    ///
736    /// // SAFETY:
737    /// // 1. `old_len..0` is empty to no elements need to be initialized.
738    /// // 2. `0 <= capacity` always holds whatever capacity is.
739    /// unsafe {
740    ///     vec.set_len(0);
741    /// }
742    /// ```
743    pub unsafe fn set_len(&mut self, new_len: usize) {
744        debug_assert!(new_len <= self.capacity());
745
746        self.len = new_len;
747    }
748
749    /// Converts the vector into `Box<[u8]>`. The returned slice is 1-aligned.
750    ///
751    /// This method reallocates and copies the underlying bytes. Any excess
752    /// capacity is dropped.
753    ///
754    /// # Examples
755    /// ```
756    /// # use rkyv::util::AlignedVec;
757    /// let mut v = AlignedVec::<16>::new();
758    /// v.extend_from_slice(&[1, 2, 3]);
759    ///
760    /// let slice = v.into_boxed_slice();
761    /// ```
762    ///
763    /// Any excess capacity is removed:
764    ///
765    /// ```
766    /// # use rkyv::util::AlignedVec;
767    /// let mut vec = AlignedVec::<16>::with_capacity(10);
768    /// vec.extend_from_slice(&[1, 2, 3]);
769    ///
770    /// assert_eq!(vec.capacity(), 10);
771    /// let slice = vec.into_boxed_slice();
772    /// assert_eq!(slice.len(), 3);
773    /// ```
774    pub fn into_boxed_slice(self) -> Box<[u8]> {
775        self.into_vec().into_boxed_slice()
776    }
777
778    /// Converts the vector into `Vec<u8>`.
779    ///
780    /// This method reallocates and copies the underlying bytes. Any excess
781    /// capacity is dropped.
782    ///
783    /// # Examples
784    /// ```
785    /// # use rkyv::util::AlignedVec;
786    /// let mut v = AlignedVec::<16>::new();
787    /// v.extend_from_slice(&[1, 2, 3]);
788    ///
789    /// let vec = v.into_vec();
790    /// assert_eq!(vec.len(), 3);
791    /// assert_eq!(vec.as_slice(), &[1, 2, 3]);
792    /// ```
793    pub fn into_vec(self) -> Vec<u8> {
794        Vec::from(self.as_ref())
795    }
796
797    /// Decompose an [`AlignedVec`] into its raw components: `(NonNull pointer,
798    /// length, capacity)`.
799    ///
800    /// The returned parts can be used to re-assemble the [`AlignedVec`] using
801    /// the [`from_parts`](AlignedVec::from_parts) function.
802    ///
803    /// After calling this function, the caller is responsible for the memory
804    /// previously managed by the [`AlignedVec`]. The only way to do this is
805    /// to convert the [`NonNull`] pointer, the length and the capacity back
806    /// into an [`AlignedVec`] using the [`from_parts`](AlignedVec::from_parts)
807    /// function, allowing the destructor to perform the cleanup.
808    ///
809    /// # Example
810    ///
811    /// ```
812    /// use rkyv::util::AlignedVec;
813    ///
814    /// let mut v: AlignedVec<16> = AlignedVec::new();
815    /// for i in 1..=5 {
816    ///     v.push(i);
817    /// }
818    ///
819    /// let (ptr, len, cap) = v.into_parts();
820    ///
821    /// let rebuilt: AlignedVec<16> =
822    ///     unsafe { AlignedVec::from_parts(ptr, len, cap) };
823    /// assert_eq!(rebuilt.as_slice(), &[1, 2, 3, 4, 5]);
824    /// ```
825    #[must_use = "losing the pointer will leak memory"]
826    pub fn into_parts(self) -> (NonNull<u8>, usize, usize) {
827        let this = ManuallyDrop::new(self);
828        (this.ptr, this.len, this.cap)
829    }
830
831    /// Create an [`AlignedVec`] directly from a [`NonNull`] pointer, a length
832    /// and a capacity.
833    ///
834    /// # Safety
835    ///
836    /// This is method is only safe to use with the parts returned from calling
837    /// [`into_parts`](AlignedVec::into_parts). The ownership of `ptr` is
838    /// transferred to the [`AlignedVec`], which may then de- or reallocate
839    /// the pointer or change the contents of the memory pointed to by the
840    /// pointer at will. Ensure that nothing else uses the pointer after
841    /// calling this function.
842    ///
843    /// # Example
844    ///
845    /// ```
846    /// use rkyv::util::AlignedVec;
847    ///
848    /// let mut v: AlignedVec<16> = AlignedVec::new();
849    /// for i in 1..=5 {
850    ///     v.push(i);
851    /// }
852    ///
853    /// let (ptr, len, cap) = v.into_parts();
854    ///
855    /// let rebuilt: AlignedVec<16> =
856    ///     unsafe { AlignedVec::from_parts(ptr, len, cap) };
857    /// assert_eq!(rebuilt.as_slice(), &[1, 2, 3, 4, 5]);
858    /// ```
859    pub unsafe fn from_parts(ptr: NonNull<u8>, len: usize, cap: usize) -> Self {
860        Self { ptr, len, cap }
861    }
862}
863
864#[cfg(feature = "std")]
865const _: () = {
866    use std::io;
867
868    impl<const A: usize> AlignedVec<A> {
869        /// Reads all bytes until EOF from `r` and appends them to this
870        /// `AlignedVec`.
871        ///
872        /// If successful, this function will return the total number of bytes
873        /// read.
874        ///
875        /// # Examples
876        /// ```
877        /// # use rkyv::util::AlignedVec;
878        ///
879        /// let source = (0..4096).map(|x| (x % 256) as u8).collect::<Vec<_>>();
880        /// let mut bytes = AlignedVec::<16>::new();
881        /// bytes.extend_from_reader(&mut source.as_slice()).unwrap();
882        ///
883        /// assert_eq!(bytes.len(), 4096);
884        /// assert_eq!(bytes[0], 0);
885        /// assert_eq!(bytes[100], 100);
886        /// assert_eq!(bytes[2945], 129);
887        /// ```
888        pub fn extend_from_reader<R: io::Read + ?Sized>(
889            &mut self,
890            r: &mut R,
891        ) -> io::Result<usize> {
892            let start_len = self.len();
893            let start_cap = self.capacity();
894
895            // Extra initialized bytes from previous loop iteration.
896            let mut initialized = 0;
897            loop {
898                if self.len() == self.capacity() {
899                    // No available capacity, reserve some space.
900                    self.reserve(32);
901                }
902
903                let read_buf_start = unsafe { self.as_mut_ptr().add(self.len) };
904                let read_buf_len = self.capacity() - self.len();
905
906                // Initialize the uninitialized portion of the available space.
907                unsafe {
908                    // The first `initialized` bytes don't need to be zeroed.
909                    // This leaves us `read_buf_len - initialized` bytes to zero
910                    // starting at `initialized`.
911                    core::ptr::write_bytes(
912                        read_buf_start.add(initialized),
913                        0,
914                        read_buf_len - initialized,
915                    );
916                }
917
918                // The entire read buffer is now initialized, so we can create a
919                // mutable slice of it.
920                let read_buf = unsafe {
921                    core::slice::from_raw_parts_mut(
922                        read_buf_start,
923                        read_buf_len,
924                    )
925                };
926
927                match r.read(read_buf) {
928                    Ok(read) => {
929                        // We filled `read` additional bytes.
930                        unsafe {
931                            self.set_len(self.len() + read);
932                        }
933                        initialized = read_buf_len - read;
934
935                        if read == 0 {
936                            return Ok(self.len() - start_len);
937                        }
938                    }
939                    Err(e) if e.kind() == io::ErrorKind::Interrupted => {
940                        continue
941                    }
942                    Err(e) => return Err(e),
943                }
944
945                if self.len() == self.capacity() && self.capacity() == start_cap
946                {
947                    // The buffer might be an exact fit. Let's read into a probe
948                    // buffer and see if it returns `Ok(0)`.
949                    // If so, we've avoided an unnecessary
950                    // doubling of the capacity. But if not, append the
951                    // probe buffer to the primary buffer and let its capacity
952                    // grow.
953                    let mut probe = [0u8; 32];
954
955                    loop {
956                        match r.read(&mut probe) {
957                            Ok(0) => return Ok(self.len() - start_len),
958                            Ok(n) => {
959                                self.extend_from_slice(&probe[..n]);
960                                break;
961                            }
962                            Err(ref e)
963                                if e.kind() == io::ErrorKind::Interrupted =>
964                            {
965                                continue
966                            }
967                            Err(e) => return Err(e),
968                        }
969                    }
970                }
971            }
972        }
973    }
974
975    impl<const A: usize> io::Write for AlignedVec<A> {
976        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
977            self.extend_from_slice(buf);
978            Ok(buf.len())
979        }
980
981        fn write_vectored(
982            &mut self,
983            bufs: &[io::IoSlice<'_>],
984        ) -> io::Result<usize> {
985            let len = bufs.iter().map(|b| b.len()).sum();
986            self.reserve(len);
987            for buf in bufs {
988                self.extend_from_slice(buf);
989            }
990            Ok(len)
991        }
992
993        fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
994            self.extend_from_slice(buf);
995            Ok(())
996        }
997
998        fn flush(&mut self) -> io::Result<()> {
999            Ok(())
1000        }
1001    }
1002};
1003
1004impl<const A: usize> From<AlignedVec<A>> for Vec<u8> {
1005    fn from(aligned: AlignedVec<A>) -> Self {
1006        aligned.to_vec()
1007    }
1008}
1009
1010impl<const A: usize> AsMut<[u8]> for AlignedVec<A> {
1011    fn as_mut(&mut self) -> &mut [u8] {
1012        self.as_mut_slice()
1013    }
1014}
1015
1016impl<const A: usize> AsRef<[u8]> for AlignedVec<A> {
1017    fn as_ref(&self) -> &[u8] {
1018        self.as_slice()
1019    }
1020}
1021
1022impl<const A: usize> Borrow<[u8]> for AlignedVec<A> {
1023    fn borrow(&self) -> &[u8] {
1024        self.as_slice()
1025    }
1026}
1027
1028impl<const A: usize> BorrowMut<[u8]> for AlignedVec<A> {
1029    fn borrow_mut(&mut self) -> &mut [u8] {
1030        self.as_mut_slice()
1031    }
1032}
1033
1034impl<const A: usize> Clone for AlignedVec<A> {
1035    fn clone(&self) -> Self {
1036        unsafe {
1037            let mut result = Self::with_capacity(self.len);
1038            result.len = self.len;
1039            core::ptr::copy_nonoverlapping(
1040                self.as_ptr(),
1041                result.as_mut_ptr(),
1042                self.len,
1043            );
1044            result
1045        }
1046    }
1047}
1048
1049impl<const A: usize> fmt::Debug for AlignedVec<A> {
1050    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1051        self.as_slice().fmt(f)
1052    }
1053}
1054
1055impl<const A: usize> Default for AlignedVec<A> {
1056    fn default() -> Self {
1057        Self::new()
1058    }
1059}
1060
1061impl<const A: usize> Deref for AlignedVec<A> {
1062    type Target = [u8];
1063
1064    fn deref(&self) -> &Self::Target {
1065        self.as_slice()
1066    }
1067}
1068
1069impl<const A: usize> DerefMut for AlignedVec<A> {
1070    fn deref_mut(&mut self) -> &mut Self::Target {
1071        self.as_mut_slice()
1072    }
1073}
1074
1075impl<const A: usize, I: slice::SliceIndex<[u8]>> Index<I> for AlignedVec<A> {
1076    type Output = <I as slice::SliceIndex<[u8]>>::Output;
1077
1078    fn index(&self, index: I) -> &Self::Output {
1079        &self.as_slice()[index]
1080    }
1081}
1082
1083impl<const A: usize, I: slice::SliceIndex<[u8]>> IndexMut<I> for AlignedVec<A> {
1084    fn index_mut(&mut self, index: I) -> &mut Self::Output {
1085        &mut self.as_mut_slice()[index]
1086    }
1087}
1088
1089// SAFETY: AlignedVec is safe to send to another thread
1090unsafe impl<const A: usize> Send for AlignedVec<A> {}
1091
1092// SAFETY: AlignedVec is safe to share between threads
1093unsafe impl<const A: usize> Sync for AlignedVec<A> {}
1094
1095impl<const A: usize> Unpin for AlignedVec<A> {}
1096
1097impl<const A: usize> ArchiveWith<AlignedVec<A>> for AsVec {
1098    type Archived = ArchivedVec<u8>;
1099    type Resolver = VecResolver;
1100
1101    fn resolve_with(
1102        field: &AlignedVec<A>,
1103        resolver: Self::Resolver,
1104        out: Place<Self::Archived>,
1105    ) {
1106        ArchivedVec::resolve_from_len(field.len(), resolver, out)
1107    }
1108}
1109
1110impl<S, const A: usize> SerializeWith<AlignedVec<A>, S> for AsVec
1111where
1112    S: Allocator + Fallible + Writer + ?Sized,
1113{
1114    fn serialize_with(
1115        field: &AlignedVec<A>,
1116        serializer: &mut S,
1117    ) -> Result<Self::Resolver, S::Error> {
1118        ArchivedVec::serialize_from_slice(field.as_slice(), serializer)
1119    }
1120}
1121
1122impl<D, const A: usize> DeserializeWith<ArchivedVec<u8>, AlignedVec<A>, D>
1123    for AsVec
1124where
1125    D: Fallible + ?Sized,
1126{
1127    fn deserialize_with(
1128        field: &ArchivedVec<u8>,
1129        _: &mut D,
1130    ) -> Result<AlignedVec<A>, D::Error> {
1131        let mut result = AlignedVec::with_capacity(field.len());
1132        result.extend_from_slice(field.as_slice());
1133        Ok(result)
1134    }
1135}