try_with_capacity for Vec, VecDeque, String

#91913

Author: kornelski

library/alloc/src/collections/vec_deque/mod.rs

@@ -559,6 +559,21 @@ impl<T> VecDeque<T> {
     pub fn with_capacity(capacity: usize) -> VecDeque<T> {
         Self::with_capacity_in(capacity, Global)
     }
+
+    /// Creates an empty deque with space for at least `capacity` elements.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use std::collections::VecDeque;
+    ///
+    /// let deque: VecDeque<u32> = VecDeque::with_capacity(10);
+    /// ```
+    #[inline]
+    #[unstable(feature = "try_with_capacity", issue = "91913")]
+    pub fn try_with_capacity(capacity: usize) -> Result<VecDeque<T>, TryReserveError> {
+        Ok(VecDeque { head: 0, len: 0, buf: RawVec::try_with_capacity_in(capacity, Global)? })
+    }
 }
 
 impl<T, A: Allocator> VecDeque<T, A> {

library/alloc/src/lib.rs

@@ -158,6 +158,7 @@
 #![feature(trusted_len)]
 #![feature(trusted_random_access)]
 #![feature(try_trait_v2)]
+#![feature(try_with_capacity)]
 #![feature(tuple_trait)]
 #![feature(unchecked_math)]
 #![feature(unicode_internals)]

library/alloc/src/raw_vec.rs

@@ -147,6 +147,13 @@ impl<T, A: Allocator> RawVec<T, A> {
         handle_reserve(Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc))
     }
 
+    /// Like `try_with_capacity`, but parameterized over the choice of
+    /// allocator for the returned `RawVec`.
+    #[inline]
+    pub fn try_with_capacity_in(capacity: usize, alloc: A) -> Result<Self, TryReserveError> {
+        Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc)
+    }
+
     /// Like `with_capacity_zeroed`, but parameterized over the choice
     /// of allocator for the returned `RawVec`.
     #[cfg(not(no_global_oom_handling))]

library/alloc/src/string.rs

@@ -492,6 +492,19 @@ impl String {
         String { vec: Vec::with_capacity(capacity) }
     }
 
+    /// Creates a new empty `String` with at least the specified capacity.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`Err`] if the capacity exceeds `isize::MAX` bytes,
+    /// or if the memory allocator reports failure.
+    ///
+    #[inline]
+    #[unstable(feature = "try_with_capacity", issue = "91913")]
+    pub fn try_with_capacity(capacity: usize) -> Result<String, TryReserveError> {
+        Ok(String { vec: Vec::try_with_capacity(capacity)? })
+    }
+
     // HACK(japaric): with cfg(test) the inherent `[T]::to_vec` method, which is
     // required for this method definition, is not available. Since we don't
     // require this method for testing purposes, I'll just stub it

library/alloc/src/vec/mod.rs

@@ -481,6 +481,22 @@ impl<T> Vec<T> {
         Self::with_capacity_in(capacity, Global)
     }
 
+    /// Constructs a new, empty `Vec<T>` with at least the specified capacity.
+    ///
+    /// The vector will be able to hold at least `capacity` elements without
+    /// reallocating. This method is allowed to allocate for more elements than
+    /// `capacity`. If `capacity` is 0, the vector will not allocate.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the capacity exceeds `isize::MAX` _bytes_,
+    /// or if the allocator reports allocation failure.
+    #[inline]
+    #[unstable(feature = "try_with_capacity", issue = "91913")]
+    pub fn try_with_capacity(capacity: usize) -> Result<Self, TryReserveError> {
+        Self::try_with_capacity_in(capacity, Global)
+    }
+
     /// Creates a `Vec<T>` directly from a pointer, a capacity, and a length.
     ///
     /// # Safety
@@ -672,6 +688,25 @@ impl<T, A: Allocator> Vec<T, A> {
         Vec { buf: RawVec::with_capacity_in(capacity, alloc), len: 0 }
     }
 
+    /// Constructs a new, empty `Vec<T, A>` with at least the specified capacity
+    /// with the provided allocator.
+    ///
+    /// The vector will be able to hold at least `capacity` elements without
+    /// reallocating. This method is allowed to allocate for more elements than
+    /// `capacity`. If `capacity` is 0, the vector will not allocate.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the capacity exceeds `isize::MAX` _bytes_,
+    /// or if the allocator reports allocation failure.
+    #[cfg(not(no_global_oom_handling))]
+    #[inline]
+    #[unstable(feature = "allocator_api", issue = "32838")]
+    // #[unstable(feature = "try_with_capacity", issue = "91913")]
+    pub fn try_with_capacity_in(capacity: usize, alloc: A) -> Result<Self, TryReserveError> {
+        Ok(Vec { buf: RawVec::try_with_capacity_in(capacity, alloc)?, len: 0 })
+    }
+
     /// Creates a `Vec<T, A>` directly from a pointer, a capacity, a length,
     /// and an allocator.
     ///

