Merge pull request #2 from SMS-Derfflinger/master

crates/eonix_mm/src/page_table/page_table.rs

@@ -87,6 +87,34 @@ where
     }
 
     pub fn iter_kernel(&self, range: VRange) -> impl Iterator<Item = &mut M::Entry> {
+        Self::iter_kernel_levels(self, range, M::LEVELS)
+    }
+
+    /// Iterates over the kernel space entries in the page table for the specified levels.
+    ///
+    /// # Parameters
+    /// - `range`: The virtual address range to iterate over.
+    /// - `levels`: A slice of `PageTableLevel` that specifies which levels of the page table
+    ///   should be included in the iteration. Each level corresponds to a level in the page
+    ///   table hierarchy, and the iterator will traverse entries at these levels.
+    ///
+    /// # Returns
+    /// An iterator over mutable references to the page table entries (`M::Entry`) within the
+    /// specified range and levels.
+    ///
+    /// # Example
+    /// ```
+    /// let range = VRange::new(0x1234000, 0x1300000);
+    /// let levels = &M::LEVELS[..2];
+    /// for pte in page_table.iter_kernel_levels(range, levels) {
+    ///     // Process each entry
+    /// }
+    /// ```
+    pub fn iter_kernel_levels(
+        &self,
+        range: VRange,
+        levels: &'static [PageTableLevel],
+    ) -> impl Iterator<Item = &mut M::Entry> {
         let alloc = self.root_table_page.allocator();
         let page_table_ptr = X::get_ptr_for_page(&self.root_table_page);
         let root_page_table = unsafe {
@@ -94,7 +122,12 @@ where
             M::RawTable::from_ptr(page_table_ptr)
         };
 
-        PageTableIterator::<M, A, X, KernelIterator>::new(root_page_table, range, alloc.clone())
+        PageTableIterator::<M, A, X, KernelIterator>::with_levels(
+            root_page_table,
+            range,
+            alloc.clone(),
+            levels,
+        )
     }
 
     fn drop_page_table_recursive(page_table: &Page<A>, levels: &[PageTableLevel]) {

crates/eonix_mm/src/page_table/paging_mode.rs

@@ -28,6 +28,10 @@ impl PageTableLevel {
         self.1
     }
 
+    pub const fn page_size(self) -> usize {
+        1 << self.nth_bit()
+    }
+
     pub const fn max_index(self) -> u16 {
         (1 << self.len()) - 1
     }

crates/eonix_mm/src/page_table/pte_iterator.rs

@@ -1,10 +1,10 @@
 use super::{
     pte::{RawAttribute, TableAttribute},
-    PagingMode, RawPageTable as _, PTE,
+    PageTableLevel, PagingMode, RawPageTable as _, PTE,
 };
 use crate::{
     address::{AddrOps as _, VRange},
-    paging::{Page, PageAccess, PageAlloc, PAGE_SIZE},
+    paging::{Page, PageAccess, PageAlloc},
 };
 use core::marker::PhantomData;
 
@@ -58,8 +58,12 @@ where
     M::Entry: 'a,
     A: PageAlloc,
     X: PageAccess,
-    K: IteratorType<M>,
 {
+    /// Specifies the hierarchy of page table levels to iterate over.
+    /// This field determines the sequence of levels in the page table
+    /// hierarchy that the iterator will traverse, starting from the
+    /// highest level and moving down to the lowest.
+    levels: &'static [PageTableLevel],
     remaining: usize,
 
     indicies: [u16; 8],
@@ -78,13 +82,11 @@ where
     K: IteratorType<M>,
 {
     fn parse_tables_starting_from(&mut self, idx_level: usize) {
-        let levels_len = M::LEVELS.len();
-
         for (idx, &pt_idx) in self
             .indicies
             .iter()
             .enumerate()
-            .take(levels_len - 1)
+            .take(self.levels.len() - 1)
             .skip(idx_level)
         {
             let [parent_table, child_table] = unsafe {
@@ -100,18 +102,30 @@ where
     }
 
     pub fn new(page_table: M::RawTable<'a>, range: VRange, alloc: A) -> Self {
+        Self::with_levels(page_table, range, alloc, M::LEVELS)
+    }
+
+    pub fn with_levels(
+        page_table: M::RawTable<'a>,
+        range: VRange,
+        alloc: A,
+        levels: &'static [PageTableLevel],
+    ) -> Self {
         let start = range.start().floor();
         let end = range.end().ceil();
 
+        let [.., last_level] = levels else { unreachable!() };
+
         let mut me = Self {
-            remaining: (end - start) / PAGE_SIZE,
+            levels,
+            remaining: (end - start) / last_level.page_size(),
             indicies: [0; 8],
             tables: [const { None }; 8],
             alloc,
             _phantom: PhantomData,
         };
 
-        for (i, level) in M::LEVELS.iter().enumerate() {
+        for (i, level) in levels.iter().enumerate() {
             me.indicies[i] = level.index_of(start);
         }
 
@@ -139,14 +153,14 @@ where
             self.remaining -= 1;
         }
 
-        let len_levels = M::LEVELS.len();
-
-        let retval = self.tables[len_levels - 1]
+        let table_level = self.levels.len() - 1;
+        let retval = self.tables[table_level]
             .as_mut()
             .unwrap()
-            .index_mut(self.indicies[len_levels - 1]);
+            .index_mut(self.indicies[table_level]);
 
-        let idx_level_start_updating = M::LEVELS
+        let idx_level_start_updating = self
+            .levels
             .iter()
             .zip(self.indicies.iter_mut())
             .enumerate()
@@ -157,7 +171,7 @@ where
             .expect("Index out of bounds");
 
         self.indicies[idx_level_start_updating] += 1;
-        self.indicies[idx_level_start_updating + 1..len_levels].fill(0);
+        self.indicies[idx_level_start_updating + 1..=table_level].fill(0);
         self.parse_tables_starting_from(idx_level_start_updating);
 
         Some(retval)