[PATCH WIP 3/8] rust: alloc: implement AllocatorWithFlags trait

[Danilo Krummrich <dakr@redhat.com>]

Signed-off-by: Danilo Krummrich <dakr@redhat.com>
---
 rust/kernel/alloc.rs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index 9bc1b48b5641..ec45ba8f77ce 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -8,6 +8,14 @@
 pub mod box_ext;
 pub mod vec_ext;
 
+use core::{
+    alloc::{Allocator, AllocError, Layout},
+    ptr,
+    ptr::NonNull,
+};
+
+use flags::*;
+
 /// Flags to be used when allocating memory.
 ///
 /// They can be combined with the operators `|`, `&`, and `!`.
@@ -68,3 +76,16 @@ pub mod flags {
     /// small allocations.
     pub const GFP_NOWAIT: Flags = Flags(bindings::GFP_NOWAIT);
 }
+
+pub unsafe trait AllocatorWithFlags: Allocator {
+    unsafe fn alloc_flags(&self, layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError>;
+    unsafe fn realloc_flags(&self, ptr: *mut u8, old_size: usize, layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError>;
+
+    fn default_allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        unsafe { self.realloc_flags(ptr::null_mut(), 0, layout, GFP_KERNEL) }
+    }
+
+    fn default_allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        unsafe { self.realloc_flags(ptr::null_mut(), 0, layout, GFP_KERNEL | __GFP_ZERO) }
+    }
+}
-- 
2.44.0