library/alloc/tests/lib.rs

@@ -19,6 +19,7 @@
 #![feature(pattern)]
 #![feature(trusted_len)]
 #![feature(try_reserve_kind)]
+#![feature(try_with_capacity)]
 #![feature(unboxed_closures)]
 #![feature(associated_type_bounds)]
 #![feature(binary_heap_into_iter_sorted)]

library/alloc/tests/string.rs

@@ -723,6 +723,17 @@ fn test_reserve_exact() {
     assert!(s.capacity() >= 33)
 }
 
+#[test]
+#[cfg_attr(miri, ignore)] // Miri does not support signalling OOM
+#[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc
+fn test_try_with_capacity() {
+    let string = String::try_with_capacity(1000).unwrap();
+    assert_eq!(0, string.len());
+    assert!(string.capacity() >= 1000 && string.capacity() <= isize::MAX as usize);
+
+    assert!(String::try_with_capacity(usize::MAX).is_err());
+}
+
 #[test]
 #[cfg_attr(miri, ignore)] // Miri does not support signalling OOM
 #[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc

library/alloc/tests/vec.rs

@@ -1694,6 +1694,18 @@ fn test_reserve_exact() {
     assert!(v.capacity() >= 33)
 }
 
+#[test]
+#[cfg_attr(miri, ignore)] // Miri does not support signalling OOM
+#[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc
+fn test_try_with_capacity() {
+    let mut vec: Vec<u32> = Vec::try_with_capacity(5).unwrap();
+    assert_eq!(0, vec.len());
+    assert!(vec.capacity() >= 5 && vec.capacity() <= isize::MAX as usize / 4);
+    assert!(vec.spare_capacity_mut().len() >= 5);
+
+    assert!(Vec::<u16>::try_with_capacity(isize::MAX as usize + 1).is_err());
+}
+
 #[test]
 #[cfg_attr(miri, ignore)] // Miri does not support signalling OOM
 #[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc

library/alloc/tests/vec_deque.rs

@@ -1182,6 +1182,17 @@ fn test_reserve_exact_2() {
     assert!(v.capacity() >= 33)
 }
 
+#[test]
+#[cfg_attr(miri, ignore)] // Miri does not support signalling OOM
+#[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc
+fn test_try_with_capacity() {
+    let vec: VecDeque<u32> = VecDeque::try_with_capacity(5).unwrap();
+    assert_eq!(0, vec.len());
+    assert!(vec.capacity() >= 5 && vec.capacity() <= isize::MAX as usize / 4);
+
+    assert!(VecDeque::<u16>::try_with_capacity(isize::MAX as usize + 1).is_err());
+}
+
 #[test]
 #[cfg_attr(miri, ignore)] // Miri does not support signalling OOM
 #[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc

tests/codegen/vec-with-capacity.rs

@@ -0,0 +1,35 @@
+// compile-flags: -O
+// ignore-debug
+// (with debug assertions turned on, `assert_unchecked` generates a real assertion)
+
+#![crate_type = "lib"]
+#![feature(try_with_capacity)]
+
+// CHECK-LABEL: @with_capacity_does_not_grow1
+#[no_mangle]
+pub fn with_capacity_does_not_grow1() -> Vec<u32> {
+    let v = Vec::with_capacity(1234);
+    // CHECK: call {{.*}}__rust_alloc(
+    // CHECK-NOT: call {{.*}}__rust_realloc
+    // CHECK-NOT: call {{.*}}capacity_overflow
+    // CHECK-NOT: call {{.*}}finish_grow
+    // CHECK-NOT: call {{.*}}reserve
+    // CHECK-NOT: memcpy
+    // CHECK-NOT: memset
+    v
+}
+
+// CHECK-LABEL: @try_with_capacity_does_not_grow2
+#[no_mangle]
+pub fn try_with_capacity_does_not_grow2() -> Option<Vec<Vec<u8>>> {
+    let v = Vec::try_with_capacity(1234).ok()?;
+    // CHECK: call {{.*}}__rust_alloc(
+    // CHECK-NOT: call {{.*}}__rust_realloc
+    // CHECK-NOT: call {{.*}}capacity_overflow
+    // CHECK-NOT: call {{.*}}finish_grow
+    // CHECK-NOT: call {{.*}}handle_alloc_error
+    // CHECK-NOT: call {{.*}}reserve
+    // CHECK-NOT: memcpy
+    // CHECK-NOT: memset
+    Some(v)
+}