hyperlight_common/arch/amd64/

vmem.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright 2025 The Hyperlight Authors.

//! x86-64 4-level page table manipulation code.
//!
//! This module implements page table setup for x86-64 long mode using 4-level paging:
//! - PML4 (Page Map Level 4) - bits 47:39 - 512 entries, each covering 512GB
//! - PDPT (Page Directory Pointer Table) - bits 38:30 - 512 entries, each covering 1GB
//! - PD (Page Directory) - bits 29:21 - 512 entries, each covering 2MB
//! - PT (Page Table) - bits 20:12 - 512 entries, each covering 4KB pages
//!
//! The code uses an iterator-based approach to walk the page table hierarchy,
//! allocating intermediate tables as needed and setting appropriate flags on leaf PTEs

use crate::vmem::{
    BasicMapping, CowMapping, MapRequest, MapResponse, Mapping, MappingKind, TableMovabilityBase,
    TableOps, TableReadOps, UpdateParent, UpdateParentNone, UpdateParentTable, Void, modify_ptes,
    write_entry_updating,
};

#[derive(Copy, Clone)]
pub(in crate::vmem) struct UpdateParentRoot {}

/// Read a PTE and return it (widened to u64) if the present bit is
/// set. The amd64 "present" encoding is a single bit (bit 0); other
/// architectures may need richer semantics, which is why this lives
/// per-arch rather than in the common module.
///
/// # Safety
/// `entry_ptr` must point to a valid page table entry.
#[inline(always)]
#[allow(clippy::useless_conversion)]
pub(super) unsafe fn read_pte_if_present<Op: TableReadOps>(
    op: &Op,
    entry_ptr: Op::TableAddr,
) -> Option<u64> {
    let pte: u64 = unsafe { op.read_entry(entry_ptr) }.into();
    if (pte & PAGE_PRESENT) != 0 {
        Some(pte)
    } else {
        None
    }
}

/// Require that a PTE is present and descend to the next-level table.
///
/// # Safety
/// `op` must provide valid page table memory.
pub(super) unsafe fn require_pte_exist<Op: TableReadOps, P: UpdateParent<Op>>(
    op: &Op,
    x: MapResponse<Op, P>,
) -> Option<MapRequest<Op, P::ChildType>>
where
    P::ChildType: UpdateParent<Op>,
{
    unsafe { read_pte_if_present(op, x.entry_ptr) }.map(|pte| MapRequest {
        #[allow(clippy::unnecessary_cast)]
        table_base: Op::from_phys((pte & PTE_ADDR_MASK) as PhysAddr),
        vmin: x.vmin,
        len: x.len,
        update_parent: x.update_parent.for_child_at_entry(x.entry_ptr),
    })
}

// Paging Flags
//
// See the following links explaining paging:
//
// * Intel® 64 and IA-32 Architectures Software Developer’s Manual, Volume 3A: System Programming Guide, Part 1
//  - Chapter 5 "Paging"
//
// https://cdrdv2.intel.com/v1/dl/getContent/671200
//
// * AMD64 Architecture Programmer’s Manual, Volume 2: System Programming, Section 5.3: Long-Mode Page Translation
//
// https://docs.amd.com/v/u/en-US/24593_3.43
//
// Or if you prefer something less formal:
//
// * Very basic description: https://stackoverflow.com/a/26945892
// * More in-depth descriptions: https://wiki.osdev.org/Paging
//

/// Page is Present
pub const PAGE_PRESENT: u64 = 1;
/// Page is Read/Write (if not set page is read only so long as the WP bit in CR0 is set to 1 - which it is in Hyperlight)
const PAGE_RW: u64 = 1 << 1;
/// Execute Disable (if this bit is set then data in the page cannot be executed)`
const PAGE_NX: u64 = 1 << 63;
/// Mask to extract the physical address from a PTE (bits 51:12)
/// This masks out the lower 12 flag bits AND the upper bits including NX (bit 63)
pub const PTE_ADDR_MASK: u64 = 0x000F_FFFF_FFFF_F000;
const PAGE_USER_ACCESS_DISABLED: u64 = 0 << 2; // U/S bit not set - supervisor mode only (no code runs in user mode for now)
const PAGE_DIRTY_SET: u64 = 1 << 6; // D - dirty bit
const PAGE_ACCESSED_SET: u64 = 1 << 5; // A - accessed bit
const PAGE_CACHE_ENABLED: u64 = 0 << 4; // PCD - page cache disable bit not set (caching enabled)
const PAGE_WRITE_BACK: u64 = 0 << 3; // PWT - page write-through bit not set (write-back caching)
const PAGE_PAT_WB: u64 = 0 << 7; // PAT - page attribute table index bit (0 for write-back memory when PCD=0, PWT=0)

// We use various patterns of the available-for-software-use bits to
// represent certain special mappings.
const PTE_AVL_MASK: u64 = 0x0000_0000_0000_0E00;
const PAGE_AVL_COW: u64 = 1 << 9;

