[PATCH v2 06/10] rust: alloc: introduce the `BoxExt` trait

[Wedson Almeida Filho <wedsonaf@gmail.com>]

From: Wedson Almeida Filho <walmeida@microsoft.com>

Make fallible versions of `new` and `new_uninit` methods available in
`Box` even though it doesn't implement them because we build `alloc`
with the `no_global_oom_handling` config.

They also have an extra `flags` parameter that allows callers to pass
flags to the allocator.

Signed-off-by: Wedson Almeida Filho <walmeida@microsoft.com>
---
 rust/kernel/alloc.rs           |  1 +
 rust/kernel/alloc/allocator.rs |  6 +++-
 rust/kernel/alloc/box_ext.rs   | 62 ++++++++++++++++++++++++++++++++++
 rust/kernel/init.rs            | 13 +++----
 rust/kernel/prelude.rs         |  2 +-
 rust/kernel/sync/arc.rs        |  3 +-
 6 files changed, 78 insertions(+), 9 deletions(-)
 create mode 100644 rust/kernel/alloc/box_ext.rs

diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index 0502bcf00bce..a4c4b1f79bc9 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -5,6 +5,7 @@
 #[cfg(not(test))]
 #[cfg(not(testlib))]
 mod allocator;
+pub mod box_ext;
 pub mod vec_ext;
 
 /// Flags to be used when allocating memory.
diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
index 01ad139e19bc..fc0439455faa 100644
--- a/rust/kernel/alloc/allocator.rs
+++ b/rust/kernel/alloc/allocator.rs
@@ -15,7 +15,11 @@
 ///
 /// - `ptr` can be either null or a pointer which has been allocated by this allocator.
 /// - `new_layout` must have a non-zero size.
