[PATCH 2/3] rust: allocator: Vmalloc: Support alignments larger than PAGE_SIZE

[Hui Zhu <hui.zhu@linux.dev>]

From: Hui Zhu <zhuhui@kylinos.cn>

This commit add code to make rust alloc Vmalloc support alignments
larger than PAGE_SIZE.

It adds a new option element to ReallocFunc. When an object supports
aligned reallocation, it can register its alignment-specific realloc
function here.
During VREALLOC initialization, it sets bindings::vrealloc_align to
this element.
When ReallocFunc::call executes, if the object supports aligned
reallocation and the alignment exceeds PAGE_SIZE, the aligned realloc
function is used to support alignment capabilities.

Co-developed-by: Geliang Tang <geliang@kernel.org>
Signed-off-by: Geliang Tang <geliang@kernel.org>
Signed-off-by: Hui Zhu <zhuhui@kylinos.cn>
---
 rust/helpers/vmalloc.c         |  7 +++++++
 rust/kernel/alloc/allocator.rs | 32 ++++++++++++++++++++------------
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/rust/helpers/vmalloc.c b/rust/helpers/vmalloc.c
index 80d34501bbc0..3290c4c4c42f 100644
--- a/rust/helpers/vmalloc.c
+++ b/rust/helpers/vmalloc.c
@@ -7,3 +7,10 @@ rust_helper_vrealloc(const void *p, size_t size, gfp_t flags)
 {
 	return vrealloc(p, size, flags);
 }
+
+void * __must_check __realloc_size(2)
+rust_helper_vrealloc_align(const void *p, size_t size, size_t align,
+			   gfp_t flags)
+{
+	return vrealloc_align(p, size, align, flags);
+}
diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
index aa2dfa9dca4c..197222e15c26 100644
--- a/rust/kernel/alloc/allocator.rs
+++ b/rust/kernel/alloc/allocator.rs
@@ -59,17 +59,25 @@ fn aligned_size(new_layout: Layout) -> usize {
 /// One of the following: `krealloc`, `vrealloc`, `kvrealloc`.
 struct ReallocFunc(
     unsafe extern "C" fn(*const crate::ffi::c_void, usize, u32) -> *mut crate::ffi::c_void,
+    Option<
+        unsafe extern "C" fn(
+            *const crate::ffi::c_void,
+            usize,
+            usize,
+            u32,
+        ) -> *mut crate::ffi::c_void,
+    >,
 );
 
 impl ReallocFunc {
     // INVARIANT: `krealloc` satisfies the type invariants.
-    const KREALLOC: Self = Self(bindings::krealloc);
+    const KREALLOC: Self = Self(bindings::krealloc, None);
 
     // INVARIANT: `vrealloc` satisfies the type invariants.
-    const VREALLOC: Self = Self(bindings::vrealloc);
+    const VREALLOC: Self = Self(bindings::vrealloc, Some(bindings::vrealloc_align));
 
     // INVARIANT: `kvrealloc` satisfies the type invariants.
-    const KVREALLOC: Self = Self(bindings::kvrealloc);
+    const KVREALLOC: Self = Self(bindings::kvrealloc, None);
 
     /// # Safety
     ///
@@ -108,9 +116,15 @@ unsafe fn call(
         // GUARANTEE:
         // - `self.0` is one of `krealloc`, `vrealloc`, `kvrealloc`.
         // - Those functions provide the guarantees of this function.
-        let raw_ptr = unsafe {
-            // If `size == 0` and `ptr != NULL` the memory behind the pointer is freed.
-            self.0(ptr.cast(), size, flags.0).cast()
+        // If `size == 0` and `ptr != NULL` the memory behind the pointer is freed.
+        let raw_ptr = if let Some(f) = self.1 {
+            if layout.align() > bindings::PAGE_SIZE {
+                unsafe { f(ptr.cast(), size, layout.align(), flags.0).cast() }
+            } else {
+                unsafe { self.0(ptr.cast(), size, flags.0).cast() }
+            }
+        } else {
+            unsafe { self.0(ptr.cast(), size, flags.0).cast() }
         };
 
         let ptr = if size == 0 {
@@ -152,12 +166,6 @@ unsafe fn realloc(
         old_layout: Layout,
         flags: Flags,
     ) -> Result<NonNull<[u8]>, AllocError> {
-        // TODO: Support alignments larger than PAGE_SIZE.
-        if layout.align() > bindings::PAGE_SIZE {
-            pr_warn!("Vmalloc does not support alignments larger than PAGE_SIZE yet.\n");
-            return Err(AllocError);
-        }
-
         // SAFETY: If not `None`, `ptr` is guaranteed to point to valid memory, which was previously
         // allocated with this `Allocator`.
         unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags) }
-- 
2.43.0