/// Returns PAGE_RW if writable is true, 0 otherwise
#[inline(always)]
const fn page_rw_flag(writable: bool) -> u64 {
    if writable { PAGE_RW } else { 0 }
}

/// Returns PAGE_NX if executable is false (NX = No Execute), 0 otherwise
#[inline(always)]
const fn page_nx_flag(executable: bool) -> u64 {
    if executable { 0 } else { PAGE_NX }
}

/// Helper function to generate a page table entry that points to another table
#[allow(clippy::identity_op)]
#[allow(clippy::precedence)]
fn pte_for_table<Op: TableOps>(table_addr: Op::TableAddr) -> u64 {
    Op::to_phys(table_addr) |
        PAGE_ACCESSED_SET | // prevent the CPU writing to the access flag
        PAGE_CACHE_ENABLED | // leave caching enabled
        PAGE_WRITE_BACK | // use write-back caching
        PAGE_USER_ACCESS_DISABLED |// dont allow user access (no code runs in user mode for now)
        PAGE_RW | // R/W - we don't use block-level permissions
        PAGE_PRESENT // P   - this entry is present
}

/// This trait is used to select appropriate implementations of
/// [`UpdateParent`] to be used, depending on whether a particular
/// implementation needs the ability to move tables.
pub(in crate::vmem) trait TableMovability<Op: TableReadOps + ?Sized, TableMoveInfo> {
    type RootUpdateParent: UpdateParent<Op, TableMoveInfo = TableMoveInfo>;
    fn root_update_parent() -> Self::RootUpdateParent;
}
impl<Op: TableOps<TableMovability = crate::vmem::MayMoveTable>> TableMovability<Op, Op::TableAddr>
    for crate::vmem::MayMoveTable
{
    type RootUpdateParent = UpdateParentRoot;
    fn root_update_parent() -> Self::RootUpdateParent {
        UpdateParentRoot {}
    }
}
impl<Op: TableReadOps> TableMovability<Op, Void> for crate::vmem::MayNotMoveTable {
    type RootUpdateParent = UpdateParentNone;
    fn root_update_parent() -> Self::RootUpdateParent {
        UpdateParentNone {}
    }
}

impl<
    Op: TableOps<TableMovability = crate::vmem::MayMoveTable>,
    P: UpdateParent<Op, TableMoveInfo = Op::TableAddr>,
> UpdateParent<Op> for UpdateParentTable<Op, P>
{
    type TableMoveInfo = Op::TableAddr;
    type ChildType = UpdateParentTable<Op, Self>;
    fn update_parent(self, op: &Op, new_ptr: Op::TableAddr) {
        let pte = pte_for_table::<Op>(new_ptr);
        unsafe {
            write_entry_updating(op, self.parent, self.entry_ptr, pte);
        }
    }
    fn for_child_at_entry(self, entry_ptr: Op::TableAddr) -> Self::ChildType {
        Self::ChildType::new(self, entry_ptr)
    }
}

impl<Op: TableOps<TableMovability = crate::vmem::MayMoveTable>> UpdateParent<Op>
    for UpdateParentRoot
{
    type TableMoveInfo = Op::TableAddr;
    type ChildType = UpdateParentTable<Op, Self>;
    fn update_parent(self, op: &Op, new_ptr: Op::TableAddr) {
        unsafe {
            op.update_root(new_ptr);
        }
    }
    fn for_child_at_entry(self, entry_ptr: Op::TableAddr) -> Self::ChildType {
        Self::ChildType::new(self, entry_ptr)
    }
}

/// Page-mapping callback to allocate a next-level page table if necessary.
/// # Safety
/// This function modifies page table data structures, and should not be called concurrently
/// with any other operations that modify the page tables.
unsafe fn alloc_pte_if_needed<
    Op: TableOps,
    P: UpdateParent<
            Op,
            TableMoveInfo = <Op::TableMovability as TableMovabilityBase<Op>>::TableMoveInfo,
        >,
>(
    op: &Op,
    x: MapResponse<Op, P>,
) -> MapRequest<Op, P::ChildType>
where
    P::ChildType: UpdateParent<Op>,
{
    let new_update_parent = x.update_parent.for_child_at_entry(x.entry_ptr);
    if let Some(pte) = unsafe { read_pte_if_present(op, x.entry_ptr) } {
        return MapRequest {
            table_base: Op::from_phys(pte & PTE_ADDR_MASK),
            vmin: x.vmin,
            len: x.len,
            update_parent: new_update_parent,
        };
    }

    let page_addr = unsafe { op.alloc_table() };

    let pte = pte_for_table::<Op>(page_addr);
    unsafe {
        write_entry_updating(op, x.update_parent, x.entry_ptr, pte);
    };
    MapRequest {
        table_base: page_addr,
        vmin: x.vmin,
        len: x.len,
        update_parent: new_update_parent,
    }
}

