iter.rs

  1// SPDX-License-Identifier: GPL-2.0
  2
  3use super::Vmalloc;
  4use crate::page;
  5use core::marker::PhantomData;
  6use core::ptr::NonNull;
  7
  8/// An [`Iterator`] of [`page::BorrowedPage`] items owned by a [`Vmalloc`] allocation.
  9///
 10/// # Guarantees
 11///
 12/// The pages iterated by the [`Iterator`] appear in the order as they are mapped in the CPU's
 13/// virtual address space ascendingly.
 14///
 15/// # Invariants
 16///
 17/// - `buf` is a valid and [`page::PAGE_SIZE`] aligned pointer into a [`Vmalloc`] allocation.
 18/// - `size` is the number of bytes from `buf` until the end of the [`Vmalloc`] allocation `buf`
 19///   points to.
 20pub struct VmallocPageIter<'a> {
 21    /// The base address of the [`Vmalloc`] buffer.
 22    buf: NonNull<u8>,
 23    /// The size of the buffer pointed to by `buf` in bytes.
 24    size: usize,
 25    /// The current page index of the [`Iterator`].
 26    index: usize,
 27    _p: PhantomData<page::BorrowedPage<'a>>,
 28}
 29
 30impl<'a> Iterator for VmallocPageIter<'a> {
 31    type Item = page::BorrowedPage<'a>;
 32
 33    fn next(&mut self) -> Option<Self::Item> {
 34        let offset = self.index.checked_mul(page::PAGE_SIZE)?;
 35
 36        // Even though `self.size()` may be smaller than `Self::page_count() * page::PAGE_SIZE`, it
 37        // is always a number between `(Self::page_count() - 1) * page::PAGE_SIZE` and
 38        // `Self::page_count() * page::PAGE_SIZE`, hence the check below is sufficient.
 39        if offset < self.size() {
 40            self.index += 1;
 41        } else {
 42            return None;
 43        }
 44
 45        // TODO: Use `NonNull::add()` instead, once the minimum supported compiler version is
 46        // bumped to 1.80 or later.
 47        //
 48        // SAFETY: `offset` is in the interval `[0, (self.page_count() - 1) * page::PAGE_SIZE]`,
 49        // hence the resulting pointer is guaranteed to be within the same allocation.
 50        let ptr = unsafe { self.buf.as_ptr().add(offset) };
 51
 52        // SAFETY: `ptr` is guaranteed to be non-null given that it is derived from `self.buf`.
 53        let ptr = unsafe { NonNull::new_unchecked(ptr) };
 54
 55        // SAFETY:
 56        // - `ptr` is a valid pointer to a `Vmalloc` allocation.
 57        // - `ptr` is valid for the duration of `'a`.
 58        Some(unsafe { Vmalloc::to_page(ptr) })
 59    }
 60
 61    fn size_hint(&self) -> (usize, Option<usize>) {
 62        let remaining = self.page_count().saturating_sub(self.index);
 63
 64        (remaining, Some(remaining))
 65    }
 66}
 67
 68impl<'a> VmallocPageIter<'a> {
 69    /// Creates a new [`VmallocPageIter`] instance.
 70    ///
 71    /// # Safety
 72    ///
 73    /// - `buf` must be a [`page::PAGE_SIZE`] aligned pointer into a [`Vmalloc`] allocation.
 74    /// - `buf` must be valid for at least the lifetime of `'a`.
 75    /// - `size` must be the number of bytes from `buf` until the end of the [`Vmalloc`] allocation
 76    ///   `buf` points to.
 77    pub unsafe fn new(buf: NonNull<u8>, size: usize) -> Self {
 78        // INVARIANT: By the safety requirements, `buf` is a valid and `page::PAGE_SIZE` aligned
 79        // pointer into a [`Vmalloc`] allocation.
 80        Self {
 81            buf,
 82            size,
 83            index: 0,
 84            _p: PhantomData,
 85        }
 86    }
 87
 88    /// Returns the size of the backing [`Vmalloc`] allocation in bytes.
 89    ///
 90    /// Note that this is the size the [`Vmalloc`] allocation has been allocated with. Hence, this
 91    /// number may be smaller than `[`Self::page_count`] * [`page::PAGE_SIZE`]`.
 92    #[inline]
 93    pub fn size(&self) -> usize {
 94        self.size
 95    }
 96
 97    /// Returns the number of pages owned by the backing [`Vmalloc`] allocation.
 98    #[inline]
 99    pub fn page_count(&self) -> usize {
100        self.size().div_ceil(page::PAGE_SIZE)
101    }
102}