change(mm): replace enum with &[PageTableLevel]

crates/eonix_mm/src/page_table/page_table.rs

@@ -7,7 +7,7 @@ use super::{
 use crate::{
     address::{PAddr, VRange},
     page_table::PageTableIterator,
-    paging::{GlobalPageAlloc, Page, PageAccess, PageAlloc, PageBlock, PageSize},
+    paging::{GlobalPageAlloc, Page, PageAccess, PageAlloc, PageBlock},
 };
 use core::{marker::PhantomData, ptr::NonNull};
 
@@ -83,11 +83,14 @@ where
             M::RawTable::from_ptr(page_table_ptr)
         };
 
-        // default 3 level
-        PageTableIterator::<M, A, X, UserIterator>::new(root_page_table, range, alloc.clone(), PageSize::_4KbPage)
+        PageTableIterator::<M, A, X, UserIterator>::new(root_page_table, range, alloc.clone())
     }
 
-    pub fn iter_kernel(&self, range: VRange, page_size: PageSize) -> impl Iterator<Item = &mut M::Entry> {
+    pub fn iter_kernel(&self, range: VRange) -> impl Iterator<Item = &mut M::Entry> {
+        Self::iter_kernel_levels(self, range, M::LEVELS)
+    }
+
+    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 {
@@ -95,7 +98,7 @@ where
             M::RawTable::from_ptr(page_table_ptr)
         };
 
-        PageTableIterator::<M, A, X, KernelIterator>::new(root_page_table, range, alloc.clone(), page_size)
+        PageTableIterator::<M, A, X, KernelIterator>::new_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/pte_iterator.rs

@@ -1,10 +1,9 @@
 use super::{
-    pte::{RawAttribute, TableAttribute},
-    PagingMode, RawPageTable as _, PTE,
+    pte::{RawAttribute, TableAttribute}, PageTableLevel, PagingMode, RawPageTable as _, PTE
 };
 use crate::{
     address::{AddrOps as _, VRange},
-    paging::{Page, PageAccess, PageAlloc, PageSize},
+    paging::{Page, PageAccess, PageAlloc, PAGE_SIZE},
 };
 use core::{marker::PhantomData};
 
@@ -60,7 +59,7 @@ where
     X: PageAccess,
     K: IteratorType<M>,
 {
-    page_size: PageSize,
+    levels: &'static [PageTableLevel],
     remaining: usize,
 
     indicies: [u16; 8],
@@ -78,22 +77,14 @@ where
     X: PageAccess,
     K: IteratorType<M>,
 {
-    fn table_level(&self) -> usize {
-        match self.page_size {
-            PageSize::_4KbPage => 3,
-            PageSize::_2MbPage => 2,
-            PageSize::_1GbPage => 1,
-        }
-    }
 
     fn parse_tables_starting_from(&mut self, idx_level: usize) {
-        let table_level = self.table_level();
 
         for (idx, &pt_idx) in self
             .indicies
             .iter()
             .enumerate()
-            .take(table_level)
+            .take(self.levels.len() - 1)
             .skip(idx_level)
         {
             let [parent_table, child_table] = unsafe {
@@ -108,20 +99,24 @@ where
         }
     }
 
-    pub fn new(page_table: M::RawTable<'a>, range: VRange, alloc: A, page_size: PageSize) -> Self {
+    pub fn new(page_table: M::RawTable<'a>, range: VRange, alloc: A) -> Self {
+        Self::new_levels(page_table, range, alloc, M::LEVELS)
+    }
+
+    pub fn new_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 mut me = Self {
-            page_size,
-            remaining: (end - start) / (page_size as usize),
+            levels,
+            remaining: (end - start) / 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);
         }
 
@@ -146,29 +141,27 @@ where
         if self.remaining == 0 {
             return None;
         } else {
-            self.remaining -= 1;
+            self.remaining -= (1 << (self.levels.last()?.nth_bit())) / PAGE_SIZE;
         }
 
-        let table_level = self.table_level();
-        let retval = self.tables[self.table_level()]
+        let table_level = self.levels.len() - 1;
+        let retval = self.tables[table_level]
             .as_mut()
             .unwrap()
             .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()
             .rev()
-            .skip_while(|(i, (level, idx))| {
-                *i >= table_level && **idx == level.max_index()
-            })
+            .skip_while(|(_, (level, idx))| **idx == level.max_index())
             .map(|(i, _)| i)
             .next()
             .expect("Index out of bounds");
 
         self.indicies[idx_level_start_updating] += 1;
-        self.indicies[idx_level_start_updating + 1..table_level].fill(0);
+        self.indicies[idx_level_start_updating + 1..=table_level].fill(0);
         self.parse_tables_starting_from(idx_level_start_updating);
 
         Some(retval)

crates/eonix_mm/src/paging.rs

@@ -3,7 +3,7 @@ mod page_alloc;
 mod pfn;
 mod raw_page;
 
-pub use page::{Page, PageAccess, PageBlock, PAGE_SIZE, PageSize, PAGE_SIZE_BITS};
+pub use page::{Page, PageAccess, PageBlock, PAGE_SIZE, PAGE_SIZE_BITS};
 pub use page_alloc::{GlobalPageAlloc, PageAlloc};
 pub use pfn::PFN;
 pub use raw_page::RawPage;

crates/eonix_mm/src/paging/page.rs

@@ -3,13 +3,6 @@ use crate::address::{AddrRange, PAddr};
 use core::{fmt, mem::ManuallyDrop, ptr::NonNull, sync::atomic::Ordering};
 
 pub const PAGE_SIZE: usize = 4096;
-
-#[derive(Clone, Copy)]
-pub enum PageSize {
-    _4KbPage = 4096,
-    _2MbPage = 2 * 1024 * 1024,
-    _1GbPage = 1 * 1024 * 1024 * 1024,
-}
 pub const PAGE_SIZE_BITS: u32 = PAGE_SIZE.trailing_zeros();
 
 /// A block of memory that is aligned to the page size and can be used for