/// Map a normal memory page
/// # Safety
/// This function modifies page table data structures, and should not be called concurrently
/// with any other operations that modify the page tables.
#[allow(clippy::identity_op)]
#[allow(clippy::precedence)]
unsafe fn map_page<
    Op: TableOps,
    P: UpdateParent<
            Op,
            TableMoveInfo = <Op::TableMovability as TableMovabilityBase<Op>>::TableMoveInfo,
        >,
>(
    op: &Op,
    mapping: &Mapping,
    r: MapResponse<Op, P>,
) {
    let pte = match &mapping.kind {
        MappingKind::Basic(bm) =>
        // TODO: Support not readable
        // NOTE: On x86-64, there is no separate "readable" bit in the page table entry.
        // This means that pages cannot be made write-only or execute-only without also being readable.
        // All pages that are mapped as writable or executable are also implicitly readable.
        // If support for "not readable" mappings is required in the future, it would need to be
        // implemented using additional mechanisms (e.g., page-fault handling or memory protection keys),
        // but for now, this architectural limitation is accepted.
        {
            (mapping.phys_base + (r.vmin - mapping.virt_base)) |
                page_nx_flag(bm.executable) | // NX - no execute unless allowed
                PAGE_PAT_WB | // PAT index bit for write-back memory
                PAGE_DIRTY_SET | // prevent the CPU writing to the dirty bit
                PAGE_ACCESSED_SET | // prevent the CPU writing to the access flag
                PAGE_CACHE_ENABLED | // leave caching enabled
                PAGE_WRITE_BACK | // use write-back caching
                PAGE_USER_ACCESS_DISABLED | // dont allow user access (no code runs in user mode for now)
                page_rw_flag(bm.writable) | // R/W - set if writable
                PAGE_PRESENT // P   - this entry is present
        }
        MappingKind::Cow(cm) => {
            (mapping.phys_base + (r.vmin - mapping.virt_base)) |
                page_nx_flag(cm.executable) | // NX - no execute unless allowed
                PAGE_AVL_COW |
                PAGE_PAT_WB | // PAT index bit for write-back memory
                PAGE_DIRTY_SET | // prevent the CPU writing to the dirty bit
                PAGE_ACCESSED_SET | // prevent the CPU writing to the access flag
                PAGE_CACHE_ENABLED | // leave caching enabled
                PAGE_WRITE_BACK | // use write-back caching
                PAGE_USER_ACCESS_DISABLED | // dont allow user access (no code runs in user mode for now)
                0 | // R/W - Cow page is never writable
                PAGE_PRESENT // P   - this entry is present
        }
        MappingKind::Unmapped => 0,
    };
    unsafe {
        write_entry_updating(op, r.update_parent, r.entry_ptr, pte);
    }
}

// There are no notable architecture-specific safety considerations
// here, and the general conditions are documented in the
// architecture-independent re-export in vmem.rs

/// Maps a contiguous virtual address range to physical memory.
///
/// This function walks the 4-level page table hierarchy (PML4 → PDPT → PD → PT),
/// allocating intermediate tables as needed via `alloc_pte_if_needed`, and finally
/// writing the leaf page table entries with the requested permissions via `map_page`.
///
/// The iterator chain processes each level:
/// 1. PML4 (47:39) - allocate PDPT if needed
/// 2. PDPT (38:30) - allocate PD if needed
/// 3. PD (29:21) - allocate PT if needed
/// 4. PT (20:12) - write final PTE with physical address and flags
///
/// Multi-space page-table walking on amd64: walks each root
/// independently and emits all leaves as `ThisSpace`. Aliased
/// intermediate-table detection is not implemented here because no
/// current embedder exercises that pattern on amd64.
#[allow(clippy::missing_safety_doc)]
pub unsafe fn walk_va_spaces<Op: TableReadOps>(
    op: &Op,
    roots: &[Op::TableAddr],
    address: u64,
    len: u64,
) -> ::alloc::vec::Vec<(
    crate::vmem::SpaceId,
    ::alloc::vec::Vec<crate::vmem::SpaceAwareMapping>,
)> {
    use ::alloc::vec::Vec;

    let mut out: Vec<(crate::vmem::SpaceId, Vec<crate::vmem::SpaceAwareMapping>)> =
        Vec::with_capacity(roots.len());

    let addr = address & ((1u64 << VA_BITS) - 1);
    let vmin = addr & !(PAGE_SIZE as u64 - 1);
    let vmax = core::cmp::min(addr + len, 1u64 << VA_BITS);

    for &root in roots {
        #[allow(clippy::unnecessary_cast)]
        let root_id: crate::vmem::SpaceId = Op::to_phys(root) as u64;
        let mut mappings: Vec<crate::vmem::SpaceAwareMapping> = Vec::new();

        let iter = modify_ptes::<47, 39, Op, _>(MapRequest {
            table_base: root,
            vmin,
            len: vmax.saturating_sub(vmin),
            update_parent: UpdateParentNone {},
        })
        .filter_map(|r| unsafe { require_pte_exist(op, r) })
        .flat_map(modify_ptes::<38, 30, Op, _>)
        .filter_map(|r| unsafe { require_pte_exist(op, r) })
        .flat_map(modify_ptes::<29, 21, Op, _>)
        .filter_map(|r| unsafe { require_pte_exist(op, r) })
        .flat_map(modify_ptes::<20, 12, Op, _>);

        for r in iter {
            let Some(pte) = (unsafe { read_pte_if_present(op, r.entry_ptr) }) else {
                continue;
            };
            let phys_addr = pte & PTE_ADDR_MASK;
            let sgn_bit = r.vmin >> (VA_BITS - 1);
            let sgn_bits = 0u64.wrapping_sub(sgn_bit) << VA_BITS;
            let virt_addr = sgn_bits | r.vmin;

            let executable = (pte & PAGE_NX) == 0;
            let avl = pte & PTE_AVL_MASK;
            let kind = if avl == PAGE_AVL_COW {
                MappingKind::Cow(CowMapping {
                    readable: true,
                    executable,
                })
            } else {
                MappingKind::Basic(BasicMapping {
                    readable: true,
                    writable: (pte & PAGE_RW) != 0,
                    executable,
                })
            };
            mappings.push(crate::vmem::SpaceAwareMapping::ThisSpace(Mapping {
                phys_base: phys_addr,
                virt_base: virt_addr,
                len: PAGE_SIZE as u64,
                kind,
            }));
        }

        out.push((root_id, mappings));
    }

    out
}

