From f228ae136e8790429e18b8b100db891a25a6aec7 Mon Sep 17 00:00:00 2001 From: Tim Diekmann Date: Sun, 8 Mar 2020 13:39:27 +0100 Subject: [PATCH] Use `AllocRef` for ZSTs in `RawVec` --- src/liballoc/raw_vec.rs | 591 +++++++++++++++------------------- src/liballoc/raw_vec/tests.rs | 1 + 2 files changed, 257 insertions(+), 335 deletions(-) diff --git a/src/liballoc/raw_vec.rs b/src/liballoc/raw_vec.rs index 345834d7daacc..827fc03cffb6c 100644 --- a/src/liballoc/raw_vec.rs +++ b/src/liballoc/raw_vec.rs @@ -4,7 +4,7 @@ use core::cmp; use core::mem; use core::ops::Drop; -use core::ptr::{self, NonNull, Unique}; +use core::ptr::Unique; use core::slice; use crate::alloc::{handle_alloc_error, AllocErr, AllocRef, Global, Layout}; @@ -19,28 +19,18 @@ mod tests; /// involved. This type is excellent for building your own data structures like Vec and VecDeque. /// In particular: /// -/// * Produces `Unique::empty()` on zero-sized types. -/// * Produces `Unique::empty()` on zero-length allocations. /// * Catches all overflows in capacity computations (promotes them to "capacity overflow" panics). /// * Guards against 32-bit systems allocating more than isize::MAX bytes. /// * Guards against overflowing your length. -/// * Aborts on OOM or calls `handle_alloc_error` as applicable. -/// * Avoids freeing `Unique::empty()`. +/// * Calls `handle_alloc_error` for fallible allocations. /// * Contains a `ptr::Unique` and thus endows the user with all related benefits. /// /// This type does not in anyway inspect the memory that it manages. When dropped it *will* /// free its memory, but it *won't* try to drop its contents. It is up to the user of `RawVec` /// to handle the actual things *stored* inside of a `RawVec`. /// -/// Note that a `RawVec` always forces its capacity to be `usize::MAX` for zero-sized types. -/// This enables you to use capacity-growing logic catch the overflows in your length -/// that might occur with zero-sized types. -/// -/// The above means that you need to be careful when round-tripping this type with a -/// `Box<[T]>`, since `capacity()` won't yield the length. However, `with_capacity`, -/// `shrink_to_fit`, and `from_box` will actually set `RawVec`'s private capacity -/// field. This allows zero-sized types to not be special-cased by consumers of -/// this type. +/// Note that a `RawVec` always returns a capacity of `usize::MAX` for zero-sized types. Beside +/// this, zero-sized types are handled like non-zero-sized types. #[allow(missing_debug_implementations)] pub struct RawVec { ptr: Unique, @@ -52,50 +42,42 @@ impl RawVec { /// Like `new`, but parameterized over the choice of allocator for /// the returned `RawVec`. pub const fn new_in(a: A) -> Self { - let cap = if mem::size_of::() == 0 { core::usize::MAX } else { 0 }; - - // `Unique::empty()` doubles as "unallocated" and "zero-sized allocation". - RawVec { ptr: Unique::empty(), cap, a } + // `cap: 0` means "unallocated". zero-sized allocations are handled by `AllocRef` + Self { ptr: Unique::empty(), cap: 0, a } } /// Like `with_capacity`, but parameterized over the choice of /// allocator for the returned `RawVec`. #[inline] pub fn with_capacity_in(capacity: usize, a: A) -> Self { - RawVec::allocate_in(capacity, false, a) + Self::allocate_in(capacity, AllocInit::Unspecified, a) } /// Like `with_capacity_zeroed`, but parameterized over the choice /// of allocator for the returned `RawVec`. #[inline] pub fn with_capacity_zeroed_in(capacity: usize, a: A) -> Self { - RawVec::allocate_in(capacity, true, a) + Self::allocate_in(capacity, AllocInit::Zero, a) } - fn allocate_in(mut capacity: usize, zeroed: bool, mut a: A) -> Self { - unsafe { - let elem_size = mem::size_of::(); + fn allocate_in(capacity: usize, init: AllocInit, mut a: A) -> Self { + let layout = Layout::array::(capacity).unwrap_or_else(|_| capacity_overflow()); + alloc_guard(layout.size()).unwrap_or_else(|_| capacity_overflow()); - let alloc_size = capacity.checked_mul(elem_size).unwrap_or_else(|| capacity_overflow()); - alloc_guard(alloc_size).unwrap_or_else(|_| capacity_overflow()); + let allocation = match init { + AllocInit::Unspecified => a.alloc(layout), + AllocInit::Zero => a.alloc_zeroed(layout), + }; + let (ptr, alloc_size) = allocation.unwrap_or_else(|_| handle_alloc_error(layout)); - // Handles ZSTs and `capacity == 0` alike. - let ptr = if alloc_size == 0 { - NonNull::::dangling() + let ptr = ptr.cast().as_ptr(); + let elem_size = mem::size_of::(); + unsafe { + if elem_size == 0 { + Self::from_raw_parts_in(ptr, capacity, a) } else { - let align = mem::align_of::(); - let layout = Layout::from_size_align(alloc_size, align).unwrap(); - let result = if zeroed { a.alloc_zeroed(layout) } else { a.alloc(layout) }; - match result { - Ok((ptr, size)) => { - capacity = size / elem_size; - ptr.cast() - } - Err(_) => handle_alloc_error(layout), - } - }; - - RawVec { ptr: ptr.into(), cap: capacity, a } + Self::from_raw_parts_in(ptr, alloc_size / elem_size, a) + } } } } @@ -140,13 +122,13 @@ impl RawVec { /// Aborts on OOM. #[inline] pub fn with_capacity(capacity: usize) -> Self { - RawVec::allocate_in(capacity, false, Global) + Self::with_capacity_in(capacity, Global) } /// Like `with_capacity`, but guarantees the buffer is zeroed. #[inline] pub fn with_capacity_zeroed(capacity: usize) -> Self { - RawVec::allocate_in(capacity, true, Global) + Self::with_capacity_zeroed_in(capacity, Global) } } @@ -158,8 +140,9 @@ impl RawVec { /// The `ptr` must be allocated (via the given allocator `a`), and with the given `capacity`. /// The `capacity` cannot exceed `isize::MAX` (only a concern on 32-bit systems). /// If the `ptr` and `capacity` come from a `RawVec` created via `a`, then this is guaranteed. + #[inline] pub unsafe fn from_raw_parts_in(ptr: *mut T, capacity: usize, a: A) -> Self { - RawVec { ptr: Unique::new_unchecked(ptr), cap: capacity, a } + Self { ptr: Unique::new_unchecked(ptr), cap: capacity, a } } } @@ -171,8 +154,9 @@ impl RawVec { /// The `ptr` must be allocated (on the system heap), and with the given `capacity`. /// The `capacity` cannot exceed `isize::MAX` (only a concern on 32-bit systems). /// If the `ptr` and `capacity` come from a `RawVec`, then this is guaranteed. + #[inline] pub unsafe fn from_raw_parts(ptr: *mut T, capacity: usize) -> Self { - RawVec { ptr: Unique::new_unchecked(ptr), cap: capacity, a: Global } + Self::from_raw_parts_in(ptr, capacity, Global) } /// Converts a `Box<[T]>` into a `RawVec`. @@ -198,7 +182,7 @@ impl RawVec { /// This will always be `usize::MAX` if `T` is zero-sized. #[inline(always)] pub fn capacity(&self) -> usize { - if mem::size_of::() == 0 { !0 } else { self.cap } + if mem::size_of::() == 0 { usize::MAX } else { self.cap } } /// Returns a shared reference to the allocator backing this `RawVec`. @@ -236,8 +220,6 @@ impl RawVec { /// /// # Panics /// - /// * Panics if `T` is zero-sized on the assumption that you managed to exhaust - /// all `usize::MAX` slots in your imaginary buffer. /// * Panics on 32-bit platforms if the requested capacity exceeds /// `isize::MAX` bytes. /// @@ -276,50 +258,10 @@ impl RawVec { #[inline(never)] #[cold] pub fn double(&mut self) { - unsafe { - let elem_size = mem::size_of::(); - - // Since we set the capacity to `usize::MAX` when `elem_size` is - // 0, getting to here necessarily means the `RawVec` is overfull. - assert!(elem_size != 0, "capacity overflow"); - - let (ptr, new_cap) = match self.current_layout() { - Some(cur) => { - // Since we guarantee that we never allocate more than - // `isize::MAX` bytes, `elem_size * self.cap <= isize::MAX` as - // a precondition, so this can't overflow. Additionally the - // alignment will never be too large as to "not be - // satisfiable", so `Layout::from_size_align` will always - // return `Some`. - // - // TL;DR, we bypass runtime checks due to dynamic assertions - // in this module, allowing us to use - // `from_size_align_unchecked`. - let new_cap = 2 * self.cap; - let new_size = new_cap * elem_size; - alloc_guard(new_size).unwrap_or_else(|_| capacity_overflow()); - let ptr_res = self.a.realloc(NonNull::from(self.ptr).cast(), cur, new_size); - match ptr_res { - Ok((ptr, new_size)) => (ptr, new_size / elem_size), - Err(_) => handle_alloc_error(Layout::from_size_align_unchecked( - new_size, - cur.align(), - )), - } - } - None => { - // Skip to 4 because tiny `Vec`'s are dumb; but not if that - // would cause overflow. - let new_cap = if elem_size > (!0) / 8 { 1 } else { 4 }; - let layout = Layout::array::(new_cap).unwrap(); - match self.a.alloc(layout) { - Ok((ptr, new_size)) => (ptr, new_size / elem_size), - Err(_) => handle_alloc_error(layout), - } - } - }; - self.ptr = ptr.cast().into(); - self.cap = new_cap; + match self.grow(Double, AllocPlacement::Unspecified, AllocInit::Unspecified) { + Err(CapacityOverflow) => capacity_overflow(), + Err(AllocError { layout, .. }) => handle_alloc_error(layout), + Ok(()) => { /* yay */ } } } @@ -331,106 +273,12 @@ impl RawVec { /// /// # Panics /// - /// * Panics if `T` is zero-sized on the assumption that you managed to exhaust - /// all `usize::MAX` slots in your imaginary buffer. /// * Panics on 32-bit platforms if the requested capacity exceeds /// `isize::MAX` bytes. #[inline(never)] #[cold] pub fn double_in_place(&mut self) -> bool { - unsafe { - let elem_size = mem::size_of::(); - let old_layout = match self.current_layout() { - Some(layout) => layout, - None => return false, // nothing to double - }; - - // Since we set the capacity to `usize::MAX` when `elem_size` is - // 0, getting to here necessarily means the `RawVec` is overfull. - assert!(elem_size != 0, "capacity overflow"); - - // Since we guarantee that we never allocate more than `isize::MAX` - // bytes, `elem_size * self.cap <= isize::MAX` as a precondition, so - // this can't overflow. - // - // Similarly to with `double` above, we can go straight to - // `Layout::from_size_align_unchecked` as we know this won't - // overflow and the alignment is sufficiently small. - let new_cap = 2 * self.cap; - let new_size = new_cap * elem_size; - alloc_guard(new_size).unwrap_or_else(|_| capacity_overflow()); - match self.a.grow_in_place(NonNull::from(self.ptr).cast(), old_layout, new_size) { - Ok(_) => { - // We can't directly divide `size`. - self.cap = new_cap; - true - } - Err(_) => false, - } - } - } - - /// The same as `reserve_exact`, but returns on errors instead of panicking or aborting. - pub fn try_reserve_exact( - &mut self, - used_capacity: usize, - needed_extra_capacity: usize, - ) -> Result<(), TryReserveError> { - self.reserve_internal(used_capacity, needed_extra_capacity, Fallible, Exact) - } - - /// Ensures that the buffer contains at least enough space to hold - /// `used_capacity + needed_extra_capacity` elements. If it doesn't already, - /// will reallocate the minimum possible amount of memory necessary. - /// Generally this will be exactly the amount of memory necessary, - /// but in principle the allocator is free to give back more than - /// we asked for. - /// - /// If `used_capacity` exceeds `self.capacity()`, this may fail to actually allocate - /// the requested space. This is not really unsafe, but the unsafe - /// code *you* write that relies on the behavior of this function may break. - /// - /// # Panics - /// - /// * Panics if the requested capacity exceeds `usize::MAX` bytes. - /// * Panics on 32-bit platforms if the requested capacity exceeds - /// `isize::MAX` bytes. - /// - /// # Aborts - /// - /// Aborts on OOM. - pub fn reserve_exact(&mut self, used_capacity: usize, needed_extra_capacity: usize) { - match self.reserve_internal(used_capacity, needed_extra_capacity, Infallible, Exact) { - Err(CapacityOverflow) => capacity_overflow(), - Err(AllocError { .. }) => unreachable!(), - Ok(()) => { /* yay */ } - } - } - - /// Calculates the buffer's new size given that it'll hold `used_capacity + - /// needed_extra_capacity` elements. This logic is used in amortized reserve methods. - /// Returns `(new_capacity, new_alloc_size)`. - fn amortized_new_size( - &self, - used_capacity: usize, - needed_extra_capacity: usize, - ) -> Result { - // Nothing we can really do about these checks, sadly. - let required_cap = - used_capacity.checked_add(needed_extra_capacity).ok_or(CapacityOverflow)?; - // Cannot overflow, because `cap <= isize::MAX`, and type of `cap` is `usize`. - let double_cap = self.cap * 2; - // `double_cap` guarantees exponential growth. - Ok(cmp::max(double_cap, required_cap)) - } - - /// The same as `reserve`, but returns on errors instead of panicking or aborting. - pub fn try_reserve( - &mut self, - used_capacity: usize, - needed_extra_capacity: usize, - ) -> Result<(), TryReserveError> { - self.reserve_internal(used_capacity, needed_extra_capacity, Fallible, Amortized) + self.grow(Double, AllocPlacement::InPlace, AllocInit::Unspecified).is_ok() } /// Ensures that the buffer contains at least enough space to hold @@ -486,12 +334,30 @@ impl RawVec { /// # } /// ``` pub fn reserve(&mut self, used_capacity: usize, needed_extra_capacity: usize) { - match self.reserve_internal(used_capacity, needed_extra_capacity, Infallible, Amortized) { + match self.try_reserve(used_capacity, needed_extra_capacity) { Err(CapacityOverflow) => capacity_overflow(), - Err(AllocError { .. }) => unreachable!(), + Err(AllocError { layout, .. }) => handle_alloc_error(layout), Ok(()) => { /* yay */ } } } + + /// The same as `reserve`, but returns on errors instead of panicking or aborting. + pub fn try_reserve( + &mut self, + used_capacity: usize, + needed_extra_capacity: usize, + ) -> Result<(), TryReserveError> { + if self.needs_to_grow(used_capacity, needed_extra_capacity) { + self.grow( + Amortized { used_capacity, needed_extra_capacity }, + AllocPlacement::Unspecified, + AllocInit::Unspecified, + ) + } else { + Ok(()) + } + } + /// Attempts to ensure that the buffer contains at least enough space to hold /// `used_capacity + needed_extra_capacity` elements. If it doesn't already have /// enough capacity, will reallocate in place enough space plus comfortable slack @@ -510,45 +376,62 @@ impl RawVec { /// * Panics on 32-bit platforms if the requested capacity exceeds /// `isize::MAX` bytes. pub fn reserve_in_place(&mut self, used_capacity: usize, needed_extra_capacity: usize) -> bool { - unsafe { - // NOTE: we don't early branch on ZSTs here because we want this - // to actually catch "asking for more than usize::MAX" in that case. - // If we make it past the first branch then we are guaranteed to - // panic. - - // Don't actually need any more capacity. If the current `cap` is 0, we can't - // reallocate in place. - // Wrapping in case they give a bad `used_capacity` - let old_layout = match self.current_layout() { - Some(layout) => layout, - None => return false, - }; - if self.capacity().wrapping_sub(used_capacity) >= needed_extra_capacity { - return false; - } + // This is more readable as putting this in one line: + // `!self.needs_to_grow(...) || self.grow(...).is_ok()` + if self.needs_to_grow(used_capacity, needed_extra_capacity) { + self.grow( + Amortized { used_capacity, needed_extra_capacity }, + AllocPlacement::InPlace, + AllocInit::Unspecified, + ) + .is_ok() + } else { + true + } + } - let new_cap = self - .amortized_new_size(used_capacity, needed_extra_capacity) - .unwrap_or_else(|_| capacity_overflow()); - - // Here, `cap < used_capacity + needed_extra_capacity <= new_cap` - // (regardless of whether `self.cap - used_capacity` wrapped). - // Therefore, we can safely call `grow_in_place`. - - let new_layout = Layout::new::().repeat(new_cap).unwrap().0; - // FIXME: may crash and burn on over-reserve - alloc_guard(new_layout.size()).unwrap_or_else(|_| capacity_overflow()); - match self.a.grow_in_place( - NonNull::from(self.ptr).cast(), - old_layout, - new_layout.size(), - ) { - Ok(_) => { - self.cap = new_cap; - true - } - Err(_) => false, - } + /// Ensures that the buffer contains at least enough space to hold + /// `used_capacity + needed_extra_capacity` elements. If it doesn't already, + /// will reallocate the minimum possible amount of memory necessary. + /// Generally this will be exactly the amount of memory necessary, + /// but in principle the allocator is free to give back more than + /// we asked for. + /// + /// If `used_capacity` exceeds `self.capacity()`, this may fail to actually allocate + /// the requested space. This is not really unsafe, but the unsafe + /// code *you* write that relies on the behavior of this function may break. + /// + /// # Panics + /// + /// * Panics if the requested capacity exceeds `usize::MAX` bytes. + /// * Panics on 32-bit platforms if the requested capacity exceeds + /// `isize::MAX` bytes. + /// + /// # Aborts + /// + /// Aborts on OOM. + pub fn reserve_exact(&mut self, used_capacity: usize, needed_extra_capacity: usize) { + match self.try_reserve_exact(used_capacity, needed_extra_capacity) { + Err(CapacityOverflow) => capacity_overflow(), + Err(AllocError { layout, .. }) => handle_alloc_error(layout), + Ok(()) => { /* yay */ } + } + } + + /// The same as `reserve_exact`, but returns on errors instead of panicking or aborting. + pub fn try_reserve_exact( + &mut self, + used_capacity: usize, + needed_extra_capacity: usize, + ) -> Result<(), TryReserveError> { + if self.needs_to_grow(used_capacity, needed_extra_capacity) { + self.grow( + Exact { used_capacity, needed_extra_capacity }, + AllocPlacement::Unspecified, + AllocInit::Unspecified, + ) + } else { + Ok(()) } } @@ -563,126 +446,169 @@ impl RawVec { /// /// Aborts on OOM. pub fn shrink_to_fit(&mut self, amount: usize) { - let elem_size = mem::size_of::(); - - // Set the `cap` because they might be about to promote to a `Box<[T]>` - if elem_size == 0 { - self.cap = amount; - return; - } - - // This check is my waterloo; it's the only thing `Vec` wouldn't have to do. - assert!(self.cap >= amount, "Tried to shrink to a larger capacity"); - - if amount == 0 { - // We want to create a new zero-length vector within the - // same allocator. We use `ptr::write` to avoid an - // erroneous attempt to drop the contents, and we use - // `ptr::read` to sidestep condition against destructuring - // types that implement Drop. - - unsafe { - let a = ptr::read(&self.a as *const A); - self.dealloc_buffer(); - ptr::write(self, RawVec::new_in(a)); - } - } else if self.cap != amount { - unsafe { - // We know here that our `amount` is greater than zero. This - // implies, via the assert above, that capacity is also greater - // than zero, which means that we've got a current layout that - // "fits" - // - // We also know that `self.cap` is greater than `amount`, and - // consequently we don't need runtime checks for creating either - // layout. - let old_size = elem_size * self.cap; - let new_size = elem_size * amount; - let align = mem::align_of::(); - let old_layout = Layout::from_size_align_unchecked(old_size, align); - match self.a.realloc(NonNull::from(self.ptr).cast(), old_layout, new_size) { - Ok((ptr, _)) => self.ptr = ptr.cast().into(), - Err(_) => { - handle_alloc_error(Layout::from_size_align_unchecked(new_size, align)) - } - } - } - self.cap = amount; + match self.shrink(amount, AllocPlacement::Unspecified) { + Err(CapacityOverflow) => capacity_overflow(), + Err(AllocError { layout, .. }) => handle_alloc_error(layout), + Ok(()) => { /* yay */ } } } } -enum Fallibility { - Fallible, - Infallible, +#[derive(Copy, Clone)] +enum Strategy { + Double, + Amortized { used_capacity: usize, needed_extra_capacity: usize }, + Exact { used_capacity: usize, needed_extra_capacity: usize }, } +use Strategy::*; -use Fallibility::*; - -enum ReserveStrategy { - Exact, - Amortized, +enum AllocInit { + Unspecified, + Zero, } -use ReserveStrategy::*; +enum AllocPlacement { + Unspecified, + InPlace, +} impl RawVec { - fn reserve_internal( + /// 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, used_capacity: usize, needed_extra_capacity: usize) -> bool { + needed_extra_capacity > self.capacity().wrapping_sub(used_capacity) + } + + /// Single method to handle all possibilities of growing the buffer. + fn grow( &mut self, - used_capacity: usize, - needed_extra_capacity: usize, - fallibility: Fallibility, - strategy: ReserveStrategy, + strategy: Strategy, + placement: AllocPlacement, + init: AllocInit, ) -> Result<(), TryReserveError> { let elem_size = mem::size_of::(); - - unsafe { - // NOTE: we don't early branch on ZSTs here because we want this - // to actually catch "asking for more than usize::MAX" in that case. - // If we make it past the first branch then we are guaranteed to - // panic. - - // Don't actually need any more capacity. - // Wrapping in case they gave a bad `used_capacity`. - if self.capacity().wrapping_sub(used_capacity) >= needed_extra_capacity { - return Ok(()); + let (new_layout, new_cap) = match strategy { + Double => unsafe { + if elem_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); + } + // Since we guarantee that we never allocate more than `isize::MAX` bytes, + // `elem_size * self.cap <= isize::MAX` as a precondition, so this can't overflow. + // Additionally the alignment will never be too large as to "not be satisfiable", + // so `Layout::from_size_align` will always return `Some`. + // + // TL;DR, we bypass runtime checks due to dynamic assertions in this module, + // allowing us to use `from_size_align_unchecked`. + let cap = if self.cap == 0 { + if elem_size > usize::MAX / 8 { 1 } else { 4 } + } else { + self.cap * 2 + }; + let layout = + Layout::from_size_align_unchecked(cap * elem_size, mem::align_of::()); + (layout, cap) + }, + Amortized { used_capacity, needed_extra_capacity } => { + // Nothing we can really do about these checks, sadly. + let required_cap = + used_capacity.checked_add(needed_extra_capacity).ok_or(CapacityOverflow)?; + // Cannot overflow, because `cap <= isize::MAX`, and type of `cap` is `usize`. + let double_cap = self.cap * 2; + // `double_cap` guarantees exponential growth. + let cap = cmp::max(double_cap, required_cap); + let layout = Layout::array::(cap).map_err(|_| CapacityOverflow)?; + (layout, cap) + } + Exact { used_capacity, needed_extra_capacity } => { + let cap = + used_capacity.checked_add(needed_extra_capacity).ok_or(CapacityOverflow)?; + let layout = Layout::array::(cap).map_err(|_| CapacityOverflow)?; + (layout, cap) } + }; - // Nothing we can really do about these checks, sadly. - let new_cap = match strategy { - Exact => { - used_capacity.checked_add(needed_extra_capacity).ok_or(CapacityOverflow)? + let allocation = if let Some(old_layout) = self.current_layout() { + debug_assert!(old_layout.align() == new_layout.align()); + debug_assert!(old_layout.size() <= new_layout.size()); + let ptr = self.ptr.cast().into(); + unsafe { + match (placement, init) { + (AllocPlacement::Unspecified, AllocInit::Unspecified) => { + self.a.realloc(ptr, old_layout, new_layout.size()) + } + (AllocPlacement::Unspecified, AllocInit::Zero) => { + self.a.realloc_zeroed(ptr, old_layout, new_layout.size()) + } + (AllocPlacement::InPlace, AllocInit::Unspecified) => self + .a + .grow_in_place(ptr, old_layout, new_layout.size()) + .map(|size| (ptr, size)) + .map_err(|_| AllocErr), + (AllocPlacement::InPlace, AllocInit::Zero) => self + .a + .grow_in_place_zeroed(ptr, old_layout, new_layout.size()) + .map(|size| (ptr, size)) + .map_err(|_| AllocErr), } - Amortized => self.amortized_new_size(used_capacity, needed_extra_capacity)?, - }; - let new_layout = Layout::array::(new_cap).map_err(|_| CapacityOverflow)?; + } + } else { + match (placement, init) { + (AllocPlacement::Unspecified, AllocInit::Unspecified) => self.a.alloc(new_layout), + (AllocPlacement::Unspecified, AllocInit::Zero) => self.a.alloc_zeroed(new_layout), + (AllocPlacement::InPlace, _) => Err(AllocErr), + } + }; + allocation + .map(|(ptr, alloc_size)| { + self.ptr = ptr.cast().into(); + if elem_size == 0 { + self.cap = new_cap; + } else { + self.cap = alloc_size / elem_size; + } + }) + .map_err(|_| TryReserveError::AllocError { layout: new_layout, non_exhaustive: () }) + } - alloc_guard(new_layout.size())?; + fn shrink(&mut self, amount: usize, placement: AllocPlacement) -> Result<(), TryReserveError> { + assert!(amount <= self.cap, "Tried to shrink to a larger capacity"); - let res = match self.current_layout() { - Some(layout) => { - debug_assert!(new_layout.align() == layout.align()); - self.a.realloc(NonNull::from(self.ptr).cast(), layout, new_layout.size()) + let elem_size = mem::size_of::(); + let old_layout = + if let Some(layout) = self.current_layout() { layout } else { return Ok(()) }; + let old_ptr = self.ptr.cast().into(); + let new_size = amount * elem_size; + + let allocation = unsafe { + match (amount, placement) { + (0, AllocPlacement::Unspecified) => { + self.dealloc_buffer(); + Ok((old_layout.dangling(), 0)) } - None => self.a.alloc(new_layout), - }; - - let (ptr, new_cap) = match (res, fallibility) { - (Err(AllocErr), Infallible) => handle_alloc_error(new_layout), - (Err(AllocErr), Fallible) => { - return Err(TryReserveError::AllocError { - layout: new_layout, - non_exhaustive: (), - }); + (_, AllocPlacement::Unspecified) => self.a.realloc(old_ptr, old_layout, new_size), + (_, AllocPlacement::InPlace) => self + .a + .shrink_in_place(old_ptr, old_layout, new_size) + .map(|size| (old_ptr, size)) + .map_err(|_| AllocErr), + } + }; + + allocation + .map(|(ptr, alloc_size)| { + self.ptr = ptr.cast().into(); + if elem_size == 0 { + self.cap = amount; + } else { + self.cap = alloc_size / elem_size; } - (Ok((ptr, new_size)), _) => (ptr, new_size / elem_size), - }; - - self.ptr = ptr.cast().into(); - self.cap = new_cap; - - Ok(()) - } + }) + .map_err(|_| TryReserveError::AllocError { + layout: unsafe { Layout::from_size_align_unchecked(new_size, old_layout.align()) }, + non_exhaustive: (), + }) } } @@ -709,11 +635,8 @@ impl RawVec { impl RawVec { /// Frees the memory owned by the `RawVec` *without* trying to drop its contents. pub unsafe fn dealloc_buffer(&mut self) { - let elem_size = mem::size_of::(); - if elem_size != 0 { - if let Some(layout) = self.current_layout() { - self.a.dealloc(NonNull::from(self.ptr).cast(), layout); - } + if let Some(layout) = self.current_layout() { + self.a.dealloc(self.ptr.cast().into(), layout); } } } @@ -721,9 +644,7 @@ impl RawVec { unsafe impl<#[may_dangle] T, A: AllocRef> Drop for RawVec { /// Frees the memory owned by the `RawVec` *without* trying to drop its contents. fn drop(&mut self) { - unsafe { - self.dealloc_buffer(); - } + unsafe { self.dealloc_buffer() } } } diff --git a/src/liballoc/raw_vec/tests.rs b/src/liballoc/raw_vec/tests.rs index 860058debe1fd..dae0797246a36 100644 --- a/src/liballoc/raw_vec/tests.rs +++ b/src/liballoc/raw_vec/tests.rs @@ -1,4 +1,5 @@ use super::*; +use core::ptr::NonNull; #[test] fn allocator_param() {