rkyv/util/alloc/
aligned_vec.rs

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