/// See [`walk_va_spaces`]: amd64 never emits `AnotherSpace`, so this
/// is unreachable in practice. It silently no-ops (rather than
/// panicking) to keep the architecture-independent re-export usable.
#[allow(clippy::missing_safety_doc)]
pub unsafe fn space_aware_map<Op: TableOps>(
    _op: &Op,
    _ref_map: crate::vmem::SpaceReferenceMapping,
    _built_roots: &::alloc::collections::BTreeMap<crate::vmem::SpaceId, Op::TableAddr>,
) {
}

#[allow(clippy::missing_safety_doc)]
pub unsafe fn map<Op: TableOps>(op: &Op, mapping: Mapping) {
    modify_ptes::<47, 39, Op, _>(MapRequest {
        table_base: op.root_table(),
        vmin: mapping.virt_base,
        len: mapping.len,
        update_parent: Op::TableMovability::root_update_parent(),
    })
    .map(|r| unsafe { alloc_pte_if_needed(op, r) })
    .flat_map(modify_ptes::<38, 30, Op, _>)
    .map(|r| unsafe { alloc_pte_if_needed(op, r) })
    .flat_map(modify_ptes::<29, 21, Op, _>)
    .map(|r| unsafe { alloc_pte_if_needed(op, r) })
    .flat_map(modify_ptes::<20, 12, Op, _>)
    .map(|r| unsafe { map_page(op, &mapping, r) })
    .for_each(drop);
}

// There are no notable architecture-specific safety considerations
// here, and the general conditions are documented in the
// architecture-independent re-export in vmem.rs

/// Translates a virtual address range to the physical address pages
/// that back it by walking the page tables.
///
/// Returns an iterator with an entry for each mapped page that
/// intersects the given range.
///
/// This takes AsRef<Op> + Copy so that on targets where the
/// operations have little state (e.g. the guest) the operations state
/// can be copied into the closure(s) in the iterator, allowing for a
/// nicer result lifetime.  On targets like the
/// building-an-original-snapshot portion of the host, where the
/// operations structure owns a large buffer, a reference can instead
/// be passed.
#[allow(clippy::missing_safety_doc)]
pub unsafe fn virt_to_phys<'a, Op: TableReadOps + 'a>(
    op: impl core::convert::AsRef<Op> + Copy + 'a,
    address: u64,
    len: u64,
) -> impl Iterator<Item = Mapping> + 'a {
    // Undo sign-extension
    let addr = address & ((1u64 << VA_BITS) - 1);
    // Mask off any sub-page bits
    let vmin = addr & !(PAGE_SIZE as u64 - 1);
    // Calculate the maximum virtual address we need to look at based on the starting
    // address and length ensuring we don't go past the end of the address space
    let vmax = core::cmp::min(addr + len, 1u64 << VA_BITS);
    modify_ptes::<47, 39, Op, _>(MapRequest {
        table_base: op.as_ref().root_table(),
        vmin,
        len: vmax - vmin,
        update_parent: UpdateParentNone {},
    })
    .filter_map(move |r| unsafe { require_pte_exist(op.as_ref(), r) })
    .flat_map(modify_ptes::<38, 30, Op, _>)
    .filter_map(move |r| unsafe { require_pte_exist(op.as_ref(), r) })
    .flat_map(modify_ptes::<29, 21, Op, _>)
    .filter_map(move |r| unsafe { require_pte_exist(op.as_ref(), r) })
    .flat_map(modify_ptes::<20, 12, Op, _>)
    .filter_map(move |r| {
        let pte = unsafe { read_pte_if_present(op.as_ref(), r.entry_ptr) }?;
        let phys_addr = pte & PTE_ADDR_MASK;
        // Re-do the sign extension
        let sgn_bit = r.vmin >> (VA_BITS - 1);
        let sgn_bits = 0u64.wrapping_sub(sgn_bit) << VA_BITS;
        let virt_addr = sgn_bits | r.vmin;

        let executable = (pte & PAGE_NX) == 0;
        let avl = pte & PTE_AVL_MASK;
        let kind = if avl == PAGE_AVL_COW {
            MappingKind::Cow(CowMapping {
                readable: true,
                executable,
            })
        } else {
            MappingKind::Basic(BasicMapping {
                readable: true,
                writable: (pte & PAGE_RW) != 0,
                executable,
            })
        };
        Some(Mapping {
            phys_base: phys_addr,
            virt_base: virt_addr,
            len: PAGE_SIZE as u64,
            kind,
        })
    })
}

