memf-core 0.2.1

//! Virtual address space and page table walking.

use lru::LruCache;
use memf_format::PhysicalMemoryProvider;
use std::cell::RefCell;
use std::num::NonZeroUsize;

use crate::pagefile::PagefileSource;
use crate::proto_pte::PrototypePteSource;
use crate::{Error, Result};

/// Translation mode for virtual-to-physical address translation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TranslationMode {
    /// x86_64 4-level paging (PML4 → PDPT → PD → PT).
    X86_64FourLevel,
    /// x86_64 5-level paging (PML5 → PML4 → PDPT → PD → PT). Linux LA57, Windows Server 2025.
    X86_645Level,
    /// AArch64 4-level page tables (4K granule, 48-bit VA). Linux, Android, macOS ARM.
    AArch64FourLevel,
}

/// A virtual address space backed by physical memory and page tables.
pub struct VirtualAddressSpace<P: PhysicalMemoryProvider> {
    physical: P,
    page_table_root: u64,
    mode: TranslationMode,
    pagefiles: Vec<Box<dyn PagefileSource>>,
    prototype_source: Option<Box<dyn PrototypePteSource>>,
    /// LRU cache: page-aligned vaddr → page-aligned paddr.
    tlb_cache: RefCell<LruCache<u64, u64>>,
}

// x86_64 page table constants
const ADDR_MASK: u64 = 0x000F_FFFF_FFFF_F000;
const PRESENT: u64 = 1;
const PS: u64 = 1 << 7;

// AArch64 4K-granule page table constants (ARMv8-A, 4K granule, 48-bit VA)
const AARCH64_VALID: u64 = 1; // bit 0: entry is valid
const AARCH64_TABLE: u64 = 1 << 1; // bit 1: 1=table/page, 0=block
const AARCH64_OA_MASK: u64 = 0x0000_FFFF_FFFF_F000; // bits [47:12]: output address

/// Number of page translations to cache per `VirtualAddressSpace` instance.
const TRANSLATION_CACHE_CAPACITY: usize = 4096;

/// Internal result of page table walk — not exposed publicly.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TranslationResult {
    /// Page is in physical memory at this address.
    Physical(u64),
    /// Page is demand-zero (all zeroes).
    DemandZero,
    /// Page is in a pagefile.
    PagefileEntry { pagefile_num: u8, page_offset: u64 },
    /// Page is a transition page (still in physical memory at this PFN-derived address).
    Transition(u64),
    /// Page uses a prototype PTE; carries the raw PTE value for resolution.
    Prototype(u64),
}

impl<P: PhysicalMemoryProvider> VirtualAddressSpace<P> {
    /// Create a new virtual address space.
    pub fn new(physical: P, page_table_root: u64, mode: TranslationMode) -> Self {
        Self {
            physical,
            page_table_root,
            mode,
            pagefiles: Vec::new(),
            prototype_source: None,
            // TRANSLATION_CACHE_CAPACITY is a non-zero compile-time constant; the
            // fallback to NonZeroUsize::MIN keeps construction infallible if it is
            // ever set to 0.
            tlb_cache: RefCell::new(LruCache::new(
                NonZeroUsize::new(TRANSLATION_CACHE_CAPACITY).unwrap_or(NonZeroUsize::MIN),
            )),
        }
    }

    /// Attach a pagefile source for resolving paged-out memory.
    pub fn with_pagefile(mut self, source: Box<dyn PagefileSource>) -> Self {
        self.pagefiles.push(source);
        self
    }

    /// Attach a prototype PTE source for resolving shared section pages.
    pub fn with_prototype_source(mut self, source: Box<dyn PrototypePteSource>) -> Self {
        self.prototype_source = Some(source);
        self
    }

    /// Translate a virtual address to a physical address.
    pub fn virt_to_phys(&self, vaddr: u64) -> Result<u64> {
        match self.mode {
            TranslationMode::X86_64FourLevel => self.walk_x86_64_4level(vaddr),
            TranslationMode::X86_645Level => match self.walk_x86_64_5level_internal(vaddr)? {
                TranslationResult::Physical(addr) | TranslationResult::Transition(addr) => Ok(addr),
                TranslationResult::DemandZero => Err(Error::PageNotPresent(vaddr)),
                TranslationResult::PagefileEntry {
                    pagefile_num,
                    page_offset,
                } => Err(Error::PagedOut {
                    vaddr,
                    pagefile_num,
                    page_offset,
                }),
                TranslationResult::Prototype(_) => Err(Error::PrototypePte(vaddr)),
            },
            TranslationMode::AArch64FourLevel => self.walk_aarch64_4level(vaddr),
        }
    }

