diff --git a/src/drop.rs b/src/drop.rs index e134c32..4ccfe26 100644 --- a/src/drop.rs +++ b/src/drop.rs @@ -25,7 +25,7 @@ impl Drop for MiniVec { } = core::ptr::read(self.buf.as_ptr().cast::
()); core::ptr::drop_in_place(core::ptr::slice_from_raw_parts_mut(self.data(), len)); - alloc::alloc::dealloc(self.buf.as_ptr(), make_layout::(cap, alignment)); + alloc::alloc::dealloc(self.buf.as_ptr(), make_layout(cap, alignment, core::mem::size_of::())); }; } } diff --git a/src/impl/helpers.rs b/src/impl/helpers.rs index 4c1ce6d..960fc55 100644 --- a/src/impl/helpers.rs +++ b/src/impl/helpers.rs @@ -22,7 +22,10 @@ pub const fn next_capacity(capacity: usize) -> usize { }; } - 2 * capacity + match capacity.checked_mul(2) { + Some(cap) => cap, + None => capacity, + } } pub fn max_align() -> usize { @@ -31,13 +34,13 @@ pub fn max_align() -> usize { core::cmp::max(align_t, header_align) } -pub fn make_layout(capacity: usize, alignment: usize) -> alloc::alloc::Layout { +pub fn make_layout(capacity: usize, alignment: usize, type_size: usize) -> alloc::alloc::Layout { let header_size = core::mem::size_of::
(); let num_bytes = if capacity == 0 { next_aligned(header_size, alignment) } else { next_aligned(header_size, alignment) - + next_aligned(capacity * core::mem::size_of::(), alignment) + + next_aligned(capacity * type_size, alignment) }; alloc::alloc::Layout::from_size_align(num_bytes, alignment).unwrap() @@ -83,14 +86,14 @@ mod tests { fn make_layout_test() { // empty // - let layout = make_layout::(0, max_align::()); + let layout = make_layout(0, max_align::(), core::mem::size_of::()); assert_eq!(layout.align(), core::mem::align_of::
()); assert_eq!(layout.size(), core::mem::size_of::
()); // non-empty, less than // - let layout = make_layout::(512, max_align::()); + let layout = make_layout(512, max_align::(), core::mem::size_of::()); assert!(core::mem::align_of::() < core::mem::align_of::
()); assert_eq!(layout.align(), core::mem::align_of::
()); assert_eq!( @@ -100,7 +103,7 @@ mod tests { // non-empty, equal // - let layout = make_layout::(512, max_align::()); + let layout = make_layout(512, max_align::(), core::mem::size_of::()); assert_eq!( core::mem::align_of::(), core::mem::align_of::
() @@ -112,7 +115,7 @@ mod tests { ); // non-empty, greater - let layout = make_layout::(512, max_align::()); + let layout = make_layout(512, max_align::(), core::mem::size_of::()); assert!(core::mem::align_of::() > core::mem::align_of::
()); assert_eq!(layout.align(), core::mem::align_of::()); assert_eq!( @@ -124,7 +127,7 @@ mod tests { ); // non-empty, over-aligned - let layout = make_layout::(512, 32); + let layout = make_layout(512, 32, core::mem::size_of::()); assert_eq!(layout.align(), 32); assert_eq!( layout.size(), diff --git a/src/lib.rs b/src/lib.rs index 47c01ec..e22610c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -77,6 +77,64 @@ use crate::r#impl::splice::make_splice_iterator; pub use crate::r#impl::{Drain, DrainFilter, IntoIter, Splice}; +#[inline] +#[allow(clippy::cast_ptr_alignment)] +fn is_default(buf: core::ptr::NonNull) -> bool { + core::ptr::eq(buf.as_ptr(), &DEFAULT_U8) +} + +#[derive(Debug, Clone, Eq, PartialEq)] +//TODO: Switch to https://doc.rust-lang.org/alloc/collections/enum.TryReserveErrorKind.html once stable +/// The error type for `try_reserve` methods. +pub struct TryReserveError { +} + +struct Grow { + buf: core::ptr::NonNull, + old_capacity: usize, + new_capacity: usize, + alignment: usize, + len: usize, + type_size: usize, +} + +impl Grow { + #[inline] + fn grow(self) -> Result>, alloc::alloc::Layout> { + if self.new_capacity == self.old_capacity { + return Ok(None); + } + + let new_layout = make_layout(self.new_capacity, self.alignment, self.type_size); + + let new_buf = if is_default(self.buf) { + unsafe { alloc::alloc::alloc(new_layout) } + } else { + let old_layout = make_layout(self.old_capacity, self.alignment, self.type_size); + + unsafe { alloc::alloc::realloc(self.buf.as_ptr(), old_layout, new_layout.size()) } + }; + + match core::ptr::NonNull::new(new_buf) { + Some(new_buf) => { + let header = Header { + len: self.len, + cap: self.new_capacity, + alignment: self.alignment, + }; + + #[allow(clippy::cast_ptr_alignment)] + unsafe { + core::ptr::write(new_buf.cast::
().as_ptr(), header); + } + + Ok(Some(new_buf)) + }, + None => Err(new_layout), + } + } +} + /// `MiniVec` is a space-optimized implementation of `alloc::vec::Vec` that is only the size of a single pointer and /// also extends portions of its API, including support for over-aligned allocations. `MiniVec` also aims to bring as /// many Nightly features from `Vec` to stable toolchains as is possible. In many cases, it is a drop-in replacement @@ -130,9 +188,8 @@ fn header_clone() { static DEFAULT_U8: u8 = 137; impl MiniVec { - #[allow(clippy::cast_ptr_alignment)] fn is_default(&self) -> bool { - core::ptr::eq(self.buf.as_ptr(), &DEFAULT_U8) + is_default(self.buf) } fn header(&self) -> &Header { @@ -167,41 +224,44 @@ impl MiniVec { fn grow(&mut self, capacity: usize, alignment: usize) { debug_assert!(capacity >= self.len()); - let old_capacity = self.capacity(); - let new_capacity = capacity; - - if new_capacity == old_capacity { - return; - } - - let new_layout = make_layout::(new_capacity, alignment); - - let len = self.len(); - - let new_buf = if self.is_default() { - unsafe { alloc::alloc::alloc(new_layout) } - } else { - let old_layout = make_layout::(old_capacity, alignment); - - unsafe { alloc::alloc::realloc(self.buf.as_ptr(), old_layout, new_layout.size()) } + let grower = Grow { + buf: self.buf, + old_capacity: self.capacity(), + new_capacity: capacity, + alignment, + len: self.len(), + type_size: core::mem::size_of::(), }; - if new_buf.is_null() { - alloc::alloc::handle_alloc_error(new_layout); + match grower.grow() { + Ok(Some(new_buf)) => { + self.buf = new_buf; + }, + Ok(None) => (), + Err(new_layout) => alloc::alloc::handle_alloc_error(new_layout), } + } - let header = Header { - len, - cap: new_capacity, - alignment, + fn try_grow(&mut self, capacity: usize, alignment: usize) -> Result<(), TryReserveError> { + debug_assert!(capacity >= self.len()); + + let grower = Grow { + buf: self.buf, + old_capacity: self.capacity(), + new_capacity: capacity, + alignment, + len: self.len(), + type_size: core::mem::size_of::(), }; - #[allow(clippy::cast_ptr_alignment)] - unsafe { - core::ptr::write(new_buf.cast::
(), header); + match grower.grow() { + Ok(Some(new_buf)) => { + self.buf = new_buf; + Ok(()) + }, + Ok(None) => Ok(()), + Err(_) => Err(TryReserveError {}), } - - self.buf = unsafe { core::ptr::NonNull::::new_unchecked(new_buf) }; } /// `append` moves every element from `other` to the back of `self`. `other.is_empty()` is `true` once this operation @@ -1056,6 +1116,78 @@ impl MiniVec { self.grow(total_required, self.alignment()); } + /// `try_reserve` ensures there is sufficient capacity for `additional` extra elements to be either + /// inserted or appended to the end of the vector. Will reallocate if needed otherwise this + /// function is a no-op. + /// + /// Guarantees that the new capacity is greater than or equal to `len() + additional`. + /// + /// # Errors + /// + /// In case of failure, old capacity is retained. + /// + /// # Example + /// + /// ``` + /// let mut vec = minivec::MiniVec::::new(); + /// + /// assert_eq!(vec.capacity(), 0); + /// + /// vec.try_reserve(128).expect("Successfully allocate extra memory"); + /// + /// assert!(vec.capacity() >= 128); + /// ``` + /// + pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> { + let capacity = self.capacity(); + let total_required = match self.len().checked_add(additional) { + Some(total_required) => total_required, + None => return Err(TryReserveError {}), + }; + + if total_required <= capacity { + return Ok(()); + } + + let mut new_capacity = next_capacity::(capacity); + while new_capacity < total_required { + new_capacity = next_capacity::(new_capacity); + } + + self.try_grow(new_capacity, self.alignment()) + } + + /// `try_reserve_exact` ensures that the capacity of the vector is exactly equal to + /// `len() + additional` unless the capacity is already sufficient in which case no operation is + /// performed. + /// + /// # Errors + /// + /// In case of failure, old capacity is retained. + /// + /// # Example + /// + /// ``` + /// let mut vec = minivec::MiniVec::::new(); + /// vec.try_reserve_exact(57).expect("Oi, no capacity"); + /// + /// assert_eq!(vec.capacity(), 57); + /// ``` + /// + pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), TryReserveError> { + let capacity = self.capacity(); + + let total_required = match self.len().checked_add(additional) { + Some(total_required) => total_required, + None => return Err(TryReserveError {}), + }; + if capacity >= total_required { + return Ok(()); + } + + self.try_grow(total_required, self.alignment()) + } + /// `resize` will clone the supplied `value` as many times as required until `len()` becomes /// `new_len`. If the current [`len()`](MiniVec::len) is greater than `new_len` then the vector /// is truncated in a way that's identical to calling `vec.truncate(new_len)`. If the `len()` diff --git a/tests/mine.rs b/tests/mine.rs index c9d0e55..81734d9 100644 --- a/tests/mine.rs +++ b/tests/mine.rs @@ -1262,3 +1262,52 @@ fn minivec_assume_minivec_init() { assert_eq!(bytes[0], 137); assert_eq!(bytes[511], 137); } + +#[test] +fn minivec_should_fail_try_reserve_impossible() { + let mut buf = minivec::mini_vec![0; 512]; + buf.try_reserve(usize::max_value()).expect_err("FAIL"); +} + +#[test] +fn minivec_should_fail_try_reserve_exact_impossible() { + let mut buf = minivec::mini_vec![0; 512]; + buf.try_reserve_exact(usize::max_value()).expect_err("FAIL"); +} + +#[test] +fn minivec_test_try_reserve() { + let mut v = MiniVec::new(); + assert_eq!(v.capacity(), 0); + + v.try_reserve(2).expect("to try 2"); + assert!(v.capacity() >= 2); + + for i in 0..16 { + v.push(i); + } + + assert!(v.capacity() >= 16); + v.try_reserve(16).expect("to try 16"); + assert!(v.capacity() >= 32); + + v.push(16); + + v.try_reserve(16).expect("to try 16 again"); + assert!(v.capacity() >= 33) +} + +#[test] +fn minivec_overaligned_allocations_with_try_reserve_exact() { + #[repr(align(256))] + struct Foo(usize); + let mut v = mini_vec![Foo(273)]; + for i in 0..0x1000 { + v.try_reserve_exact(i).expect("try reserve exact"); + assert!(v[0].0 == 273); + assert!(v.as_ptr() as usize & 0xff == 0); + v.shrink_to_fit(); + assert!(v[0].0 == 273); + assert!(v.as_ptr() as usize & 0xff == 0); + } +}