const VA_BITS: usize = 48; // We use 48-bit virtual addresses at the moment.

pub const PAGE_SIZE: usize = 4096;
pub const PAGE_TABLE_SIZE: usize = 4096;
pub type PageTableEntry = u64;
pub type VirtAddr = u64;
pub type PhysAddr = u64;

#[cfg(test)]
mod tests {
    use alloc::vec;
    use alloc::vec::Vec;
    use core::cell::RefCell;

    use super::*;
    use crate::vmem::{
        BasicMapping, Mapping, MappingKind, MayNotMoveTable, PAGE_TABLE_ENTRIES_PER_TABLE,
        TableOps, TableReadOps, Void, bits,
    };

    /// A mock TableOps implementation for testing that stores page tables in memory
    /// needed because the `GuestPageTableBuffer` is in hyperlight_host which would cause a circular dependency
    struct MockTableOps {
        tables: RefCell<Vec<[u64; PAGE_TABLE_ENTRIES_PER_TABLE]>>,
    }

    // for virt_to_phys
    impl core::convert::AsRef<MockTableOps> for MockTableOps {
        fn as_ref(&self) -> &Self {
            self
        }
    }

    impl MockTableOps {
        fn new() -> Self {
            // Start with one table (the root/PML4)
            Self {
                tables: RefCell::new(vec![[0u64; PAGE_TABLE_ENTRIES_PER_TABLE]]),
            }
        }

        fn table_count(&self) -> usize {
            self.tables.borrow().len()
        }

        fn get_entry(&self, table_idx: usize, entry_idx: usize) -> u64 {
            self.tables.borrow()[table_idx][entry_idx]
        }
    }

    impl TableReadOps for MockTableOps {
        type TableAddr = (usize, usize); // (table_index, entry_index)

        fn entry_addr(addr: Self::TableAddr, entry_offset: u64) -> Self::TableAddr {
            // Convert to physical address, add offset, convert back
            let phys = Self::to_phys(addr) + entry_offset;
            Self::from_phys(phys)
        }

        unsafe fn read_entry(&self, addr: Self::TableAddr) -> u64 {
            self.tables.borrow()[addr.0][addr.1]
        }

        fn to_phys(addr: Self::TableAddr) -> PhysAddr {
            // Each table is 4KB, entries are 8 bytes
            (addr.0 as u64 * PAGE_TABLE_SIZE as u64) + (addr.1 as u64 * 8)
        }

        fn from_phys(addr: PhysAddr) -> Self::TableAddr {
            let table_idx = (addr / PAGE_TABLE_SIZE as u64) as usize;
            let entry_idx = ((addr % PAGE_TABLE_SIZE as u64) / 8) as usize;
            (table_idx, entry_idx)
        }

        fn root_table(&self) -> Self::TableAddr {
            (0, 0)
        }
    }

    impl TableOps for MockTableOps {
        type TableMovability = MayNotMoveTable;

        unsafe fn alloc_table(&self) -> Self::TableAddr {
            let mut tables = self.tables.borrow_mut();
            let idx = tables.len();
            tables.push([0u64; PAGE_TABLE_ENTRIES_PER_TABLE]);
            (idx, 0)
        }

        unsafe fn write_entry(&self, addr: Self::TableAddr, entry: u64) -> Option<Void> {
            self.tables.borrow_mut()[addr.0][addr.1] = entry;
            None
        }

        unsafe fn update_root(&self, impossible: Void) {
            match impossible {}
        }
    }

    // ==================== bits() function tests ====================

    #[test]
    fn test_bits_extracts_pml4_index() {
        // PML4 uses bits 47:39
        // Address 0x0000_0080_0000_0000 should have PML4 index 1
        let addr: u64 = 0x0000_0080_0000_0000;
        assert_eq!(bits::<47, 39>(addr), 1);
    }

