From: Hui Zhu <zhuhui@kylinos.cn>

Add a sample to the samples memory allocator usage.

Co-developed-by: Geliang Tang <geliang@kernel.org>
Signed-off-by: Geliang Tang <geliang@kernel.org>
Signed-off-by: Hui Zhu <zhuhui@kylinos.cn>
---
 samples/rust/Kconfig           |  10 ++++
 samples/rust/Makefile          |   1 +
 samples/rust/rust_allocator.rs | 104 +++++++++++++++++++++++++++++++++
 3 files changed, 115 insertions(+)
 create mode 100644 samples/rust/rust_allocator.rs

diff --git a/samples/rust/Kconfig b/samples/rust/Kconfig
index 7f7371a004ee..79c73f6c5216 100644
--- a/samples/rust/Kconfig
+++ b/samples/rust/Kconfig
@@ -105,6 +105,16 @@ config SAMPLE_RUST_DRIVER_AUXILIARY
 
 	  If unsure, say N.
 
+config SAMPLE_RUST_ALLOCATOR
+	tristate "Allocator Test Driver"
+	help
+	  This option builds the Rust allocator Test driver sample.
+
+	  To compile this as a module, choose M here:
+	  the module will be called rust_dma.
+
+	  If unsure, say N.
+
 config SAMPLE_RUST_HOSTPROGS
 	bool "Host programs"
 	help
diff --git a/samples/rust/Makefile b/samples/rust/Makefile
index bd2faad63b4f..b378959eab19 100644
--- a/samples/rust/Makefile
+++ b/samples/rust/Makefile
@@ -10,6 +10,7 @@ obj-$(CONFIG_SAMPLE_RUST_DRIVER_PLATFORM)	+= rust_driver_platform.o
 obj-$(CONFIG_SAMPLE_RUST_DRIVER_FAUX)		+= rust_driver_faux.o
 obj-$(CONFIG_SAMPLE_RUST_DRIVER_AUXILIARY)	+= rust_driver_auxiliary.o
 obj-$(CONFIG_SAMPLE_RUST_CONFIGFS)		+= rust_configfs.o
+obj-$(CONFIG_SAMPLE_RUST_ALLOCATOR)		+= rust_allocator.o
 
 rust_print-y := rust_print_main.o rust_print_events.o
 
diff --git a/samples/rust/rust_allocator.rs b/samples/rust/rust_allocator.rs
new file mode 100644
index 000000000000..13d23cc9d682
--- /dev/null
+++ b/samples/rust/rust_allocator.rs
@@ -0,0 +1,104 @@
+// SPDX-License-Identifier: GPL-2.0
+
+// Copyright (c) 2025, Kylin Software
+
+//! Rust allocator sample.
+
+use core::{alloc::Layout, ptr::NonNull};
+use kernel::alloc::allocator;
+use kernel::alloc::Allocator;
+use kernel::bindings;
+use kernel::prelude::*;
+
+module! {
+    type: RustAllocator,
+    name: "rust_allocator",
+    authors: ["Rust for Linux Contributors"],
+    description: "Rust allocator sample",
+    license: "GPL",
+}
+
+const VMALLOC_ARG: [(usize, usize); 2] = [
+    (bindings::PAGE_SIZE * 4, bindings::PAGE_SIZE * 2),
+    (1024, 128),
+];
+
+struct RustAllocator {
+    vmalloc_vec: KVec<(usize, Layout)>,
+}
+
+fn vmalloc_align(size: usize, align: usize) -> Result<(NonNull<[u8]>, Layout)> {
+    let layout = Layout::from_size_align(size, align).map_err(|_| EINVAL)?;
+
+    Ok((
+        <allocator::Vmalloc as Allocator>::alloc(layout, GFP_KERNEL).map_err(|_| EINVAL)?,
+        layout,
+    ))
+}
+
+fn vfree(addr: usize, layout: Layout) {
+    let vmalloc_ptr = NonNull::new(addr as *mut u8);
+    if let Some(ptr) = vmalloc_ptr {
+        unsafe {
+            <allocator::Vmalloc as Allocator>::free(ptr, layout);
+        }
+    } else {
+        pr_err!("Failed to vfree: pointer is null\n");
+    }
+}
+
+fn check_ptr(ptr: NonNull<[u8]>, size: usize, align: usize) -> (usize, bool) {
+    let current_size = unsafe { ptr.as_ref().len() };
+    if current_size != size {
+        pr_err!(
+            "The length to be allocated is {}, and the actually allocated memory length is {}.\n",
+            size,
+            current_size
+        );
+        return (0, false);
+    }
+
+    let addr = ptr.cast::<u8>().as_ptr() as usize;
+    debug_assert!(align.is_power_of_two());
+    if addr & (align - 1) != 0 {
+        pr_err!("Address {:#x} is not aligned with {:#x}.\n", addr, align);
+        return (0, false);
+    }
+
+    (addr, true)
+}
+
+fn clear_vmalloc_vec(v: &KVec<(usize, Layout)>) {
+    for (addr, layout) in v {
+        vfree(*addr, *layout);
+    }
+}
+
+impl kernel::Module for RustAllocator {
+    fn init(_module: &'static ThisModule) -> Result<Self> {
+        pr_info!("Rust allocator sample (init)\n");
+
+        let mut vmalloc_vec = KVec::new();
+        for (size, align) in VMALLOC_ARG {
+            let (ptr, layout) = vmalloc_align(size, align)?;
+
+            let (addr, is_ok) = check_ptr(ptr, size, align);
+            if !is_ok {
+                clear_vmalloc_vec(&vmalloc_vec);
+                return Err(EINVAL);
+            }
+
+            vmalloc_vec.push((addr, layout), GFP_KERNEL)?;
+        }
+
+        Ok(RustAllocator { vmalloc_vec })
+    }
+}
+
+impl Drop for RustAllocator {
+    fn drop(&mut self) {
+        pr_info!("Rust allocator sample (exit)\n");
+
+        clear_vmalloc_vec(&self.vmalloc_vec);
+    }
+}
-- 
2.43.0

