rust: alloc: introduce the `BoxExt` trait

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>
Reviewed-by: Boqun Feng <boqun.feng@gmail.com>
Reviewed-by: Benno Lossin <benno.lossin@proton.me>
Link: https://lore.kernel.org/r/20240328013603.206764-7-wedsonaf@gmail.com
[ Used `Box::write()` to avoid one `unsafe` block as suggested by Boqun. ]
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>

1 files changed, 58 insertions, 0 deletions

diff --git a/rust/kernel/alloc/box_ext.rs b/rust/kernel/alloc/box_ext.rs
new file mode 100644
index 000000000000..76653d6f4257
--- /dev/null
+++ b/rust/kernel/alloc/box_ext.rs
@@ -0,0 +1,58 @@
+// 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> {
+    fn new(x: T, flags: Flags) -> Result<Self, AllocError> {
+        let b = <Self as BoxExt<_>>::new_uninit(flags)?;
+        Ok(Box::write(b, x))
+    }
+
+    #[cfg(any(test, testlib))]
+    fn new_uninit(_flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError> {
+        Ok(Box::new_uninit())
+    }
+
+    #[cfg(not(any(test, testlib)))]
+    fn new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError> {
+        let ptr = if core::mem::size_of::<MaybeUninit<T>>() == 0 {
+            core::ptr::NonNull::<_>::dangling().as_ptr()
+        } else {
+            let layout = core::alloc::Layout::new::<MaybeUninit<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) };
+            if ptr.is_null() {
+                return Err(AllocError);
+            }
+
+            ptr.cast::<MaybeUninit<T>>()
+        };
+
+        // 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) })
+    }
+}