    #[test]
    fn test_bits_extracts_pdpt_index() {
        // PDPT uses bits 38:30
        // Address with PDPT index 1: bit 30 set = 0x4000_0000 (1GB)
        let addr: u64 = 0x4000_0000;
        assert_eq!(bits::<38, 30>(addr), 1);
    }

    #[test]
    fn test_bits_extracts_pd_index() {
        // PD uses bits 29:21
        // Address 0x0000_0000_0020_0000 (2MB) should have PD index 1
        let addr: u64 = 0x0000_0000_0020_0000;
        assert_eq!(bits::<29, 21>(addr), 1);
    }

    #[test]
    fn test_bits_extracts_pt_index() {
        // PT uses bits 20:12
        // Address 0x0000_0000_0000_1000 (4KB) should have PT index 1
        let addr: u64 = 0x0000_0000_0000_1000;
        assert_eq!(bits::<20, 12>(addr), 1);
    }

    #[test]
    fn test_bits_max_index() {
        // Maximum 9-bit index is 511
        // PML4 index 511 = bits 47:39 all set = 0x0000_FF80_0000_0000
        let addr: u64 = 0x0000_FF80_0000_0000;
        assert_eq!(bits::<47, 39>(addr), 511);
    }

    // ==================== PTE flag tests ====================

    #[test]
    fn test_page_rw_flag_writable() {
        assert_eq!(page_rw_flag(true), PAGE_RW);
    }

    #[test]
    fn test_page_rw_flag_readonly() {
        assert_eq!(page_rw_flag(false), 0);
    }

    #[test]
    fn test_page_nx_flag_executable() {
        assert_eq!(page_nx_flag(true), 0); // Executable = no NX bit
    }

    #[test]
    fn test_page_nx_flag_not_executable() {
        assert_eq!(page_nx_flag(false), PAGE_NX);
    }

    // ==================== map() function tests ====================

    #[test]
    fn test_map_single_page() {
        let ops = MockTableOps::new();
        let mapping = Mapping {
            phys_base: 0x1000,
            virt_base: 0x1000,
            len: PAGE_SIZE as u64,
            kind: MappingKind::Basic(BasicMapping {
                readable: true,
                writable: true,
                executable: false,
            }),
        };

        unsafe { map(&ops, mapping) };

        // Should have allocated: PML4(exists) + PDPT + PD + PT = 4 tables
        assert_eq!(ops.table_count(), 4);

        // Check PML4 entry 0 points to PDPT (table 1) with correct flags
        let pml4_entry = ops.get_entry(0, 0);
        assert_ne!(pml4_entry & PAGE_PRESENT, 0, "PML4 entry should be present");
        assert_ne!(pml4_entry & PAGE_RW, 0, "PML4 entry should be writable");

        // Check the leaf PTE has correct flags
        // PT is table 3, entry 1 (for virt_base 0x1000)
        let pte = ops.get_entry(3, 1);
        assert_ne!(pte & PAGE_PRESENT, 0, "PTE should be present");
        assert_ne!(pte & PAGE_RW, 0, "PTE should be writable");
        assert_ne!(pte & PAGE_NX, 0, "PTE should have NX set (not executable)");
        assert_eq!(pte & PTE_ADDR_MASK, 0x1000, "PTE should map to phys 0x1000");
    }

    #[test]
    fn test_map_executable_page() {
        let ops = MockTableOps::new();
        let mapping = Mapping {
            phys_base: 0x2000,
            virt_base: 0x2000,
            len: PAGE_SIZE as u64,
            kind: MappingKind::Basic(BasicMapping {
                readable: true,
                writable: false,
                executable: true,
            }),
        };

        unsafe { map(&ops, mapping) };

        // PT is table 3, entry 2 (for virt_base 0x2000)
        let pte = ops.get_entry(3, 2);
        assert_ne!(pte & PAGE_PRESENT, 0, "PTE should be present");
        assert_eq!(pte & PAGE_RW, 0, "PTE should be read-only");
        assert_eq!(pte & PAGE_NX, 0, "PTE should NOT have NX set (executable)");
    }

    #[test]
    fn test_map_multiple_pages() {
        let ops = MockTableOps::new();
        let mapping = Mapping {
            phys_base: 0x10000,
            virt_base: 0x10000,
            len: 4 * PAGE_SIZE as u64, // 4 pages = 16KB
            kind: MappingKind::Basic(BasicMapping {
                readable: true,
                writable: true,
                executable: false,
            }),
        };

        unsafe { map(&ops, mapping) };

        // Check all 4 PTEs are present
        for i in 0..4 {
            let entry_idx = 16 + i; // 0x10000 / 0x1000 = 16
            let pte = ops.get_entry(3, entry_idx);
            assert_ne!(pte & PAGE_PRESENT, 0, "PTE {} should be present", i);
            let expected_phys = 0x10000 + (i as u64 * PAGE_SIZE as u64);
            assert_eq!(
                pte & PTE_ADDR_MASK,
                expected_phys,
                "PTE {} should map to correct phys addr",
                i
            );
        }
    }

