Introduce try_reserve friends

src/drop.rs

@@ -25,7 +25,7 @@ impl<T> Drop for MiniVec<T> {
       } = core::ptr::read(self.buf.as_ptr().cast::<Header>());
 
       core::ptr::drop_in_place(core::ptr::slice_from_raw_parts_mut(self.data(), len));
-      alloc::alloc::dealloc(self.buf.as_ptr(), make_layout::<T>(cap, alignment));
+      alloc::alloc::dealloc(self.buf.as_ptr(), make_layout(cap, alignment, core::mem::size_of::<T>()));
     };
   }
 }

src/impl/helpers.rs

@@ -22,7 +22,10 @@ pub const fn next_capacity<T>(capacity: usize) -> usize {
     };
   }
 
-  2 * capacity
+  match capacity.checked_mul(2) {
+      Some(cap) => cap,
+      None => capacity,
+  }
 }
 
 pub fn max_align<T>() -> usize {
@@ -31,13 +34,13 @@ pub fn max_align<T>() -> usize {
   core::cmp::max(align_t, header_align)
 }
 
-pub fn make_layout<T>(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::<Header>();
   let num_bytes = if capacity == 0 {
     next_aligned(header_size, alignment)
   } else {
     next_aligned(header_size, alignment)
-      + next_aligned(capacity * core::mem::size_of::<T>(), 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::<i32>(0, max_align::<i32>());
+    let layout = make_layout(0, max_align::<i32>(), core::mem::size_of::<i32>());
 
     assert_eq!(layout.align(), core::mem::align_of::<Header>());
     assert_eq!(layout.size(), core::mem::size_of::<Header>());
 
     // non-empty, less than
     //
-    let layout = make_layout::<i32>(512, max_align::<i32>());
+    let layout = make_layout(512, max_align::<i32>(), core::mem::size_of::<i32>());
     assert!(core::mem::align_of::<i32>() < core::mem::align_of::<Header>());
     assert_eq!(layout.align(), core::mem::align_of::<Header>());
     assert_eq!(
@@ -100,7 +103,7 @@ mod tests {
 
     // non-empty, equal
     //
-    let layout = make_layout::<i64>(512, max_align::<i64>());
+    let layout = make_layout(512, max_align::<i64>(), core::mem::size_of::<i64>());
     assert_eq!(
       core::mem::align_of::<i64>(),
       core::mem::align_of::<Header>()
@@ -112,7 +115,7 @@ mod tests {
     );
 
     // non-empty, greater
-    let layout = make_layout::<OverAligned>(512, max_align::<OverAligned>());
+    let layout = make_layout(512, max_align::<OverAligned>(), core::mem::size_of::<OverAligned>());
     assert!(core::mem::align_of::<OverAligned>() > core::mem::align_of::<Header>());
     assert_eq!(layout.align(), core::mem::align_of::<OverAligned>());
     assert_eq!(
@@ -124,7 +127,7 @@ mod tests {
     );
 
     // non-empty, over-aligned
-    let layout = make_layout::<i32>(512, 32);
+    let layout = make_layout(512, 32, core::mem::size_of::<i32>());
     assert_eq!(layout.align(), 32);
     assert_eq!(
       layout.size(),

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<u8>) -> 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<u8>,
+    old_capacity: usize,
+    new_capacity: usize,
+    alignment: usize,
+    len: usize,
+    type_size: usize,
+}
+
+impl Grow {
+    #[inline]
+    fn grow(self) -> Result<Option<core::ptr::NonNull<u8>>, 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::<Header>().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<T> MiniVec<T> {
-  #[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<T> MiniVec<T> {
   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::<T>(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::<T>(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::<T>(),
     };
 
-    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::<T>(),
     };
 
-    #[allow(clippy::cast_ptr_alignment)]
-    unsafe {
-      core::ptr::write(new_buf.cast::<Header>(), 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::<u8>::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<T> MiniVec<T> {
     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::<i32>::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::<T>(capacity);
+    while new_capacity < total_required {
+      new_capacity = next_capacity::<T>(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::<i32>::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()`

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);
+  }
+}