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