    /// Read `buf.len()` bytes from virtual address `vaddr`, handling page boundary crossings.
    ///
    /// Uses `walk_x86_64_4level_internal()` to resolve each 4K chunk, transparently
    /// handling physical, transition, demand-zero, and pagefile pages.
    pub fn read_virt(&self, vaddr: u64, buf: &mut [u8]) -> Result<()> {
        if buf.is_empty() {
            return Ok(());
        }

        let mut offset = 0usize;
        let mut current_vaddr = vaddr;

        while offset < buf.len() {
            let page_off = (current_vaddr & 0xFFF) as usize;
            let remaining_in_page = 0x1000 - page_off;
            let remaining_to_read = buf.len() - offset;
            let chunk = remaining_to_read.min(remaining_in_page);

            let result = match self.mode {
                TranslationMode::X86_64FourLevel => {
                    self.walk_x86_64_4level_internal(current_vaddr)?
                }
                TranslationMode::X86_645Level => self.walk_x86_64_5level_internal(current_vaddr)?,
                TranslationMode::AArch64FourLevel => self.walk_aarch64_internal(current_vaddr)?,
            };

            match result {
                TranslationResult::Physical(paddr) | TranslationResult::Transition(paddr) => {
                    let n = self
                        .physical
                        .read_phys(paddr, &mut buf[offset..offset + chunk])?;
                    if n == 0 {
                        return Err(Error::PartialRead {
                            addr: vaddr,
                            requested: buf.len(),
                            got: offset,
                        });
                    }
                    offset += n;
                    current_vaddr = current_vaddr.wrapping_add(n as u64);
                }
                TranslationResult::DemandZero => {
                    buf[offset..offset + chunk].fill(0);
                    offset += chunk;
                    current_vaddr = current_vaddr.wrapping_add(chunk as u64);
                }
                TranslationResult::PagefileEntry {
                    pagefile_num,
                    page_offset,
                } => {
                    let page = self.read_pagefile_page(current_vaddr, pagefile_num, page_offset)?;
                    buf[offset..offset + chunk].copy_from_slice(&page[page_off..page_off + chunk]);
                    offset += chunk;
                    current_vaddr = current_vaddr.wrapping_add(chunk as u64);
                }
                TranslationResult::Prototype(raw_pte) => {
                    if let Some(ref source) = self.prototype_source {
                        if let Some(phys_base) = source.resolve(raw_pte) {
                            let paddr = phys_base + (current_vaddr & 0xFFF);
                            let n = self
                                .physical
                                .read_phys(paddr, &mut buf[offset..offset + chunk])?;
                            if n == 0 {
                                return Err(Error::PartialRead {
                                    addr: vaddr,
                                    requested: buf.len(),
                                    got: offset,
                                });
                            }
                            offset += n;
                            current_vaddr = current_vaddr.wrapping_add(n as u64);
                        } else {
                            return Err(Error::PrototypePte(current_vaddr));
                        }
                    } else {
                        return Err(Error::PrototypePte(current_vaddr));
                    }
                }
            }
        }

        Ok(())
    }

    fn read_pagefile_page(
        &self,
        vaddr: u64,
        pagefile_num: u8,
        page_offset: u64,
    ) -> Result<[u8; 4096]> {
        for source in &self.pagefiles {
            if source.pagefile_number() == pagefile_num {
                if let Some(page) = source.read_page(page_offset)? {
                    return Ok(page);
                }
                break;
            }
        }
        Err(Error::PagedOut {
            vaddr,
            pagefile_num,
            page_offset,
        })
    }

    /// Return a reference to the underlying physical memory provider.
    pub fn physical(&self) -> &P {
        &self.physical
    }

    /// Return the translation mode.
    pub fn mode(&self) -> TranslationMode {
        self.mode
    }

    /// Reverse-map: return a kernel-half virtual address that resolves to the
    /// page at `target_phys` (page granularity), or `None`.
    ///
    /// Enumerates the x86-64 four-level page tables under the current root over
    /// the kernel half (PML4 indices 256..512), checking 2 MiB / 4 KiB leaves and
    /// **transition** PTEs — a resident page whose PTE is not marked valid, which
    /// is exactly the state of a live-acquired kernel image whose header reads as
    /// not-present through a stale/guessed VA. This is how the kernel image base
    /// is recovered when the low-stub hint is wrong: physically locate ntoskrnl,
    /// then ask which kernel VA maps to it. `max_leaves` bounds the work so a
    /// crafted page-table tree cannot turn this into a denial of service.
    ///
    /// Only meaningful for [`TranslationMode::X86_64FourLevel`]; returns `None`
    /// for other modes. 1 GiB pages are skipped (the kernel image is mapped with
    /// 2 MiB / 4 KiB pages, never inside a 1 GiB page).
    #[must_use]
    pub fn find_kernel_va_for_phys(&self, target_phys: u64, max_leaves: usize) -> Option<u64> {
        if !matches!(self.mode, TranslationMode::X86_64FourLevel) {
            return None;
        }
        let target = target_phys & !0xFFF;
        let root = self.page_table_root;
        let mut budget = max_leaves;
        // p4 >= 256 ⇒ bit 47 is always set, so canonicalization always
        // sign-extends into the kernel half.
        let canon = |p4: u64, p3: u64, p2: u64, p1: u64| -> u64 {
            ((p4 << 39) | (p3 << 30) | (p2 << 21) | (p1 << 12)) | 0xFFFF_0000_0000_0000
        };
        for p4 in 256u64..512 {
            let Ok(pml4e) = self.read_pte(root + p4 * 8) else {
                continue;
            };
            if pml4e & PRESENT == 0 {
                continue;
            }
            let pdpt = pml4e & ADDR_MASK;
            for p3 in 0u64..512 {
                let Ok(pdpte) = self.read_pte(pdpt + p3 * 8) else {
                    continue;
                };
                if pdpte & PRESENT == 0 {
                    continue;
                }
                // 1 GiB pages are not used to map the kernel image; skip them.
                if pdpte & PS != 0 {
                    continue;
                }
                let pd = pdpte & ADDR_MASK;
                for p2 in 0u64..512 {
                    let Ok(pde) = self.read_pte(pd + p2 * 8) else {
                        continue;
                    };
                    if pde & PRESENT == 0 {
                        continue;
                    }
                    if pde & PS != 0 {
                        let base = pde & 0x000F_FFFF_FFE0_0000;
                        if target >= base && target < base + (1 << 21) {
                            return Some(canon(p4, p3, p2, 0) | (target - base));
                        }
                        budget = budget.checked_sub(1)?;
                        continue;
                    }
                    let pt = pde & ADDR_MASK;
                    for p1 in 0u64..512 {
                        let Ok(pte) = self.read_pte(pt + p1 * 8) else {
                            continue;
                        };
                        // Match present OR transition leaves: a transition PTE
                        // (PRESENT=0, bit 11 set) still points at a resident page.
                        let leaf = if pte & PRESENT != 0 {
                            Some(pte & ADDR_MASK)
                        } else if pte & (1 << 11) != 0 {
                            Some(((pte >> 12) & 0xF_FFFF_FFFF) * 0x1000)
                        } else {
                            None
                        };
                        match leaf {
                            Some(phys) if phys == target => return Some(canon(p4, p3, p2, p1)),
                            Some(_) => budget = budget.checked_sub(1)?,
                            None => {}
                        }
                    }
                }
            }
        }
        None
    }

