aligned_size() dates back to when Rust did support kmalloc() only, but
is now used in ReallocFunc::call() and hence for all allocators.

However, the additional padding applied by aligned_size() is only
required by the kmalloc() allocator backend.

Hence, replace aligned_size() with Kmalloc::aligned_layout() and use it
for the affected allocators, i.e. kmalloc() and kvmalloc(), only.

While at it, make Kmalloc::aligned_layout() public, such that Rust
abstractions, which have to call subsystem specific kmalloc() based
allocation primitives directly, can make use of it.

Fixes: 8a799831fc63 ("rust: alloc: implement `ReallocFunc`")
Signed-off-by: Danilo Krummrich <dakr@kernel.org>
---
 rust/kernel/alloc/allocator.rs | 30 ++++++++++++++++++------------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
index aa2dfa9dca4c..3331ef338f3b 100644
--- a/rust/kernel/alloc/allocator.rs
+++ b/rust/kernel/alloc/allocator.rs
@@ -43,17 +43,6 @@
 /// For more details see [self].
 pub struct KVmalloc;
 
-/// Returns a proper size to alloc a new object aligned to `new_layout`'s alignment.
-fn aligned_size(new_layout: Layout) -> usize {
-    // Customized layouts from `Layout::from_size_align()` can have size < align, so pad first.
-    let layout = new_layout.pad_to_align();
-
-    // Note that `layout.size()` (after padding) is guaranteed to be a multiple of `layout.align()`
-    // which together with the slab guarantees means the `krealloc` will return a properly aligned
-    // object (see comments in `kmalloc()` for more information).
-    layout.size()
-}
-
 /// # Invariants
 ///
 /// One of the following: `krealloc`, `vrealloc`, `kvrealloc`.
