rust: alloc: add some try_* methods we need

In preparation for enabling `no_global_oom_handling` for `alloc`,
we need to add some new methods.

They are all marked as:

    #[stable(feature = "kernel", since = "1.0.0")]

for easy identification.

Signed-off-by: Miguel Ojeda <ojeda@kernel.org>

Author: ojeda

rust/alloc/raw_vec.rs

@@ -20,7 +20,7 @@ use crate::collections::TryReserveError::{self, *};
 #[cfg(test)]
 mod tests;
 
-#[cfg(not(no_global_oom_handling))]
+#[allow(dead_code)]
 enum AllocInit {
     /// The contents of the new memory are uninitialized.
     Uninitialized,
@@ -93,6 +93,16 @@ impl<T> RawVec<T, Global> {
         Self::with_capacity_in(capacity, Global)
     }
 
+    /// Tries to create a `RawVec` (on the system heap) with exactly the
+    /// capacity and alignment requirements for a `[T; capacity]`. This is
+    /// equivalent to calling `RawVec::new` when `capacity` is `0` or `T` is
+    /// zero-sized. Note that if `T` is zero-sized this means you will
+    /// *not* get a `RawVec` with the requested capacity.
+    #[inline]
+    pub fn try_with_capacity(capacity: usize) -> Result<Self, TryReserveError> {
+        Self::try_with_capacity_in(capacity, Global)
+    }
+
     /// Like `with_capacity`, but guarantees the buffer is zeroed.
     #[cfg(not(no_global_oom_handling))]
     #[inline]
@@ -144,6 +154,13 @@ impl<T, A: Allocator> RawVec<T, A> {
         Self::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))]
@@ -218,6 +235,29 @@ impl<T, A: Allocator> RawVec<T, A> {
         }
     }
 
+    fn try_allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Result<Self, TryReserveError> {
+        if mem::size_of::<T>() == 0 {
+            return Ok(Self::new_in(alloc));
+        }
+
+        let layout = Layout::array::<T>(capacity)?;
+        alloc_guard(layout.size())?;
+        let result = match init {
+            AllocInit::Uninitialized => alloc.allocate(layout),
+            AllocInit::Zeroed => alloc.allocate_zeroed(layout),
+        };
+        let ptr = match result {
+            Ok(ptr) => ptr,
+            Err(_) => return Err(TryReserveError::AllocError { layout, non_exhaustive: () }),
+        };
+
+        Ok(Self {
+            ptr: unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) },
+            cap: Self::capacity_from_bytes(ptr.len()),
+            alloc,
+        })
+    }
+
     /// Reconstitutes a `RawVec` from a pointer, capacity, and allocator.
     ///
     /// # Safety