    fn read_pte(&self, addr: u64) -> Result<u64> {
        let mut buf = [0u8; 8];
        let n = self.physical.read_phys(addr, &mut buf)?;
        if n < 8 {
            return Err(Error::PartialRead {
                addr,
                requested: 8,
                got: n,
            });
        }
        Ok(u64::from_le_bytes(buf))
    }

    fn walk_x86_64_4level(&self, vaddr: u64) -> Result<u64> {
        let result = self.walk_x86_64_4level_internal(vaddr)?;
        match result {
            TranslationResult::Physical(addr) | TranslationResult::Transition(addr) => Ok(addr),
            TranslationResult::DemandZero => Err(Error::PageNotPresent(vaddr)),
            TranslationResult::PagefileEntry {
                pagefile_num,
                page_offset,
            } => Err(Error::PagedOut {
                vaddr,
                pagefile_num,
                page_offset,
            }),
            TranslationResult::Prototype(_) => Err(Error::PrototypePte(vaddr)),
        }
    }

    fn walk_x86_64_4level_internal(&self, vaddr: u64) -> Result<TranslationResult> {
        self.walk_4level_from(self.page_table_root, vaddr)
    }

    fn walk_4level_from(&self, pml4_root: u64, vaddr: u64) -> Result<TranslationResult> {
        let page_vaddr = vaddr & !0xFFF;
        // peek avoids promoting on read; no mut borrow needed on the hot path.
        if let Some(&paddr_base) = self.tlb_cache.borrow().peek(&page_vaddr) {
            return Ok(TranslationResult::Physical(paddr_base | (vaddr & 0xFFF)));
        }

        let pml4_idx = (vaddr >> 39) & 0x1FF;
        let pdpt_idx = (vaddr >> 30) & 0x1FF;
        let pd_idx = (vaddr >> 21) & 0x1FF;
        let pt_idx = (vaddr >> 12) & 0x1FF;
        let page_offset = vaddr & 0xFFF;

        // PML4
        let pml4e = self.read_pte(pml4_root + pml4_idx * 8)?;
        if pml4e & PRESENT == 0 {
            return Err(Error::PageNotPresent(vaddr));
        }

        // PDPT
        let pdpt_base = pml4e & ADDR_MASK;
        let pdpte = self.read_pte(pdpt_base + pdpt_idx * 8)?;
        if pdpte & PRESENT == 0 {
            return Err(Error::PageNotPresent(vaddr));
        }

        // 1GB huge page check
        if pdpte & PS != 0 {
            let phys_base = pdpte & 0x000F_FFFF_C000_0000;
            let offset_1g = vaddr & 0x3FFF_FFFF;
            let phys = phys_base | offset_1g;
            self.tlb_cache.borrow_mut().put(page_vaddr, phys & !0xFFF);
            return Ok(TranslationResult::Physical(phys));
        }

        // PD
        let pd_base = pdpte & ADDR_MASK;
        let pde = self.read_pte(pd_base + pd_idx * 8)?;
        if pde & PRESENT == 0 {
            return Err(Error::PageNotPresent(vaddr));
        }

        // 2MB large page check
        if pde & PS != 0 {
            let phys_base = pde & 0x000F_FFFF_FFE0_0000;
            let offset_2m = vaddr & 0x1F_FFFF;
            let phys = phys_base | offset_2m;
            self.tlb_cache.borrow_mut().put(page_vaddr, phys & !0xFFF);
            return Ok(TranslationResult::Physical(phys));
        }

        // PT (4K page)
        let pt_base = pde & ADDR_MASK;
        let pte = self.read_pte(pt_base + pt_idx * 8)?;

        if pte & PRESENT != 0 {
            let phys_base = pte & ADDR_MASK;
            let phys = phys_base | page_offset;
            self.tlb_cache.borrow_mut().put(page_vaddr, phys_base);
            return Ok(TranslationResult::Physical(phys));
        }

        // Non-present PTE decoding (PT level only)
        Ok(Self::decode_non_present_pte(pte, page_offset))
    }

    fn walk_x86_64_5level_internal(&self, vaddr: u64) -> Result<TranslationResult> {
        let pml5_idx = (vaddr >> 48) & 0x1FF;
        let pml5e = self.read_pte(self.page_table_root + pml5_idx * 8)?;
        if pml5e & PRESENT == 0 {
            return Err(Error::PageNotPresent(vaddr));
        }
        let pml4_root = pml5e & ADDR_MASK;
        self.walk_4level_from(pml4_root, vaddr)
    }

    fn walk_aarch64_4level(&self, vaddr: u64) -> Result<u64> {
        match self.walk_aarch64_internal(vaddr)? {
            TranslationResult::Physical(addr) | TranslationResult::Transition(addr) => Ok(addr),
            TranslationResult::DemandZero
            | TranslationResult::PagefileEntry { .. }
            | TranslationResult::Prototype(_) => Err(Error::PageNotPresent(vaddr)),
        }
    }

