[PATCH WIP 8/8] rust: alloc: implement vmalloc allocator

[Danilo Krummrich <dakr@redhat.com>]

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

diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
index 173347ff018a..d3d0bbc41ab2 100644
--- a/rust/kernel/alloc/allocator.rs
+++ b/rust/kernel/alloc/allocator.rs
@@ -11,6 +11,7 @@
 use crate::alloc::AllocatorWithFlags;
 
 pub struct KernelAllocator;
+pub struct VmAllocator;
 
 pub(crate) fn aligned_size(new_layout: Layout) -> usize {
     // Customized layouts from `Layout::from_size_align()` can have size < align, so pad first.
@@ -123,6 +124,42 @@ unsafe fn realloc_flags(&self, ptr: *mut u8, _old_size: usize, layout: Layout, f
     }
 }
 
+unsafe impl Allocator for VmAllocator {
+    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 { bindings::vfree(ptr.as_ptr() as *const core::ffi::c_void) }
+    }
+}
+
+unsafe impl AllocatorWithFlags for VmAllocator {
+    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, src: *mut u8, old_size: usize, layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError>
+    {
+        let size = aligned_size(layout);
+
+        unsafe {
+            let dst = bindings::__vmalloc(size as u64, flags.0);
+            if !src.is_null() {
+                core::ptr::copy_nonoverlapping(src, dst as *mut u8, old_size);
+                bindings::vfree(src as *const core::ffi::c_void);
+            }
+            let ptr = NonNull::new(dst as *mut u8).ok_or(AllocError)?;
+            Ok(NonNull::slice_from_raw_parts(ptr, size))
+        }
+    }
+}
+
 #[global_allocator]
 static ALLOCATOR: KernelAllocator = KernelAllocator;
 
-- 
2.44.0