Rollup merge of #120504 - kornelski:try_with_capacity, r=Amanieu

Vec::try_with_capacity

Related to #91913

Implements try_with_capacity for `Vec`, `VecDeque`, and `String`. I can follow it up with more collections if desired.

`Vec::try_with_capacity()` is functionally equivalent to the current stable:

```rust
let mut v = Vec::new();
v.try_reserve_exact(n)?
```

However, `try_reserve` calls non-inlined `finish_grow`, which requires old and new `Layout`, and is designed to reallocate memory. There is benefit to using `try_with_capacity`, besides syntax convenience, because it generates much smaller code at the call site with a direct call to the allocator. There's codegen test included.

It's also a very desirable functionality for users of `no_global_oom_handling` (Rust-for-Linux), since it makes a very commonly used function available in that environment (`with_capacity` is used much more frequently than all `(try_)reserve(_exact)`).

14 files changed, 189 insertions, 31 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 4d9694d4beb..ace73c0fdaa 100644
--- a/library/alloc/src/raw_vec.rs
+++ b/library/alloc/src/raw_vec.rs
@@ -17,10 +17,19 @@ use crate::collections::TryReserveErrorKind::*;
 #[cfg(test)]
 mod tests;
 
+// One central function responsible for reporting capacity overflows. This'll
+// ensure that the code generation related to these panics is minimal as there's
+// only one location which panics rather than a bunch throughout the module.
 #[cfg(not(no_global_oom_handling))]
+#[cfg_attr(not(feature = "panic_immediate_abort"), inline(never))]
+fn capacity_overflow() -> ! {
+    panic!("capacity overflow");
+}
+
 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,
 }