    fn walk_aarch64_internal(&self, vaddr: u64) -> Result<TranslationResult> {
        let page_vaddr = vaddr & !0xFFF;
        // peek avoids promoting on read; no mut borrow needed on the hot path.
        if let Some(&paddr_base) = self.tlb_cache.borrow().peek(&page_vaddr) {
            return Ok(TranslationResult::Physical(paddr_base | (vaddr & 0xFFF)));
        }

        let l0_idx = (vaddr >> 39) & 0x1FF;
        let l1_idx = (vaddr >> 30) & 0x1FF;
        let l2_idx = (vaddr >> 21) & 0x1FF;
        let l3_idx = (vaddr >> 12) & 0x1FF;
        let page_off = vaddr & 0xFFF;

        // Level 0 (PGD)
        let l0e = self.read_pte(self.page_table_root + l0_idx * 8)?;
        if l0e & AARCH64_VALID == 0 {
            return Err(Error::PageNotPresent(vaddr));
        }
        let l1_base = l0e & AARCH64_OA_MASK;

        // Level 1 (PUD)
        let l1e = self.read_pte(l1_base + l1_idx * 8)?;
        if l1e & AARCH64_VALID == 0 {
            return Err(Error::PageNotPresent(vaddr));
        }
        if l1e & AARCH64_TABLE == 0 {
            // 1GB block entry
            let phys_base = l1e & 0x0000_FFFF_C000_0000;
            let phys = phys_base | (vaddr & 0x3FFF_FFFF);
            self.tlb_cache.borrow_mut().put(page_vaddr, phys & !0xFFF);
            return Ok(TranslationResult::Physical(phys));
        }
        let l2_base = l1e & AARCH64_OA_MASK;

        // Level 2 (PMD)
        let l2e = self.read_pte(l2_base + l2_idx * 8)?;
        if l2e & AARCH64_VALID == 0 {
            return Err(Error::PageNotPresent(vaddr));
        }
        if l2e & AARCH64_TABLE == 0 {
            // 2MB block entry
            let phys_base = l2e & 0x0000_FFFF_FFE0_0000;
            let phys = phys_base | (vaddr & 0x001F_FFFF);
            self.tlb_cache.borrow_mut().put(page_vaddr, phys & !0xFFF);
            return Ok(TranslationResult::Physical(phys));
        }
        let l3_base = l2e & AARCH64_OA_MASK;

        // Level 3 (PTE) — bit1=1 means "page" at L3
        let l3e = self.read_pte(l3_base + l3_idx * 8)?;
        if l3e & AARCH64_VALID == 0 {
            return Err(Error::PageNotPresent(vaddr));
        }
        let phys_base = l3e & AARCH64_OA_MASK;
        let phys = phys_base | page_off;
        self.tlb_cache.borrow_mut().put(page_vaddr, phys_base);
        Ok(TranslationResult::Physical(phys))
    }

    fn decode_non_present_pte(pte: u64, page_offset: u64) -> TranslationResult {
        if pte == 0 {
            return TranslationResult::DemandZero;
        }
        if pte & (1 << 11) != 0 {
            let pfn = (pte >> 12) & 0xF_FFFF_FFFF;
            return TranslationResult::Transition(pfn * 0x1000 + page_offset);
        }
        if pte & (1 << 10) != 0 {
            return TranslationResult::Prototype(pte);
        }
        let pagefile_num = ((pte >> 1) & 0xF) as u8;
        let pf_page_offset = (pte >> 12) & 0xF_FFFF_FFFF;
        TranslationResult::PagefileEntry {
            pagefile_num,
            page_offset: pf_page_offset,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_builders::{flags, PageTableBuilder, SyntheticPhysMem};

    /// Build a minimal 4-level page table (CR3=0) mapping `vaddr`'s leaf PTE to
    /// the raw `leaf_pte` value, so tests can install present/transition/large
    /// entries directly. Tables live at fixed physical pages.
    fn manual_pt(vaddr: u64, leaf_pte: u64, large_2m: bool) -> SyntheticPhysMem {
        let mut mem = SyntheticPhysMem::new(16 * 1024 * 1024);
        let p4 = (vaddr >> 39) & 0x1FF;
        let p3 = (vaddr >> 30) & 0x1FF;
        let p2 = (vaddr >> 21) & 0x1FF;
        let p1 = (vaddr >> 12) & 0x1FF;
        let (pdpt, pd, pt) = (0x1000u64, 0x2000u64, 0x3000u64);
        mem.write_u64(p4 * 8, pdpt | PRESENT); // PML4 base is physical 0
        mem.write_u64(pdpt + p3 * 8, pd | PRESENT);
        if large_2m {
            mem.write_u64(pd + p2 * 8, leaf_pte);
        } else {
            mem.write_u64(pd + p2 * 8, pt | PRESENT);
            mem.write_u64(pt + p1 * 8, leaf_pte);
        }
        mem
    }

    #[test]
    fn reverse_map_finds_4k_present_page() {
        let vaddr = 0xFFFF_F802_4020_1000u64;
        let phys = 0x6000u64;
        let mem = manual_pt(vaddr, phys | PRESENT, false);
        let vas = VirtualAddressSpace::new(mem, 0, TranslationMode::X86_64FourLevel);
        assert_eq!(vas.find_kernel_va_for_phys(phys, 100_000), Some(vaddr));
    }

    #[test]
    fn reverse_map_finds_transition_page() {
        // Transition PTE: PRESENT=0, bit 11 set, PFN in bits 12.. — a resident
        // page (e.g. a paged-image header) that a stale VA reads as not-present.
        let vaddr = 0xFFFF_F802_4020_2000u64;
        let pfn = 0x7u64; // phys 0x7000
        let trans_pte = (pfn << 12) | (1 << 11);
        let mem = manual_pt(vaddr, trans_pte, false);
        let vas = VirtualAddressSpace::new(mem, 0, TranslationMode::X86_64FourLevel);
        assert_eq!(vas.find_kernel_va_for_phys(0x7000, 100_000), Some(vaddr));
    }

    #[test]
    fn reverse_map_finds_2m_large_page_with_offset() {
        let vaddr = 0xFFFF_F802_4040_0000u64; // 2 MiB aligned
        let base = 0x20_0000u64;
        let mem = manual_pt(vaddr, base | PRESENT | PS, true);
        let vas = VirtualAddressSpace::new(mem, 0, TranslationMode::X86_64FourLevel);
        // A phys 0x3000 into the 2 MiB page maps to vaddr + 0x3000.
        assert_eq!(
            vas.find_kernel_va_for_phys(base + 0x3000, 100_000),
            Some(vaddr + 0x3000)
        );
    }

    #[test]
    fn reverse_map_returns_none_when_absent() {
        let mem = manual_pt(0xFFFF_F802_4020_1000u64, 0x6000 | PRESENT, false);
        let vas = VirtualAddressSpace::new(mem, 0, TranslationMode::X86_64FourLevel);
        assert_eq!(vas.find_kernel_va_for_phys(0xDEAD_0000, 100_000), None);
    }

    #[test]
    fn reverse_map_respects_max_leaves_budget() {
        // Two present leaves, neither matching; a budget of 1 must give up
        // (the second non-matching leaf exhausts the budget) and return None.
        let vaddr = 0xFFFF_F802_4020_1000u64;
        let mut mem = manual_pt(vaddr, 0x6000 | PRESENT, false);
        let p2 = (vaddr >> 21) & 0x1FF;
        let pt = 0x3000u64;
        // Add a second present leaf in the same PT at index +1.
        let p1b = ((vaddr >> 12) & 0x1FF) + 1;
        let _ = p2;
        mem.write_u64(pt + p1b * 8, 0x9000 | PRESENT);
        let vas = VirtualAddressSpace::new(mem, 0, TranslationMode::X86_64FourLevel);
        assert_eq!(vas.find_kernel_va_for_phys(0xABCD_0000, 1), None);
    }

    #[test]
    fn reverse_map_skips_1g_huge_page() {
        // A present 1 GiB PDPTE is skipped (the kernel image is never mapped in
        // a 1 GiB page); the search returns None even if target falls inside it.
        let vaddr = 0xFFFF_F802_4000_0000u64;
        let mut mem = SyntheticPhysMem::new(16 * 1024 * 1024);
        let p4 = (vaddr >> 39) & 0x1FF;
        let p3 = (vaddr >> 30) & 0x1FF;
        mem.write_u64(p4 * 8, 0x1000 | PRESENT);
        mem.write_u64(0x1000 + p3 * 8, 0x4000_0000 | PRESENT | PS); // 1 GiB page
        let vas = VirtualAddressSpace::new(mem, 0, TranslationMode::X86_64FourLevel);
        assert_eq!(vas.find_kernel_va_for_phys(0x4000_1000, 100_000), None);
    }

    #[test]
    fn reverse_map_skips_read_error_branch() {
        // A PML4E pointing past the synthetic image makes the PDPT read fail;
        // the walk must skip that subtree, then find the valid mapping elsewhere.
        let vaddr = 0xFFFF_F802_4020_1000u64;
        let mut mem = manual_pt(vaddr, 0x6000 | PRESENT, false);
        // Point PML4 index 256 at an out-of-range PDPT (16 MiB image, this is 64 MiB).
        mem.write_u64(256 * 8, 0x400_0000 | PRESENT);
        let vas = VirtualAddressSpace::new(mem, 0, TranslationMode::X86_64FourLevel);
        assert_eq!(vas.find_kernel_va_for_phys(0x6000, 100_000), Some(vaddr));
    }

    #[test]
    fn reverse_map_skips_non_matching_2m_page() {
        // A present 2 MiB page that does not contain the target is skipped; the
        // search then finds nothing (covers the large-page non-match branch).
        let vaddr = 0xFFFF_F802_4040_0000u64;
        let mem = manual_pt(vaddr, 0x20_0000 | PRESENT | PS, true);
        let vas = VirtualAddressSpace::new(mem, 0, TranslationMode::X86_64FourLevel);
        assert_eq!(vas.find_kernel_va_for_phys(0x40_0000, 100_000), None);
    }

    #[test]
    fn reverse_map_none_for_non_4level_mode() {
        let mem = manual_pt(0xFFFF_F802_4020_1000u64, 0x6000 | PRESENT, false);
        let vas = VirtualAddressSpace::new(mem, 0, TranslationMode::AArch64FourLevel);
        assert_eq!(vas.find_kernel_va_for_phys(0x6000, 100_000), None);
    }

    #[test]
    fn translate_4k_page() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let paddr: u64 = 0x0080_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_4k(vaddr, paddr, flags::WRITABLE)
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let result = vas.virt_to_phys(vaddr).unwrap();
        assert_eq!(result, paddr);
    }

    #[test]
    fn translate_4k_with_offset() {
        let vaddr: u64 = 0xFFFF_8000_0010_0ABC;
        let paddr_base: u64 = 0x0080_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_4k(vaddr & !0xFFF, paddr_base, flags::WRITABLE)
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let result = vas.virt_to_phys(vaddr).unwrap();
        assert_eq!(result, paddr_base + 0xABC);
    }

    #[test]
    fn read_virt_4k() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let paddr: u64 = 0x0080_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_4k(vaddr, paddr, flags::WRITABLE)
            .write_phys(paddr, &[0xDE, 0xAD, 0xBE, 0xEF])
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let mut buf = [0u8; 4];
        vas.read_virt(vaddr, &mut buf).unwrap();
        assert_eq!(buf, [0xDE, 0xAD, 0xBE, 0xEF]);
    }

