[PATCH WIP 5/8] rust: alloc: implement AllocatorWithFlags for KernelAllocator

[Danilo Krummrich <dakr@redhat.com>]

Signed-off-by: Danilo Krummrich <dakr@redhat.com>
---
 rust/kernel/alloc.rs           |  2 +-
 rust/kernel/alloc/allocator.rs | 38 ++++++++++++++++++++++++++++++++--
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index ec45ba8f77ce..12c6c4e423f5 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -4,7 +4,7 @@
 
 #[cfg(not(test))]
 #[cfg(not(testlib))]
-mod allocator;
+pub mod allocator;
 pub mod box_ext;
 pub mod vec_ext;
 
diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
index 236100649ae1..173347ff018a 100644
--- a/rust/kernel/alloc/allocator.rs
+++ b/rust/kernel/alloc/allocator.rs
@@ -3,12 +3,14 @@
 //! Allocator support.
 
 use super::{flags::*, Flags};
-use core::alloc::{GlobalAlloc, Layout};
+use core::alloc::{GlobalAlloc, Layout, Allocator, AllocError};
 use core::ptr;
+use core::ptr::NonNull;
 
 use crate::bindings;
+use crate::alloc::AllocatorWithFlags;
 
-struct KernelAllocator;
+pub struct KernelAllocator;
 
 pub(crate) fn aligned_size(new_layout: Layout) -> usize {
     // Customized layouts from `Layout::from_size_align()` can have size < align, so pad first.
@@ -89,6 +91,38 @@ unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
     }
 }
 
+unsafe impl Allocator for KernelAllocator {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        self.default_allocate(layout)
+    }
+
+    fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        self.default_allocate_zeroed(layout)
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        unsafe { self.dealloc(ptr.as_ptr(), layout) }
+    }
+}
+
+unsafe impl AllocatorWithFlags for KernelAllocator {
+    unsafe fn alloc_flags(&self, layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError>
+    {
+        unsafe { self.realloc_flags(ptr::null_mut(), 0, layout, flags) }
+    }
+
+    unsafe fn realloc_flags(&self, ptr: *mut u8, _old_size: usize, layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError>
+    {
+        let size = aligned_size(layout);
+
+        unsafe {
+            let raw_ptr = krealloc(ptr, size, flags);
+            let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
+            Ok(NonNull::slice_from_raw_parts(ptr, size))
+        }
+    }
+}
+
 #[global_allocator]
 static ALLOCATOR: KernelAllocator = KernelAllocator;
 
-- 
2.44.0