@@ -93,6 +102,8 @@ impl<T> RawVec<T, Global> {
     /// zero-sized. Note that if `T` is zero-sized this means you will
     /// *not* get a `RawVec` with the requested capacity.
     ///
+    /// Non-fallible version of `try_with_capacity`
+    ///
     /// # Panics
     ///
     /// Panics if the requested capacity exceeds `isize::MAX` bytes.
@@ -104,7 +115,7 @@ impl<T> RawVec<T, Global> {
     #[must_use]
     #[inline]
     pub fn with_capacity(capacity: usize) -> Self {
-        Self::with_capacity_in(capacity, Global)
+        handle_reserve(Self::try_allocate_in(capacity, AllocInit::Uninitialized, Global))
     }
 
     /// Like `with_capacity`, but guarantees the buffer is zeroed.
@@ -142,7 +153,14 @@ impl<T, A: Allocator> RawVec<T, A> {
     #[cfg(not(no_global_oom_handling))]
     #[inline]
     pub fn with_capacity_in(capacity: usize, alloc: A) -> Self {
-        Self::allocate_in(capacity, AllocInit::Uninitialized, alloc)
+        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
@@ -150,7 +168,7 @@ impl<T, A: Allocator> RawVec<T, A> {
     #[cfg(not(no_global_oom_handling))]
     #[inline]
     pub fn with_capacity_zeroed_in(capacity: usize, alloc: A) -> Self {
-        Self::allocate_in(capacity, AllocInit::Zeroed, alloc)
+        handle_reserve(Self::try_allocate_in(capacity, AllocInit::Zeroed, alloc))
     }
 
     /// Converts the entire buffer into `Box<[MaybeUninit<T>]>` with the specified `len`.
@@ -179,35 +197,41 @@ impl<T, A: Allocator> RawVec<T, A> {
         }
     }
 
-    #[cfg(not(no_global_oom_handling))]
-    fn allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Self {
+    fn try_allocate_in(
+        capacity: usize,
+        init: AllocInit,
+        alloc: A,
+    ) -> Result<Self, TryReserveError> {
         // Don't allocate here because `Drop` will not deallocate when `capacity` is 0.
+
         if T::IS_ZST || capacity == 0 {
-            Self::new_in(alloc)
+            Ok(Self::new_in(alloc))
         } else {
             // We avoid `unwrap_or_else` here because it bloats the amount of
             // LLVM IR generated.
             let layout = match Layout::array::<T>(capacity) {
                 Ok(layout) => layout,
-                Err(_) => capacity_overflow(),
+                Err(_) => return Err(CapacityOverflow.into()),
             };
-            match alloc_guard(layout.size()) {
-                Ok(_) => {}
-                Err(_) => capacity_overflow(),
+
+            if let Err(err) = alloc_guard(layout.size()) {
+                return Err(err);
             }
+
             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 {
                 Ok(ptr) => ptr,
-                Err(_) => handle_alloc_error(layout),
+                Err(_) => return Err(AllocError { layout, non_exhaustive: () }.into()),
             };
 
             // Allocators currently return a `NonNull<[u8]>` whose length
             // matches the size requested. If that ever changes, the capacity
             // here should change to `ptr.len() / mem::size_of::<T>()`.
-            Self { ptr: Unique::from(ptr.cast()), cap: unsafe { Cap(capacity) }, alloc }
+            Ok(Self { ptr: Unique::from(ptr.cast()), cap: unsafe { Cap(capacity) }, alloc })
         }
     }
 
@@ -537,11 +561,11 @@ unsafe impl<#[may_dangle] T, A: Allocator> Drop for RawVec<T, A> {
 // Central function for reserve error handling.
 #[cfg(not(no_global_oom_handling))]
 #[inline]
-fn handle_reserve(result: Result<(), TryReserveError>) {
+fn handle_reserve<T>(result: Result<T, TryReserveError>) -> T {
     match result.map_err(|e| e.kind()) {
+        Ok(res) => res,
         Err(CapacityOverflow) => capacity_overflow(),
         Err(AllocError { layout, .. }) => handle_alloc_error(layout),
-        Ok(()) => { /* yay */ }
     }
 }
 
@@ -561,12 +585,3 @@ fn alloc_guard(alloc_size: usize) -> Result<(), TryReserveError> {
         Ok(())
     }
 }
-
-// One central function responsible for reporting capacity overflows. This'll
-// ensure that the code generation related to these panics is minimal as there's
-// only one location which panics rather than a bunch throughout the module.
-#[cfg(not(no_global_oom_handling))]
-#[cfg_attr(not(feature = "panic_immediate_abort"), inline(never))]
-fn capacity_overflow() -> ! {
-    panic!("capacity overflow");
-}
diff --git a/library/alloc/src/raw_vec/tests.rs b/library/alloc/src/raw_vec/tests.rs
index f8cada01c03..4194be53061 100644
--- a/library/alloc/src/raw_vec/tests.rs
+++ b/library/alloc/src/raw_vec/tests.rs
@@ -105,13 +105,14 @@ fn zst() {
     let v: RawVec<ZST> = RawVec::with_capacity_in(100, Global);
     zst_sanity(&v);
 
-    let v: RawVec<ZST> = RawVec::allocate_in(0, AllocInit::Uninitialized, Global);
+    let v: RawVec<ZST> = RawVec::try_allocate_in(0, AllocInit::Uninitialized, Global).unwrap();
     zst_sanity(&v);
 
-    let v: RawVec<ZST> = RawVec::allocate_in(100, AllocInit::Uninitialized, Global);
+    let v: RawVec<ZST> = RawVec::try_allocate_in(100, AllocInit::Uninitialized, Global).unwrap();
     zst_sanity(&v);
 
-    let mut v: RawVec<ZST> = RawVec::allocate_in(usize::MAX, AllocInit::Uninitialized, Global);
+    let mut v: RawVec<ZST> =
+        RawVec::try_allocate_in(usize::MAX, AllocInit::Uninitialized, Global).unwrap();
     zst_sanity(&v);
 
     // Check all these operations work as expected with zero-sized elements.
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 c0e934b3b1f..f2f42e63d6b 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 ed928994ad6..e8496989bcf 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/hygiene/panic-location.run.stderr b/tests/ui/hygiene/panic-location.run.stderr
index 5c552411da7..ec0ce18c3df 100644
--- a/tests/ui/hygiene/panic-location.run.stderr
+++ b/tests/ui/hygiene/panic-location.run.stderr
@@ -1,3 +1,3 @@
-thread 'main' panicked at library/alloc/src/raw_vec.rs:571:5:
+thread 'main' panicked at library/alloc/src/raw_vec.rs:26:5:
 capacity overflow
 note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
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
    |