rust: add fallible allocation methods

These are all fallible alternatives to all currently used methods that
may perform infallible allocation.

`assume_fallible` is a semantically marker function to denote that we
have proved that the closure passed to it will not perform infallible
allocation.

Signed-off-by: Gary Guo <[email protected]>

Author: nbdd0121

rust/kernel/io_buffer.rs

@@ -2,6 +2,7 @@
 
 //! Buffers used in IO.
 
+use crate::traits::VecExt;
 use crate::Result;
 use alloc::vec::Vec;
 use core::mem::{size_of, MaybeUninit};
@@ -29,9 +30,8 @@ pub trait IoBufferReader {
     ///
     /// Returns `EFAULT` if the address does not currently point to mapped, readable memory.
     fn read_all(&mut self) -> Result<Vec<u8>> {
-        let mut data = Vec::<u8>::new();
-        data.try_reserve_exact(self.len())?;
-        data.resize(self.len(), 0);
+        let mut data = Vec::try_with_capacity(self.len())?;
+        data.try_resize(self.len(), 0)?;
 
         // SAFETY: The output buffer is valid as we just allocated it.
         unsafe { self.read_raw(data.as_mut_ptr(), data.len())? };

rust/kernel/module_param.rs

@@ -5,6 +5,7 @@
 //! C header: [`include/linux/moduleparam.h`](../../../include/linux/moduleparam.h)
 
 use crate::str::CStr;
+use crate::traits::TryToOwned;
 use core::fmt::Write;
 
 /// Types that can be used for module parameters.
@@ -475,9 +476,7 @@ impl ModuleParam for StringParam {
         let slab_available = unsafe { crate::bindings::slab_is_available() };
         arg.and_then(|arg| {
             if slab_available {
-                let mut vec = alloc::vec::Vec::new();
-                vec.try_reserve_exact(arg.len()).ok()?;
-                vec.extend_from_slice(arg);
+                let vec = arg.try_to_owned().ok()?;
                 Some(StringParam::Owned(vec))
             } else {
                 Some(StringParam::Ref(arg))

rust/kernel/prelude.rs

@@ -23,4 +23,4 @@ pub use super::static_assert;
 
 pub use super::{KernelModule, Result};
 
-pub use crate::traits::TryPin;
+pub use crate::traits::{TryPin, TryToOwned, VecExt as _};

rust/kernel/traits.rs

@@ -4,8 +4,19 @@
 
 use core::{ops::Deref, pin::Pin};
 
+use alloc::collections::TryReserveError;
+use alloc::string::String;
+use alloc::vec::Vec;
 use alloc::{alloc::AllocError, sync::Arc};
 
+type AllocResult<T = ()> = core::result::Result<T, AllocError>;
+type CollectionResult<T = ()> = core::result::Result<T, TryReserveError>;
+
+#[inline]
+fn assume_fallible<T, F: FnOnce() -> T>(f: F) -> T {
+    f()
+}
+
 /// Trait which provides a fallible version of `pin()` for pointer types.
 ///
 /// Common pointer types which implement a `pin()` method include [`Box`](alloc::boxed::Box) and [`Arc`].
@@ -24,3 +35,97 @@ impl<T> TryPin<Arc<T>> for Arc<T> {
         Ok(unsafe { Pin::new_unchecked(Arc::try_new(data)?) })
     }
 }
+
+pub trait TryToOwned {
+    type Owned: core::borrow::Borrow<Self>;
+
+    #[must_use = "cloning is often expensive and is not expected to have side effects"]
+    fn try_to_owned(&self) -> AllocResult<Self::Owned>;
+
+    fn try_clone_into(&self, target: &mut Self::Owned) -> AllocResult {
+        *target = self.try_to_owned()?;
+        Ok(())
+    }
+}
+
+impl<T: Clone> TryToOwned for [T] {
+    type Owned = Vec<T>;
+
+    fn try_to_owned(&self) -> AllocResult<Vec<T>> {
+        let mut vec = Vec::new();
+        self.try_clone_into(&mut vec)?;
+        Ok(vec)
+    }
+
+    fn try_clone_into(&self, target: &mut Vec<T>) -> AllocResult {
+        // Ensure target has enough capacity
+        target
+            .try_reserve_exact(self.len().saturating_sub(target.len()))
+            .map_err(|_| AllocError)?;
+
+        target.clear();
+        assume_fallible(|| target.extend_from_slice(self));
+        Ok(())
+    }
+}
+
+impl TryToOwned for str {
+    type Owned = String;
+
+    fn try_to_owned(&self) -> AllocResult<String> {
+        let mut vec = String::new();
+        self.try_clone_into(&mut vec)?;
+        Ok(vec)
+    }
+
+    fn try_clone_into(&self, target: &mut String) -> AllocResult {
+        // Ensure target has enough capacity
+        target
+            .try_reserve_exact(self.len().saturating_sub(target.len()))
+            .map_err(|_| AllocError)?;
+
+        target.clear();
+        assume_fallible(|| target.push_str(self));
+        Ok(())
+    }
+}
+
+pub trait VecExt<T> {
+    fn try_with_capacity(capacity: usize) -> CollectionResult<Self>
+    where
+        Self: Sized;
+
+    fn try_extend_from_slice(&mut self, other: &[T]) -> CollectionResult
+    where
+        T: Clone;
+
+    fn try_resize(&mut self, new_len: usize, value: T) -> CollectionResult
+    where
+        T: Clone;
+}
+
+impl<T> VecExt<T> for alloc::vec::Vec<T> {
+    fn try_with_capacity(capacity: usize) -> CollectionResult<Self> {
+        let mut vec = Self::new();
+        vec.try_reserve_exact(capacity)?;
+        Ok(vec)
+    }
+
+    fn try_extend_from_slice(&mut self, other: &[T]) -> CollectionResult
+    where
+        T: Clone,
+    {
+        self.try_reserve(other.len())?;
+        assume_fallible(|| self.extend_from_slice(other));
+        Ok(())
+    }
+
+    fn try_resize(&mut self, new_len: usize, value: T) -> CollectionResult
+    where
+        T: Clone,
+    {
+        self.try_reserve(new_len.saturating_sub(self.len()))?;
+        assume_fallible(|| self.resize(new_len, value));
+        Ok(())
+    }
+}

samples/rust/rust_minimal.rs

@@ -25,7 +25,7 @@ impl KernelModule for RustMinimal {
         pr_info!("Am I built-in? {}\n", !cfg!(MODULE));
 
         Ok(RustMinimal {
-            message: "on the heap!".to_owned(),
+            message: "on the heap!".try_to_owned()?,
         })
     }
 }