[PATCH v2 5/6] rust: alloc: kbox: implement PageOwner for VBox

[Danilo Krummrich <dakr@kernel.org>]

Implement the PageOwner trait for VBox; it allows to iterate and borrow
the backing pages of a VBox. This, for instance, is useful in
combination with VBox backing a scatterlist.

Signed-off-by: Danilo Krummrich <dakr@kernel.org>
---
 rust/kernel/alloc/kbox.rs | 42 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/rust/kernel/alloc/kbox.rs b/rust/kernel/alloc/kbox.rs
index 856d05aa60f1..bcca798d466b 100644
--- a/rust/kernel/alloc/kbox.rs
+++ b/rust/kernel/alloc/kbox.rs
@@ -18,6 +18,7 @@
 
 use crate::ffi::c_void;
 use crate::init::InPlaceInit;
+use crate::page;
 use crate::types::ForeignOwnable;
 use pin_init::{InPlaceWrite, Init, PinInit, ZeroableOption};
 
@@ -598,3 +599,44 @@ fn drop(&mut self) {
         unsafe { A::free(self.0.cast(), layout) };
     }
 }
+
+/// # Examples
+///
+/// ```
+/// # use kernel::{page, prelude::*};
+/// use kernel::page::PageOwner;
+///
+/// let mut vbox = VBox::<[u8; page::PAGE_SIZE]>::new_uninit(GFP_KERNEL)?;
+///
+/// let page = vbox.page_iter().next().expect("At least one page should be available.\n");
+///
+/// // SAFETY: There is no concurrent read or write to the same page.
+/// unsafe { page.fill_zero_raw(0, page::PAGE_SIZE)? };
+/// # Ok::<(), Error>(())
+/// ```
+impl<T> page::PageOwner for VBox<T> {
+    fn page_iter<'a>(&'a mut self) -> impl Iterator<Item = page::BorrowedPage<'a>> {
+        (0..self.page_count()).map(|idx| {
+            let ptr = self.0.as_ptr().cast::<u8>();
+
+            // TODO: Use `NonNull::add()` instead, once the minimum supported compiler version is
+            // bumped to 1.80 or later.
+            //
+            // SAFETY: `idx` is in the interval `[0, self.page_count())`, hence the resulting
+            // pointer is guaranteed to be within the same allocation.
+            let ptr = unsafe { ptr.add(idx * page::PAGE_SIZE) };
+
+            // SAFETY: `ptr` is guaranteed to be non-null given that it is derived from `self.0`.
+            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`.
+            unsafe { Vmalloc::to_page(ptr) }
+        })
+    }
+
+    fn page_count(&self) -> usize {
+        core::mem::size_of::<T>().div_ceil(page::PAGE_SIZE)
+    }
+}
-- 
2.50.1