From 40fa8279eec2481ce18deb86eeba34061138e2ee Mon Sep 17 00:00:00 2001 From: Simon Sapin Date: Wed, 22 Apr 2020 23:24:17 +0200 Subject: [PATCH 1/4] Stop using `Vec` as an allocator --- lib.rs | 38 ++++++++++++++++++++++++++++++++------ 1 file changed, 32 insertions(+), 6 deletions(-) diff --git a/lib.rs b/lib.rs index ed83a75..9307f6d 100644 --- a/lib.rs +++ b/lib.rs @@ -40,7 +40,7 @@ //! //! When this feature is enabled, `SmallVec` works with any arrays of any size, not just a fixed //! list of sizes. -//! +//! //! ### `specialization` //! //! **This feature is unstable and requires a nightly build of the Rust toolchain.** @@ -71,6 +71,7 @@ pub extern crate alloc; #[cfg(any(test, feature = "write"))] extern crate std; +use alloc::alloc::{Layout, LayoutErr}; use alloc::boxed::Box; use alloc::{vec, vec::Vec}; use core::borrow::{Borrow, BorrowMut}; @@ -200,9 +201,31 @@ impl ExtendFromSlice for Vec { } } +#[derive(Debug)] +enum CollectionAllocErr { + CapacityOverflow, +} + +impl From for CollectionAllocErr { + fn from(_: LayoutErr) -> Self { + CollectionAllocErr::CapacityOverflow + } +} + +/// FIXME: use `Layout::array` when we require a Rust version where it’s stable +/// https://github.com/rust-lang/rust/issues/55724 +fn layout_array(n: usize) -> Result { + let size = mem::size_of::().checked_mul(n) + .ok_or(CollectionAllocErr::CapacityOverflow)?; + let align = mem::align_of::(); + Layout::from_size_align(size, align) + .map_err(|_| CollectionAllocErr::CapacityOverflow) +} + unsafe fn deallocate(ptr: *mut T, capacity: usize) { - let _vec: Vec = Vec::from_raw_parts(ptr, 0, capacity); - // Let it drop. + // This unwrap should succeed since the same did when allocating. + let layout = layout_array::(capacity).unwrap(); + alloc::alloc::dealloc(ptr as *mut u8, layout) } /// An iterator that removes the items from a `SmallVec` and yields them by value. @@ -705,9 +728,12 @@ impl SmallVec { ptr::copy_nonoverlapping(ptr, self.data.inline_mut(), len); self.capacity = len; } else if new_cap != cap { - let mut vec = Vec::with_capacity(new_cap); - let new_alloc = vec.as_mut_ptr(); - mem::forget(vec); + // Panic on overflow + let layout = layout_array::(new_cap).unwrap(); + let new_alloc = NonNull::new(alloc::alloc::alloc(layout)) + .unwrap_or_else(|| alloc::alloc::handle_alloc_error(layout)) + .cast() + .as_ptr(); ptr::copy_nonoverlapping(ptr, new_alloc, len); self.data = SmallVecData::from_heap(new_alloc, len); self.capacity = new_cap; From 62a36229789b3781d62e38a2b510dff3541131e5 Mon Sep 17 00:00:00 2001 From: Simon Sapin Date: Wed, 22 Apr 2020 23:58:05 +0200 Subject: [PATCH 2/4] Add `try_reserve` and friends --- Cargo.toml | 2 +- lib.rs | 71 +++++++++++++++++++++++++++++++++++++++++------------- 2 files changed, 55 insertions(+), 18 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 7fe519a..0e8f13a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "smallvec" -version = "1.3.0" +version = "1.4.0" edition = "2018" authors = ["Simon Sapin "] license = "MIT/Apache-2.0" diff --git a/lib.rs b/lib.rs index 9307f6d..88a1edc 100644 --- a/lib.rs +++ b/lib.rs @@ -201,9 +201,16 @@ impl ExtendFromSlice for Vec { } } +/// Error type for APIs with fallible heap allocation #[derive(Debug)] -enum CollectionAllocErr { +pub enum CollectionAllocErr { + /// Overflow `usize::MAX` or other error during size computation CapacityOverflow, + /// The allocator return an error + AllocErr { + /// The layout that was passed to the allocator + layout: Layout, + }, } impl From for CollectionAllocErr { @@ -212,6 +219,14 @@ impl From for CollectionAllocErr { } } +fn infallible(result: Result) -> T { + match result { + Ok(x) => x, + Err(CollectionAllocErr::CapacityOverflow) => panic!("capacity overflow"), + Err(CollectionAllocErr::AllocErr { layout }) => alloc::alloc::handle_alloc_error(layout), + } +} + /// FIXME: use `Layout::array` when we require a Rust version where it’s stable /// https://github.com/rust-lang/rust/issues/55724 fn layout_array(n: usize) -> Result { @@ -714,36 +729,44 @@ impl SmallVec { /// Re-allocate to set the capacity to `max(new_cap, inline_size())`. /// - /// Panics if `new_cap` is less than the vector's length. + /// Panics if `new_cap` is less than the vector's length + /// or if the capacity computation overflows `usize`. pub fn grow(&mut self, new_cap: usize) { + infallible(self.try_grow(new_cap)) + } + + /// Re-allocate to set the capacity to `max(new_cap, inline_size())`. + /// + /// Panics if `new_cap` is less than the vector's length + pub fn try_grow(&mut self, new_cap: usize) -> Result<(), CollectionAllocErr> { unsafe { let (ptr, &mut len, cap) = self.triple_mut(); let unspilled = !self.spilled(); assert!(new_cap >= len); if new_cap <= self.inline_size() { if unspilled { - return; + return Ok(()); } self.data = SmallVecData::from_inline(MaybeUninit::uninit()); ptr::copy_nonoverlapping(ptr, self.data.inline_mut(), len); self.capacity = len; } else if new_cap != cap { - // Panic on overflow - let layout = layout_array::(new_cap).unwrap(); + let layout = layout_array::(new_cap)?; let new_alloc = NonNull::new(alloc::alloc::alloc(layout)) - .unwrap_or_else(|| alloc::alloc::handle_alloc_error(layout)) + .ok_or(CollectionAllocErr::AllocErr { layout })? .cast() .as_ptr(); ptr::copy_nonoverlapping(ptr, new_alloc, len); self.data = SmallVecData::from_heap(new_alloc, len); self.capacity = new_cap; if unspilled { - return; + return Ok(()); } } else { - return; + return Ok(()); } deallocate(ptr, cap); + Ok(()) } } @@ -751,11 +774,16 @@ impl SmallVec { /// /// May reserve more space to avoid frequent reallocations. /// - /// If the new capacity would overflow `usize` then it will be set to `usize::max_value()` - /// instead. (This means that inserting `additional` new elements is not guaranteed to be - /// possible after calling this function.) + /// Panics if the capacity computation overflows `usize`. #[inline] pub fn reserve(&mut self, additional: usize) { + infallible(self.try_reserve(additional)) + } + + /// Reserve capacity for `additional` more elements to be inserted. + /// + /// May reserve more space to avoid frequent reallocations. + pub fn try_reserve(&mut self, additional: usize) -> Result<(), CollectionAllocErr> { // prefer triple_mut() even if triple() would work // so that the optimizer removes duplicated calls to it // from callers like insert() @@ -764,8 +792,10 @@ impl SmallVec { let new_cap = len .checked_add(additional) .and_then(usize::checked_next_power_of_two) - .unwrap_or(usize::max_value()); - self.grow(new_cap); + .ok_or(CollectionAllocErr::CapacityOverflow)?; + self.try_grow(new_cap) + } else { + Ok(()) } } @@ -773,12 +803,19 @@ impl SmallVec { /// /// Panics if the new capacity overflows `usize`. pub fn reserve_exact(&mut self, additional: usize) { + infallible(self.try_reserve_exact(additional)) + } + + /// Reserve the minimum capacity for `additional` more elements to be inserted. + pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), CollectionAllocErr> { let (_, &mut len, cap) = self.triple_mut(); if cap - len < additional { - match len.checked_add(additional) { - Some(cap) => self.grow(cap), - None => panic!("reserve_exact overflow"), - } + let new_cap = len + .checked_add(additional) + .ok_or(CollectionAllocErr::CapacityOverflow)?; + self.try_grow(new_cap) + } else { + Ok(()) } } From bdfc4296d70112a17ef6077fb6653414b0979698 Mon Sep 17 00:00:00 2001 From: Simon Sapin Date: Thu, 23 Apr 2020 00:06:13 +0200 Subject: [PATCH 3/4] Use `realloc` when possible --- lib.rs | 30 +++++++++++++++++++----------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/lib.rs b/lib.rs index 88a1edc..3772763 100644 --- a/lib.rs +++ b/lib.rs @@ -750,22 +750,30 @@ impl SmallVec { self.data = SmallVecData::from_inline(MaybeUninit::uninit()); ptr::copy_nonoverlapping(ptr, self.data.inline_mut(), len); self.capacity = len; + deallocate(ptr, cap); } else if new_cap != cap { let layout = layout_array::(new_cap)?; - let new_alloc = NonNull::new(alloc::alloc::alloc(layout)) - .ok_or(CollectionAllocErr::AllocErr { layout })? - .cast() - .as_ptr(); - ptr::copy_nonoverlapping(ptr, new_alloc, len); - self.data = SmallVecData::from_heap(new_alloc, len); - self.capacity = new_cap; + let new_alloc; if unspilled { - return Ok(()); + new_alloc = NonNull::new(alloc::alloc::alloc(layout)) + .ok_or(CollectionAllocErr::AllocErr { layout })? + .cast() + .as_ptr(); + ptr::copy_nonoverlapping(ptr, new_alloc, len); + } else { + // This should never fail since the same succeeded + // when previously allocating `ptr`. + let old_layout = layout_array::(cap)?; + + let new_ptr = alloc::alloc::realloc(ptr as *mut u8, old_layout, layout.size()); + new_alloc = NonNull::new(new_ptr) + .ok_or(CollectionAllocErr::AllocErr { layout })? + .cast() + .as_ptr(); } - } else { - return Ok(()); + self.data = SmallVecData::from_heap(new_alloc, len); + self.capacity = new_cap; } - deallocate(ptr, cap); Ok(()) } } From 3ec778f1d0ffe79c3d6e2e8cda4de5600affc7f9 Mon Sep 17 00:00:00 2001 From: Simon Sapin Date: Thu, 23 Apr 2020 13:58:43 +0200 Subject: [PATCH 4/4] Early returns --- lib.rs | 28 +++++++++++++--------------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/lib.rs b/lib.rs index 3772763..caea125 100644 --- a/lib.rs +++ b/lib.rs @@ -796,15 +796,14 @@ impl SmallVec { // so that the optimizer removes duplicated calls to it // from callers like insert() let (_, &mut len, cap) = self.triple_mut(); - if cap - len < additional { - let new_cap = len - .checked_add(additional) - .and_then(usize::checked_next_power_of_two) - .ok_or(CollectionAllocErr::CapacityOverflow)?; - self.try_grow(new_cap) - } else { - Ok(()) + if cap - len >= additional { + return Ok(()); } + let new_cap = len + .checked_add(additional) + .and_then(usize::checked_next_power_of_two) + .ok_or(CollectionAllocErr::CapacityOverflow)?; + self.try_grow(new_cap) } /// Reserve the minimum capacity for `additional` more elements to be inserted. @@ -817,14 +816,13 @@ impl SmallVec { /// Reserve the minimum capacity for `additional` more elements to be inserted. pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), CollectionAllocErr> { let (_, &mut len, cap) = self.triple_mut(); - if cap - len < additional { - let new_cap = len - .checked_add(additional) - .ok_or(CollectionAllocErr::CapacityOverflow)?; - self.try_grow(new_cap) - } else { - Ok(()) + if cap - len >= additional { + return Ok(()); } + let new_cap = len + .checked_add(additional) + .ok_or(CollectionAllocErr::CapacityOverflow)?; + self.try_grow(new_cap) } /// Shrink the capacity of the vector as much as possible.