-unsafe fn krealloc_aligned(ptr: *mut u8, new_layout: Layout, flags: bindings::gfp_t) -> *mut u8 {
+pub(crate) unsafe fn krealloc_aligned(
+    ptr: *mut u8,
+    new_layout: Layout,
+    flags: bindings::gfp_t,
+) -> *mut u8 {
     // Customized layouts from `Layout::from_size_align()` can have size < align, so pad first.
     let layout = new_layout.pad_to_align();
 
diff --git a/rust/kernel/alloc/box_ext.rs b/rust/kernel/alloc/box_ext.rs
new file mode 100644
index 000000000000..aaa6a20d12b6
--- /dev/null
+++ b/rust/kernel/alloc/box_ext.rs
@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Extensions to [`Box`] for fallible allocations.
+
+use super::Flags;
+use alloc::boxed::Box;
+use core::alloc::AllocError;
+use core::mem::MaybeUninit;
+use core::result::Result;
+
+/// Extensions to [`Box`].
+pub trait BoxExt<T>: Sized {
+    /// Allocates a new box.
+    ///
+    /// The allocation may fail, in which case an error is returned.
+    fn new(x: T, flags: Flags) -> Result<Self, AllocError>;
+
+    /// Allocates a new uninitialised box.
+    ///
+    /// The allocation may fail, in which case an error is returned.
+    fn new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError>;
+}
+
+impl<T> BoxExt<T> for Box<T> {
+    #[cfg(any(test, testlib))]
+    fn new(x: T, _flags: Flags) -> Result<Self, AllocError> {
+        Ok(Box::new(x))
+    }
+
+    #[cfg(not(any(test, testlib)))]
+    fn new(x: T, flags: Flags) -> Result<Self, AllocError> {
+        let ptr = if core::mem::size_of::<T>() == 0 {
+            core::ptr::NonNull::<T>::dangling().as_ptr()
+        } else {
+            let layout = core::alloc::Layout::new::<T>();
+
+            // SAFETY: Memory is being allocated (first arg is null). The only other source of
+            // safety issues is sleeping on atomic context, which is addressed by klint. Lastly,
+            // the type is not a SZT (checked above).
+            let ptr = unsafe {
+                super::allocator::krealloc_aligned(core::ptr::null_mut(), layout, flags.0)
+            };
+            if ptr.is_null() {
+                return Err(AllocError);
+            }
+
+            let ptr = ptr.cast::<T>();
+
+            // SAFETY: We just allocated the memory above, it is valid for write.
+            unsafe { ptr.write(x) };
+            ptr
+        };
+
+        // SAFETY: For non-zero-sized types, we allocate above using the global allocator. For
+        // zero-sized types, we use `NonNull::dangling`.
+        Ok(unsafe { Box::from_raw(ptr) })
+    }
+
+    fn new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError> {
+        <Box<_> as BoxExt<_>>::new(MaybeUninit::<T>::uninit(), flags)
+    }
+}
diff --git a/rust/kernel/init.rs b/rust/kernel/init.rs
index 424257284d16..9c798cffc8e4 100644
--- a/rust/kernel/init.rs
+++ b/rust/kernel/init.rs
@@ -210,6 +210,7 @@
 //! [`pin_init!`]: crate::pin_init!
 
 use crate::{
+    alloc::{box_ext::BoxExt, flags::*},
     error::{self, Error},
     sync::UniqueArc,
     types::{Opaque, ScopeGuard},
@@ -305,9 +306,9 @@ macro_rules! stack_pin_init {
 ///
 /// stack_try_pin_init!(let foo: Result<Pin<&mut Foo>, AllocError> = pin_init!(Foo {
 ///     a <- new_mutex!(42),
-///     b: Box::try_new(Bar {
+///     b: Box::new(Bar {
 ///         x: 64,
-///     })?,
+///     }, GFP_KERNEL)?,
 /// }));
 /// let foo = foo.unwrap();
 /// pr_info!("a: {}", &*foo.a.lock());
@@ -331,9 +332,9 @@ macro_rules! stack_pin_init {
 ///
 /// stack_try_pin_init!(let foo: Pin<&mut Foo> =? pin_init!(Foo {
 ///     a <- new_mutex!(42),
-///     b: Box::try_new(Bar {
+///     b: Box::new(Bar {
 ///         x: 64,
-///     })?,
+///     }, GFP_KERNEL)?,
 /// }));
 /// pr_info!("a: {}", &*foo.a.lock());
 /// # Ok::<_, AllocError>(())
@@ -1158,7 +1159,7 @@ fn try_pin_init<E>(init: impl PinInit<T, E>) -> Result<Pin<Self>, E>
     where
         E: From<AllocError>,
     {
-        let mut this = Box::try_new_uninit()?;
+        let mut this = <Box<_> as BoxExt<_>>::new_uninit(GFP_KERNEL)?;
         let slot = this.as_mut_ptr();
         // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
         // slot is valid and will not be moved, because we pin it later.
@@ -1172,7 +1173,7 @@ fn try_init<E>(init: impl Init<T, E>) -> Result<Self, E>
     where
         E: From<AllocError>,
     {
-        let mut this = Box::try_new_uninit()?;
+        let mut this = <Box<_> as BoxExt<_>>::new_uninit(GFP_KERNEL)?;
         let slot = this.as_mut_ptr();
         // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
         // slot is valid.
diff --git a/rust/kernel/prelude.rs b/rust/kernel/prelude.rs
index 827e4dfd77df..b37a0b3180fb 100644
--- a/rust/kernel/prelude.rs
+++ b/rust/kernel/prelude.rs
@@ -14,7 +14,7 @@
 #[doc(no_inline)]
 pub use core::pin::Pin;
 
-pub use crate::alloc::{flags::*, vec_ext::VecExt};
+pub use crate::alloc::{box_ext::BoxExt, flags::*, vec_ext::VecExt};
 
 #[doc(no_inline)]
 pub use alloc::{boxed::Box, vec::Vec};
diff --git a/rust/kernel/sync/arc.rs b/rust/kernel/sync/arc.rs
index 7d4c4bf58388..1252a1b630ed 100644
--- a/rust/kernel/sync/arc.rs
+++ b/rust/kernel/sync/arc.rs
@@ -16,6 +16,7 @@
 //! [`Arc`]: https://doc.rust-lang.org/std/sync/struct.Arc.html
 
 use crate::{
+    alloc::{box_ext::BoxExt, flags::*},
     bindings,
     error::{self, Error},
     init::{self, InPlaceInit, Init, PinInit},
@@ -170,7 +171,7 @@ pub fn try_new(contents: T) -> Result<Self, AllocError> {
             data: contents,
         };
 
-        let inner = Box::try_new(value)?;
+        let inner = <Box<_> as BoxExt<_>>::new(value, GFP_KERNEL)?;
 
         // SAFETY: We just created `inner` with a reference count of 1, which is owned by the new
         // `Arc` object.
-- 
2.34.1