@@ -88,7 +77,7 @@ unsafe fn call(
         old_layout: Layout,
         flags: Flags,
     ) -> Result<NonNull<[u8]>, AllocError> {
-        let size = aligned_size(layout);
+        let size = layout.size();
         let ptr = match ptr {
             Some(ptr) => {
                 if old_layout.size() == 0 {
@@ -123,6 +112,17 @@ unsafe fn call(
     }
 }
 
+impl Kmalloc {
+    /// Returns a [`Layout`] that makes [`Kmalloc`] fulfill the requested size and alignment of
+    /// `layout`.
+    pub const fn aligned_layout(layout: Layout) -> Layout {
+        // Note that `layout.size()` (after padding) is guaranteed to be a multiple of
+        // `layout.align()` which together with the slab guarantees means that `Kmalloc` will return
+        // a properly aligned object (see comments in `kmalloc()` for more information).
+        layout.pad_to_align()
+    }
+}
+
 // SAFETY: `realloc` delegates to `ReallocFunc::call`, which guarantees that
 // - memory remains valid until it is explicitly freed,
 // - passing a pointer to a valid memory allocation is OK,
@@ -135,6 +135,8 @@ unsafe fn realloc(
         old_layout: Layout,
         flags: Flags,
     ) -> Result<NonNull<[u8]>, AllocError> {
+        let layout = Kmalloc::aligned_layout(layout);
+
         // SAFETY: `ReallocFunc::call` has the same safety requirements as `Allocator::realloc`.
         unsafe { ReallocFunc::KREALLOC.call(ptr, layout, old_layout, flags) }
     }
@@ -176,6 +178,10 @@ unsafe fn realloc(
         old_layout: Layout,
         flags: Flags,
     ) -> Result<NonNull<[u8]>, AllocError> {
+        // `KVmalloc` may use the `Kmalloc` backend, hence we have to enforce a `Kmalloc`
+        // compatible layout.
+        let layout = Kmalloc::aligned_layout(layout);
+
         // TODO: Support alignments larger than PAGE_SIZE.
         if layout.align() > bindings::PAGE_SIZE {
             pr_warn!("KVmalloc does not support alignments larger than PAGE_SIZE yet.\n");
-- 
2.50.0

On Thu, Jul 31, 2025 at 05:48:06PM +0200, Danilo Krummrich wrote:
> aligned_size() dates back to when Rust did support kmalloc() only, but
> is now used in ReallocFunc::call() and hence for all allocators.
> 
> However, the additional padding applied by aligned_size() is only
> required by the kmalloc() allocator backend.
> 
> Hence, replace aligned_size() with Kmalloc::aligned_layout() and use it
> for the affected allocators, i.e. kmalloc() and kvmalloc(), only.
> 
> While at it, make Kmalloc::aligned_layout() public, such that Rust
> abstractions, which have to call subsystem specific kmalloc() based
> allocation primitives directly, can make use of it.
> 
> Fixes: 8a799831fc63 ("rust: alloc: implement `ReallocFunc`")
> Signed-off-by: Danilo Krummrich <dakr@kernel.org>

I guess vmalloc handles alignment in a different way ... ok makes sense
to me.

Reviewed-by: Alice Ryhl <aliceryhl@google.com>

On Thu, Jul 31, 2025 at 5:49 PM Danilo Krummrich <dakr@kernel.org> wrote:
>
> +impl Kmalloc {
> +    /// Returns a [`Layout`] that makes [`Kmalloc`] fulfill the requested size and alignment of
> +    /// `layout`.
> +    pub const fn aligned_layout(layout: Layout) -> Layout {

I think this `const fn` here was removed when applying to make it work
with older compilers, right?

I was fixing another `rusttest` thing and noticed while applying
these. I had a patch to fix it, since we can actually just use the
feature, and then I noticed it wasn't in the tree. Since I have it, I
am attaching it for reference in case the now-stable feature is
needed, e.g. if you want to make that `const fn` again.

Cheers,
Miguel

On Sat Aug 16, 2025 at 10:40 PM CEST, Miguel Ojeda wrote:
> On Thu, Jul 31, 2025 at 5:49 PM Danilo Krummrich <dakr@kernel.org> wrote:
>>
>> +impl Kmalloc {
>> +    /// Returns a [`Layout`] that makes [`Kmalloc`] fulfill the requested size and alignment of
>> +    /// `layout`.
>> +    pub const fn aligned_layout(layout: Layout) -> Layout {
>
> I think this `const fn` here was removed when applying to make it work
> with older compilers, right?

Yes, that's correct.

> I was fixing another `rusttest` thing and noticed while applying
> these. I had a patch to fix it, since we can actually just use the
> feature, and then I noticed it wasn't in the tree. Since I have it, I
> am attaching it for reference in case the now-stable feature is
> needed, e.g. if you want to make that `const fn` again.

I think it doesn't add much value for this to be a const function anyways, but
maybe it is more useful elsewhere? In that case, I think it also doesn't hurt
to Kmalloc::aligned_layout() const again.

On Thu, Jul 31, 2025 at 5:49 PM Danilo Krummrich <dakr@kernel.org> wrote:
>
> aligned_size() dates back to when Rust did support kmalloc() only, but
> is now used in ReallocFunc::call() and hence for all allocators.
>
> However, the additional padding applied by aligned_size() is only
> required by the kmalloc() allocator backend.
>
> Hence, replace aligned_size() with Kmalloc::aligned_layout() and use it
> for the affected allocators, i.e. kmalloc() and kvmalloc(), only.
>
> While at it, make Kmalloc::aligned_layout() public, such that Rust
> abstractions, which have to call subsystem specific kmalloc() based
> allocation primitives directly, can make use of it.
>
> Fixes: 8a799831fc63 ("rust: alloc: implement `ReallocFunc`")
> Signed-off-by: Danilo Krummrich <dakr@kernel.org>

Did this need Cc: stable or was it skipped since it is just extra padding?

(i.e. not sure what you usually do for DRM, but I was looking at this
since it is alloc since I would normally pick it.)

Thanks!

Cheers,
Miguel

On Tue Aug 12, 2025 at 9:52 PM CEST, Miguel Ojeda wrote:
> Did this need Cc: stable or was it skipped since it is just extra padding?

I don't think so, it just lead to pad to the alignment for Vmalloc too.

Technically, this makes no difference, since Vmalloc is always PAGE_SIZE aligned
and the size always a multiple of PAGE_SIZE.

However, the patch is a prerequisite for the DRM device fix in patch 2.

- Danilo

On Tue, Aug 12, 2025 at 10:00 PM Danilo Krummrich <dakr@kernel.org> wrote:
>
> I don't think so, it just lead to pad to the alignment for Vmalloc too.
>
> Technically, this makes no difference, since Vmalloc is always PAGE_SIZE aligned
> and the size always a multiple of PAGE_SIZE.

Got it, thanks for the quick reply! Then I guess we could have skipped
the Fixes in this one, but it is not a big deal and as usual it
depends on how one defines "fix".

Cheers,
Miguel

On Tue Aug 12, 2025 at 10:12 PM CEST, Miguel Ojeda wrote:
> On Tue, Aug 12, 2025 at 10:00 PM Danilo Krummrich <dakr@kernel.org> wrote:
>>
>> I don't think so, it just lead to pad to the alignment for Vmalloc too.
>>
>> Technically, this makes no difference, since Vmalloc is always PAGE_SIZE aligned
>> and the size always a multiple of PAGE_SIZE.
>
> Got it, thanks for the quick reply! Then I guess we could have skipped
> the Fixes in this one, but it is not a big deal and as usual it
> depends on how one defines "fix".

Yeah, in the past I was more on the "'Fixes:' for actual bugs only" side of
things, but I changed my mind a bit; I find it useful to have this as a
reference even for minor issues that might not be actual bugs, such as this
one.