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