    #[test]
    fn translate_2mb_page() {
        let vaddr: u64 = 0xFFFF_8000_0020_0000;
        let paddr: u64 = 0x0100_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_2m(vaddr, paddr, flags::WRITABLE)
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let result = vas.virt_to_phys(vaddr).unwrap();
        assert_eq!(result, paddr);

        // Test with offset within the 2MB page
        let result_offset = vas.virt_to_phys(vaddr + 0x1234).unwrap();
        assert_eq!(result_offset, paddr + 0x1234);
    }

    #[test]
    fn translate_1gb_page() {
        let vaddr: u64 = 0xFFFF_8000_4000_0000;
        let paddr: u64 = 0x4000_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_1g(vaddr, paddr, flags::WRITABLE)
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let result = vas.virt_to_phys(vaddr).unwrap();
        assert_eq!(result, paddr);

        // Test with offset within the 1GB page
        let result_offset = vas.virt_to_phys(vaddr + 0x12_3456).unwrap();
        assert_eq!(result_offset, paddr + 0x12_3456);
    }

    #[test]
    fn non_present_page_returns_error() {
        let (cr3, mem) = PageTableBuilder::new().build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let result = vas.virt_to_phys(0xFFFF_8000_0010_0000);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::PageNotPresent(addr) => assert_eq!(addr, 0xFFFF_8000_0010_0000),
            other => panic!("unexpected error: {other}"),
        }
    }

    #[test]
    fn read_virt_cross_page_boundary() {
        // Map two adjacent virtual pages to different physical pages
        let vaddr_page1: u64 = 0xFFFF_8000_0010_0000;
        let vaddr_page2: u64 = 0xFFFF_8000_0010_1000;
        let paddr1: u64 = 0x0080_0000;
        let paddr2: u64 = 0x0090_0000;

        // Write data at end of page1 and start of page2
        let (cr3, mem) = PageTableBuilder::new()
            .map_4k(vaddr_page1, paddr1, flags::WRITABLE)
            .map_4k(vaddr_page2, paddr2, flags::WRITABLE)
            .write_phys(paddr1 + 0xFFC, &[0xAA, 0xBB, 0xCC, 0xDD])
            .write_phys(paddr2, &[0x11, 0x22, 0x33, 0x44])
            .build();

        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let mut buf = [0u8; 8];
        vas.read_virt(vaddr_page1 + 0xFFC, &mut buf).unwrap();
        assert_eq!(buf, [0xAA, 0xBB, 0xCC, 0xDD, 0x11, 0x22, 0x33, 0x44]);
    }

    #[test]
    fn read_virt_empty_buffer() {
        let (cr3, mem) = PageTableBuilder::new().build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let mut buf = [];
        vas.read_virt(0xFFFF_8000_0010_0000, &mut buf).unwrap();
    }

    #[test]
    fn virt_to_phys_4k_direct() {
        // Test virt_to_phys as the public API (not just internal translate)
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let paddr: u64 = 0x0080_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_4k(vaddr, paddr, flags::WRITABLE)
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        assert_eq!(vas.virt_to_phys(vaddr).unwrap(), paddr);
        assert_eq!(vas.virt_to_phys(vaddr + 0x42).unwrap(), paddr + 0x42);
    }

    #[test]
    fn physical_accessor() {
        let (cr3, mem) = PageTableBuilder::new()
            .write_phys(0x5000, &[0xAB; 8])
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let phys = vas.physical();
        let mut buf = [0u8; 4];
        let n = phys.read_phys(0x5000, &mut buf).unwrap();
        assert_eq!(n, 4);
        assert_eq!(buf, [0xAB; 4]);
    }

    #[test]
    fn demand_zero_pte_returns_page_not_present() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let (cr3, mem) = PageTableBuilder::new().map_demand_zero(vaddr).build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let result = vas.virt_to_phys(vaddr);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::PageNotPresent(addr) => assert_eq!(addr, vaddr),
            other => panic!("expected PageNotPresent, got: {other}"),
        }
    }

    #[test]
    fn transition_pte_resolves_to_physical() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let pfn: u64 = 0x800;
        let (cr3, mem) = PageTableBuilder::new()
            .map_transition(vaddr, pfn)
            .write_phys(pfn * 0x1000, &[0xDE, 0xAD, 0xBE, 0xEF])
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let paddr = vas.virt_to_phys(vaddr).unwrap();
        assert_eq!(paddr, pfn * 0x1000);
    }

    #[test]
    fn transition_pte_with_offset() {
        let vaddr_base: u64 = 0xFFFF_8000_0010_0000;
        let vaddr: u64 = vaddr_base + 0x42;
        let pfn: u64 = 0x800;
        let (cr3, mem) = PageTableBuilder::new()
            .map_transition(vaddr_base, pfn)
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let paddr = vas.virt_to_phys(vaddr).unwrap();
        assert_eq!(paddr, pfn * 0x1000 + 0x42);
    }

    #[test]
    fn prototype_pte_returns_error() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let (cr3, mem) = PageTableBuilder::new().map_prototype(vaddr).build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let result = vas.virt_to_phys(vaddr);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::PrototypePte(addr) => assert_eq!(addr, vaddr),
            other => panic!("expected PrototypePte, got: {other}"),
        }
    }

    #[test]
    fn pagefile_pte_returns_paged_out() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_pagefile(vaddr, 0, 0x1234)
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let result = vas.virt_to_phys(vaddr);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::PagedOut {
                vaddr: v,
                pagefile_num,
                page_offset,
            } => {
                assert_eq!(v, vaddr);
                assert_eq!(pagefile_num, 0);
                assert_eq!(page_offset, 0x1234);
            }
            other => panic!("expected PagedOut, got: {other}"),
        }
    }

    #[test]
    fn pagefile_pte_number_routing() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_pagefile(vaddr, 2, 0xABCD)
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let result = vas.virt_to_phys(vaddr);
        match result.unwrap_err() {
            Error::PagedOut {
                pagefile_num,
                page_offset,
                ..
            } => {
                assert_eq!(pagefile_num, 2);
                assert_eq!(page_offset, 0xABCD);
            }
            other => panic!("expected PagedOut, got: {other}"),
        }
    }

    use crate::test_builders::{MockPagefileSource, MockPrototypePteSource};

    #[test]
    fn read_virt_demand_zero_returns_zeroes() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let (cr3, mem) = PageTableBuilder::new().map_demand_zero(vaddr).build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let mut buf = [0xFFu8; 4096];
        vas.read_virt(vaddr, &mut buf).unwrap();
        assert!(
            buf.iter().all(|&b| b == 0),
            "demand-zero page must be all zeroes"
        );
    }

    #[test]
    fn read_virt_transition_reads_physical() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let pfn: u64 = 0x800;
        let (cr3, mem) = PageTableBuilder::new()
            .map_transition(vaddr, pfn)
            .write_phys(pfn * 0x1000, &[0xCA, 0xFE, 0xBA, 0xBE])
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let mut buf = [0u8; 4];
        vas.read_virt(vaddr, &mut buf).unwrap();
        assert_eq!(buf, [0xCA, 0xFE, 0xBA, 0xBE]);
    }

    #[test]
    fn read_virt_pagefile_with_provider() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let page_offset: u64 = 0x10;
        let mut page_data = [0u8; 4096];
        page_data[0..4].copy_from_slice(&[0xDE, 0xAD, 0xBE, 0xEF]);

        let (cr3, mem) = PageTableBuilder::new()
            .map_pagefile(vaddr, 0, page_offset)
            .build();

        let mock = MockPagefileSource::new(0, vec![(page_offset, page_data)]);
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel)
            .with_pagefile(Box::new(mock));

        let mut buf = [0u8; 4];
        vas.read_virt(vaddr, &mut buf).unwrap();
        assert_eq!(buf, [0xDE, 0xAD, 0xBE, 0xEF]);
    }

    #[test]
    fn read_virt_pagefile_without_provider_errors() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let (cr3, mem) = PageTableBuilder::new().map_pagefile(vaddr, 0, 0x10).build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let mut buf = [0u8; 4];
        let result = vas.read_virt(vaddr, &mut buf);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::PagedOut {
                pagefile_num: 0,
                page_offset: 0x10,
                ..
            } => {}
            other => panic!("expected PagedOut, got: {other}"),
        }
    }

    #[test]
    fn read_virt_prototype_pte_errors() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let (cr3, mem) = PageTableBuilder::new().map_prototype(vaddr).build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let mut buf = [0u8; 4];
        let result = vas.read_virt(vaddr, &mut buf);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::PrototypePte(addr) => assert_eq!(addr, vaddr),
            other => panic!("expected PrototypePte, got: {other}"),
        }
    }

    #[test]
    fn read_virt_pagefile_number_routing() {
        let vaddr1: u64 = 0xFFFF_8000_0010_0000;
        let vaddr2: u64 = 0xFFFF_8000_0010_1000;

        let mut page0_data = [0u8; 4096];
        page0_data[0..4].copy_from_slice(&[0x11, 0x22, 0x33, 0x44]);
        let mut page1_data = [0u8; 4096];
        page1_data[0..4].copy_from_slice(&[0xAA, 0xBB, 0xCC, 0xDD]);

        let (cr3, mem) = PageTableBuilder::new()
            .map_pagefile(vaddr1, 0, 0x10)
            .map_pagefile(vaddr2, 1, 0x20)
            .build();

        let mock0 = MockPagefileSource::new(0, vec![(0x10, page0_data)]);
        let mock1 = MockPagefileSource::new(1, vec![(0x20, page1_data)]);

        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel)
            .with_pagefile(Box::new(mock0))
            .with_pagefile(Box::new(mock1));

        let mut buf1 = [0u8; 4];
        vas.read_virt(vaddr1, &mut buf1).unwrap();
        assert_eq!(buf1, [0x11, 0x22, 0x33, 0x44]);

        let mut buf2 = [0u8; 4];
        vas.read_virt(vaddr2, &mut buf2).unwrap();
        assert_eq!(buf2, [0xAA, 0xBB, 0xCC, 0xDD]);
    }

    #[test]
    fn read_virt_pagefile_out_of_range() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_pagefile(vaddr, 0, 0x9999)
            .build();
        let mock = MockPagefileSource::new(0, vec![]);
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel)
            .with_pagefile(Box::new(mock));
        let mut buf = [0u8; 4];
        let result = vas.read_virt(vaddr, &mut buf);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::PagedOut {
                page_offset: 0x9999,
                ..
            } => {}
            other => panic!("expected PagedOut, got: {other}"),
        }
    }

    #[test]
    fn read_virt_mixed_pages_cross_boundary() {
        let vaddr1: u64 = 0xFFFF_8000_0010_0000;
        let vaddr2: u64 = 0xFFFF_8000_0010_1000;
        let vaddr3: u64 = 0xFFFF_8000_0010_2000;
        let paddr1: u64 = 0x0080_0000;

        let mut pf_page = [0u8; 4096];
        pf_page[0..4].copy_from_slice(&[0xBB; 4]);

        let (cr3, mem) = PageTableBuilder::new()
            .map_4k(vaddr1, paddr1, flags::WRITABLE)
            .write_phys(paddr1 + 0xFFC, &[0xAA; 4])
            .map_pagefile(vaddr2, 0, 0x10)
            .map_demand_zero(vaddr3)
            .build();

        let mock = MockPagefileSource::new(0, vec![(0x10, pf_page)]);
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel)
            .with_pagefile(Box::new(mock));

        // Read spanning: last 4 bytes of phys page + first 4 bytes of pagefile page
        let mut buf = [0u8; 8];
        vas.read_virt(vaddr1 + 0xFFC, &mut buf).unwrap();
        assert_eq!(buf, [0xAA, 0xAA, 0xAA, 0xAA, 0xBB, 0xBB, 0xBB, 0xBB]);

        // Read spanning: last 4 bytes of pagefile page + first 4 bytes of demand-zero page
        let mut buf2 = [0u8; 8];
        vas.read_virt(vaddr2 + 0xFFC, &mut buf2).unwrap();
        assert_eq!(buf2, [0u8; 8]);
    }

    #[test]
    fn read_virt_prototype_pte_resolves_when_source_provided() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let resolved_paddr: u64 = 0x00A0_0000;
        // Raw PTE: bit 10 set + some upper bits to identify this PTE
        let raw_pte: u64 = (1 << 10) | (0xABCu64 << 12);

        let (cr3, mem) = PageTableBuilder::new()
            .map_prototype_raw(vaddr, raw_pte)
            .write_phys(resolved_paddr, &[0xDE, 0xAD, 0xBE, 0xEF])
            .build();

        let mock = MockPrototypePteSource::new(vec![(raw_pte, resolved_paddr)]);
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel)
            .with_prototype_source(Box::new(mock));

        let mut buf = [0u8; 4];
        vas.read_virt(vaddr, &mut buf).unwrap();
        assert_eq!(buf, [0xDE, 0xAD, 0xBE, 0xEF]);
    }

    #[test]
    fn read_virt_prototype_pte_errors_when_source_returns_none() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let raw_pte: u64 = (1 << 10) | (0xDEFu64 << 12);

        let (cr3, mem) = PageTableBuilder::new()
            .map_prototype_raw(vaddr, raw_pte)
            .build();

        // Source returns None for this PTE value
        let mock = MockPrototypePteSource::new(vec![]);
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel)
            .with_prototype_source(Box::new(mock));

        let mut buf = [0u8; 4];
        let result = vas.read_virt(vaddr, &mut buf);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::PrototypePte(addr) => assert_eq!(addr, vaddr),
            other => panic!("expected PrototypePte, got: {other}"),
        }
    }

    #[test]
    fn translation_cache_hit_returns_same_result() {
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let paddr: u64 = 0x0080_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_4k(vaddr, paddr, flags::WRITABLE)
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let first = vas.virt_to_phys(vaddr).unwrap();
        let second = vas.virt_to_phys(vaddr).unwrap();
        assert_eq!(first, second);
        assert_eq!(first, paddr);
    }

    #[test]
    fn translate_5level_variant_compiles() {
        // Verify X86_645Level variant exists and is distinct from 4-level.
        assert_ne!(
            std::mem::discriminant(&TranslationMode::X86_645Level),
            std::mem::discriminant(&TranslationMode::X86_64FourLevel),
        );
    }

    #[test]
    fn translate_5level_4k_page() {
        // For an address that fits within the first PML5 entry's PML4 table,
        // 5-level mode must produce the same result as 4-level mode
        // (PageTableBuilder builds x86_64 4-level tables; a 5-level walk with
        // PML5[0] pointing to CR3's PML4 table is equivalent to a 4-level walk).
        // The test verifies the new mode dispatches without panicking.
        let vaddr: u64 = 0xFFFF_8000_0010_0000;
        let paddr: u64 = 0x0090_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_4k(vaddr, paddr, flags::WRITABLE)
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_645Level);
        // With a 5-level page table where PML5[idx] = cr3's PML4 root,
        // the walk should succeed. For now, just verify it doesn't panic.
        let _ = vas.virt_to_phys(vaddr); // may return Ok or Err — must not panic
    }

    #[test]
    fn translation_cache_capacity_100_distinct_pages() {
        use std::collections::HashSet;
        let base_vaddr: u64 = 0xFFFF_8000_0000_0000;
        let mut builder = PageTableBuilder::new();
        for i in 0..200u64 {
            builder = builder.map_4k(
                base_vaddr + i * 0x1000,
                0x1000 + i * 0x1000,
                flags::WRITABLE,
            );
        }
        let (cr3, mem) = builder.build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);
        let mut results = HashSet::new();
        for i in 0..200u64 {
            let paddr = vas.virt_to_phys(base_vaddr + i * 0x1000).unwrap();
            results.insert(paddr);
        }
        assert_eq!(
            results.len(),
            200,
            "each page must map to a distinct physical address"
        );
    }

    #[test]
    fn read_virt_prototype_pte_with_page_offset() {
        let vaddr_base: u64 = 0xFFFF_8000_0010_0000;
        let vaddr: u64 = vaddr_base + 0x100;
        let resolved_paddr: u64 = 0x00B0_0000;
        let raw_pte: u64 = (1 << 10) | (0x123u64 << 12);

        let (cr3, mem) = PageTableBuilder::new()
            .map_prototype_raw(vaddr_base, raw_pte)
            .write_phys(resolved_paddr + 0x100, &[0xCA, 0xFE])
            .build();

        let mock = MockPrototypePteSource::new(vec![(raw_pte, resolved_paddr)]);
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel)
            .with_prototype_source(Box::new(mock));

        let mut buf = [0u8; 2];
        vas.read_virt(vaddr, &mut buf).unwrap();
        assert_eq!(buf, [0xCA, 0xFE]);
    }

    #[test]
    fn aarch64_mode_distinct_from_x86() {
        assert_ne!(
            std::mem::discriminant(&TranslationMode::AArch64FourLevel),
            std::mem::discriminant(&TranslationMode::X86_64FourLevel),
        );
    }

    #[test]
    fn aarch64_non_present_returns_error() {
        // An empty page table (zeroed PGD entry) should return PageNotPresent.
        let (cr3, mem) = PageTableBuilder::new().build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::AArch64FourLevel);
        let result = vas.virt_to_phys(0x1000);
        assert!(matches!(result, Err(Error::PageNotPresent(_))));
    }

    #[test]
    fn aarch64_translate_4k_page() {
        // AArch64 and x86_64 4-level page table structures are identical at the
        // structural level (9-bit indices, 8-byte PTEs, 4K granule).
        // PageTableBuilder sets PRESENT+WRITABLE bits = 0b11, which maps to
        // AArch64 VALID(bit0)=1 + TABLE(bit1)=1 for table entries.
        // So an x86_64-built table is also a valid AArch64 L0–L3 table.
        let vaddr: u64 = 0x0000_0000_0010_0000; // canonical user-space AArch64 vaddr
        let paddr: u64 = 0x0080_0000;
        let (cr3, mem) = PageTableBuilder::new()
            .map_4k(vaddr, paddr, flags::WRITABLE)
            .build();
        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::AArch64FourLevel);
        let result = vas.virt_to_phys(vaddr);
        // x86_64 PTEs use bits[51:12] for address and bit0 for PRESENT.
        // AArch64 PTEs use bits[47:12] for OA and bit0 for VALID, bit1 for TABLE.
        // PageTableBuilder sets bit0+bit1 (PRESENT+WRITABLE), so AArch64 sees VALID+TABLE.
        // The physical address bits overlap, so the walk should succeed.
        assert!(
            result.is_ok() || matches!(result, Err(Error::PageNotPresent(_))),
            "must not panic; got {result:?}"
        );
    }

    #[test]
    fn multiple_mappings_same_pml4() {
        let vaddr1: u64 = 0xFFFF_8000_0010_0000;
        let vaddr2: u64 = 0xFFFF_8000_0010_1000;
        let paddr1: u64 = 0x0080_0000;
        let paddr2: u64 = 0x0090_0000;

        let (cr3, mem) = PageTableBuilder::new()
            .map_4k(vaddr1, paddr1, flags::WRITABLE)
            .map_4k(vaddr2, paddr2, flags::WRITABLE)
            .write_phys(paddr1, &[0x11; 8])
            .write_phys(paddr2, &[0x22; 8])
            .build();

        let vas = VirtualAddressSpace::new(mem, cr3, TranslationMode::X86_64FourLevel);

        let mut buf1 = [0u8; 8];
        vas.read_virt(vaddr1, &mut buf1).unwrap();
        assert_eq!(buf1, [0x11; 8]);

        let mut buf2 = [0u8; 8];
        vas.read_virt(vaddr2, &mut buf2).unwrap();
        assert_eq!(buf2, [0x22; 8]);
    }
}