rust: add atomic allocated box

Signed-off-by: Andreas Hindborg <a.hindborg@samsung.com>

2 files changed, 30 insertions, 0 deletions

diff --git a/rust/kernel/box_ext.rs b/rust/kernel/box_ext.rs
new file mode 100644
index 000000000000..0bd3808e1d89
--- /dev/null
+++ b/rust/kernel/box_ext.rs
@@ -0,0 +1,29 @@
+use alloc::alloc::AllocError;
+use alloc::boxed::Box;
+use core::mem::MaybeUninit;
+
+use crate::bindings;
+
+pub trait BoxExt<T: ?Sized> {
+    fn try_new_atomic(x: T) -> Result<Self, AllocError>
+    where
+        Self: Sized;
+}
+
+#[cfg(not(test))]
+#[cfg(not(testlib))]
+impl<T> BoxExt<T> for Box<T> {
+    fn try_new_atomic(x: T) -> Result<Box<T>, AllocError> {
+        let layout = core::alloc::Layout::new::<core::mem::MaybeUninit<T>>();
+        let ptr = crate::allocator::ALLOCATOR
+            .allocate_with_flags(layout, bindings::GFP_ATOMIC)?
+            .cast();
+        let mut boxed: Box<MaybeUninit<T>> =
+            unsafe { Box::from_raw_in(ptr.as_ptr(), alloc::alloc::Global) };
+
+        unsafe {
+            boxed.as_mut_ptr().write(x);
+            Ok(boxed.assume_init())
+        }
+    }
+}
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 8673c792e5db..cc4551f27e11 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -33,6 +33,7 @@ extern crate self as kernel;
 #[cfg(not(test))]
 #[cfg(not(testlib))]
 mod allocator;
+pub mod box_ext;
 mod build_assert;
 pub mod device;
 pub mod dma;