Re: [PATCH v4 3/7] rust: alloc: implement VmallocPageIter

[Alice Ryhl <aliceryhl@google.com>]

On Thu, Aug 14, 2025 at 11:33:42AM +0200, Danilo Krummrich wrote:
> Introduce the VmallocPageIter type; an instance of VmallocPageIter may
> be exposed by owners of vmalloc allocations to provide borrowed access
> to the backing pages of the vmalloc allocation.
> 
> For instance, this is useful to access and borrow the backing pages of
> allocation primitives, such as Box and Vec, backing a scatterlist.
> 
> Suggested-by: Alice Ryhl <aliceryhl@google.com>
> Signed-off-by: Danilo Krummrich <dakr@kernel.org>

Some nits below, but:

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

>  rust/kernel/alloc/allocator.rs | 98 ++++++++++++++++++++++++++++++++++
>  1 file changed, 98 insertions(+)
> 
> diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
> index 2315f5063011..5cff3069bff4 100644
> --- a/rust/kernel/alloc/allocator.rs
> +++ b/rust/kernel/alloc/allocator.rs
> @@ -10,6 +10,7 @@
>  
>  use super::Flags;
>  use core::alloc::Layout;
> +use core::marker::PhantomData;
>  use core::ptr;
>  use core::ptr::NonNull;
>  
> @@ -236,3 +237,100 @@ unsafe fn realloc(
>          unsafe { ReallocFunc::KVREALLOC.call(ptr, layout, old_layout, flags) }
>      }
>  }
> +
> +/// An [`Iterator`] of [`page::BorrowedPage`] items owned by a [`Vmalloc`] allocation.
> +///
> +/// # Guarantees
> +///
> +/// The pages iterated by the [`Iterator`] appear in the order as they are mapped in the CPU's
> +/// virtual address space ascendingly.
> +///
> +/// # Invariants
> +///
> +/// - `buf` is a valid pointer to the beginning of a [`Vmalloc`] allocation.
> +/// - `size` is the size of the [`Vmalloc`] allocation `buf` points to.
> +pub struct VmallocPageIter<'a> {
> +    /// The base address of the [`Vmalloc`] buffer.
> +    buf: NonNull<u8>,
> +    /// The size of the buffer pointed to by `buf` in bytes.
> +    size: usize,
> +    /// The current page index of the [`Iterator`].
> +    index: usize,
> +    _p: PhantomData<page::BorrowedPage<'a>>,
> +}
> +
> +impl<'a> Iterator for VmallocPageIter<'a> {
> +    type Item = page::BorrowedPage<'a>;
> +
> +    fn next(&mut self) -> Option<Self::Item> {
> +        let offset = self.index.checked_mul(page::PAGE_SIZE)?;
> +
> +        // Even though `self.size()` may be smaller than `Self::page_count() * page::PAGE_SIZE`, it
> +        // is always a number between `(Self::page_count() - 1) * page::PAGE_SIZE` and
> +        // `Self::page_count() * page::PAGE_SIZE`, hence the check below is sufficient.
> +        if offset < self.size() {
> +            self.index += 1;
> +        } else {
> +            return None;
> +        }
> +
> +        // TODO: Use `NonNull::add()` instead, once the minimum supported compiler version is
> +        // bumped to 1.80 or later.
> +        //
> +        // SAFETY: `offset` is in the interval `[0, (self.page_count() - 1) * page::PAGE_SIZE]`,
> +        // hence the resulting pointer is guaranteed to be within the same allocation.
> +        let ptr = unsafe { self.buf.as_ptr().add(offset) };
> +
> +        // SAFETY: `ptr` is guaranteed to be non-null given that it is derived from `self.buf`.
> +        let ptr = unsafe { NonNull::new_unchecked(ptr) };
> +
> +        // SAFETY:
> +        // - `ptr` is a valid pointer to a `Vmalloc` allocation.
> +        // - `ptr` is valid for the duration of `'a`.
> +        Some(unsafe { Vmalloc::to_page(ptr) })
> +    }
> +
> +    fn size_hint(&self) -> (usize, Option<usize>) {
> +        let remaining = self.page_count().saturating_sub(self.index);
> +
> +        (remaining, Some(remaining))
> +    }
> +}
> +
> +impl<'a> VmallocPageIter<'a> {
> +    /// Creates a new [`VmallocPageIter`] instance.
> +    ///
> +    /// # Safety
> +    ///
> +    /// - `buf` must be a pointer to the beginning of a [`Vmalloc`] allocation.

It has to be the beginning of the vmalloc allocation?

> +    /// - `buf` points to must at least be valid for the lifetime of `'a`.

This grammar doesn't quite compute for me.

> +    /// - `size` must be the size of the [`Vmalloc`] allocation `buf` points to.
> +    pub unsafe fn new(buf: NonNull<u8>, size: usize) -> Self {
> +        // INVARIANT: `buf` is a valid pointer to the beginning of a `Vmalloc` allocation by the
> +        // safety requirements of this function.
> +        Self {
> +            buf,
> +            size,
> +            index: 0,
> +            _p: PhantomData,
> +        }
> +    }
> +
> +    /// Returns the base address of the backing [`Vmalloc`] allocation.
> +    pub fn base_address(&self) -> NonNull<u8> {
> +        self.buf
> +    }
> +
> +    /// Returns the size of the backing [`Vmalloc`] allocation in bytes.
> +    ///
> +    /// Note that this is the size the [`Vmalloc`] allocation has been allocated with. Hence, this
> +    /// number may be smaller than `[`Self::page_count`] * [`page::PAGE_SIZE`]`.
> +    pub fn size(&self) -> usize {

How about #[inline] on all these?

> +        self.size
> +    }
> +
> +    /// Returns the number of pages owned by the backing [`Vmalloc`] allocation.
> +    pub fn page_count(&self) -> usize {
> +        self.size().div_ceil(page::PAGE_SIZE)
> +    }
> +}
> -- 
> 2.50.1
>