From 89ae3169c7ba848feee34af456a2b740ecd7e079 Mon Sep 17 00:00:00 2001
From: Kornel <kornel@geekhood.net>
Date: Tue, 30 Jan 2024 15:58:08 +0000
Subject: [PATCH] try_with_capacity for RawVec

---
 library/alloc/src/raw_vec.rs       | 36 +++++++++++++++++-------------
 library/alloc/src/raw_vec/tests.rs |  7 +++---
 2 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs
index dd8d6f6c7e634..2d86fa377f16d 100644
--- a/library/alloc/src/raw_vec.rs
+++ b/library/alloc/src/raw_vec.rs
@@ -17,7 +17,6 @@ use crate::collections::TryReserveErrorKind::*;
 #[cfg(test)]
 mod tests;
 
-#[cfg(not(no_global_oom_handling))]
 enum AllocInit {
     /// The contents of the new memory are uninitialized.
     Uninitialized,
@@ -93,6 +92,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 +105,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 +143,7 @@ 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 `with_capacity_zeroed`, but parameterized over the choice
@@ -150,7 +151,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 +180,40 @@ 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),
                 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 +543,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 */ }
     }
 }
 
diff --git a/library/alloc/src/raw_vec/tests.rs b/library/alloc/src/raw_vec/tests.rs
index f8cada01c0309..4194be530612d 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.