@@ -394,6 +434,16 @@ impl<T, A: Allocator> RawVec<T, A> {
     pub fn shrink_to_fit(&mut self, amount: usize) {
         handle_reserve(self.shrink(amount));
     }
+
+    /// Tries to shrink the allocation down to the specified amount. If the given amount
+    /// is 0, actually completely deallocates.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the given amount is *larger* than the current capacity.
+    pub fn try_shrink_to_fit(&mut self, amount: usize) -> Result<(), TryReserveError> {
+        self.shrink(amount)
+    }
 }
 
 impl<T, A: Allocator> RawVec<T, A> {
@@ -465,7 +515,6 @@ impl<T, A: Allocator> RawVec<T, A> {
         Ok(())
     }
 
-    #[cfg(not(no_global_oom_handling))]
     fn shrink(&mut self, amount: usize) -> Result<(), TryReserveError> {
         assert!(amount <= self.capacity(), "Tried to shrink to a larger capacity");
 

rust/alloc/slice.rs

@@ -95,11 +95,11 @@ use core::mem::size_of;
 use core::ptr;
 
 use crate::alloc::Allocator;
-#[cfg(not(no_global_oom_handling))]
 use crate::alloc::Global;
 #[cfg(not(no_global_oom_handling))]
 use crate::borrow::ToOwned;
 use crate::boxed::Box;
+use crate::collections::TryReserveError;
 use crate::vec::Vec;
 
 #[unstable(feature = "slice_range", issue = "76393")]
@@ -155,6 +155,7 @@ mod hack {
     use core::alloc::Allocator;
 
     use crate::boxed::Box;
+    use crate::collections::TryReserveError;
     use crate::vec::Vec;
 
     // We shouldn't add inline attribute to this since this is used in
@@ -174,13 +175,24 @@ mod hack {
         T::to_vec(s, alloc)
     }
 
+    #[inline]
+    pub fn try_to_vec<T: TryConvertVec, A: Allocator>(s: &[T], alloc: A) -> Result<Vec<T, A>, TryReserveError> {
+        T::try_to_vec(s, alloc)
+    }
+
     #[cfg(not(no_global_oom_handling))]
     pub trait ConvertVec {
         fn to_vec<A: Allocator>(s: &[Self], alloc: A) -> Vec<Self, A>
         where
             Self: Sized;
     }
 
+    pub trait TryConvertVec {
+        fn try_to_vec<A: Allocator>(s: &[Self], alloc: A) -> Result<Vec<Self, A>, TryReserveError>
+        where
+            Self: Sized;
+    }
+
     #[cfg(not(no_global_oom_handling))]
     impl<T: Clone> ConvertVec for T {
         #[inline]
@@ -233,6 +245,42 @@ mod hack {
             v
         }
     }
+
+    impl<T: Clone> TryConvertVec for T {
+        #[inline]
+        default fn try_to_vec<A: Allocator>(s: &[Self], alloc: A) -> Result<Vec<Self, A>, TryReserveError> {
+            struct DropGuard<'a, T, A: Allocator> {
+                vec: &'a mut Vec<T, A>,
+                num_init: usize,
+            }
+            impl<'a, T, A: Allocator> Drop for DropGuard<'a, T, A> {
+                #[inline]
+                fn drop(&mut self) {
+                    // SAFETY:
+                    // items were marked initialized in the loop below
+                    unsafe {
+                        self.vec.set_len(self.num_init);
+                    }
+                }
+            }
+            let mut vec = Vec::try_with_capacity_in(s.len(), alloc)?;
+            let mut guard = DropGuard { vec: &mut vec, num_init: 0 };
+            let slots = guard.vec.spare_capacity_mut();
+            // .take(slots.len()) is necessary for LLVM to remove bounds checks
+            // and has better codegen than zip.
+            for (i, b) in s.iter().enumerate().take(slots.len()) {
+                guard.num_init = i;
+                slots[i].write(b.clone());
+            }
+            core::mem::forget(guard);
+            // SAFETY:
+            // the vec was allocated and initialized above to at least this length.
+            unsafe {
+                vec.set_len(s.len());
+            }
+            Ok(vec)
+        }
+    }
 }
 
 #[lang = "slice_alloc"]
@@ -472,6 +520,24 @@ impl<T> [T] {
         self.to_vec_in(Global)
     }
 
+    /// Tries to copy `self` into a new `Vec`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let s = [10, 40, 30];
+    /// let x = s.try_to_vec().unwrap();
+    /// // Here, `s` and `x` can be modified independently.
+    /// ```
+    #[inline]
+    #[stable(feature = "kernel", since = "1.0.0")]
+    pub fn try_to_vec(&self) -> Result<Vec<T>, TryReserveError>
+    where
+        T: Clone,
+    {
+        self.try_to_vec_in(Global)
+    }
+
     /// Copies `self` into a new `Vec` with an allocator.
     ///
     /// # Examples
@@ -496,6 +562,29 @@ impl<T> [T] {
         hack::to_vec(self, alloc)
     }
 
+    /// Tries to copy `self` into a new `Vec` with an allocator.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(allocator_api)]
+    ///
+    /// use std::alloc::System;
+    ///
+    /// let s = [10, 40, 30];
+    /// let x = s.try_to_vec_in(System).unwrap();
+    /// // Here, `s` and `x` can be modified independently.
+    /// ```
+    #[inline]
+    #[stable(feature = "kernel", since = "1.0.0")]
+    pub fn try_to_vec_in<A: Allocator>(&self, alloc: A) -> Result<Vec<T, A>, TryReserveError>
+    where
+        T: Clone,
+    {
+        // N.B., see the `hack` module in this file for more details.
+        hack::try_to_vec(self, alloc)
+    }
+
     /// Converts `self` into a vector without clones or allocation.
     ///
     /// The resulting vector can be converted back into a box via

rust/alloc/str.rs

@@ -38,6 +38,7 @@ use core::unicode::conversions;
 
 use crate::borrow::ToOwned;
 use crate::boxed::Box;
+use crate::collections::TryReserveError;
 use crate::slice::{Concat, Join, SliceIndex};
 use crate::string::String;
 use crate::vec::Vec;
@@ -575,6 +576,22 @@ impl str {
         // make_ascii_lowercase() preserves the UTF-8 invariant.
         unsafe { String::from_utf8_unchecked(bytes) }
     }
+
+    /// Tries to create a `String`.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// let s: &str = "a";
+    /// let ss: String = s.try_to_owned().unwrap();
+    /// ```
+    #[inline]
+    #[stable(feature = "kernel", since = "1.0.0")]
+    pub fn try_to_owned(&self) -> Result<String, TryReserveError> {
+        unsafe { Ok(String::from_utf8_unchecked(self.as_bytes().try_to_vec()?)) }
+    }
 }
 
 /// Converts a boxed slice of bytes to a boxed string slice without checking

