page-table-generic 0.2.2

use crate::{page_table_entry::Pte, Access, PTEArch, PTEGeneric, PTEInfo, PageTableRef};

pub struct TableIter<'a, 'b: 'a, P: PTEArch, A: Access> {
    access: &'b A,
    idx_stack: [usize; 12],
    table_stack: [Option<PageTableRef<'a, P>>; 12],
    level: usize,
    max_level: usize,
    start_vaddr: *const u8,
    table_size: usize,
}

impl<'a, 'b: 'a, P: PTEArch, A: Access> TableIter<'a, 'b, P, A> {
    pub fn new(va: *const u8, root: PageTableRef<'a, P>, access: &'b A) -> Self {
        let mut table_stack = [None; 12];
        let max_level = root.level();
        let table_size = root.table_size();
        table_stack[max_level - 1] = Some(root);

        TableIter {
            idx_stack: [0; 12],
            table_stack,
            table_size,
            level: max_level,
            access,
            start_vaddr: va,
            max_level,
        }
    }

    fn idx(&self) -> usize {
        self.idx_stack[self.level - 1]
    }

    fn table(&self) -> PageTableRef<'a, P> {
        self.table_stack[self.level - 1].unwrap()
    }

    fn entries(&self) -> &'a [Pte<P>] {
        self.table().as_slice(self.access)
    }

    fn pte(&self) -> Option<PTEGeneric> {
        let idx = self.idx();
        let entries = self.entries();
        if idx >= entries.len() {
            return None;
        }
        Some(entries[idx].read())
    }

    fn idx_next(&mut self, pte: &PTEGeneric) {
        if pte.is_block || self.level == 1 || !pte.valid() {
            let table_size = self.table_size;
            self.idx_stack[self.level - 1] += 1;
            if self.level < self.max_level && self.idx() >= table_size {
                self.level += 1;
                self.idx_stack[self.level - 1] += 1;
            }
        } else {
            self.level -= 1;
            self.idx_stack[self.level - 1] = 0;
            self.table_stack[self.level - 1] = Some(PageTableRef::from_addr(pte.paddr, self.level));
        }
    }

    fn vaddr(&self) -> *const u8 {
        unsafe { self.start_vaddr.add(self.idx() * self.table().entry_size()) }
    }
}

impl<'a, 'b: 'a, P: PTEArch, A: Access> Iterator for TableIter<'a, 'b, P, A> {
    type Item = PTEInfo;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.pte() {
                Some(pte) => {
                    let out = if pte.valid() {
                        let vaddr = self.vaddr();
                        Some(PTEInfo {
                            level: self.level,
                            vaddr,
                            pte: pte.clone(),
                        })
                    } else {
                        None
                    };
                    self.idx_next(&pte);
                    if let Some(out) = out {
                        return Some(out);
                    }
                }
                None => {
                    if self.level == self.max_level {
                        return None;
                    } else {
                        self.level += 1;
                        self.idx_stack[self.level - 1] += 1;
                        continue;
                    }
                }
            };
        }
    }
}