mod.rs - source
Description: Source of the Rust file `library/alloc/src/raw_vec/mod.rs`.

#
alloc/raw_vec/
mod.rs

    1#![unstable(feature = "raw_vec_internals", reason = "unstable const warnings", issue = "none")]
    2#![cfg_attr(test, allow(dead_code))]
    3
    4// Note: This module is also included in the alloctests crate using #[path] to
    5// run the tests. See the comment there for an explanation why this is the case.
    6
    7use core::marker::{Destruct, PhantomData};
    8use core::mem::{ManuallyDrop, MaybeUninit, SizedTypeProperties};
    9use core::ptr::{self, Alignment, NonNull, Unique};
    10use core::{cmp, hint};
    11
    12#[cfg(not(no_global_oom_handling))]
    13use crate::alloc::handle_alloc_error;
    14use crate::alloc::{Allocator, Global, Layout};
    15use crate::boxed::Box;
    16use crate::collections::TryReserveError;
    17use crate::collections::TryReserveErrorKind::*;
    18
    19#[cfg(test)]
    20mod tests;
    21
    22// One central function responsible for reporting capacity overflows. This'll
    23// ensure that the code generation related to these panics is minimal as there's
    24// only one location which panics rather than a bunch throughout the module.
    25#[cfg(not(no_global_oom_handling))]
    26#[cfg_attr(not(panic = "immediate-abort"), inline(never))]
    27const fn capacity_overflow() -> ! {
    28    panic!("capacity overflow");
    29}
    30
    31enum AllocInit {
    32    /// The contents of the new memory are uninitialized.
    33    Uninitialized,
    34    #[cfg(not(no_global_oom_handling))]
    35    /// The new memory is guaranteed to be zeroed.
    36    Zeroed,
    37}
    38
    39type Cap = core::num::niche_types::UsizeNoHighBit;
    40
    41const ZERO_CAP: Cap = unsafe { Cap::new_unchecked(0) };
    42
    43/// `Cap(cap)`, except if `T` is a ZST then `Cap::ZERO`.
    44///
    45/// # Safety: cap must be <= `isize::MAX`.
    46unsafe fn new_cap<T>(cap: usize) -> Cap {
    47    if T::IS_ZST { ZERO_CAP } else { unsafe { Cap::new_unchecked(cap) } }
    48}
    49
    50/// A low-level utility for more ergonomically allocating, reallocating, and deallocating
    51/// a buffer of memory on the heap without having to worry about all the corner cases
    52/// involved. This type is excellent for building your own data structures like Vec and VecDeque.
    53/// In particular:
    54///
    55/// * Produces `Unique::dangling()` on zero-sized types.
    56/// * Produces `Unique::dangling()` on zero-length allocations.
    57/// * Avoids freeing `Unique::dangling()`.
    58/// * Catches all overflows in capacity computations (promotes them to "capacity overflow" panics).
    59/// * Guards against 32-bit systems allocating more than `isize::MAX` bytes.
    60/// * Guards against overflowing your length.
    61/// * Calls `handle_alloc_error` for fallible allocations.
    62/// * Contains a `ptr::Unique` and thus endows the user with all related benefits.
    63/// * Uses the excess returned from the allocator to use the largest available capacity.
    64///
    65/// This type does not in anyway inspect the memory that it manages. When dropped it *will*
    66/// free its memory, but it *won't* try to drop its contents. It is up to the user of `RawVec`
    67/// to handle the actual things *stored* inside of a `RawVec`.
    68///
    69/// Note that the excess of a zero-sized types is always infinite, so `capacity()` always returns
    70/// `usize::MAX`. This means that you need to be careful when round-tripping this type with a
    71/// `Box<[T]>`, since `capacity()` won't yield the length.
    72#[allow(missing_debug_implementations)]
    73pub(crate) struct RawVec<T, A: Allocator = Global> {
    74    inner: RawVecInner<A>,
    75    _marker: PhantomData<T>,
    76}
    77
    78/// Like a `RawVec`, but only generic over the allocator, not the type.
    79///
    80/// As such, all the methods need the layout passed-in as a parameter.
    81///
    82/// Having this separation reduces the amount of code we need to monomorphize,
    83/// as most operations don't need the actual type, just its layout.
    84#[allow(missing_debug_implementations)]
    85struct RawVecInner<A: Allocator = Global> {
    86    ptr: Unique<u8>,
    87    /// Never used for ZSTs; it's `capacity()`'s responsibility to return usize::MAX in that case.
    88    ///
    89    /// # Safety
    90    ///
    91    /// `cap` must be in the `0..=isize::MAX` range.
    92    cap: Cap,
    93    alloc: A,
    94}
    95
    96impl<T> RawVec<T, Global> {
    97    /// Creates the biggest possible `RawVec` (on the system heap)
    98    /// without allocating. If `T` has positive size, then this makes a
    99    /// `RawVec` with capacity `0`. If `T` is zero-sized, then it makes a
    100    /// `RawVec` with capacity `usize::MAX`. Useful for implementing
    101    /// delayed allocation.
    102    #[must_use]
    103    pub(crate) const fn new() -> Self {
    104        Self::new_in(Global)
    105    }
    106
    107    /// Creates a `RawVec` (on the system heap) with exactly the
    108    /// capacity and alignment requirements for a `[T; capacity]`. This is
    109    /// equivalent to calling `RawVec::new` when `capacity` is `0` or `T` is
    110    /// zero-sized. Note that if `T` is zero-sized this means you will
    111    /// *not* get a `RawVec` with the requested capacity.
    112    ///
    113    /// Non-fallible version of `try_with_capacity`
    114    ///
    115    /// # Panics
    116    ///
    117    /// Panics if the requested capacity exceeds `isize::MAX` bytes.
    118    ///
    119    /// # Aborts
    120    ///
    121    /// Aborts on OOM.
    122    #[cfg(not(any(no_global_oom_handling, test)))]
    123    #[must_use]
    124    #[inline]
    125    pub(crate) fn with_capacity(capacity: usize) -> Self {
    126        Self { inner: RawVecInner::with_capacity(capacity, T::LAYOUT), _marker: PhantomData }
    127    }
    128
    129    /// Like `with_capacity`, but guarantees the buffer is zeroed.
    130    #[cfg(not(any(no_global_oom_handling, test)))]
    131    #[must_use]
    132    #[inline]
    133    pub(crate) fn with_capacity_zeroed(capacity: usize) -> Self {
    134        Self {
    135            inner: RawVecInner::with_capacity_zeroed_in(capacity, Global, T::LAYOUT),
    136            _marker: PhantomData,
    137        }
    138    }
    139}
    140
    141impl RawVecInner<Global> {
    142    #[cfg(not(any(no_global_oom_handling, test)))]
    143    #[must_use]
    144    #[inline]
    145    fn with_capacity(capacity: usize, elem_layout: Layout) -> Self {
    146        match Self::try_allocate_in(capacity, AllocInit::Uninitialized, Global, elem_layout) {
    147            Ok(res) => res,
    148            Err(err) => handle_error(err),
    149        }
    150    }
    151}
    152
    153// Tiny Vecs are dumb. Skip to:
    154// - 8 if the element size is 1, because any heap allocator is likely
    155//   to round up a request of less than 8 bytes to at least 8 bytes.
    156// - 4 if elements are moderate-sized (<= 1 KiB).
    157// - 1 otherwise, to avoid wasting too much space for very short Vecs.
    158const fn min_non_zero_cap(size: usize) -> usize {
    159    if size == 1 {
    160        8
    161    } else if size <= 1024 {
    162        4
    163    } else {
    164        1
    165    }
    166}
    167
    168#[rustc_const_unstable(feature = "const_heap", issue = "79597")]
    169#[rustfmt::skip] // FIXME(fee1-dead): temporary measure before rustfmt is bumped
    170const impl<T, A: [const] Allocator + [const] Destruct> RawVec<T, A> {
    171    /// Like `with_capacity`, but parameterized over the choice of
    172    /// allocator for the returned `RawVec`.
    173    #[cfg(not(no_global_oom_handling))]
    174    #[inline]
    175    pub(crate) fn with_capacity_in(capacity: usize, alloc: A) -> Self {
    176        Self {
    177            inner: RawVecInner::with_capacity_in(capacity, alloc, T::LAYOUT),
    178            _marker: PhantomData,
    179        }
    180    }
    181
    182    /// A specialized version of `self.reserve(len, 1)` which requires the
    183    /// caller to ensure `len == self.capacity()`.
    184    #[cfg(not(no_global_oom_handling))]
    185    #[inline(never)]
    186    pub(crate) fn grow_one(&mut self) {
    187        // SAFETY: All calls on self.inner pass T::LAYOUT as the elem_layout
    188        unsafe { self.inner.grow_one(T::LAYOUT) }
    189    }
    190}
    191
    192impl<T, A: Allocator> RawVec<T, A> {
    193    #[cfg(not(no_global_oom_handling))]
    194    pub(crate) const MIN_NON_ZERO_CAP: usize = min_non_zero_cap(size_of::<T>());
    195
    196    /// Like `new`, but parameterized over the choice of allocator for
    197    /// the returned `RawVec`.
    198    #[inline]
    199    pub(crate) const fn new_in(alloc: A) -> Self {
    200        // Check assumption made in `current_memory`
    201        const { assert!(T::LAYOUT.size() % T::LAYOUT.align() == 0) };
    202        Self { inner: RawVecInner::new_in(alloc, Alignment::of::<T>()), _marker: PhantomData }
    203    }
    204
    205    /// Like `try_with_capacity`, but parameterized over the choice of
    206    /// allocator for the returned `RawVec`.
    207    #[inline]
    208    pub(crate) fn try_with_capacity_in(capacity: usize, alloc: A) -> Result<Self, TryReserveError> {
    209        match RawVecInner::try_with_capacity_in(capacity, alloc, T::LAYOUT) {
    210            Ok(inner) => Ok(Self { inner, _marker: PhantomData }),
    211            Err(e) => Err(e),
    212        }
    213    }
    214
    215    /// Like `with_capacity_zeroed`, but parameterized over the choice
    216    /// of allocator for the returned `RawVec`.
    217    #[cfg(not(no_global_oom_handling))]
    218    #[inline]
    219    pub(crate) fn with_capacity_zeroed_in(capacity: usize, alloc: A) -> Self {
    220        Self {
    221            inner: RawVecInner::with_capacity_zeroed_in(capacity, alloc, T::LAYOUT),
    222            _marker: PhantomData,
    223        }
    224    }
    225
    226    /// Converts the entire buffer into `Box<[MaybeUninit<T>]>` with the specified `len`.
    227    ///
    228    /// Note that this will correctly reconstitute any `cap` changes
    229    /// that may have been performed. (See description of type for details.)
    230    ///
    231    /// # Safety
    232    ///
    233    /// * `len` must be greater than or equal to the most recently requested capacity, and
    234    /// * `len` must be less than or equal to `self.capacity()`.
    235    ///
    236    /// Note, that the requested capacity and `self.capacity()` could differ, as
    237    /// an allocator could overallocate and return a greater memory block than requested.
    238    pub(crate) unsafe fn into_box(self, len: usize) -> Box<[MaybeUninit<T>], A> {
    239        // Sanity-check one half of the safety requirement (we cannot check the other half).
    240        debug_assert!(
    241            len <= self.capacity(),
    242            "`len` must be smaller than or equal to `self.capacity()`"
    243        );
    244
    245        let me = ManuallyDrop::new(self);
    246        unsafe {
    247            let slice = ptr::slice_from_raw_parts_mut(me.ptr() as *mut MaybeUninit<T>, len);
    248            Box::from_raw_in(slice, ptr::read(&me.inner.alloc))
    249        }
    250    }
    251
    252    /// Reconstitutes a `RawVec` from a pointer, capacity, and allocator.
    253    ///
    254    /// # Safety
    255    ///
    256    /// The `ptr` must be allocated (via the given allocator `alloc`), and with the given
    257    /// `capacity`.
    258    /// The `capacity` cannot exceed `isize::MAX` for sized types. (only a concern on 32-bit
    259    /// systems). For ZSTs capacity is ignored.
    260    /// If the `ptr` and `capacity` come from a `RawVec` created via `alloc`, then this is
    261    /// guaranteed.
    262    #[inline]
    263    pub(crate) unsafe fn from_raw_parts_in(ptr: *mut T, capacity: usize, alloc: A) -> Self {
    264        // SAFETY: Precondition passed to the caller
    265        unsafe {
    266            let ptr = ptr.cast();
    267            let capacity = new_cap::<T>(capacity);
    268            Self {
    269                inner: RawVecInner::from_raw_parts_in(ptr, capacity, alloc),
    270                _marker: PhantomData,
    271            }
    272        }
    273    }
    274
    275    /// A convenience method for hoisting the non-null precondition out of [`RawVec::from_raw_parts_in`].
    276    ///
    277    /// # Safety
    278    ///
    279    /// See [`RawVec::from_raw_parts_in`].
    280    #[inline]
    281    pub(crate) unsafe fn from_nonnull_in(ptr: NonNull<T>, capacity: usize, alloc: A) -> Self {
    282        // SAFETY: Precondition passed to the caller
    283        unsafe {
    284            let ptr = ptr.cast();
    285            let capacity = new_cap::<T>(capacity);
    286            Self { inner: RawVecInner::from_nonnull_in(ptr, capacity, alloc), _marker: PhantomData }
    287        }
    288    }
    289
    290    /// Gets a raw pointer to the start of the allocation. Note that this is
    291    /// `Unique::dangling()` if `capacity == 0` or `T` is zero-sized. In the former case, you must
    292    /// be careful.
    293    #[inline]
    294    pub(crate) const fn ptr(&self) -> *mut T {
    295        self.inner.ptr()
    296    }
    297
    298    #[inline]
    299    pub(crate) const fn non_null(&self) -> NonNull<T> {
    300        self.inner.non_null()
    301    }
    302
    303    /// Gets the capacity of the allocation.
    304    ///
    305    /// This will always be `usize::MAX` if `T` is zero-sized.
    306    #[inline]
    307    pub(crate) const fn capacity(&self) -> usize {
    308        self.inner.capacity(size_of::<T>())
    309    }
    310
    311    /// Returns a shared reference to the allocator backing this `RawVec`.
    312    #[inline]
    313    pub(crate) fn allocator(&self) -> &A {
    314        self.inner.allocator()
    315    }
    316
    317    /// Ensures that the buffer contains at least enough space to hold `len +
    318    /// additional` elements. If it doesn't already have enough capacity, will
    319    /// reallocate enough space plus comfortable slack space to get amortized
    320    /// *O*(1) behavior. Will limit this behavior if it would needlessly cause
    321    /// itself to panic.
    322    ///
    323    /// If `len` exceeds `self.capacity()`, this may fail to actually allocate
    324    /// the requested space. This is not really unsafe, but the unsafe
    325    /// code *you* write that relies on the behavior of this function may break.
    326    ///
    327    /// This is ideal for implementing a bulk-push operation like `extend`.
    328    ///
    329    /// # Panics
    330    ///
    331    /// Panics if the new capacity exceeds `isize::MAX` _bytes_.
    332    ///
    333    /// # Aborts
    334    ///
    335    /// Aborts on OOM.
    336    #[cfg(not(no_global_oom_handling))]
    337    #[inline]
    338    pub(crate) fn reserve(&mut self, len: usize, additional: usize) {
    339        // SAFETY: All calls on self.inner pass T::LAYOUT as the elem_layout
    340        unsafe { self.inner.reserve(len, additional, T::LAYOUT) }
    341    }
    342
    343    /// The same as `reserve`, but returns on errors instead of panicking or aborting.
    344    pub(crate) fn try_reserve(
    345        &mut self,
    346        len: usize,
    347        additional: usize,
    348    ) -> Result<(), TryReserveError> {
    349        // SAFETY: All calls on self.inner pass T::LAYOUT as the elem_layout
    350        unsafe { self.inner.try_reserve(len, additional, T::LAYOUT) }
    351    }
    352
    353    /// Ensures that the buffer contains at least enough space to hold `len +
    354    /// additional` elements. If it doesn't already, will reallocate the
    355    /// minimum possible amount of memory necessary. Generally this will be
    356    /// exactly the amount of memory necessary, but in principle the allocator
    357    /// is free to give back more than we asked for.
    358    ///
    359    /// If `len` exceeds `self.capacity()`, this may fail to actually allocate
    360    /// the requested space. This is not really unsafe, but the unsafe code
    361    /// *you* write that relies on the behavior of this function may break.
    362    ///
    363    /// # Panics
    364    ///
    365    /// Panics if the new capacity exceeds `isize::MAX` _bytes_.
    366    ///
    367    /// # Aborts
    368    ///
    369    /// Aborts on OOM.
    370    #[cfg(not(no_global_oom_handling))]
    371    pub(crate) fn reserve_exact(&mut self, len: usize, additional: usize) {
    372        // SAFETY: All calls on self.inner pass T::LAYOUT as the elem_layout
    373        unsafe { self.inner.reserve_exact(len, additional, T::LAYOUT) }
    374    }
    375
    376    /// The same as `reserve_exact`, but returns on errors instead of panicking or aborting.
    377    pub(crate) fn try_reserve_exact(
    378        &mut self,
    379        len: usize,
    380        additional: usize,
    381    ) -> Result<(), TryReserveError> {
    382        // SAFETY: All calls on self.inner pass T::LAYOUT as the elem_layout
    383        unsafe { self.inner.try_reserve_exact(len, additional, T::LAYOUT) }
    384    }
    385
    386    /// Shrinks the buffer down to the specified capacity. If the given amount
    387    /// is 0, actually completely deallocates.
    388    ///
    389    /// # Panics
    390    ///
    391    /// Panics if the given amount is *larger* than the current capacity.
    392    ///
    393    /// # Aborts
    394    ///
    395    /// Aborts on OOM.
    396    #[cfg(not(no_global_oom_handling))]
    397    #[inline]
    398    pub(crate) fn shrink_to_fit(&mut self, cap: usize) {
    399        // SAFETY: All calls on self.inner pass T::LAYOUT as the elem_layout
    400        unsafe { self.inner.shrink_to_fit(cap, T::LAYOUT) }
    401    }
    402}
    403
    404unsafe impl<#[may_dangle] T, A: Allocator> Drop for RawVec<T, A> {
    405    /// Frees the memory owned by the `RawVec` *without* trying to drop its contents.
    406    fn drop(&mut self) {
    407        // SAFETY: We are in a Drop impl, self.inner will not be used again.
    408        unsafe { self.inner.deallocate(T::LAYOUT) }
    409    }
    410}
    411
    412#[rustc_const_unstable(feature = "const_heap", issue = "79597")]
    413#[rustfmt::skip] // FIXME(fee1-dead): temporary measure before rustfmt is bumped
    414const impl<A: [const] Allocator + [const] Destruct> RawVecInner<A> {
    415    #[cfg(not(no_global_oom_handling))]
    416    #[inline]
    417    fn with_capacity_in(capacity: usize, alloc: A, elem_layout: Layout) -> Self {
    418        match Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc, elem_layout) {
    419            Ok(this) => {
    420                unsafe {
    421                    // Make it more obvious that a subsequent Vec::reserve(capacity) will not allocate.
    422                    hint::assert_unchecked(!this.needs_to_grow(0, capacity, elem_layout));
    423                }
    424                this
    425            }
    426            Err(err) => handle_error(err),
    427        }
    428    }
    429
    430    fn try_allocate_in(
    431        capacity: usize,
    432        init: AllocInit,
    433        alloc: A,
    434        elem_layout: Layout,
    435    ) -> Result<Self, TryReserveError> {
    436        // We avoid `unwrap_or_else` here because it bloats the amount of
    437        // LLVM IR generated.
    438        let layout = match layout_array(capacity, elem_layout) {
    439            Ok(layout) => layout,
    440            Err(_) => return Err(CapacityOverflow.into()),
    441        };
    442
    443        // Don't allocate here because `Drop` will not deallocate when `capacity` is 0.
    444        if layout.size() == 0 {
    445            return Ok(Self::new_in(alloc, elem_layout.alignment()));
    446        }
    447
    448        let result = match init {
    449            AllocInit::Uninitialized => alloc.allocate(layout),
    450            #[cfg(not(no_global_oom_handling))]
    451            AllocInit::Zeroed => alloc.allocate_zeroed(layout),
    452        };
    453        let ptr = match result {
    454            Ok(ptr) => ptr,
    455            Err(_) => return Err(AllocError { layout, non_exhaustive: () }.into()),
    456        };
    457
    458        // Allocators currently return a `NonNull<[u8]>` whose length
    459        // matches the size requested. If that ever changes, the capacity
    460        // here should change to `ptr.len() / size_of::<T>()`.
    461        Ok(Self {
    462            ptr: Unique::from(ptr.cast()),
    463            cap: unsafe { Cap::new_unchecked(capacity) },
    464            alloc,
    465        })
    466    }
    467
    468    /// # Safety
    469    /// - `elem_layout` must be valid for `self`, i.e. it must be the same `elem_layout` used to
    470    ///   initially construct `self`
    471    /// - `elem_layout`'s size must be a multiple of its alignment
    472    #[cfg(not(no_global_oom_handling))]
    473    #[inline]
    474    unsafe fn grow_one(&mut self, elem_layout: Layout) {
    475        // SAFETY: Precondition passed to caller
    476        if let Err(err) = unsafe { self.grow_amortized(self.cap.as_inner(), 1, elem_layout) } {
    477            handle_error(err);
    478        }
    479    }
    480
    481    /// # Safety
    482    /// - `elem_layout` must be valid for `self`, i.e. it must be the same `elem_layout` used to
    483    ///   initially construct `self`
    484    /// - `elem_layout`'s size must be a multiple of its alignment
    485    /// - The sum of `len` and `additional` must be greater than the current capacity
    486    unsafe fn grow_amortized(
    487        &mut self,
    488        len: usize,
    489        additional: usize,
    490        elem_layout: Layout,
    491    ) -> Result<(), TryReserveError> {
    492        // This is ensured by the calling contexts.
    493        debug_assert!(additional > 0);
    494
    495        if elem_layout.size() == 0 {
    496            // Since we return a capacity of `usize::MAX` when `elem_size` is
    497            // 0, getting to here necessarily means the `RawVec` is overfull.
    498            return Err(CapacityOverflow.into());
    499        }
    500
    501        // Nothing we can really do about these checks, sadly.
    502        let required_cap = len.checked_add(additional).ok_or(CapacityOverflow)?;
    503
    504        // This guarantees exponential growth. The doubling cannot overflow
    505        // because `cap <= isize::MAX` and the type of `cap` is `usize`.
    506        let cap = cmp::max(self.cap.as_inner() * 2, required_cap);
    507        let cap = cmp::max(min_non_zero_cap(elem_layout.size()), cap);
    508
    509        // SAFETY:
    510        // - cap >= len + additional
    511        // - other preconditions passed to caller
    512        let ptr = unsafe { self.finish_grow(cap, elem_layout)? };
    513
    514        // SAFETY: `finish_grow` would have failed if `cap > isize::MAX`
    515        unsafe { self.set_ptr_and_cap(ptr, cap) };
    516        Ok(())
    517    }
    518
    519    /// # Safety
    520    /// - `elem_layout` must be valid for `self`, i.e. it must be the same `elem_layout` used to
    521    ///   initially construct `self`
    522    /// - `elem_layout`'s size must be a multiple of its alignment
    523    /// - `cap` must be greater than the current capacity
    524    // not marked inline(never) since we want optimizers to be able to observe the specifics of this
    525    // function, see tests/codegen-llvm/vec-reserve-extend.rs.
    526    #[cold]
    527    unsafe fn finish_grow(
    528        &self,
    529        cap: usize,
    530        elem_layout: Layout,
    531    ) -> Result<NonNull<[u8]>, TryReserveError> {
    532        let new_layout = layout_array(cap, elem_layout)?;
    533
    534        let memory = if let Some((ptr, old_layout)) = unsafe { self.current_memory(elem_layout) } {
    535            // FIXME(const-hack): switch to `debug_assert_eq`
    536            debug_assert!(old_layout.align() == new_layout.align());
    537            unsafe {
    538                // The allocator checks for alignment equality
    539                hint::assert_unchecked(old_layout.align() == new_layout.align());
    540                self.alloc.grow(ptr, old_layout, new_layout)
    541            }
    542        } else {
    543            self.alloc.allocate(new_layout)
    544        };
    545
    546        // FIXME(const-hack): switch back to `map_err`
    547        match memory {
    548            Ok(memory) => Ok(memory),
    549            Err(_) => Err(AllocError { layout: new_layout, non_exhaustive: () }.into()),
    550        }
    551    }
    552}
    553
    554impl<A: Allocator> RawVecInner<A> {
    555    #[inline]
    556    const fn new_in(alloc: A, align: Alignment) -> Self {
    557        let ptr = Unique::from_non_null(NonNull::without_provenance(align.as_nonzero()));
    558        // `cap: 0` means "unallocated". zero-sized types are ignored.
    559        Self { ptr, cap: ZERO_CAP, alloc }
    560    }
    561
    562    #[inline]
    563    fn try_with_capacity_in(
    564        capacity: usize,
    565        alloc: A,
    566        elem_layout: Layout,
    567    ) -> Result<Self, TryReserveError> {
    568        Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc, elem_layout)
    569    }
    570
    571    #[cfg(not(no_global_oom_handling))]
    572    #[inline]
    573    fn with_capacity_zeroed_in(capacity: usize, alloc: A, elem_layout: Layout) -> Self {
    574        match Self::try_allocate_in(capacity, AllocInit::Zeroed, alloc, elem_layout) {
    575            Ok(res) => res,
    576            Err(err) => handle_error(err),
    577        }
    578    }
    579
    580    #[inline]
    581    unsafe fn from_raw_parts_in(ptr: *mut u8, cap: Cap, alloc: A) -> Self {
    582        Self { ptr: unsafe { Unique::new_unchecked(ptr) }, cap, alloc }
    583    }
    584
    585    #[inline]
    586    unsafe fn from_nonnull_in(ptr: NonNull<u8>, cap: Cap, alloc: A) -> Self {
    587        Self { ptr: Unique::from(ptr), cap, alloc }
    588    }
    589
    590    #[inline]
    591    const fn ptr<T>(&self) -> *mut T {
    592        self.non_null::<T>().as_ptr()
    593    }
    594
    595    #[inline]
    596    const fn non_null<T>(&self) -> NonNull<T> {
    597        self.ptr.cast().as_non_null_ptr()
    598    }
    599
    600    #[inline]
    601    const fn capacity(&self, elem_size: usize) -> usize {
    602        if elem_size == 0 { usize::MAX } else { self.cap.as_inner() }
    603    }
    604
    605    #[inline]
    606    fn allocator(&self) -> &A {
    607        &self.alloc
    608    }
    609
    610    /// # Safety
    611    /// - `elem_layout` must be valid for `self`, i.e. it must be the same `elem_layout` used to
    612    ///   initially construct `self`
    613    /// - `elem_layout`'s size must be a multiple of its alignment
    614    #[inline]
    615    #[rustc_const_unstable(feature = "const_heap", issue = "79597")]
    616    const unsafe fn current_memory(&self, elem_layout: Layout) -> Option<(NonNull<u8>, Layout)> {
    617        if elem_layout.size() == 0 || self.cap.as_inner() == 0 {
    618            None
    619        } else {
    620            // We could use Layout::array here which ensures the absence of isize and usize overflows
    621            // and could hypothetically handle differences between stride and size, but this memory
    622            // has already been allocated so we know it can't overflow and currently Rust does not
    623            // support such types. So we can do better by skipping some checks and avoid an unwrap.
    624            unsafe {
    625                let alloc_size = elem_layout.size().unchecked_mul(self.cap.as_inner());
    626                let layout = Layout::from_size_align_unchecked(alloc_size, elem_layout.align());
    627                Some((self.ptr.into(), layout))
    628            }
    629        }
    630    }
    631
    632    /// # Safety
    633    /// - `elem_layout` must be valid for `self`, i.e. it must be the same `elem_layout` used to
    634    ///   initially construct `self`
    635    /// - `elem_layout`'s size must be a multiple of its alignment
    636    #[cfg(not(no_global_oom_handling))]
    637    #[inline]
    638    unsafe fn reserve(&mut self, len: usize, additional: usize, elem_layout: Layout) {
    639        // Callers expect this function to be very cheap when there is already sufficient capacity.
    640        // Therefore, we move all the resizing and error-handling logic from grow_amortized and
    641        // handle_reserve behind a call, while making sure that this function is likely to be
    642        // inlined as just a comparison and a call if the comparison fails.
    643        #[cold]
    644        unsafe fn do_reserve_and_handle<A: Allocator>(
    645            slf: &mut RawVecInner<A>,
    646            len: usize,
    647            additional: usize,
    648            elem_layout: Layout,
    649        ) {
    650            // SAFETY: Precondition passed to caller
    651            if let Err(err) = unsafe { slf.grow_amortized(len, additional, elem_layout) } {
    652                handle_error(err);
    653            }
    654        }
    655
    656        if self.needs_to_grow(len, additional, elem_layout) {
    657            unsafe {
    658                do_reserve_and_handle(self, len, additional, elem_layout);
    659            }
    660        }
    661    }
    662
    663    /// # Safety
    664    /// - `elem_layout` must be valid for `self`, i.e. it must be the same `elem_layout` used to
    665    ///   initially construct `self`
    666    /// - `elem_layout`'s size must be a multiple of its alignment
    667    unsafe fn try_reserve(
    668        &mut self,
    669        len: usize,
    670        additional: usize,
    671        elem_layout: Layout,
    672    ) -> Result<(), TryReserveError> {
    673        if self.needs_to_grow(len, additional, elem_layout) {
    674            // SAFETY: Precondition passed to caller
    675            unsafe {
    676                self.grow_amortized(len, additional, elem_layout)?;
    677            }
    678        }
    679        unsafe {
    680            // Inform the optimizer that the reservation has succeeded or wasn't needed
    681            hint::assert_unchecked(!self.needs_to_grow(len, additional, elem_layout));
    682        }
    683        Ok(())
    684    }
    685
    686    /// # Safety
    687    /// - `elem_layout` must be valid for `self`, i.e. it must be the same `elem_layout` used to
    688    ///   initially construct `self`
    689    /// - `elem_layout`'s size must be a multiple of its alignment
    690    #[cfg(not(no_global_oom_handling))]
    691    unsafe fn reserve_exact(&mut self, len: usize, additional: usize, elem_layout: Layout) {
    692        // SAFETY: Precondition passed to caller
    693        if let Err(err) = unsafe { self.try_reserve_exact(len, additional, elem_layout) } {
    694            handle_error(err);
    695        }
    696    }
    697
    698    /// # Safety
    699    /// - `elem_layout` must be valid for `self`, i.e. it must be the same `elem_layout` used to
    700    ///   initially construct `self`
    701    /// - `elem_layout`'s size must be a multiple of its alignment
    702    unsafe fn try_reserve_exact(
    703        &mut self,
    704        len: usize,
    705        additional: usize,
    706        elem_layout: Layout,
    707    ) -> Result<(), TryReserveError> {
    708        if self.needs_to_grow(len, additional, elem_layout) {
    709            // SAFETY: Precondition passed to caller
    710            unsafe {
    711                self.grow_exact(len, additional, elem_layout)?;
    712            }
    713        }
    714        unsafe {
    715            // Inform the optimizer that the reservation has succeeded or wasn't needed
    716            hint::assert_unchecked(!self.needs_to_grow(len, additional, elem_layout));
    717        }
    718        Ok(())
    719    }
    720
    721    /// # Safety
    722    /// - `elem_layout` must be valid for `self`, i.e. it must be the same `elem_layout` used to
    723    ///   initially construct `self`
    724    /// - `elem_layout`'s size must be a multiple of its alignment
    725    /// - `cap` must be less than or equal to `self.capacity(elem_layout.size())`
    726    #[cfg(not(no_global_oom_handling))]
    727    #[inline]
    728    unsafe fn shrink_to_fit(&mut self, cap: usize, elem_layout: Layout) {
    729        if let Err(err) = unsafe { self.shrink(cap, elem_layout) } {
    730            handle_error(err);
    731        }
    732    }
    733
    734    #[inline]
    735    const fn needs_to_grow(&self, len: usize, additional: usize, elem_layout: Layout) -> bool {
    736        additional > self.capacity(elem_layout.size()).wrapping_sub(len)
    737    }
    738
    739    #[inline]
    740    #[rustc_const_unstable(feature = "const_heap", issue = "79597")]
    741    const unsafe fn set_ptr_and_cap(&mut self, ptr: NonNull<[u8]>, cap: usize) {
    742        // Allocators currently return a `NonNull<[u8]>` whose length matches
    743        // the size requested. If that ever changes, the capacity here should
    744        // change to `ptr.len() / size_of::<T>()`.
    745        self.ptr = Unique::from(ptr.cast());
    746        self.cap = unsafe { Cap::new_unchecked(cap) };
    747    }
    748
    749    /// # Safety
    750    /// - `elem_layout` must be valid for `self`, i.e. it must be the same `elem_layout` used to
    751    ///   initially construct `self`
    752    /// - `elem_layout`'s size must be a multiple of its alignment
    753    /// - The sum of `len` and `additional` must be greater than the current capacity
    754    unsafe fn grow_exact(
    755        &mut self,
    756        len: usize,
    757        additional: usize,
    758        elem_layout: Layout,
    759    ) -> Result<(), TryReserveError> {
    760        if elem_layout.size() == 0 {
    761            // Since we return a capacity of `usize::MAX` when the type size is
    762            // 0, getting to here necessarily means the `RawVec` is overfull.
    763            return Err(CapacityOverflow.into());
    764        }
    765
    766        let cap = len.checked_add(additional).ok_or(CapacityOverflow)?;
    767
    768        // SAFETY: preconditions passed to caller
    769        let ptr = unsafe { self.finish_grow(cap, elem_layout)? };
    770
    771        // SAFETY: `finish_grow` would have failed if `cap > isize::MAX`
    772        unsafe { self.set_ptr_and_cap(ptr, cap) };
    773        Ok(())
    774    }
    775
    776    /// # Safety
    777    /// - `elem_layout` must be valid for `self`, i.e. it must be the same `elem_layout` used to
    778    ///   initially construct `self`
    779    /// - `elem_layout`'s size must be a multiple of its alignment
    780    /// - `cap` must be less than or equal to `self.capacity(elem_layout.size())`
    781    #[cfg(not(no_global_oom_handling))]
    782    #[inline]
    783    unsafe fn shrink(&mut self, cap: usize, elem_layout: Layout) -> Result<(), TryReserveError> {
    784        assert!(cap <= self.capacity(elem_layout.size()), "Tried to shrink to a larger capacity");
    785        // SAFETY: Just checked this isn't trying to grow
    786        unsafe { self.shrink_unchecked(cap, elem_layout) }
    787    }
    788
    789    /// `shrink`, but without the capacity check.
    790    ///
    791    /// This is split out so that `shrink` can inline the check, since it
    792    /// optimizes out in things like `shrink_to_fit`, without needing to
    793    /// also inline all this code, as doing that ends up failing the
    794    /// `vec-shrink-panic` codegen test when `shrink_to_fit` ends up being too
    795    /// big for LLVM to be willing to inline.
    796    ///
    797    /// # Safety
    798    /// `cap <= self.capacity()`
    799    #[cfg(not(no_global_oom_handling))]
    800    unsafe fn shrink_unchecked(
    801        &mut self,
    802        cap: usize,
    803        elem_layout: Layout,
    804    ) -> Result<(), TryReserveError> {
    805        // SAFETY: Precondition passed to caller
    806        let Some((ptr, layout)) = (unsafe { self.current_memory(elem_layout) }) else {
    807            return Ok(());
    808        };
    809
    810        // If shrinking to 0, deallocate the buffer. We don't reach this point
    811        // for the T::IS_ZST case since current_memory() will have returned
    812        // None.
    813        if cap == 0 {
    814            unsafe { self.alloc.deallocate(ptr, layout) };
    815            self.ptr =
    816                unsafe { Unique::new_unchecked(ptr::without_provenance_mut(elem_layout.align())) };
    817            self.cap = ZERO_CAP;
    818        } else {
    819            let ptr = unsafe {
    820                // Layout cannot overflow here because it would have
    821                // overflowed earlier when capacity was larger.
    822                let new_size = elem_layout.size().unchecked_mul(cap);
    823                let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
    824                self.alloc
    825                    .shrink(ptr, layout, new_layout)
    826                    .map_err(|_| AllocError { layout: new_layout, non_exhaustive: () })?
    827            };
    828            // SAFETY: if the allocation is valid, then the capacity is too
    829            unsafe {
    830                self.set_ptr_and_cap(ptr, cap);
    831            }
    832        }
    833        Ok(())
    834    }
    835
    836    /// # Safety
    837    ///
    838    /// This function deallocates the owned allocation, but does not update `ptr` or `cap` to
    839    /// prevent double-free or use-after-free. Essentially, do not do anything with the caller
    840    /// after this function returns.
    841    /// Ideally this function would take `self` by move, but it cannot because it exists to be
    842    /// called from a `Drop` impl.
    843    unsafe fn deallocate(&mut self, elem_layout: Layout) {
    844        // SAFETY: Precondition passed to caller
    845        if let Some((ptr, layout)) = unsafe { self.current_memory(elem_layout) } {
    846            unsafe {
    847                self.alloc.deallocate(ptr, layout);
    848            }
    849        }
    850    }
    851}
    852
    853// Central function for reserve error handling.
    854#[cfg(not(no_global_oom_handling))]
    855#[cold]
    856#[optimize(size)]
    857#[rustc_const_unstable(feature = "const_heap", issue = "79597")]
    858const fn handle_error(e: TryReserveError) -> ! {
    859    match e.kind() {
    860        CapacityOverflow => capacity_overflow(),
    861        AllocError { layout, .. } => handle_alloc_error(layout),
    862    }
    863}
    864
    865#[inline]
    866#[rustc_const_unstable(feature = "const_heap", issue = "79597")]
    867const fn layout_array(cap: usize, elem_layout: Layout) -> Result<Layout, TryReserveError> {
    868    // FIXME(const-hack) return to using `map` and `map_err` once `const_closures` is implemented
    869    match elem_layout.repeat(cap) {
    870        Ok((layout, _pad)) => Ok(layout),
    871        Err(_) => Err(CapacityOverflow.into()),
    872    }
    873}