rust/alloc/sync.rs

@@ -36,10 +36,10 @@ use crate::alloc::{box_free, WriteCloneIntoRaw};
 use crate::alloc::{AllocError, Allocator, Global, Layout};
 use crate::borrow::{Cow, ToOwned};
 use crate::boxed::Box;
+use crate::collections::TryReserveError;
 use crate::rc::is_dangling;
 #[cfg(not(no_global_oom_handling))]
 use crate::string::String;
-#[cfg(not(no_global_oom_handling))]
 use crate::vec::Vec;
 
 #[cfg(test)]
@@ -1188,6 +1188,18 @@ impl<T> Arc<[T]> {
         }
     }
 
+    /// Tries to allocate an `ArcInner<[T]>` with the given length.
+    unsafe fn try_allocate_for_slice(len: usize) -> Result<*mut ArcInner<[T]>, TryReserveError> {
+        unsafe {
+            let layout = Layout::array::<T>(len)?;
+            Self::try_allocate_for_layout(
+                layout,
+                |l| Global.allocate(l),
+                |mem| ptr::slice_from_raw_parts_mut(mem as *mut T, len) as *mut ArcInner<[T]>,
+            ).map_err(|_| TryReserveError::AllocError { layout, non_exhaustive: () })
+        }
+    }
+
     /// Copy elements from slice into newly allocated Arc<\[T\]>
     ///
     /// Unsafe because the caller must either take ownership or bind `T: Copy`.
@@ -1202,6 +1214,19 @@ impl<T> Arc<[T]> {
         }
     }
 
+    /// Tries to copy elements from slice into newly allocated Arc<\[T\]>
+    ///
+    /// Unsafe because the caller must either take ownership or bind `T: Copy`.
+    unsafe fn try_copy_from_slice(v: &[T]) -> Result<Arc<[T]>, TryReserveError> {
+        unsafe {
+            let ptr = Self::try_allocate_for_slice(v.len())?;
+
+            ptr::copy_nonoverlapping(v.as_ptr(), &mut (*ptr).data as *mut [T] as *mut T, v.len());
+
+            Ok(Self::from_ptr(ptr))
+        }
+    }
+
     /// Constructs an `Arc<[T]>` from an iterator known to be of a certain size.
     ///
     /// Behavior is undefined should the size be wrong.
@@ -2415,6 +2440,32 @@ impl<T> From<Vec<T>> for Arc<[T]> {
     }
 }
 
+// Avoid `error: specializing impl repeats parameter` implementing `TryFrom`.
+impl<T> Arc<[T]> {
+    /// Tries to allocate a reference-counted slice and move `v`'s items into it.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use std::sync::Arc;
+    /// let unique: Vec<i32> = vec![1, 2, 3];
+    /// let shared: Arc<[i32]> = Arc::try_from(unique).unwrap();
+    /// assert_eq!(&[1, 2, 3], &shared[..]);
+    /// ```
+    #[stable(feature = "kernel", since = "1.0.0")]
+    #[inline]
+    pub fn try_from_vec(mut v: Vec<T>) -> Result<Self, TryReserveError> {
+        unsafe {
+            let arc = Arc::try_copy_from_slice(&v)?;
+
+            // Allow the Vec to free its memory, but not destroy its contents
+            v.set_len(0);
+
+            Ok(arc)
+        }
+    }
+}
+
 #[stable(feature = "shared_from_cow", since = "1.45.0")]
 impl<'a, B> From<Cow<'a, B>> for Arc<B>
 where