[PATCH v3 3/7] rust: alloc: implement VmallocPageIter

[Danilo Krummrich <dakr@kernel.org>]

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>
---
 rust/kernel/alloc/allocator.rs | 92 ++++++++++++++++++++++++++++++++++
 1 file changed, 92 insertions(+)

diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
index 2315f5063011..3ed2d5ff1e37 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,94 @@ 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<&'a u8>,
+}
+
+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) })
+    }
+}
+
+impl<'a> VmallocPageIter<'a> {
+    /// Creates a new [`VmallocPageIter`] instance.
+    ///
+    /// # Safety
+    ///
+    /// - `buf` must be a pointer to the beginning of a [`Vmalloc`] allocation.
+    /// - `buf` points to must at least be valid for the lifetime of `'a`.
+    /// - `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 {
+        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