    #[test]
    fn test_map_reuses_existing_tables() {
        let ops = MockTableOps::new();

        // Map first region
        let mapping1 = Mapping {
            phys_base: 0x1000,
            virt_base: 0x1000,
            len: PAGE_SIZE as u64,
            kind: MappingKind::Basic(BasicMapping {
                readable: true,
                writable: true,
                executable: false,
            }),
        };
        unsafe { map(&ops, mapping1) };
        let tables_after_first = ops.table_count();

        // Map second region in same PT (different page)
        let mapping2 = Mapping {
            phys_base: 0x5000,
            virt_base: 0x5000,
            len: PAGE_SIZE as u64,
            kind: MappingKind::Basic(BasicMapping {
                readable: true,
                writable: true,
                executable: false,
            }),
        };
        unsafe { map(&ops, mapping2) };

        // Should NOT allocate new tables (reuses existing hierarchy)
        assert_eq!(
            ops.table_count(),
            tables_after_first,
            "Should reuse existing page tables"
        );
    }

    // ==================== virt_to_phys() tests ====================

    #[test]
    fn test_virt_to_phys_mapped_address() {
        let ops = MockTableOps::new();
        let mapping = Mapping {
            phys_base: 0x1000,
            virt_base: 0x1000,
            len: PAGE_SIZE as u64,
            kind: MappingKind::Basic(BasicMapping {
                readable: true,
                writable: true,
                executable: false,
            }),
        };

        unsafe { map(&ops, mapping) };

        let result = unsafe { virt_to_phys(&ops, 0x1000, 1).next() };
        assert!(result.is_some(), "Should find mapped address");
        let mapping = result.unwrap();
        assert_eq!(mapping.phys_base, 0x1000);
    }

    #[test]
    fn test_virt_to_phys_unaligned_virt() {
        let ops = MockTableOps::new();
        let mapping = Mapping {
            phys_base: 0x1000,
            virt_base: 0x1000,
            len: PAGE_SIZE as u64,
            kind: MappingKind::Basic(BasicMapping {
                readable: true,
                writable: true,
                executable: false,
            }),
        };

        unsafe { map(&ops, mapping) };

        let result = unsafe { virt_to_phys(&ops, 0x1234, 1).next() };
        assert!(result.is_some(), "Should find mapped address");
        let mapping = result.unwrap();
        assert_eq!(mapping.phys_base, 0x1000);
    }

    #[test]
    fn test_virt_to_phys_unaligned_virt_and_across_pages_len() {
        let ops = MockTableOps::new();
        let mapping = Mapping {
            phys_base: 0x1000,
            virt_base: 0x1000,
            len: 2 * PAGE_SIZE as u64, // 2 page
            kind: MappingKind::Basic(BasicMapping {
                readable: true,
                writable: true,
                executable: false,
            }),
        };

        unsafe { map(&ops, mapping) };

        let mappings = unsafe { virt_to_phys(&ops, 0x1F00, 0x300).collect::<Vec<_>>() };
        assert_eq!(mappings.len(), 2, "Should return 2 mappings for 2 pages");
        assert_eq!(mappings[0].phys_base, 0x1000);
        assert_eq!(mappings[1].phys_base, 0x2000);
    }

    #[test]
    fn test_virt_to_phys_unaligned_virt_and_multiple_page_len() {
        let ops = MockTableOps::new();
        let mapping = Mapping {
            phys_base: 0x1000,
            virt_base: 0x1000,
            len: PAGE_SIZE as u64 * 2 + 0x200, // 2 page + 512 bytes
            kind: MappingKind::Basic(BasicMapping {
                readable: true,
                writable: true,
                executable: false,
            }),
        };

        unsafe { map(&ops, mapping) };

        let mappings =
            unsafe { virt_to_phys(&ops, 0x1234, PAGE_SIZE as u64 * 2 + 0x10).collect::<Vec<_>>() };
        assert_eq!(mappings.len(), 3, "Should return 3 mappings for 3 pages");
        assert_eq!(mappings[0].phys_base, 0x1000);
        assert_eq!(mappings[1].phys_base, 0x2000);
        assert_eq!(mappings[2].phys_base, 0x3000);
    }

    #[test]
    fn test_virt_to_phys_perms() {
        let test = |kind| {
            let ops = MockTableOps::new();
            let mapping = Mapping {
                phys_base: 0x1000,
                virt_base: 0x1000,
                len: PAGE_SIZE as u64,
                kind,
            };
            unsafe { map(&ops, mapping) };
            let result = unsafe { virt_to_phys(&ops, 0x1000, 1).next() };
            let mapping = result.unwrap();
            assert_eq!(mapping.kind, kind);
        };
        test(MappingKind::Basic(BasicMapping {
            readable: true,
            writable: false,
            executable: false,
        }));
        test(MappingKind::Basic(BasicMapping {
            readable: true,
            writable: false,
            executable: true,
        }));
        test(MappingKind::Basic(BasicMapping {
            readable: true,
            writable: true,
            executable: false,
        }));
        test(MappingKind::Basic(BasicMapping {
            readable: true,
            writable: true,
            executable: true,
        }));
        test(MappingKind::Cow(CowMapping {
            readable: true,
            executable: false,
        }));
        test(MappingKind::Cow(CowMapping {
            readable: true,
            executable: true,
        }));
    }

