try_with_capacity for Vec, VecDeque, String

#91913

12 files changed, 155 insertions, 4 deletions

diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs
index c35bab5ef66..41adc2e79dc 100644
--- a/library/alloc/src/collections/vec_deque/mod.rs
+++ b/library/alloc/src/collections/vec_deque/mod.rs
@@ -559,6 +559,30 @@ 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.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the capacity exceeds `isize::MAX` _bytes_,
+    /// or if the allocator reports allocation failure.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(try_with_capacity)]
+    /// # #[allow(unused)]
+    /// # fn example() -> Result<(), std::collections::TryReserveError> {
+    /// use std::collections::VecDeque;
+    ///
+    /// let deque: VecDeque<u32> = VecDeque::try_with_capacity(10)?;
+    /// # Ok(()) }
+    /// ```
+    #[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> {
diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs
index 28695ade5bf..ca504b05a96 100644
--- a/library/alloc/src/lib.rs
+++ b/library/alloc/src/lib.rs
@@ -163,6 +163,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)]
diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs
index 2d86fa377f1..c5cf12209d9 100644
--- a/library/alloc/src/raw_vec.rs
+++ b/library/alloc/src/raw_vec.rs
@@ -20,6 +20,7 @@ mod tests;
 enum AllocInit {
     /// The contents of the new memory are uninitialized.
     Uninitialized,
+    #[cfg(not(no_global_oom_handling))]
     /// The new memory is guaranteed to be zeroed.
     Zeroed,
 }
@@ -146,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))]
@@ -203,6 +211,7 @@ impl<T, A: Allocator> RawVec<T, A> {
 
             let result = match init {
                 AllocInit::Uninitialized => alloc.allocate(layout),
+                #[cfg(not(no_global_oom_handling))]
                 AllocInit::Zeroed => alloc.allocate_zeroed(layout),
             };
             let ptr = match result {
diff --git a/library/alloc/src/string.rs b/library/alloc/src/string.rs
index 98ded7f6cdf..c4dcff1b1c4 100644
--- a/library/alloc/src/string.rs
+++ b/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
diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
index 7bd19875584..4b8b095c752 100644
--- a/library/alloc/src/vec/mod.rs
+++ b/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 length, and a capacity.
     ///
     /// # Safety
@@ -672,6 +688,24 @@ 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.
+    #[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 length, a capacity,
     /// and an allocator.
     ///
diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs
index c4e89a58a05..e4e1a02fd83 100644
--- a/library/alloc/tests/lib.rs
+++ b/library/alloc/tests/lib.rs
@@ -20,6 +20,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)]
diff --git a/library/alloc/tests/string.rs b/library/alloc/tests/string.rs
index 711e4eef2e7..e20ceae87b0 100644
--- a/library/alloc/tests/string.rs
+++ b/library/alloc/tests/string.rs
@@ -726,6 +726,17 @@ fn test_reserve_exact() {
 #[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
 fn test_try_reserve() {
     // These are the interesting cases:
     // * exactly isize::MAX should never trigger a CapacityOverflow (can be OOM)
diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs
index 15ee4d65205..aa95b4e9770 100644
--- a/library/alloc/tests/vec.rs
+++ b/library/alloc/tests/vec.rs
@@ -1697,6 +1697,18 @@ fn test_reserve_exact() {
 #[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
 fn test_try_reserve() {
     // These are the interesting cases:
     // * exactly isize::MAX should never trigger a CapacityOverflow (can be OOM)
diff --git a/library/alloc/tests/vec_deque.rs b/library/alloc/tests/vec_deque.rs
index eda2f8bb812..cea5de4dd59 100644
--- a/library/alloc/tests/vec_deque.rs
+++ b/library/alloc/tests/vec_deque.rs
@@ -1185,6 +1185,17 @@ fn test_reserve_exact_2() {
 #[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
 fn test_try_reserve() {
     // These are the interesting cases:
     // * exactly isize::MAX should never trigger a CapacityOverflow (can be OOM)
diff --git a/tests/codegen/vec-with-capacity.rs b/tests/codegen/vec-with-capacity.rs
new file mode 100644
index 00000000000..47051f2eef8
--- /dev/null
+++ b/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)
+}
diff --git a/tests/ui/suggestions/deref-path-method.stderr b/tests/ui/suggestions/deref-path-method.stderr
index a2b68fa966f..b27d9aef066 100644
--- a/tests/ui/suggestions/deref-path-method.stderr
+++ b/tests/ui/suggestions/deref-path-method.stderr
@@ -7,9 +7,9 @@ LL |     Vec::contains(&vec, &0);
 note: if you're trying to build a new `Vec<_, _>` consider using one of the following associated functions:
       Vec::<T>::new
       Vec::<T>::with_capacity
+      Vec::<T>::try_with_capacity
       Vec::<T>::from_raw_parts
-      Vec::<T, A>::new_in
-      and 2 others
+      and 4 others
   --> $SRC_DIR/alloc/src/vec/mod.rs:LL:COL
 help: the function `contains` is implemented on `[_]`
    |
diff --git a/tests/ui/ufcs/bad-builder.stderr b/tests/ui/ufcs/bad-builder.stderr
index e1c5e45b3eb..9cfeb7a5d09 100644
--- a/tests/ui/ufcs/bad-builder.stderr
+++ b/tests/ui/ufcs/bad-builder.stderr
@@ -7,9 +7,9 @@ LL |     Vec::<Q>::mew()
 note: if you're trying to build a new `Vec<Q>` consider using one of the following associated functions:
       Vec::<T>::new
       Vec::<T>::with_capacity
+      Vec::<T>::try_with_capacity
       Vec::<T>::from_raw_parts
-      Vec::<T, A>::new_in
-      and 2 others
+      and 4 others
   --> $SRC_DIR/alloc/src/vec/mod.rs:LL:COL
 help: there is an associated function `new` with a similar name
    |