xref: /linux/rust/kernel/alloc/allocator/iter.rs (revision 07fdad3a93756b872da7b53647715c48d0f4a2d0)
1 // SPDX-License-Identifier: GPL-2.0
2 
3 use super::Vmalloc;
4 use crate::page;
5 use core::marker::PhantomData;
6 use 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.
20 pub 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 
30 impl<'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 
68 impl<'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 }
103