Skip to content

Commit

Permalink
Add try_reserve and friends
Browse files Browse the repository at this point in the history
  • Loading branch information
SimonSapin committed Apr 22, 2020
1 parent 40fa827 commit 62a3622
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 18 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
@@ -1,6 +1,6 @@
[package]
name = "smallvec"
version = "1.3.0"
version = "1.4.0"
edition = "2018"
authors = ["Simon Sapin <simon.sapin@exyr.org>"]
license = "MIT/Apache-2.0"
Expand Down
71 changes: 54 additions & 17 deletions lib.rs
Expand Up @@ -201,9 +201,16 @@ impl<T: Clone> ExtendFromSlice<T> for Vec<T> {
}
}

/// 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<LayoutErr> for CollectionAllocErr {
Expand All @@ -212,6 +219,14 @@ impl From<LayoutErr> for CollectionAllocErr {
}
}

fn infallible<T>(result: Result<T, CollectionAllocErr>) -> 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<T>(n: usize) -> Result<Layout, CollectionAllocErr> {
Expand Down Expand Up @@ -714,48 +729,61 @@ impl<A: Array> SmallVec<A> {

/// 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::<A::Item>(new_cap).unwrap();
let layout = layout_array::<A::Item>(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(())
}
}

/// Reserve capacity for `additional` more elements to be inserted.
///
/// 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()
Expand All @@ -764,21 +792,30 @@ impl<A: Array> SmallVec<A> {
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(())
}
}

/// Reserve the minimum capacity for `additional` more elements to be inserted.
///
/// 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(())
}
}

Expand Down

0 comments on commit 62a3622

Please sign in to comment.