On Tue Jul 15, 2025 at 11:59 AM CEST, Hui Zhu wrote:
> +impl kernel::Module for RustAllocator {
> +    fn init(_module: &'static ThisModule) -> Result<Self> {
> +        pr_info!("Rust allocator sample (init)\n");
> +
> +        let mut vmalloc_vec = KVec::new();
> +        for (size, align) in VMALLOC_ARG {
> +            let (ptr, layout) = vmalloc_align(size, align)?;

Ok, I think I get the idea, you want to demonstrate how to use the Allocator
trait for raw memory allocations.

However, doing so is discouraged unless there's really no other way. One obvious
example are Rust's own memory allocation primitives, such as Box and Vec.

So, instead of this raw allocation, you can just use VBox::new() or
VBox::new_uninit() in the following way.

	[repr(align(ALIGN))]
	struct Blob([u8; SIZE]);

	// Creates a vmalloc allocation of size `SIZE` with an alignment of
	// `ALIGN`. The allocation is freed once `b` is dropped.
	let b = VBox::<Blob>::new_uninit(GFP_KERNEL)?;

This way you don't have to handle the layout and the Allocator type yourself and
you also don't have to care about explicitly calling vfree(), VBox does all this
for you.

> +
> +            let (addr, is_ok) = check_ptr(ptr, size, align);
> +            if !is_ok {
> +                clear_vmalloc_vec(&vmalloc_vec);
> +                return Err(EINVAL);
> +            }
> +
> +            vmalloc_vec.push((addr, layout), GFP_KERNEL)?;
> +        }
> +
> +        Ok(RustAllocator { vmalloc_vec })
> +    }
> +}

Hi Danilo,

Thanks for your help.

On Tue, Jul 15, 2025 at 12:37:52PM +0200, Danilo Krummrich wrote:
> On Tue Jul 15, 2025 at 11:59 AM CEST, Hui Zhu wrote:
> > +impl kernel::Module for RustAllocator {
> > +    fn init(_module: &'static ThisModule) -> Result<Self> {
> > +        pr_info!("Rust allocator sample (init)\n");
> > +
> > +        let mut vmalloc_vec = KVec::new();
> > +        for (size, align) in VMALLOC_ARG {
> > +            let (ptr, layout) = vmalloc_align(size, align)?;
> 
> Ok, I think I get the idea, you want to demonstrate how to use the Allocator
> trait for raw memory allocations.
> 
> However, doing so is discouraged unless there's really no other way. One obvious
> example are Rust's own memory allocation primitives, such as Box and Vec.
> 
> So, instead of this raw allocation, you can just use VBox::new() or
> VBox::new_uninit() in the following way.
> 
> 	[repr(align(ALIGN))]
> 	struct Blob([u8; SIZE]);
> 
> 	// Creates a vmalloc allocation of size `SIZE` with an alignment of
> 	// `ALIGN`. The allocation is freed once `b` is dropped.
> 	let b = VBox::<Blob>::new_uninit(GFP_KERNEL)?;
> 
> This way you don't have to handle the layout and the Allocator type yourself and
> you also don't have to care about explicitly calling vfree(), VBox does all this
> for you.
> 
> > +
> > +            let (addr, is_ok) = check_ptr(ptr, size, align);
> > +            if !is_ok {
> > +                clear_vmalloc_vec(&vmalloc_vec);
> > +                return Err(EINVAL);
> > +            }
> > +
> > +            vmalloc_vec.push((addr, layout), GFP_KERNEL)?;
> > +        }
> > +
> > +        Ok(RustAllocator { vmalloc_vec })
> > +    }
> > +}
>

I sent version v2.
It only included the sample code and updated to use VBox according to your comments.

Best,
Hui