    #[test]
    fn test_virt_to_phys_unmapped_address() {
        let ops = MockTableOps::new();
        // Don't map anything

        let result = unsafe { virt_to_phys(&ops, 0x1000, 1).next() };
        assert!(result.is_none(), "Should return None for unmapped address");
    }

    #[test]
    fn test_virt_to_phys_partially_mapped() {
        let ops = MockTableOps::new();
        let mapping = Mapping {
            phys_base: 0x1000,
            virt_base: 0x1000,
            len: PAGE_SIZE as u64,
            kind: MappingKind::Basic(BasicMapping {
                readable: true,
                writable: true,
                executable: false,
            }),
        };

        unsafe { map(&ops, mapping) };

        // Query an address in a different PT entry (unmapped)
        let result = unsafe { virt_to_phys(&ops, 0x5000, 1).next() };
        assert!(
            result.is_none(),
            "Should return None for unmapped address in same PT"
        );
    }

    // ==================== ModifyPteIterator tests ====================

    #[test]
    fn test_modify_pte_iterator_single_page() {
        let ops = MockTableOps::new();
        let request = MapRequest {
            table_base: ops.root_table(),
            vmin: 0x1000,
            len: PAGE_SIZE as u64,
            update_parent: UpdateParentNone {},
        };

        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps, _>(request).collect();
        assert_eq!(responses.len(), 1, "Single page should yield one response");
        assert_eq!(responses[0].vmin, 0x1000);
        assert_eq!(responses[0].len, PAGE_SIZE as u64);
    }

    #[test]
    fn test_modify_pte_iterator_multiple_pages() {
        let ops = MockTableOps::new();
        let request = MapRequest {
            table_base: ops.root_table(),
            vmin: 0x1000,
            len: 3 * PAGE_SIZE as u64,
            update_parent: UpdateParentNone {},
        };

        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps, _>(request).collect();
        assert_eq!(responses.len(), 3, "3 pages should yield 3 responses");
    }

    #[test]
    fn test_modify_pte_iterator_zero_length() {
        let ops = MockTableOps::new();
        let request = MapRequest {
            table_base: ops.root_table(),
            vmin: 0x1000,
            len: 0,
            update_parent: UpdateParentNone {},
        };

        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps, _>(request).collect();
        assert_eq!(responses.len(), 0, "Zero length should yield no responses");
    }

    #[test]
    fn test_modify_pte_iterator_unaligned_start() {
        let ops = MockTableOps::new();
        // Start at 0x1800 (mid-page), map 0x1000 bytes
        // Should cover 0x1800-0x1FFF (first page) and 0x2000-0x27FF (second page)
        let request = MapRequest {
            table_base: ops.root_table(),
            vmin: 0x1800,
            len: 0x1000,
            update_parent: UpdateParentNone {},
        };

        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps, _>(request).collect();
        assert_eq!(
            responses.len(),
            2,
            "Unaligned mapping spanning 2 pages should yield 2 responses"
        );
        assert_eq!(responses[0].vmin, 0x1800);
        assert_eq!(responses[0].len, 0x800); // Remaining in first page
        assert_eq!(responses[1].vmin, 0x2000);
        assert_eq!(responses[1].len, 0x800); // Continuing in second page
    }

    // ==================== TableOps entry_addr tests ====================

    #[test]
    fn test_entry_addr_from_table_base() {
        // entry_addr is called with a table base (entry_index = 0) and a byte offset
        // offset = entry_index * 8, so offset 40 means entry 5
        let result = MockTableOps::entry_addr((2, 0), 40);
        assert_eq!(result, (2, 5), "Should return (table 2, entry 5)");
    }

    #[test]
    fn test_entry_addr_with_nonzero_base_entry() {
        // Even though entry_addr is typically called with entry_index=0,
        // it should handle non-zero base correctly by adding the offset
        // Base: table 1, entry 10 (phys = 1*4096 + 10*8 = 4176)
        // Offset: 16 bytes (2 entries)
        // Result phys: 4176 + 16 = 4192 = 1*4096 + 12*8 → (1, 12)
        let result = MockTableOps::entry_addr((1, 10), 16);
        assert_eq!(result, (1, 12), "Should add offset to base entry");
    }

    #[test]
    fn test_to_phys_from_phys_roundtrip() {
        // Verify to_phys and from_phys are inverses
        let addr = (3, 42);
        let phys = MockTableOps::to_phys(addr);
        let back = MockTableOps::from_phys(phys);
        assert_eq!(back, addr, "to_phys/from_phys should roundtrip");
    }
}