diff --git a/alloc/src/raw_vec.rs b/alloc/src/raw_vec.rs index 5b84df9ecef30..9c8fa7ceff4e5 100644 --- a/alloc/src/raw_vec.rs +++ b/alloc/src/raw_vec.rs @@ -1,7 +1,7 @@ #![unstable(feature = "raw_vec_internals", reason = "unstable const warnings", issue = "none")] -use core::alloc::LayoutError; -use core::mem::{self, ManuallyDrop, MaybeUninit, SizedTypeProperties}; +use core::marker::PhantomData; +use core::mem::{ManuallyDrop, MaybeUninit, SizedTypeProperties}; use core::ptr::{self, NonNull, Unique}; use core::{cmp, hint}; @@ -40,6 +40,13 @@ struct Cap(usize); impl Cap { const ZERO: Cap = unsafe { Cap(0) }; + + /// `Cap(cap)`, except if `T` is a ZST then `Cap::ZERO`. + /// + /// # Safety: cap must be <= `isize::MAX`. + unsafe fn new(cap: usize) -> Self { + if T::IS_ZST { Cap::ZERO } else { unsafe { Self(cap) } } + } } /// A low-level utility for more ergonomically allocating, reallocating, and deallocating @@ -66,7 +73,19 @@ impl Cap { /// `Box<[T]>`, since `capacity()` won't yield the length. #[allow(missing_debug_implementations)] pub(crate) struct RawVec { - ptr: Unique, + inner: RawVecInner, + _marker: PhantomData, +} + +/// Like a `RawVec`, but only generic over the allocator, not the type. +/// +/// As such, all the methods need the layout passed-in as a parameter. +/// +/// Having this separation reduces the amount of code we need to monomorphize, +/// as most operations don't need the actual type, just its layout. +#[allow(missing_debug_implementations)] +struct RawVecInner { + ptr: Unique, /// Never used for ZSTs; it's `capacity()`'s responsibility to return usize::MAX in that case. /// /// # Safety @@ -90,8 +109,9 @@ impl RawVec { /// `RawVec` with capacity `usize::MAX`. Useful for implementing /// delayed allocation. #[must_use] + #[rustc_const_stable(feature = "raw_vec_internals_const", since = "1.81")] pub const fn new() -> Self { - Self::new_in(Global) + Self { inner: RawVecInner::new::(), _marker: PhantomData } } /// Creates a `RawVec` (on the system heap) with exactly the @@ -113,10 +133,7 @@ impl RawVec { #[must_use] #[inline] pub fn with_capacity(capacity: usize) -> Self { - match Self::try_allocate_in(capacity, AllocInit::Uninitialized, Global) { - Ok(res) => res, - Err(err) => handle_error(err), - } + Self { inner: RawVecInner::with_capacity(capacity, T::LAYOUT), _marker: PhantomData } } /// Like `with_capacity`, but guarantees the buffer is zeroed. @@ -124,29 +141,56 @@ impl RawVec { #[must_use] #[inline] pub fn with_capacity_zeroed(capacity: usize) -> Self { - Self::with_capacity_zeroed_in(capacity, Global) + Self { + inner: RawVecInner::with_capacity_zeroed_in(capacity, Global, T::LAYOUT), + _marker: PhantomData, + } } } -impl RawVec { - // Tiny Vecs are dumb. Skip to: - // - 8 if the element size is 1, because any heap allocators is likely - // to round up a request of less than 8 bytes to at least 8 bytes. - // - 4 if elements are moderate-sized (<= 1 KiB). - // - 1 otherwise, to avoid wasting too much space for very short Vecs. - pub(crate) const MIN_NON_ZERO_CAP: usize = if mem::size_of::() == 1 { +impl RawVecInner { + #[must_use] + #[rustc_const_stable(feature = "raw_vec_internals_const", since = "1.81")] + const fn new() -> Self { + Self::new_in(Global, core::mem::align_of::()) + } + + #[cfg(not(any(no_global_oom_handling, test)))] + #[must_use] + #[inline] + fn with_capacity(capacity: usize, elem_layout: Layout) -> Self { + match Self::try_allocate_in(capacity, AllocInit::Uninitialized, Global, elem_layout) { + Ok(res) => res, + Err(err) => handle_error(err), + } + } +} + +// Tiny Vecs are dumb. Skip to: +// - 8 if the element size is 1, because any heap allocators is likely +// to round up a request of less than 8 bytes to at least 8 bytes. +// - 4 if elements are moderate-sized (<= 1 KiB). +// - 1 otherwise, to avoid wasting too much space for very short Vecs. +const fn min_non_zero_cap(size: usize) -> usize { + if size == 1 { 8 - } else if mem::size_of::() <= 1024 { + } else if size <= 1024 { 4 } else { 1 - }; + } +} + +impl RawVec { + #[cfg(not(no_global_oom_handling))] + pub(crate) const MIN_NON_ZERO_CAP: usize = min_non_zero_cap(size_of::()); /// Like `new`, but parameterized over the choice of allocator for /// the returned `RawVec`. + #[inline] + #[rustc_const_stable(feature = "raw_vec_internals_const", since = "1.81")] pub const fn new_in(alloc: A) -> Self { - // `cap: 0` means "unallocated". zero-sized types are ignored. - Self { ptr: Unique::dangling(), cap: Cap::ZERO, alloc } + Self { inner: RawVecInner::new_in(alloc, align_of::()), _marker: PhantomData } } /// Like `with_capacity`, but parameterized over the choice of @@ -154,9 +198,9 @@ impl RawVec { #[cfg(not(no_global_oom_handling))] #[inline] pub fn with_capacity_in(capacity: usize, alloc: A) -> Self { - match Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc) { - Ok(res) => res, - Err(err) => handle_error(err), + Self { + inner: RawVecInner::with_capacity_in(capacity, alloc, T::LAYOUT), + _marker: PhantomData, } } @@ -164,7 +208,10 @@ impl RawVec { /// allocator for the returned `RawVec`. #[inline] pub fn try_with_capacity_in(capacity: usize, alloc: A) -> Result { - Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc) + match RawVecInner::try_with_capacity_in(capacity, alloc, T::LAYOUT) { + Ok(inner) => Ok(Self { inner, _marker: PhantomData }), + Err(e) => Err(e), + } } /// Like `with_capacity_zeroed`, but parameterized over the choice @@ -172,9 +219,9 @@ impl RawVec { #[cfg(not(no_global_oom_handling))] #[inline] pub fn with_capacity_zeroed_in(capacity: usize, alloc: A) -> Self { - match Self::try_allocate_in(capacity, AllocInit::Zeroed, alloc) { - Ok(res) => res, - Err(err) => handle_error(err), + Self { + inner: RawVecInner::with_capacity_zeroed_in(capacity, alloc, T::LAYOUT), + _marker: PhantomData, } } @@ -200,45 +247,7 @@ impl RawVec { let me = ManuallyDrop::new(self); unsafe { let slice = ptr::slice_from_raw_parts_mut(me.ptr() as *mut MaybeUninit, len); - Box::from_raw_in(slice, ptr::read(&me.alloc)) - } - } - - fn try_allocate_in( - capacity: usize, - init: AllocInit, - alloc: A, - ) -> Result { - // Don't allocate here because `Drop` will not deallocate when `capacity` is 0. - - if T::IS_ZST || capacity == 0 { - Ok(Self::new_in(alloc)) - } else { - // We avoid `unwrap_or_else` here because it bloats the amount of - // LLVM IR generated. - let layout = match Layout::array::(capacity) { - Ok(layout) => layout, - Err(_) => return Err(CapacityOverflow.into()), - }; - - if let Err(err) = alloc_guard(layout.size()) { - return Err(err); - } - - let result = match init { - AllocInit::Uninitialized => alloc.allocate(layout), - #[cfg(not(no_global_oom_handling))] - AllocInit::Zeroed => alloc.allocate_zeroed(layout), - }; - let ptr = match result { - Ok(ptr) => ptr, - Err(_) => return Err(AllocError { layout, non_exhaustive: () }.into()), - }; - - // Allocators currently return a `NonNull<[u8]>` whose length - // matches the size requested. If that ever changes, the capacity - // here should change to `ptr.len() / mem::size_of::()`. - Ok(Self { ptr: Unique::from(ptr.cast()), cap: unsafe { Cap(capacity) }, alloc }) + Box::from_raw_in(slice, ptr::read(&me.inner.alloc)) } } @@ -254,8 +263,15 @@ impl RawVec { /// guaranteed. #[inline] pub unsafe fn from_raw_parts_in(ptr: *mut T, capacity: usize, alloc: A) -> Self { - let cap = if T::IS_ZST { Cap::ZERO } else { unsafe { Cap(capacity) } }; - Self { ptr: unsafe { Unique::new_unchecked(ptr) }, cap, alloc } + // SAFETY: Precondition passed to the caller + unsafe { + let ptr = ptr.cast(); + let capacity = Cap::new::(capacity); + Self { + inner: RawVecInner::from_raw_parts_in(ptr, capacity, alloc), + _marker: PhantomData, + } + } } /// A convenience method for hoisting the non-null precondition out of [`RawVec::from_raw_parts_in`]. @@ -264,9 +280,13 @@ impl RawVec { /// /// See [`RawVec::from_raw_parts_in`]. #[inline] - pub(crate) unsafe fn from_nonnull_in(ptr: NonNull, capacity: usize, alloc: A) -> Self { - let cap = if T::IS_ZST { Cap::ZERO } else { unsafe { Cap(capacity) } }; - Self { ptr: Unique::from(ptr), cap, alloc } + pub unsafe fn from_nonnull_in(ptr: NonNull, capacity: usize, alloc: A) -> Self { + // SAFETY: Precondition passed to the caller + unsafe { + let ptr = ptr.cast(); + let capacity = Cap::new::(capacity); + Self { inner: RawVecInner::from_nonnull_in(ptr, capacity, alloc), _marker: PhantomData } + } } /// Gets a raw pointer to the start of the allocation. Note that this is @@ -274,43 +294,26 @@ impl RawVec { /// be careful. #[inline] pub fn ptr(&self) -> *mut T { - self.ptr.as_ptr() + self.inner.ptr() } #[inline] pub fn non_null(&self) -> NonNull { - NonNull::from(self.ptr) + self.inner.non_null() } /// Gets the capacity of the allocation. /// /// This will always be `usize::MAX` if `T` is zero-sized. - #[inline(always)] + #[inline] pub fn capacity(&self) -> usize { - if T::IS_ZST { usize::MAX } else { self.cap.0 } + self.inner.capacity(size_of::()) } /// Returns a shared reference to the allocator backing this `RawVec`. + #[inline] pub fn allocator(&self) -> &A { - &self.alloc - } - - fn current_memory(&self) -> Option<(NonNull, Layout)> { - if T::IS_ZST || self.cap.0 == 0 { - None - } else { - // We could use Layout::array here which ensures the absence of isize and usize overflows - // and could hypothetically handle differences between stride and size, but this memory - // has already been allocated so we know it can't overflow and currently Rust does not - // support such types. So we can do better by skipping some checks and avoid an unwrap. - const { assert!(mem::size_of::() % mem::align_of::() == 0) }; - unsafe { - let align = mem::align_of::(); - let size = mem::size_of::().unchecked_mul(self.cap.0); - let layout = Layout::from_size_align_unchecked(size, align); - Some((self.ptr.cast().into(), layout)) - } - } + self.inner.allocator() } /// Ensures that the buffer contains at least enough space to hold `len + @@ -335,24 +338,7 @@ impl RawVec { #[cfg(not(no_global_oom_handling))] #[inline] pub fn reserve(&mut self, len: usize, additional: usize) { - // Callers expect this function to be very cheap when there is already sufficient capacity. - // Therefore, we move all the resizing and error-handling logic from grow_amortized and - // handle_reserve behind a call, while making sure that this function is likely to be - // inlined as just a comparison and a call if the comparison fails. - #[cold] - fn do_reserve_and_handle( - slf: &mut RawVec, - len: usize, - additional: usize, - ) { - if let Err(err) = slf.grow_amortized(len, additional) { - handle_error(err); - } - } - - if self.needs_to_grow(len, additional) { - do_reserve_and_handle(self, len, additional); - } + self.inner.reserve(len, additional, T::LAYOUT) } /// A specialized version of `self.reserve(len, 1)` which requires the @@ -360,21 +346,12 @@ impl RawVec { #[cfg(not(no_global_oom_handling))] #[inline(never)] pub fn grow_one(&mut self) { - if let Err(err) = self.grow_amortized(self.cap.0, 1) { - handle_error(err); - } + self.inner.grow_one(T::LAYOUT) } /// The same as `reserve`, but returns on errors instead of panicking or aborting. pub fn try_reserve(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> { - if self.needs_to_grow(len, additional) { - self.grow_amortized(len, additional)?; - } - unsafe { - // Inform the optimizer that the reservation has succeeded or wasn't needed - hint::assert_unchecked(!self.needs_to_grow(len, additional)); - } - Ok(()) + self.inner.try_reserve(len, additional, T::LAYOUT) } /// Ensures that the buffer contains at least enough space to hold `len + @@ -396,9 +373,7 @@ impl RawVec { /// Aborts on OOM. #[cfg(not(no_global_oom_handling))] pub fn reserve_exact(&mut self, len: usize, additional: usize) { - if let Err(err) = self.try_reserve_exact(len, additional) { - handle_error(err); - } + self.inner.reserve_exact(len, additional, T::LAYOUT) } /// The same as `reserve_exact`, but returns on errors instead of panicking or aborting. @@ -407,14 +382,7 @@ impl RawVec { len: usize, additional: usize, ) -> Result<(), TryReserveError> { - if self.needs_to_grow(len, additional) { - self.grow_exact(len, additional)?; - } - unsafe { - // Inform the optimizer that the reservation has succeeded or wasn't needed - hint::assert_unchecked(!self.needs_to_grow(len, additional)); - } - Ok(()) + self.inner.try_reserve_exact(len, additional, T::LAYOUT) } /// Shrinks the buffer down to the specified capacity. If the given amount @@ -430,22 +398,230 @@ impl RawVec { #[cfg(not(no_global_oom_handling))] #[inline] pub fn shrink_to_fit(&mut self, cap: usize) { - if let Err(err) = self.shrink(cap) { + self.inner.shrink_to_fit(cap, T::LAYOUT) + } +} + +unsafe impl<#[may_dangle] T, A: Allocator> Drop for RawVec { + /// Frees the memory owned by the `RawVec` *without* trying to drop its contents. + fn drop(&mut self) { + // SAFETY: We are in a Drop impl, self.inner will not be used again. + unsafe { self.inner.deallocate(T::LAYOUT) } + } +} + +impl RawVecInner { + #[inline] + #[rustc_const_stable(feature = "raw_vec_internals_const", since = "1.81")] + const fn new_in(alloc: A, align: usize) -> Self { + let ptr = unsafe { core::mem::transmute(align) }; + // `cap: 0` means "unallocated". zero-sized types are ignored. + Self { ptr, cap: Cap::ZERO, alloc } + } + + #[cfg(not(no_global_oom_handling))] + #[inline] + fn with_capacity_in(capacity: usize, alloc: A, elem_layout: Layout) -> Self { + match Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc, elem_layout) { + Ok(this) => { + unsafe { + // Make it more obvious that a subsquent Vec::reserve(capacity) will not allocate. + hint::assert_unchecked(!this.needs_to_grow(0, capacity, elem_layout)); + } + this + } + Err(err) => handle_error(err), + } + } + + #[inline] + fn try_with_capacity_in( + capacity: usize, + alloc: A, + elem_layout: Layout, + ) -> Result { + Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc, elem_layout) + } + + #[cfg(not(no_global_oom_handling))] + #[inline] + fn with_capacity_zeroed_in(capacity: usize, alloc: A, elem_layout: Layout) -> Self { + match Self::try_allocate_in(capacity, AllocInit::Zeroed, alloc, elem_layout) { + Ok(res) => res, + Err(err) => handle_error(err), + } + } + + fn try_allocate_in( + capacity: usize, + init: AllocInit, + alloc: A, + elem_layout: Layout, + ) -> Result { + // We avoid `unwrap_or_else` here because it bloats the amount of + // LLVM IR generated. + let layout = match layout_array(capacity, elem_layout) { + Ok(layout) => layout, + Err(_) => return Err(CapacityOverflow.into()), + }; + + // Don't allocate here because `Drop` will not deallocate when `capacity` is 0. + if layout.size() == 0 { + return Ok(Self::new_in(alloc, elem_layout.align())); + } + + if let Err(err) = alloc_guard(layout.size()) { + return Err(err); + } + + let result = match init { + AllocInit::Uninitialized => alloc.allocate(layout), + #[cfg(not(no_global_oom_handling))] + AllocInit::Zeroed => alloc.allocate_zeroed(layout), + }; + let ptr = match result { + Ok(ptr) => ptr, + Err(_) => return Err(AllocError { layout, non_exhaustive: () }.into()), + }; + + // Allocators currently return a `NonNull<[u8]>` whose length + // matches the size requested. If that ever changes, the capacity + // here should change to `ptr.len() / mem::size_of::()`. + Ok(Self { ptr: Unique::from(ptr.cast()), cap: unsafe { Cap(capacity) }, alloc }) + } + + #[inline] + unsafe fn from_raw_parts_in(ptr: *mut u8, cap: Cap, alloc: A) -> Self { + Self { ptr: unsafe { Unique::new_unchecked(ptr) }, cap, alloc } + } + + #[inline] + unsafe fn from_nonnull_in(ptr: NonNull, cap: Cap, alloc: A) -> Self { + Self { ptr: Unique::from(ptr), cap, alloc } + } + + #[inline] + fn ptr(&self) -> *mut T { + self.non_null::().as_ptr() + } + + #[inline] + fn non_null(&self) -> NonNull { + self.ptr.cast().into() + } + + #[inline] + fn capacity(&self, elem_size: usize) -> usize { + if elem_size == 0 { usize::MAX } else { self.cap.0 } + } + + #[inline] + fn allocator(&self) -> &A { + &self.alloc + } + + #[inline] + fn current_memory(&self, elem_layout: Layout) -> Option<(NonNull, Layout)> { + if elem_layout.size() == 0 || self.cap.0 == 0 { + None + } else { + // We could use Layout::array here which ensures the absence of isize and usize overflows + // and could hypothetically handle differences between stride and size, but this memory + // has already been allocated so we know it can't overflow and currently Rust does not + // support such types. So we can do better by skipping some checks and avoid an unwrap. + unsafe { + let alloc_size = elem_layout.size().unchecked_mul(self.cap.0); + let layout = Layout::from_size_align_unchecked(alloc_size, elem_layout.align()); + Some((self.ptr.into(), layout)) + } + } + } + + #[cfg(not(no_global_oom_handling))] + #[inline] + fn reserve(&mut self, len: usize, additional: usize, elem_layout: Layout) { + // Callers expect this function to be very cheap when there is already sufficient capacity. + // Therefore, we move all the resizing and error-handling logic from grow_amortized and + // handle_reserve behind a call, while making sure that this function is likely to be + // inlined as just a comparison and a call if the comparison fails. + #[cold] + fn do_reserve_and_handle( + slf: &mut RawVecInner, + len: usize, + additional: usize, + elem_layout: Layout, + ) { + if let Err(err) = slf.grow_amortized(len, additional, elem_layout) { + handle_error(err); + } + } + + if self.needs_to_grow(len, additional, elem_layout) { + do_reserve_and_handle(self, len, additional, elem_layout); + } + } + + #[cfg(not(no_global_oom_handling))] + #[inline] + fn grow_one(&mut self, elem_layout: Layout) { + if let Err(err) = self.grow_amortized(self.cap.0, 1, elem_layout) { handle_error(err); } } -} -impl RawVec { - /// Returns if the buffer needs to grow to fulfill the needed extra capacity. - /// Mainly used to make inlining reserve-calls possible without inlining `grow`. - fn needs_to_grow(&self, len: usize, additional: usize) -> bool { - additional > self.capacity().wrapping_sub(len) + fn try_reserve( + &mut self, + len: usize, + additional: usize, + elem_layout: Layout, + ) -> Result<(), TryReserveError> { + if self.needs_to_grow(len, additional, elem_layout) { + self.grow_amortized(len, additional, elem_layout)?; + } + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + hint::assert_unchecked(!self.needs_to_grow(len, additional, elem_layout)); + } + Ok(()) } - /// # Safety: - /// - /// `cap` must not exceed `isize::MAX`. + #[cfg(not(no_global_oom_handling))] + fn reserve_exact(&mut self, len: usize, additional: usize, elem_layout: Layout) { + if let Err(err) = self.try_reserve_exact(len, additional, elem_layout) { + handle_error(err); + } + } + + fn try_reserve_exact( + &mut self, + len: usize, + additional: usize, + elem_layout: Layout, + ) -> Result<(), TryReserveError> { + if self.needs_to_grow(len, additional, elem_layout) { + self.grow_exact(len, additional, elem_layout)?; + } + unsafe { + // Inform the optimizer that the reservation has succeeded or wasn't needed + hint::assert_unchecked(!self.needs_to_grow(len, additional, elem_layout)); + } + Ok(()) + } + + #[cfg(not(no_global_oom_handling))] + #[inline] + fn shrink_to_fit(&mut self, cap: usize, elem_layout: Layout) { + if let Err(err) = self.shrink(cap, elem_layout) { + handle_error(err); + } + } + + #[inline] + fn needs_to_grow(&self, len: usize, additional: usize, elem_layout: Layout) -> bool { + additional > self.capacity(elem_layout.size()).wrapping_sub(len) + } + + #[inline] unsafe fn set_ptr_and_cap(&mut self, ptr: NonNull<[u8]>, cap: usize) { // Allocators currently return a `NonNull<[u8]>` whose length matches // the size requested. If that ever changes, the capacity here should @@ -454,18 +630,16 @@ impl RawVec { self.cap = unsafe { Cap(cap) }; } - // This method is usually instantiated many times. So we want it to be as - // small as possible, to improve compile times. But we also want as much of - // its contents to be statically computable as possible, to make the - // generated code run faster. Therefore, this method is carefully written - // so that all of the code that depends on `T` is within it, while as much - // of the code that doesn't depend on `T` as possible is in functions that - // are non-generic over `T`. - fn grow_amortized(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> { + fn grow_amortized( + &mut self, + len: usize, + additional: usize, + elem_layout: Layout, + ) -> Result<(), TryReserveError> { // This is ensured by the calling contexts. debug_assert!(additional > 0); - if T::IS_ZST { + if elem_layout.size() == 0 { // Since we return a capacity of `usize::MAX` when `elem_size` is // 0, getting to here necessarily means the `RawVec` is overfull. return Err(CapacityOverflow.into()); @@ -477,33 +651,34 @@ impl RawVec { // This guarantees exponential growth. The doubling cannot overflow // because `cap <= isize::MAX` and the type of `cap` is `usize`. let cap = cmp::max(self.cap.0 * 2, required_cap); - let cap = cmp::max(Self::MIN_NON_ZERO_CAP, cap); + let cap = cmp::max(min_non_zero_cap(elem_layout.size()), cap); - let new_layout = Layout::array::(cap); + let new_layout = layout_array(cap, elem_layout)?; - // `finish_grow` is non-generic over `T`. - let ptr = finish_grow(new_layout, self.current_memory(), &mut self.alloc)?; + let ptr = finish_grow(new_layout, self.current_memory(elem_layout), &mut self.alloc)?; // SAFETY: finish_grow would have resulted in a capacity overflow if we tried to allocate more than `isize::MAX` items + unsafe { self.set_ptr_and_cap(ptr, cap) }; Ok(()) } - // The constraints on this method are much the same as those on - // `grow_amortized`, but this method is usually instantiated less often so - // it's less critical. - fn grow_exact(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> { - if T::IS_ZST { + fn grow_exact( + &mut self, + len: usize, + additional: usize, + elem_layout: Layout, + ) -> Result<(), TryReserveError> { + if elem_layout.size() == 0 { // Since we return a capacity of `usize::MAX` when the type size is // 0, getting to here necessarily means the `RawVec` is overfull. return Err(CapacityOverflow.into()); } let cap = len.checked_add(additional).ok_or(CapacityOverflow)?; - let new_layout = Layout::array::(cap); + let new_layout = layout_array(cap, elem_layout)?; - // `finish_grow` is non-generic over `T`. - let ptr = finish_grow(new_layout, self.current_memory(), &mut self.alloc)?; - // SAFETY: `finish_grow` would have resulted in a capacity overflow if we tried to allocate more than `isize::MAX` items + let ptr = finish_grow(new_layout, self.current_memory(elem_layout), &mut self.alloc)?; + // SAFETY: finish_grow would have resulted in a capacity overflow if we tried to allocate more than `isize::MAX` items unsafe { self.set_ptr_and_cap(ptr, cap); } @@ -512,10 +687,10 @@ impl RawVec { #[cfg(not(no_global_oom_handling))] #[inline] - fn shrink(&mut self, cap: usize) -> Result<(), TryReserveError> { - assert!(cap <= self.capacity(), "Tried to shrink to a larger capacity"); + fn shrink(&mut self, cap: usize, elem_layout: Layout) -> Result<(), TryReserveError> { + assert!(cap <= self.capacity(elem_layout.size()), "Tried to shrink to a larger capacity"); // SAFETY: Just checked this isn't trying to grow - unsafe { self.shrink_unchecked(cap) } + unsafe { self.shrink_unchecked(cap, elem_layout) } } /// `shrink`, but without the capacity check. @@ -529,23 +704,27 @@ impl RawVec { /// # Safety /// `cap <= self.capacity()` #[cfg(not(no_global_oom_handling))] - unsafe fn shrink_unchecked(&mut self, cap: usize) -> Result<(), TryReserveError> { - let (ptr, layout) = if let Some(mem) = self.current_memory() { mem } else { return Ok(()) }; - // See current_memory() why this assert is here - const { assert!(mem::size_of::() % mem::align_of::() == 0) }; + unsafe fn shrink_unchecked( + &mut self, + cap: usize, + elem_layout: Layout, + ) -> Result<(), TryReserveError> { + let (ptr, layout) = + if let Some(mem) = self.current_memory(elem_layout) { mem } else { return Ok(()) }; // If shrinking to 0, deallocate the buffer. We don't reach this point // for the T::IS_ZST case since current_memory() will have returned // None. if cap == 0 { unsafe { self.alloc.deallocate(ptr, layout) }; - self.ptr = Unique::dangling(); + self.ptr = + unsafe { Unique::new_unchecked(ptr::without_provenance_mut(elem_layout.align())) }; self.cap = Cap::ZERO; } else { let ptr = unsafe { - // `Layout::array` cannot overflow here because it would have + // Layout cannot overflow here because it would have // overflowed earlier when capacity was larger. - let new_size = mem::size_of::().unchecked_mul(cap); + let new_size = elem_layout.size().unchecked_mul(cap); let new_layout = Layout::from_size_align_unchecked(new_size, layout.align()); self.alloc .shrink(ptr, layout, new_layout) @@ -558,24 +737,32 @@ impl RawVec { } Ok(()) } + + /// # Safety + /// + /// This function deallocates the owned allocation, but does not update `ptr` or `cap` to + /// prevent double-free or use-after-free. Essentially, do not do anything with the caller + /// after this function returns. + /// Ideally this function would take `self` by move, but it cannot because it exists to be + /// called from a `Drop` impl. + unsafe fn deallocate(&mut self, elem_layout: Layout) { + if let Some((ptr, layout)) = self.current_memory(elem_layout) { + unsafe { + self.alloc.deallocate(ptr, layout); + } + } + } } -// This function is outside `RawVec` to minimize compile times. See the comment -// above `RawVec::grow_amortized` for details. (The `A` parameter isn't -// significant, because the number of different `A` types seen in practice is -// much smaller than the number of `T` types.) #[inline(never)] fn finish_grow( - new_layout: Result, + new_layout: Layout, current_memory: Option<(NonNull, Layout)>, alloc: &mut A, ) -> Result, TryReserveError> where A: Allocator, { - // Check for the error here to minimize the size of `RawVec::grow_*`. - let new_layout = new_layout.map_err(|_| CapacityOverflow)?; - alloc_guard(new_layout.size())?; let memory = if let Some((ptr, old_layout)) = current_memory { @@ -592,15 +779,6 @@ where memory.map_err(|_| AllocError { layout: new_layout, non_exhaustive: () }.into()) } -unsafe impl<#[may_dangle] T, A: Allocator> Drop for RawVec { - /// Frees the memory owned by the `RawVec` *without* trying to drop its contents. - fn drop(&mut self) { - if let Some((ptr, layout)) = self.current_memory() { - unsafe { self.alloc.deallocate(ptr, layout) } - } - } -} - // Central function for reserve error handling. #[cfg(not(no_global_oom_handling))] #[cold] @@ -627,3 +805,8 @@ fn alloc_guard(alloc_size: usize) -> Result<(), TryReserveError> { Ok(()) } } + +#[inline] +fn layout_array(cap: usize, elem_layout: Layout) -> Result { + elem_layout.repeat(cap).map(|(layout, _pad)| layout).map_err(|_| CapacityOverflow.into()) +} diff --git a/alloc/src/raw_vec/tests.rs b/alloc/src/raw_vec/tests.rs index 48c6e5f46f8db..d78ded104fb09 100644 --- a/alloc/src/raw_vec/tests.rs +++ b/alloc/src/raw_vec/tests.rs @@ -43,9 +43,9 @@ fn allocator_param() { let a = BoundedAlloc { fuel: Cell::new(500) }; let mut v: RawVec = RawVec::with_capacity_in(50, a); - assert_eq!(v.alloc.fuel.get(), 450); + assert_eq!(v.inner.alloc.fuel.get(), 450); v.reserve(50, 150); // (causes a realloc, thus using 50 + 150 = 200 units of fuel) - assert_eq!(v.alloc.fuel.get(), 250); + assert_eq!(v.inner.alloc.fuel.get(), 250); } #[test] @@ -86,7 +86,7 @@ struct ZST; fn zst_sanity(v: &RawVec) { assert_eq!(v.capacity(), usize::MAX); assert_eq!(v.ptr(), core::ptr::Unique::::dangling().as_ptr()); - assert_eq!(v.current_memory(), None); + assert_eq!(v.inner.current_memory(T::LAYOUT), None); } #[test] @@ -106,22 +106,11 @@ fn zst() { let v: RawVec = RawVec::with_capacity_in(100, Global); zst_sanity(&v); - let v: RawVec = RawVec::try_allocate_in(0, AllocInit::Uninitialized, Global).unwrap(); - zst_sanity(&v); - - let v: RawVec = RawVec::try_allocate_in(100, AllocInit::Uninitialized, Global).unwrap(); - zst_sanity(&v); - - let mut v: RawVec = - RawVec::try_allocate_in(usize::MAX, AllocInit::Uninitialized, Global).unwrap(); + let mut v: RawVec = RawVec::with_capacity_in(usize::MAX, Global); zst_sanity(&v); // Check all these operations work as expected with zero-sized elements. - assert!(!v.needs_to_grow(100, usize::MAX - 100)); - assert!(v.needs_to_grow(101, usize::MAX - 100)); - zst_sanity(&v); - v.reserve(100, usize::MAX - 100); //v.reserve(101, usize::MAX - 100); // panics, in `zst_reserve_panic` below zst_sanity(&v); @@ -138,12 +127,12 @@ fn zst() { assert_eq!(v.try_reserve_exact(101, usize::MAX - 100), cap_err); zst_sanity(&v); - assert_eq!(v.grow_amortized(100, usize::MAX - 100), cap_err); - assert_eq!(v.grow_amortized(101, usize::MAX - 100), cap_err); + assert_eq!(v.inner.grow_amortized(100, usize::MAX - 100, ZST::LAYOUT), cap_err); + assert_eq!(v.inner.grow_amortized(101, usize::MAX - 100, ZST::LAYOUT), cap_err); zst_sanity(&v); - assert_eq!(v.grow_exact(100, usize::MAX - 100), cap_err); - assert_eq!(v.grow_exact(101, usize::MAX - 100), cap_err); + assert_eq!(v.inner.grow_exact(100, usize::MAX - 100, ZST::LAYOUT), cap_err); + assert_eq!(v.inner.grow_exact(101, usize::MAX - 100, ZST::LAYOUT), cap_err); zst_sanity(&v); } diff --git a/core/src/mem/mod.rs b/core/src/mem/mod.rs index ea2dcdce6e89e..7a9ca4011be84 100644 --- a/core/src/mem/mod.rs +++ b/core/src/mem/mod.rs @@ -5,6 +5,7 @@ #![stable(feature = "rust1", since = "1.0.0")] +use crate::alloc::Layout; use crate::marker::DiscriminantKind; use crate::{clone, cmp, fmt, hash, intrinsics, ptr}; @@ -1238,6 +1239,10 @@ pub trait SizedTypeProperties: Sized { #[doc(hidden)] #[unstable(feature = "sized_type_properties", issue = "none")] const IS_ZST: bool = size_of::() == 0; + + #[doc(hidden)] + #[unstable(feature = "sized_type_properties", issue = "none")] + const LAYOUT: Layout = Layout::new::(); } #[doc(hidden)] #[unstable(feature = "sized_type_properties", issue = "none")]