strat9_kernel/memory/

address_space.rs

//! Per-process address spaces for Strat9-OS.
//!
//! Each task owns an `AddressSpace` backed by a PML4 page table.
//! Kernel tasks share a single kernel address space. User tasks get a fresh
//! PML4 with the kernel half (entries 256..512) cloned from the kernel's table.
//!
//! PML4 is the Page Map Level 4, the top-level page table in x86_64's 4-level paging scheme. It
//! contains 512 entries, each covering a 512 GiB region of the virtual address space. By cloning the
//! kernel half of the PML4, each user address space automatically shares the kernel mappings without needing to duplicate them.
//! Source : https://wiki.osdev.org/Memory_Management
//!
//!
//! x86_64 virtual address space layout:
//!   - PML4[0..256]   => user space (per-process, zeroed for new AS)
//!   - PML4[256..512] => kernel space (shared, cloned from kernel L4)
//!

use alloc::{collections::BTreeMap, sync::Arc, vec::Vec};
use core::sync::atomic::{AtomicU32, Ordering};

use spin::Once;
use x86_64::{
    registers::control::{Cr3, Cr3Flags},
    structures::paging::{
        mapper::TranslateResult, Mapper, OffsetPageTable, Page, PageTable, PageTableFlags,
        PhysFrame as X86PhysFrame, Size2MiB, Size4KiB, Translate,
    },
    PhysAddr, VirtAddr,
};

use crate::{
    capability::CapId,
    memory::{
        allocate_mapping_cap_id, mapping_index, paging::BuddyFrameAllocator, release_owned_block,
        resolve_handle, try_register_mapping_identity, unregister_mapping_identity, BlockHandle,
        MappingRef,
    },
    process::task::Pid,
    sync::SpinLock,
};

/// Flags describing permissions for a virtual memory region.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct VmaFlags {
    pub readable: bool,
    pub writable: bool,
    pub executable: bool,
    pub user_accessible: bool,
}

impl VmaFlags {
    /// Convert to x86_64 page table flags.
    pub fn to_page_flags(self) -> PageTableFlags {
        let mut flags = PageTableFlags::PRESENT;
        if self.writable {
            flags |= PageTableFlags::WRITABLE;
        }
        if !self.executable {
            flags |= PageTableFlags::NO_EXECUTE;
        }
        if self.user_accessible {
            flags |= PageTableFlags::USER_ACCESSIBLE;
        }
        flags
    }
}

/// Type/purpose of a virtual memory region.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VmaType {
    /// Zero-filled anonymous memory (heap, mmap).
    Anonymous,
    /// Stack region (grows downward).
    Stack,
    /// Code/text segment (typically RX).
    Code,
    /// Kernel-internal mapping.
    Kernel,
}

/// Supported page sizes for VMAs.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VmaPageSize {
    /// Standard 4 KiB page.
    Small,
    /// Huge 2 MiB page.
    Huge,
}

impl VmaPageSize {
    /// Performs the bytes operation.
    pub fn bytes(self) -> u64 {
        match self {
            VmaPageSize::Small => 4096,
            VmaPageSize::Huge => 2 * 1024 * 1024,
        }
    }
}

/// A tracked virtual memory region within an address space.
#[derive(Debug, Clone)]
pub struct VirtualMemoryRegion {
    /// Start virtual address (page-aligned).
    pub start: u64,
    /// Number of pages in this region (size depends on `page_size`).
    pub page_count: usize,
    /// Access permissions.
    pub flags: VmaFlags,
    /// Purpose of this region.
    pub vma_type: VmaType,
    /// Size of each page in this region.
    pub page_size: VmaPageSize,
}

/// An effective mapping currently installed in the page tables.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EffectiveMapping {
    /// Start virtual address of the mapping.
    pub start: u64,
    /// Internal capability identifier associated with this mapping.
    pub cap_id: CapId,
    /// Physical block currently backing this mapping.
    pub handle: BlockHandle,
    /// Hardware page-table flags currently installed for this mapping.
    pub flags: PageTableFlags,
    /// Page size of the mapping.
    pub page_size: VmaPageSize,
}

/// A per-process address space backed by a PML4 page table.
///
/// Kernel tasks share a single `AddressSpace` (the kernel AS).
/// User tasks each get their own, with kernel entries (PML4[256..512]) cloned
/// so that the kernel is always mapped regardless of which AS is active.
///
/// # Lock ordering
///
/// `regions` and `effective_mappings` are `SpinLock<BTreeMap<...>>`.
/// Both BTreeMap insert/remove may allocate from the heap allocator.
/// This is safe because:
///
///   1. **No IRQ-reachable path.** These locks are only taken from process
///      context (syscall handlers, fork, munmap, page-fault). No interrupt
///      handler or softirq acquires either lock, so IRQ-disabling is not
///      required and cannot deadlock with the heap allocator.
///
///   2. **Heap lock order is consistent.** The implicit heap lock is always
///      innermost (address_space => heap), never outermost. No other lock is
///      acquired between the SpinLock hold and the potential heap alloc,
///      so no ABBA cycle exists.
///
///   3. **Bounded contention.** Each lock protects a per-process map; only
///      one task at a time operates on a given address space, so contention
///      is low and hold times are short.
///
/// If an IRQ or NMI path ever needs to query these maps in the future,
/// the `SpinLock` must be replaced with a lock-free snapshot or a `Mutex`
/// (sleepable, IRQ-safe via `spin::Mutex`).
pub struct AddressSpace {
    /// Physical address of the PML4 table (loaded into CR3).
    cr3_phys: PhysAddr,
    /// Virtual address of the PML4 table (via HHDM, for reading/modifying).
    l4_table_virt: VirtAddr,
    /// Whether this is the kernel address space (never freed).
    is_kernel: bool,
    /// Tracked virtual memory regions (key = start address).
    regions: SpinLock<BTreeMap<u64, VirtualMemoryRegion>>,
    /// Tracked effective mappings (key = mapping start address).
    effective_mappings: SpinLock<BTreeMap<u64, EffectiveMapping>>,
    /// Process identifier owning this address space, when bound to a process.
    owner_pid: AtomicU32,
}

// SAFETY: AddressSpace is protected by the scheduler lock and per-task ownership.
// The PML4 table is accessed through HHDM virtual addresses which are valid on all CPUs.
unsafe impl Send for AddressSpace {}
unsafe impl Sync for AddressSpace {}

impl AddressSpace {
    /// Create the kernel address space by wrapping the current (boot) CR3.
    ///
    /// # Safety
    /// Must be called exactly once, during single-threaded init, after paging is initialized.
    pub unsafe fn new_kernel() -> Self {
        let (level_4_frame, _flags) = Cr3::read();
        let cr3_phys = level_4_frame.start_address();
        let l4_table_virt = VirtAddr::new(crate::memory::phys_to_virt(cr3_phys.as_u64()));

        log::info!(
            "Kernel address space initialized: CR3={:#x}",
            cr3_phys.as_u64()
        );

        AddressSpace {
            cr3_phys,
            l4_table_virt,
            is_kernel: true,
            regions: SpinLock::new(BTreeMap::new()),
            effective_mappings: SpinLock::new(BTreeMap::new()),
            owner_pid: AtomicU32::new(0),
        }
    }

    /// Create a new user address space with the kernel half cloned.
    ///
    /// Allocates a fresh PML4 frame, zeroes it, then copies entries 256..512
    /// from the kernel PML4. This shares the kernel's L3/L2/L1 subtrees so
    /// kernel mapping changes propagate automatically.
    pub fn new_user() -> Result<Self, &'static str> {
        // Allocate a frame for the new PML4 table.
        let new_l4_phys =
            crate::sync::with_irqs_disabled(|token| crate::memory::allocate_frame(token))
                .map_err(|_| "Failed to allocate PML4 frame")?
                .start_address;

        let new_l4_virt = VirtAddr::new(crate::memory::phys_to_virt(new_l4_phys.as_u64()));

        // Zero the entire table first (clears user-half entries 0..256).
        // SAFETY: new_l4_virt points to a freshly allocated, HHDM-mapped frame.
        unsafe {
            core::ptr::write_bytes(new_l4_virt.as_mut_ptr::<u8>(), 0, 4096);
        }

        // Clone kernel entries (PML4[256..512]) from the kernel's L4 table.
        let kernel_l4_phys = crate::memory::paging::kernel_l4_phys();
        let kernel_l4_virt = VirtAddr::new(crate::memory::phys_to_virt(kernel_l4_phys.as_u64()));

        // SAFETY: Both pointers are valid HHDM-mapped page tables. We only read
        // from the kernel table and write to the freshly allocated table.
        unsafe {
            let kernel_l4 = &*(kernel_l4_virt.as_ptr::<PageTable>());
            let new_l4 = &mut *(new_l4_virt.as_mut_ptr::<PageTable>());
            for i in 256..512 {
                new_l4[i] = kernel_l4[i].clone();
            }
        }

        // ---------- LAPIC low-half mapping (HHDM=0 workaround) ----------
        //
        // When Limine provides a non-zero HHDM offset the LAPIC is mapped in
        // PML4[256..512] (kernel half) and is already shared above.
        //
        // When HHDM=0 the LAPIC is identity-mapped at its physical address
        // (0xFEE00000) in the low half (PML4[0]).  Every Ring-0 interrupt
        // handler calls apic::eoi() which writes to this address.  If the
        // handler fires while a user CR3 is active the write faults because
        // PML4[0] is absent in the user page tables.
        //
        // Fix: map just the LAPIC 4KiB MMIO page into every new user AS using
        // a fresh private L3/L2/L1 hierarchy (no sharing with the kernel's
        // page table subtrees at the LAPIC virtual address).
        {
            let lapic_phys = crate::arch::x86_64::apic::lapic_phys();
            if lapic_phys != 0 {
                let lapic_virt = crate::memory::phys_to_virt(lapic_phys);
                // Only needed when LAPIC is in the low half.
                if lapic_virt < 0xFFFF_8000_0000_0000 {
                    let phys_offset = VirtAddr::new(crate::memory::hhdm_offset());
                    // SAFETY: new_l4_virt is the freshly allocated user PML4.
                    let l4 = unsafe { &mut *new_l4_virt.as_mut_ptr::<PageTable>() };
                    let mut mapper = unsafe { OffsetPageTable::new(l4, phys_offset) };
                    let mut buddy = crate::memory::paging::BuddyFrameAllocator;
                    let mmio_flags = PageTableFlags::PRESENT
                        | PageTableFlags::WRITABLE
                        | PageTableFlags::NO_CACHE;
                    let lapic_page =
                        Page::<Size4KiB>::containing_address(VirtAddr::new(lapic_virt));
                    let lapic_frame =
                        X86PhysFrame::<Size4KiB>::containing_address(PhysAddr::new(lapic_phys));
                    // Use map_to_with_table_flags to avoid USER_ACCESSIBLE on
                    // intermediate tables so user code cannot reach LAPIC MMIO.
                    match unsafe { mapper.map_to(lapic_page, lapic_frame, mmio_flags, &mut buddy) }
                    {
                        Ok(flush) => flush.flush(),
                        Err(e) => {
                            crate::serial_println!(
                                "[as] WARN: failed to map LAPIC ({:#x}) in user AS: {:?}",
                                lapic_phys,
                                e
                            );
                        }
                    }
                }
            }
        }

        log::debug!(
            "User address space created: CR3={:#x} (kernel entries cloned from {:#x})",
            new_l4_phys.as_u64(),
            kernel_l4_phys.as_u64()
        );

        Ok(AddressSpace {
            cr3_phys: new_l4_phys,
            l4_table_virt: new_l4_virt,
            is_kernel: false,
            regions: SpinLock::new(BTreeMap::new()),
            effective_mappings: SpinLock::new(BTreeMap::new()),
            owner_pid: AtomicU32::new(0),
        })
    }

    /// Registers an effective mapping in the address space tracking table.
    pub fn register_effective_mapping(
        &self,
        mapping: EffectiveMapping,
    ) -> Result<(), &'static str> {
        let previous_at_start = self.effective_mapping_by_start(mapping.start);
        if let Some(previous) = previous_at_start {
            if previous.handle == mapping.handle && previous.cap_id == mapping.cap_id {
                self.effective_mappings
                    .lock()
                    .insert(mapping.start, mapping);
                if let Some(pid) = self.owner_pid() {
                    mapping_index().unregister(mapping.cap_id, pid, VirtAddr::new(mapping.start));
                    mapping_index().register(
                        mapping.cap_id,
                        MappingRef {
                            pid,
                            vaddr: VirtAddr::new(mapping.start),
                            page_size: mapping.page_size,
                        },
                    );
                }
                return Ok(());
            }
        }

        if let Err(error) = try_register_mapping_identity(mapping.handle, mapping.cap_id) {
            if error != crate::memory::OwnerError::CapAlreadyPresent {
                log::warn!(
                    "memory: failed to register effective mapping identity cap={} block={:#x}/{} vaddr={:#x}: {:?}",
                    mapping.cap_id.as_u64(),
                    mapping.handle.base.as_u64(),
                    mapping.handle.order,
                    mapping.start,
                    error
                );
                return Err("Failed to register effective mapping identity");
            }
        }

        let replaced = self
            .effective_mappings
            .lock()
            .insert(mapping.start, mapping);
        if let Some(previous) = replaced {
            if let Some(block) = unregister_mapping_identity(previous.handle, previous.cap_id) {
                release_owned_block(block);
            }
            if let Some(pid) = self.owner_pid() {
                mapping_index().unregister(previous.cap_id, pid, VirtAddr::new(previous.start));
            }
        }

        if let Some(pid) = self.owner_pid() {
            mapping_index().register(
                mapping.cap_id,
                MappingRef {
                    pid,
                    vaddr: VirtAddr::new(mapping.start),
                    page_size: mapping.page_size,
                },
            );
        }
        Ok(())
    }

    /// Removes an effective mapping from the address space tracking table.
    pub fn unregister_effective_mapping(&self, start: u64) -> Option<EffectiveMapping> {
        let mapping = self.effective_mappings.lock().remove(&start);
        if let Some(mapping) = mapping {
            if let Some(block) = unregister_mapping_identity(mapping.handle, mapping.cap_id) {
                release_owned_block(block);
            }
            if let Some(pid) = self.owner_pid() {
                mapping_index().unregister(mapping.cap_id, pid, VirtAddr::new(mapping.start));
            }
            Some(mapping)
        } else {
            None
        }
    }

    /// Updates the hardware flags recorded for an effective mapping.
    pub fn update_effective_mapping_flags(&self, start: u64, flags: PageTableFlags) -> bool {
        if let Some(mapping) = self.effective_mappings.lock().get_mut(&start) {
            mapping.flags = flags;
            true
        } else {
            false
        }
    }

    /// Returns the effective mapping that starts exactly at `start`.
    pub fn effective_mapping_by_start(&self, start: u64) -> Option<EffectiveMapping> {
        self.effective_mappings.lock().get(&start).copied()
    }

    /// Unmaps the effective mapping that starts at `start`.
    pub fn unmap_effective_mapping(&self, start: u64) -> Result<(), &'static str> {
        let mapping = self
            .effective_mapping_by_start(start)
            .ok_or("Mapping not found")?;
        self.unmap_range(start, mapping.page_size.bytes())
    }

    /// Returns the effective mapping covering `addr`, if any.
    pub fn effective_mapping_containing(&self, addr: u64) -> Option<EffectiveMapping> {
        let mappings = self.effective_mappings.lock();
        if let Some(mapping) = mappings.get(&(addr & !(VmaPageSize::Small.bytes() - 1))) {
            if mapping.page_size == VmaPageSize::Small {
                return Some(*mapping);
            }
        }
        mappings
            .get(&(addr & !(VmaPageSize::Huge.bytes() - 1)))
            .copied()
    }

    /// Binds this address space to the given process identifier.
    pub fn set_owner_pid(&self, pid: Pid) {
        let previous = self.owner_pid.swap(pid, Ordering::Relaxed);
        let mappings: Vec<EffectiveMapping> = {
            let guard = self.effective_mappings.lock();
            guard.values().copied().collect()
        };

        if previous != 0 && previous != pid {
            for mapping in mappings.iter().copied() {
                mapping_index().unregister(mapping.cap_id, previous, VirtAddr::new(mapping.start));
            }
        }

        if pid != 0 {
            for mapping in mappings {
                mapping_index().register(
                    mapping.cap_id,
                    MappingRef {
                        pid,
                        vaddr: VirtAddr::new(mapping.start),
                        page_size: mapping.page_size,
                    },
                );
            }
        }
    }

    /// Returns the owning process identifier, if one has been assigned.
    pub fn owner_pid(&self) -> Option<Pid> {
        match self.owner_pid.load(Ordering::Relaxed) {
            0 => None,
            pid => Some(pid),
        }
    }

    /// Construct a temporary `OffsetPageTable` mapper for this address space.
    ///
    /// # Safety
    /// The caller must ensure exclusive access to the page tables (e.g. via
    /// the scheduler lock or single-threaded context).
    pub(crate) unsafe fn mapper(&self) -> OffsetPageTable<'_> {
        let phys_offset = VirtAddr::new(crate::memory::hhdm_offset());
        // SAFETY: l4_table_virt is the HHDM-mapped address of our PML4.
        // The caller guarantees exclusive access.
        unsafe {
            OffsetPageTable::new(
                &mut *self.l4_table_virt.as_mut_ptr::<PageTable>(),
                phys_offset,
            )
        }
    }

    /// Reserve a contiguous region of virtual pages without allocating physical frames.
    ///
    /// The pages will be mapped lazily during page faults (Demand Paging).
    pub fn reserve_region(
        &self,
        start: u64,
        page_count: usize,
        flags: VmaFlags,
        vma_type: VmaType,
        page_size: VmaPageSize,
    ) -> Result<(), &'static str> {
        let page_bytes = page_size.bytes();
        if page_count == 0 || start % page_bytes != 0 {
            return Err("Invalid region arguments");
        }
        let len = (page_count as u64)
            .checked_mul(page_bytes)
            .ok_or("Region length overflow")?;
        let end = start.checked_add(len).ok_or("Region end overflow")?;
        const USER_SPACE_END: u64 = 0x0000_8000_0000_0000;
        if end > USER_SPACE_END {
            return Err("Region out of user-space range");
        }

        // Reject overlapping VMAs
        {
            let regions = self.regions.lock();
            if regions.iter().any(|(&vma_start, vma)| {
                let vma_end = vma_start
                    .saturating_add((vma.page_count as u64).saturating_mul(vma.page_size.bytes()));
                vma_start < end && vma_end > start
            }) {
                return Err("Region overlaps existing mapping");
            }
        }

        // Enforce per-silo memory quota (best effort; non-silo tasks are ignored).
        crate::silo::charge_current_task_memory(len).map_err(|_| "Silo memory quota exceeded")?;

        // Track the region, attempting to merge with previous.
        let mut regions = self.regions.lock();
        let mut merged = false;

        if let Some((&prev_start, prev_vma)) = regions.range(..start).next_back() {
            let prev_end = prev_start + (prev_vma.page_count as u64) * prev_vma.page_size.bytes();
            if prev_end == start
                && prev_vma.flags == flags
                && prev_vma.vma_type == vma_type
                && prev_vma.page_size == page_size
            {
                let new_count = prev_vma
                    .page_count
                    .checked_add(page_count)
                    .ok_or("Region page_count overflow")?;
                let updated_vma = VirtualMemoryRegion {
                    start: prev_start,
                    page_count: new_count,
                    flags,
                    vma_type,
                    page_size,
                };
                regions.insert(prev_start, updated_vma);
                merged = true;
            }
        }

        if !merged {
            let region = VirtualMemoryRegion {
                start,
                page_count,
                flags,
                vma_type,
                page_size,
            };
            regions.insert(start, region);
        }

        log::trace!(
            "Reserved lazy region: {:#x} ({} pages, size={:?})",
            start,
            page_count,
            page_size
        );
        Ok(())
    }

    /// Handle a page fault by checking if the address falls within a reserved VMA.
    ///
    /// If it does, allocates a physical frame and maps it.
    pub fn handle_fault(&self, fault_addr: u64) -> Result<(), &'static str> {
        use x86_64::structures::paging::mapper::MapToError;

        // 1. Find the VMA covering this address
        let vma = {
            let regions = self.regions.lock();
            let mut iter = regions.range(..=fault_addr);
            let (&start, vma) = iter.next_back().ok_or("No VMA found for address")?;
            let end = start + (vma.page_count as u64) * vma.page_size.bytes();
            if fault_addr >= end {
                return Err("Address outside VMA bounds");
            }
            vma.clone()
        };

        // Align fault address to the page size used by this VMA.
        let page_bytes = vma.page_size.bytes();
        let page_addr = fault_addr & !(page_bytes - 1);

        // 2. Only Anonymous/Stack regions support demand paging for now
        match vma.vma_type {
            VmaType::Anonymous | VmaType::Stack | VmaType::Code => {}
            _ => return Err("VMA type does not support demand paging"),
        }

        // 3. If already mapped (race/re-fault), treat as handled.
        if self.translate(VirtAddr::new(page_addr)).is_some() {
            return Ok(());
        }

        // 4. Allocate and map a single page of the required size.
        //
        // IMPORTANT: `allocate_frame` (order-0) now goes through
        // `FrameAllocOptions::new()` which zeroes by default.  For order > 0
        // (huge pages) we still need a manual zero via the HHDM.
        //
        // The zero MUST go through phys_to_virt (HHDM), NOT through the user
        // virtual address, because the user address space is not necessarily
        // the currently active CR3.  Writing through `page_addr as *mut u8`
        // would either write into a different process's memory or fault.
        let mut frame_allocator = crate::memory::paging::BuddyFrameAllocator;
        let order = match vma.page_size {
            VmaPageSize::Small => 0,
            VmaPageSize::Huge => 9,
        };

        let frame = crate::sync::with_irqs_disabled(|token| {
            if order == 0 {
                crate::memory::allocate_frame(token)
            } else {
                let f = crate::memory::allocate_phys_contiguous(token, order)?;
                // SAFETY: phys_to_virt gives a valid HHDM pointer for this
                // frame; we have exclusive ownership from the buddy allocator.
                unsafe {
                    core::ptr::write_bytes(
                        crate::memory::phys_to_virt(f.start_address.as_u64()) as *mut u8,
                        0,
                        page_bytes as usize,
                    );
                }
                Ok(f)
            }
        })
        .map_err(|_| "OOM during demand paging")?;

        let mut page_flags = vma.flags.to_page_flags();

        // SAFETY: We own the address space.
        unsafe {
            let mut mapper = self.mapper();
            match vma.page_size {
                VmaPageSize::Small => {
                    let page =
                        Page::<Size4KiB>::from_start_address(VirtAddr::new(page_addr)).unwrap();
                    let phys_frame =
                        x86_64::structures::paging::PhysFrame::<Size4KiB>::containing_address(
                            frame.start_address,
                        );
                    match mapper.map_to(page, phys_frame, page_flags, &mut frame_allocator) {
                        Ok(flush) => {
                            flush.flush();
                        }
                        Err(MapToError::PageAlreadyMapped(_)) => {
                            crate::sync::with_irqs_disabled(|token| {
                                crate::memory::free_phys_contiguous(token, frame, order);
                            });
                            return Ok(());
                        }
                        Err(_) => {
                            crate::sync::with_irqs_disabled(|token| {
                                crate::memory::free_phys_contiguous(token, frame, order);
                            });
                            return Err("Failed to map demand page (4K)");
                        }
                    }
                }
                VmaPageSize::Huge => {
                    let page =
                        Page::<Size2MiB>::from_start_address(VirtAddr::new(page_addr)).unwrap();
                    let phys_frame =
                        x86_64::structures::paging::PhysFrame::<Size2MiB>::containing_address(
                            frame.start_address,
                        );
                    page_flags |= PageTableFlags::HUGE_PAGE;
                    match mapper.map_to(page, phys_frame, page_flags, &mut frame_allocator) {
                        Ok(flush) => {
                            flush.flush();
                        }
                        Err(MapToError::PageAlreadyMapped(_)) => {
                            crate::sync::with_irqs_disabled(|token| {
                                crate::memory::free_phys_contiguous(token, frame, order);
                            });
                            return Ok(());
                        }
                        Err(_) => {
                            crate::sync::with_irqs_disabled(|token| {
                                crate::memory::free_phys_contiguous(token, frame, order);
                            });
                            return Err("Failed to map demand page (2M)");
                        }
                    }
                }
            }
        }

        if self
            .register_effective_mapping(EffectiveMapping {
                start: page_addr,
                cap_id: allocate_mapping_cap_id(),
                handle: resolve_handle(frame.start_address),
                flags: page_flags,
                page_size: vma.page_size,
            })
            .is_err()
        {
            unsafe {
                let mut mapper = self.mapper();
                match vma.page_size {
                    VmaPageSize::Small => {
                        let page =
                            Page::<Size4KiB>::from_start_address(VirtAddr::new(page_addr)).unwrap();
                        if let Ok((_, flush)) = mapper.unmap(page) {
                            flush.flush();
                        }
                    }
                    VmaPageSize::Huge => {
                        let page =
                            Page::<Size2MiB>::from_start_address(VirtAddr::new(page_addr)).unwrap();
                        if let Ok((_, flush)) = mapper.unmap(page) {
                            flush.flush();
                        }
                    }
                }
            }
            crate::sync::with_irqs_disabled(|token| {
                crate::memory::free_phys_contiguous(token, frame, order);
            });
            return Err("Failed to track demand page mapping");
        }

        // Initialize COW refcount.
        //
        // Order-0 frames come from FrameAllocOptions which stamps refcount=1
        // via CAS(REFCOUNT_UNUSED → 1) : the frame is already "sole owner".
        // Huge pages (order > 0) are raw-allocated with REFCOUNT_UNUSED still
        // in the metadata; initialise explicitly to 1 here.
        //
        // Do NOT call frame_inc_ref for fresh allocations: that would push the
        // count to 2, breaking the COW semantics (refcount==1 means sole owner).
        // frame_inc_ref is correct only when sharing an existing frame (fork).
        if order != 0 {
            crate::memory::cow::handle_init_ref(resolve_handle(frame.start_address));
        }

        Ok(())
    }

    /// Map a contiguous region of pages backed by newly allocated physical frames.
    ///
    /// Frames are allocated from the buddy allocator and zero-filled.
    /// The region is tracked in the VMA list.
    pub fn map_region(
        &self,
        start: u64,
        page_count: usize,
        flags: VmaFlags,
        vma_type: VmaType,
        page_size: VmaPageSize,
    ) -> Result<(), &'static str> {
        let page_bytes = page_size.bytes();
        if page_count == 0 || start % page_bytes != 0 {
            return Err("Invalid region arguments");
        }
        let len = (page_count as u64)
            .checked_mul(page_bytes)
            .ok_or("Region length overflow")?;
        let end = start.checked_add(len).ok_or("Region end overflow")?;
        const USER_SPACE_END: u64 = 0x0000_8000_0000_0000;
        if end > USER_SPACE_END {
            return Err("Region out of user-space range");
        }

        // Reject overlapping VMAs early
        {
            let regions = self.regions.lock();
            if regions.iter().any(|(&vma_start, vma)| {
                let vma_end = vma_start
                    .saturating_add((vma.page_count as u64).saturating_mul(vma.page_size.bytes()));
                vma_start < end && vma_end > start
            }) {
                return Err("Region overlaps existing mapping");
            }
        }

        // Enforce per-silo memory quota for eagerly mapped regions.
        crate::silo::charge_current_task_memory(len).map_err(|_| "Silo memory quota exceeded")?;

        let page_flags = flags.to_page_flags();
        let mut frame_allocator = BuddyFrameAllocator;

        // SAFETY: we have logical ownership of this address space.
        let mut mapper = unsafe { self.mapper() };
        let mut mapped_pages = 0usize;

        for i in 0..page_count {
            let page_addr = start
                .checked_add((i as u64).saturating_mul(page_bytes))
                .ok_or("Page address overflow")?;

            // Allocate a physical frame of appropriate size.
            //
            // order-0 frames go through FrameAllocOptions (zeroed + metadata
            // stamped).  order > 0 (huge pages) are zeroed manually via HHDM.
            let order = match page_size {
                VmaPageSize::Small => 0,
                VmaPageSize::Huge => 9,
            };

            let frame = crate::sync::with_irqs_disabled(|token| {
                if order == 0 {
                    crate::memory::allocate_frame(token)
                } else {
                    let f = crate::memory::allocate_phys_contiguous(token, order)?;
                    unsafe {
                        let virt = crate::memory::phys_to_virt(f.start_address.as_u64());
                        core::ptr::write_bytes(virt as *mut u8, 0, page_bytes as usize);
                    }
                    Ok(f)
                }
            })
            .map_err(|_| "Failed to allocate frame")?;

            // Map the page.
            let map_ok = match page_size {
                VmaPageSize::Small => {
                    use x86_64::structures::paging::Size4KiB;
                    let page = Page::<Size4KiB>::from_start_address(VirtAddr::new(page_addr))
                        .map_err(|_| "Map 4K: invalid page address")?;
                    let phys_frame =
                        x86_64::structures::paging::PhysFrame::<Size4KiB>::containing_address(
                            frame.start_address,
                        );
                    unsafe {
                        mapper
                            .map_to(page, phys_frame, page_flags, &mut frame_allocator)
                            .map(|flush| flush.flush())
                            .is_ok()
                    }
                }
                VmaPageSize::Huge => {
                    use x86_64::structures::paging::Size2MiB;
                    let page = Page::<Size2MiB>::from_start_address(VirtAddr::new(page_addr))
                        .map_err(|_| "Map 2M: invalid page address")?;
                    let phys_frame =
                        x86_64::structures::paging::PhysFrame::<Size2MiB>::containing_address(
                            frame.start_address,
                        );
                    let mut huge_flags = page_flags;
                    huge_flags |= PageTableFlags::HUGE_PAGE;
                    unsafe {
                        mapper
                            .map_to(page, phys_frame, huge_flags, &mut frame_allocator)
                            .map(|flush| flush.flush())
                            .is_ok()
                    }
                }
            };

            if !map_ok {
                log::error!(
                    "map_region: map_to failed at page {} vaddr={:#x} size={:?}",
                    i,
                    page_addr,
                    page_size
                );
                // Free frame for this page that failed to map.
                crate::sync::with_irqs_disabled(|token| {
                    crate::memory::free_phys_contiguous(token, frame, order);
                });

                // Roll back already mapped pages to keep state consistent.
                for j in (0..mapped_pages).rev() {
                    let rb_addr = start + (j as u64) * page_bytes;
                    match page_size {
                        VmaPageSize::Small => {
                            use x86_64::structures::paging::Size4KiB;
                            let rb_page =
                                Page::<Size4KiB>::from_start_address(VirtAddr::new(rb_addr))
                                    .map_err(|_| "Rollback: invalid 4K page address")?;
                            if let Ok((_, rb_flush)) = mapper.unmap(rb_page) {
                                rb_flush.flush();
                                let _ = self.unregister_effective_mapping(rb_addr);
                            }
                        }
                        VmaPageSize::Huge => {
                            use x86_64::structures::paging::Size2MiB;
                            let rb_page =
                                Page::<Size2MiB>::from_start_address(VirtAddr::new(rb_addr))
                                    .map_err(|_| "Rollback: invalid 2M page address")?;
                            if let Ok((_, rb_flush)) = mapper.unmap(rb_page) {
                                rb_flush.flush();
                                let _ = self.unregister_effective_mapping(rb_addr);
                            }
                        }
                    }
                }

                crate::silo::release_current_task_memory(len);
                return Err("Failed to map page");
            }

            // Initialize COW refcount (same logic as demand_page above).
            let effective_flags = match page_size {
                VmaPageSize::Small => page_flags,
                VmaPageSize::Huge => page_flags | PageTableFlags::HUGE_PAGE,
            };
            if self
                .register_effective_mapping(EffectiveMapping {
                    start: page_addr,
                    cap_id: allocate_mapping_cap_id(),
                    handle: resolve_handle(frame.start_address),
                    flags: effective_flags,
                    page_size,
                })
                .is_err()
            {
                match page_size {
                    VmaPageSize::Small => {
                        use x86_64::structures::paging::Size4KiB;
                        let page = Page::<Size4KiB>::from_start_address(VirtAddr::new(page_addr))
                            .map_err(|_| "Rollback: invalid 4K page address")?;
                        if let Ok((_, flush)) = mapper.unmap(page) {
                            flush.flush();
                        }
                    }
                    VmaPageSize::Huge => {
                        use x86_64::structures::paging::Size2MiB;
                        let page = Page::<Size2MiB>::from_start_address(VirtAddr::new(page_addr))
                            .map_err(|_| "Rollback: invalid 2M page address")?;
                        if let Ok((_, flush)) = mapper.unmap(page) {
                            flush.flush();
                        }
                    }
                }
                crate::sync::with_irqs_disabled(|token| {
                    crate::memory::free_phys_contiguous(token, frame, order);
                });
                for j in (0..mapped_pages).rev() {
                    let rb_addr = start + (j as u64) * page_bytes;
                    match page_size {
                        VmaPageSize::Small => {
                            use x86_64::structures::paging::Size4KiB;
                            let rb_page =
                                Page::<Size4KiB>::from_start_address(VirtAddr::new(rb_addr))
                                    .map_err(|_| "Rollback: invalid 4K page address")?;
                            if let Ok((_, rb_flush)) = mapper.unmap(rb_page) {
                                rb_flush.flush();
                                let _ = self.unregister_effective_mapping(rb_addr);
                            }
                        }
                        VmaPageSize::Huge => {
                            use x86_64::structures::paging::Size2MiB;
                            let rb_page =
                                Page::<Size2MiB>::from_start_address(VirtAddr::new(rb_addr))
                                    .map_err(|_| "Rollback: invalid 2M page address")?;
                            if let Ok((_, rb_flush)) = mapper.unmap(rb_page) {
                                rb_flush.flush();
                                let _ = self.unregister_effective_mapping(rb_addr);
                            }
                        }
                    }
                }
                crate::silo::release_current_task_memory(len);
                return Err("Failed to track mapped region page");
            }

            mapped_pages += 1;
        }

        // Track the region
        let mut regions = self.regions.lock();
        let region = VirtualMemoryRegion {
            start,
            page_count,
            flags,
            vma_type,
            page_size,
        };
        regions.insert(start, region);

        let end = start + (page_count as u64) * page_bytes;
        crate::trace_mem!(
            crate::trace::category::MEM_MAP,
            crate::trace::TraceKind::MemMap,
            page_size.bytes(),
            crate::trace::TraceTaskCtx {
                task_id: 0,
                pid: 0,
                tid: 0,
                cr3: self.cr3_phys.as_u64(),
            },
            0,
            start,
            end,
            page_count as u64
        );

        Ok(())
    }

    /// Maps shared frames.
    pub fn map_shared_frames(
        &self,
        start: u64,
        frame_phys_addrs: &[u64],
        flags: VmaFlags,
        vma_type: VmaType,
    ) -> Result<(), &'static str> {
        self.map_shared_frames_with_cap_ids(start, frame_phys_addrs, None, flags, vma_type)
    }

    /// Maps shared physical blocks with optional stable mapping identities.
    pub fn map_shared_handles_with_cap_ids(
        &self,
        start: u64,
        handles: &[BlockHandle],
        mapping_cap_ids: Option<&[CapId]>,
        flags: VmaFlags,
        vma_type: VmaType,
        page_size: VmaPageSize,
    ) -> Result<(), &'static str> {
        let page_count = handles.len();
        let page_bytes = page_size.bytes();
        if page_count == 0 || start % page_bytes != 0 {
            return Err("Invalid shared region arguments");
        }
        if mapping_cap_ids.is_some_and(|cap_ids| cap_ids.len() != page_count) {
            return Err("Shared mapping identity count mismatch");
        }
        let len = (page_count as u64)
            .checked_mul(page_bytes)
            .ok_or("Shared region length overflow")?;
        let end = start.checked_add(len).ok_or("Shared region end overflow")?;
        const USER_SPACE_END: u64 = 0x0000_8000_0000_0000;
        if end > USER_SPACE_END {
            return Err("Shared region out of user-space range");
        }

        {
            let regions = self.regions.lock();
            if regions.iter().any(|(&vma_start, vma)| {
                let vma_end = vma_start
                    .saturating_add((vma.page_count as u64).saturating_mul(vma.page_size.bytes()));
                vma_start < end && vma_end > start
            }) {
                return Err("Shared region overlaps existing mapping");
            }
        }

        let mut page_flags = flags.to_page_flags();
        if page_size == VmaPageSize::Huge {
            page_flags |= PageTableFlags::HUGE_PAGE;
        }
        let mut frame_allocator = BuddyFrameAllocator;
        let mut mapper = unsafe { self.mapper() };
        let mut mapped_pages = 0usize;

        for (index, handle) in handles.iter().copied().enumerate() {
            let page_addr = start
                .checked_add((index as u64) * page_bytes)
                .ok_or("Shared page address overflow")?;

            let map_ok = match page_size {
                VmaPageSize::Small => {
                    let page = Page::<Size4KiB>::from_start_address(VirtAddr::new(page_addr))
                        .map_err(|_| "Map shared: invalid 4K page address")?;
                    let frame = X86PhysFrame::<Size4KiB>::containing_address(handle.base);
                    unsafe {
                        mapper
                            .map_to(page, frame, page_flags, &mut frame_allocator)
                            .map(|flush| flush.flush())
                            .is_ok()
                    }
                }
                VmaPageSize::Huge => {
                    let page = Page::<Size2MiB>::from_start_address(VirtAddr::new(page_addr))
                        .map_err(|_| "Map shared: invalid 2M page address")?;
                    let frame = X86PhysFrame::<Size2MiB>::containing_address(handle.base);
                    unsafe {
                        mapper
                            .map_to(page, frame, page_flags, &mut frame_allocator)
                            .map(|flush| flush.flush())
                            .is_ok()
                    }
                }
            };

            if !map_ok {
                for rollback in (0..mapped_pages).rev() {
                    let rb_addr = start + (rollback as u64) * page_bytes;
                    match page_size {
                        VmaPageSize::Small => {
                            if let Ok(rb_page) =
                                Page::<Size4KiB>::from_start_address(VirtAddr::new(rb_addr))
                            {
                                if let Ok((_, rb_flush)) = mapper.unmap(rb_page) {
                                    rb_flush.flush();
                                    let _ = self.unregister_effective_mapping(rb_addr);
                                }
                            }
                        }
                        VmaPageSize::Huge => {
                            if let Ok(rb_page) =
                                Page::<Size2MiB>::from_start_address(VirtAddr::new(rb_addr))
                            {
                                if let Ok((_, rb_flush)) = mapper.unmap(rb_page) {
                                    rb_flush.flush();
                                    let _ = self.unregister_effective_mapping(rb_addr);
                                }
                            }
                        }
                    }
                }
                return Err("Failed to map shared page");
            }

            if self
                .register_effective_mapping(EffectiveMapping {
                    start: page_addr,
                    cap_id: mapping_cap_ids
                        .and_then(|cap_ids| cap_ids.get(index).copied())
                        .unwrap_or_else(allocate_mapping_cap_id),
                    handle,
                    flags: page_flags,
                    page_size,
                })
                .is_err()
            {
                match page_size {
                    VmaPageSize::Small => {
                        if let Ok(page) =
                            Page::<Size4KiB>::from_start_address(VirtAddr::new(page_addr))
                        {
                            if let Ok((_, flush)) = mapper.unmap(page) {
                                flush.flush();
                            }
                        }
                    }
                    VmaPageSize::Huge => {
                        if let Ok(page) =
                            Page::<Size2MiB>::from_start_address(VirtAddr::new(page_addr))
                        {
                            if let Ok((_, flush)) = mapper.unmap(page) {
                                flush.flush();
                            }
                        }
                    }
                }
                for rollback in (0..mapped_pages).rev() {
                    let rb_addr = start + (rollback as u64) * page_bytes;
                    match page_size {
                        VmaPageSize::Small => {
                            if let Ok(rb_page) =
                                Page::<Size4KiB>::from_start_address(VirtAddr::new(rb_addr))
                            {
                                if let Ok((_, rb_flush)) = mapper.unmap(rb_page) {
                                    rb_flush.flush();
                                    let _ = self.unregister_effective_mapping(rb_addr);
                                }
                            }
                        }
                        VmaPageSize::Huge => {
                            if let Ok(rb_page) =
                                Page::<Size2MiB>::from_start_address(VirtAddr::new(rb_addr))
                            {
                                if let Ok((_, rb_flush)) = mapper.unmap(rb_page) {
                                    rb_flush.flush();
                                    let _ = self.unregister_effective_mapping(rb_addr);
                                }
                            }
                        }
                    }
                }
                return Err("Failed to track shared mapping");
            }
            mapped_pages += 1;
        }

        self.regions.lock().insert(
            start,
            VirtualMemoryRegion {
                start,
                page_count,
                flags,
                vma_type,
                page_size,
            },
        );
        Ok(())
    }

    /// Maps shared frames with optional stable mapping identities.
    pub fn map_shared_frames_with_cap_ids(
        &self,
        start: u64,
        frame_phys_addrs: &[u64],
        mapping_cap_ids: Option<&[CapId]>,
        flags: VmaFlags,
        vma_type: VmaType,
    ) -> Result<(), &'static str> {
        let handles = frame_phys_addrs
            .iter()
            .copied()
            .map(|phys_addr| resolve_handle(PhysAddr::new(phys_addr)))
            .collect::<Vec<_>>();
        self.map_shared_handles_with_cap_ids(
            start,
            &handles,
            mapping_cap_ids,
            flags,
            vma_type,
            VmaPageSize::Small,
        )
    }

    /// Unmap a previously mapped region and free the backing frames.
    pub fn unmap_region(
        &self,
        start: u64,
        page_count: usize,
        page_size: VmaPageSize,
    ) -> Result<(), &'static str> {
        let page_bytes = page_size.bytes();
        // SAFETY: We have logical ownership of this address space.
        let mut mapper = unsafe { self.mapper() };

        for i in 0..page_count {
            let page_addr = start + (i as u64) * page_bytes;

            let _frame_addr = match page_size {
                VmaPageSize::Small => {
                    use x86_64::structures::paging::Size4KiB;
                    let page = Page::<Size4KiB>::from_start_address(VirtAddr::new(page_addr))
                        .map_err(|_| "Failed to unmap: invalid 4K page address")?;
                    let (frame, flush) =
                        mapper.unmap(page).map_err(|_| "Failed to unmap 4K page")?;
                    flush.flush();
                    frame.start_address()
                }
                VmaPageSize::Huge => {
                    use x86_64::structures::paging::Size2MiB;
                    let page = Page::<Size2MiB>::from_start_address(VirtAddr::new(page_addr))
                        .map_err(|_| "Failed to unmap: invalid 2M page address")?;
                    let (frame, flush) =
                        mapper.unmap(page).map_err(|_| "Failed to unmap 2M page")?;
                    flush.flush();
                    frame.start_address()
                }
            };

            // COW-aware refcount decrement: free only when last mapping disappears.
            let _ = self.unregister_effective_mapping(page_addr);
        }

        // Remove from VMA tracking.
        self.regions.lock().remove(&start);

        log::trace!(
            "Unmapped region: {:#x}..{:#x} ({} pages, size={:?})",
            start,
            start + (page_count as u64) * page_bytes,
            page_count,
            page_size
        );

        let end = start + (page_count as u64) * page_bytes;
        crate::trace_mem!(
            crate::trace::category::MEM_UNMAP,
            crate::trace::TraceKind::MemUnmap,
            page_size.bytes(),
            crate::trace::TraceTaskCtx {
                task_id: 0,
                pid: 0,
                tid: 0,
                cr3: self.cr3_phys.as_u64(),
            },
            0,
            start,
            end,
            page_count as u64
        );

        let released = (page_count as u64).saturating_mul(page_bytes);
        crate::silo::release_current_task_memory(released);

        Ok(())
    }

    /// Find a free virtual address range of `n_pages` pages of `page_size` starting at or after `hint`.
    pub fn find_free_vma_range(
        &self,
        hint: u64,
        n_pages: usize,
        page_size: VmaPageSize,
    ) -> Option<u64> {
        if n_pages == 0 {
            return None;
        }
        let page_bytes = page_size.bytes();
        let length = (n_pages as u64).checked_mul(page_bytes)?;
        let upper_limit: u64 = 0x0000_8000_0000_0000; // USER_SPACE_END

        // Round hint up to a page boundary
        let mut candidate = (hint.saturating_add(page_bytes - 1)) & !(page_bytes - 1);
        if candidate == 0 {
            candidate = page_bytes;
        }

        let regions = self.regions.lock();
        for (&vma_start, vma) in regions.iter() {
            let vma_end = vma_start + vma.page_count as u64 * vma.page_size.bytes();

            // A gap exists before this VMA : candidate fits.
            if candidate.saturating_add(length) <= vma_start {
                break;
            }

            // Candidate overlaps this VMA; skip past it.
            if vma_end > candidate {
                candidate = (vma_end.saturating_add(page_bytes - 1)) & !(page_bytes - 1);
            }
        }

        // Final bounds check.
        if candidate.checked_add(length)? <= upper_limit {
            Some(candidate)
        } else {
            None
        }
    }

    /// Return true if any tracked VMA overlaps `[addr, addr + len)`.
    pub fn has_mapping_in_range(&self, addr: u64, len: u64) -> bool {
        let end = match addr.checked_add(len) {
            Some(v) => v,
            None => return true,
        };
        let regions = self.regions.lock();
        regions.iter().any(|(&vma_start, vma)| {
            let vma_end = vma_start
                .saturating_add((vma.page_count as u64).saturating_mul(vma.page_size.bytes()));
            vma_start < end && vma_end > addr
        })
    }

    /// Return the tracked VMA that starts exactly at `start`.
    pub fn region_by_start(&self, start: u64) -> Option<VirtualMemoryRegion> {
        let regions = self.regions.lock();
        regions.get(&start).cloned()
    }

    /// Returns true if any page in `[addr, addr + len)` is currently mapped.
    pub fn any_mapped_in_range(
        &self,
        addr: u64,
        len: u64,
        page_size: VmaPageSize,
    ) -> Result<bool, &'static str> {
        if len == 0 {
            return Ok(false);
        }
        let end = addr
            .checked_add(len)
            .ok_or("any_mapped_in_range: address overflow")?;
        let step = page_size.bytes();
        let mut cur = addr;
        while cur < end {
            if self.translate(VirtAddr::new(cur)).is_some() {
                return Ok(true);
            }
            cur = cur
                .checked_add(step)
                .ok_or("any_mapped_in_range: loop overflow")?;
        }
        Ok(false)
    }

    /// Performs the protect range operation.
    pub fn protect_range(&self, addr: u64, len: u64, flags: VmaFlags) -> Result<(), &'static str> {
        if len == 0 {
            return Ok(());
        }
        let end = addr
            .checked_add(len)
            .ok_or("protect_range: address overflow")?;
        let mut cursor = addr;

        {
            let regions = self.regions.lock();
            for (&vma_start, vma) in regions.iter() {
                let vma_end = vma_start + vma.page_count as u64 * vma.page_size.bytes();
                if vma_start >= end || vma_end <= addr {
                    continue;
                }
                if vma.page_size == VmaPageSize::Huge {
                    let range_start = core::cmp::max(vma_start, addr);
                    let range_end = core::cmp::min(vma_end, end);
                    if range_start % vma.page_size.bytes() != 0
                        || range_end % vma.page_size.bytes() != 0
                    {
                        return Err(
                            "protect_range: partial mprotect of 2MiB pages is not supported",
                        );
                    }
                }
            }
        }

        let mut touched = false;
        while cursor < end {
            let region_info = {
                let regions = self.regions.lock();
                regions
                    .iter()
                    .find(|(&vma_start, vma)| {
                        let vma_end = vma_start + vma.page_count as u64 * vma.page_size.bytes();
                        vma_start < end && vma_end > cursor
                    })
                    .map(|(&k, v)| (k, v.clone()))
            };

            let Some((vma_start, vma)) = region_info else {
                break;
            };
            touched = true;

            let vma_end = vma_start + vma.page_count as u64 * vma.page_size.bytes();
            let range_start = core::cmp::max(vma_start, cursor);
            let range_end = core::cmp::min(vma_end, end);
            let page_bytes = vma.page_size.bytes();
            let new_pt_flags = flags.to_page_flags();

            let mut mapper = unsafe { self.mapper() };
            let mut page_addr = range_start;
            while page_addr < range_end {
                if mapper.translate_addr(VirtAddr::new(page_addr)).is_none() {
                    page_addr += page_bytes;
                    continue;
                }
                unsafe {
                    match vma.page_size {
                        VmaPageSize::Small => {
                            let page =
                                Page::<Size4KiB>::from_start_address(VirtAddr::new(page_addr))
                                    .map_err(|_| "protect_range: invalid 4K page address")?;
                            mapper
                                .update_flags(page, new_pt_flags)
                                .map(|f| f.ignore())
                                .map_err(|_| "protect_range: update 4K flags failed")?;
                            let _ = self.update_effective_mapping_flags(page_addr, new_pt_flags);
                        }
                        VmaPageSize::Huge => {
                            let mut huge_flags = new_pt_flags;
                            huge_flags |= PageTableFlags::HUGE_PAGE;
                            let page =
                                Page::<Size2MiB>::from_start_address(VirtAddr::new(page_addr))
                                    .map_err(|_| "protect_range: invalid 2M page address")?;
                            mapper
                                .update_flags(page, huge_flags)
                                .map(|f| f.ignore())
                                .map_err(|_| "protect_range: update 2M flags failed")?;
                            let _ = self.update_effective_mapping_flags(page_addr, huge_flags);
                        }
                    }
                }
                page_addr += page_bytes;
            }

            {
                let mut regions = self.regions.lock();
                regions.remove(&vma_start);

                if range_start > vma_start {
                    let leading_pages = ((range_start - vma_start) / page_bytes) as usize;
                    regions.insert(
                        vma_start,
                        VirtualMemoryRegion {
                            start: vma_start,
                            page_count: leading_pages,
                            flags: vma.flags,
                            vma_type: vma.vma_type,
                            page_size: vma.page_size,
                        },
                    );
                }

                let middle_pages = ((range_end - range_start) / page_bytes) as usize;
                if middle_pages > 0 {
                    regions.insert(
                        range_start,
                        VirtualMemoryRegion {
                            start: range_start,
                            page_count: middle_pages,
                            flags,
                            vma_type: vma.vma_type,
                            page_size: vma.page_size,
                        },
                    );
                }

                if range_end < vma_end {
                    let trailing_pages = ((vma_end - range_end) / page_bytes) as usize;
                    regions.insert(
                        range_end,
                        VirtualMemoryRegion {
                            start: range_end,
                            page_count: trailing_pages,
                            flags: vma.flags,
                            vma_type: vma.vma_type,
                            page_size: vma.page_size,
                        },
                    );
                }
            }

            cursor = range_end;
        }

        if !touched {
            return Err("protect_range: no mapped region in range");
        }
        Ok(())
    }

    /// Unmaps range.
    pub fn unmap_range(&self, addr: u64, len: u64) -> Result<(), &'static str> {
        if len == 0 {
            return Ok(());
        }
        let end = addr
            .checked_add(len)
            .ok_or("unmap_range: address overflow")?;

        // Pre-validate huge-page overlaps: partial unmap of 2MiB mappings is
        // not supported yet. Callers must unmap on huge-page boundaries.
        {
            let regions = self.regions.lock();
            for (&vma_start, vma) in regions.iter() {
                let vma_end = vma_start + vma.page_count as u64 * vma.page_size.bytes();
                if vma_start >= end || vma_end <= addr {
                    continue;
                }
                if vma.page_size == VmaPageSize::Huge {
                    let range_start = core::cmp::max(vma_start, addr);
                    let range_end = core::cmp::min(vma_end, end);
                    if range_start % vma.page_size.bytes() != 0
                        || range_end % vma.page_size.bytes() != 0
                    {
                        return Err("unmap_range: partial unmap of 2MiB pages is not supported");
                    }
                }
            }
        }

        // Process regions one by one to avoid heap allocation (Vec)
        let mut released_bytes = 0u64;
        loop {
            // Find the first overlapping region
            let region_info = {
                let regions = self.regions.lock();
                regions
                    .iter()
                    .find(|(&vma_start, vma)| {
                        let vma_end = vma_start + vma.page_count as u64 * vma.page_size.bytes();
                        vma_start < end && vma_end > addr
                    })
                    .map(|(&k, v)| (k, v.clone()))
            };

            let Some((vma_start, vma)) = region_info else {
                break; // No more overlapping regions
            };

            let vma_end = vma_start + vma.page_count as u64 * vma.page_size.bytes();
            let range_start = core::cmp::max(vma_start, addr);
            let range_end = core::cmp::min(vma_end, end);
            released_bytes = released_bytes.saturating_add(range_end.saturating_sub(range_start));

            // 1. Hardware unmap
            // SAFETY: Logical ownership of address space.
            let mut mapper = unsafe { self.mapper() };
            let mut page_addr = range_start;
            let page_bytes = vma.page_size.bytes();
            while page_addr < range_end {
                // Lazy VMAs can contain unfaulted pages (no PTE). In that case
                // there is nothing to unmap in hardware; just update VMA metadata.
                if mapper.translate_addr(VirtAddr::new(page_addr)).is_none() {
                    page_addr += page_bytes;
                    continue;
                }

                let _frame_addr = match vma.page_size {
                    VmaPageSize::Small => {
                        use x86_64::structures::paging::Size4KiB;
                        let page = Page::<Size4KiB>::from_start_address(VirtAddr::new(page_addr))
                            .map_err(|_| "unmap_range: invalid 4K page address")?;
                        let (frame, flush) = mapper
                            .unmap(page)
                            .map_err(|_| "unmap_range: unmap 4K failed")?;
                        flush.flush();
                        frame.start_address()
                    }
                    VmaPageSize::Huge => {
                        use x86_64::structures::paging::Size2MiB;
                        let page = Page::<Size2MiB>::from_start_address(VirtAddr::new(page_addr))
                            .map_err(|_| "unmap_range: invalid 2M page address")?;
                        let (frame, flush) = mapper
                            .unmap(page)
                            .map_err(|_| "unmap_range: unmap 2M failed")?;
                        flush.flush();
                        frame.start_address()
                    }
                };

                let _ = self.unregister_effective_mapping(page_addr);
                page_addr += page_bytes;
            }

            // 2. Update tracking: remove and re-insert fragments
            {
                let mut regions = self.regions.lock();
                regions.remove(&vma_start);

                if range_start > vma_start {
                    let leading_pages =
                        ((range_start - vma_start) / vma.page_size.bytes()) as usize;
                    regions.insert(
                        vma_start,
                        VirtualMemoryRegion {
                            start: vma_start,
                            page_count: leading_pages,
                            flags: vma.flags,
                            vma_type: vma.vma_type,
                            page_size: vma.page_size,
                        },
                    );
                }

                if range_end < vma_end {
                    let trailing_pages = ((vma_end - range_end) / vma.page_size.bytes()) as usize;
                    regions.insert(
                        range_end,
                        VirtualMemoryRegion {
                            start: range_end,
                            page_count: trailing_pages,
                            flags: vma.flags,
                            vma_type: vma.vma_type,
                            page_size: vma.page_size,
                        },
                    );
                }
            }
        }

        crate::silo::release_current_task_memory(released_bytes);
        Ok(())
    }

    /// Translate a virtual address to its mapped physical address.
    pub fn translate(&self, vaddr: VirtAddr) -> Option<PhysAddr> {
        // SAFETY: Read-only access to the page tables.
        let mapper = unsafe { self.mapper() };
        mapper.translate_addr(vaddr)
    }

    /// Translate a virtual address to the current block handle and page-table flags.
    pub fn translate_to_handle(&self, vaddr: VirtAddr) -> Option<(BlockHandle, PageTableFlags)> {
        // SAFETY: Read-only access to the page tables.
        let mapper = unsafe { self.mapper() };
        let translated = mapper.translate(vaddr);
        match translated {
            TranslateResult::Mapped { frame, flags, .. } => {
                Some((resolve_handle(frame.start_address()), flags))
            }
            TranslateResult::NotMapped | TranslateResult::InvalidFrameAddress(_) => None,
        }
    }

    /// Get the physical address of this address space's PML4 table.
    pub fn cr3(&self) -> PhysAddr {
        self.cr3_phys
    }

    /// Switch the CPU to this address space by writing CR3.
    ///
    /// Skips the write if CR3 already points to this address space (avoids
    /// unnecessary TLB flush).
    ///
    /// # Safety
    /// The caller must ensure this address space's page tables are valid and
    /// that the kernel half is correctly mapped.
    pub unsafe fn switch_to(&self) {
        let (current_frame, _) = Cr3::read();
        if current_frame.start_address() == self.cr3_phys {
            return; // Already active : skip to avoid TLB flush.
        }

        // SAFETY: cr3_phys points to a valid, 4KiB-aligned PML4 table with
        // the kernel half correctly populated.
        unsafe {
            let frame =
                X86PhysFrame::from_start_address(self.cr3_phys).expect("CR3 address not aligned");
            crate::e9_println!("C");
            Cr3::write(frame, Cr3Flags::empty());
            crate::e9_println!("c");
        }
    }

    /// Whether this is the kernel address space.
    pub fn is_kernel(&self) -> bool {
        self.is_kernel
    }

    /// Check if this address space has any user-space memory mappings.
    pub fn has_user_mappings(&self) -> bool {
        if self.is_kernel {
            return false;
        }
        let regions = self.regions.lock();
        // Check for any non-kernel mappings.
        regions.values().any(|vma| vma.vma_type != VmaType::Kernel)
    }

    fn teardown_effective_mapping_for_drop(&self, mapping: EffectiveMapping) {
        // SAFETY: the address space is being torn down; any remaining user mapping
        // must be detached from ownership tracking before page-table reclamation.
        unsafe {
            let mut mapper = self.mapper();
            if mapper
                .translate_addr(VirtAddr::new(mapping.start))
                .is_some()
            {
                match mapping.page_size {
                    VmaPageSize::Small => {
                        let page =
                            Page::<Size4KiB>::from_start_address(VirtAddr::new(mapping.start))
                                .unwrap();
                        if let Err(error) = mapper.unmap(page) {
                            log::warn!(
                                "memory: drop cleanup failed to unmap 4K page at {:#x}: {:?}",
                                mapping.start,
                                error
                            );
                        }
                    }
                    VmaPageSize::Huge => {
                        let page =
                            Page::<Size2MiB>::from_start_address(VirtAddr::new(mapping.start))
                                .unwrap();
                        if let Err(error) = mapper.unmap(page) {
                            log::warn!(
                                "memory: drop cleanup failed to unmap 2M page at {:#x}: {:?}",
                                mapping.start,
                                error
                            );
                        }
                    }
                }
            }
        }

        let _ = self.unregister_effective_mapping(mapping.start);
    }

    fn teardown_region_for_drop(&self, start: u64, region: &VirtualMemoryRegion) {
        let len = (region.page_count as u64).saturating_mul(region.page_size.bytes());
        let mut page_addr = start;
        let page_bytes = region.page_size.bytes();

        while page_addr < start.saturating_add(len) {
            if let Some(mapping) = self.effective_mapping_by_start(page_addr) {
                self.teardown_effective_mapping_for_drop(mapping);
            }
            page_addr += page_bytes;
        }

        let _ = self.regions.lock().remove(&start);
        crate::silo::release_current_task_memory(len);
    }

    /// Unmap all tracked user regions (best-effort).
    ///
    /// This frees user frames and clears the VMA list. Kernel mappings are untouched.
    /// Does not allocate memory.
    pub fn unmap_all_user_regions(&self) {
        if self.is_kernel {
            return;
        }

        loop {
            let first = {
                let guard = self.regions.lock();
                guard
                    .iter()
                    .next()
                    .map(|(&start, region)| (start, region.clone()))
            };

            let Some((start, region)) = first else {
                break;
            };

            let len = (region.page_count as u64).saturating_mul(region.page_size.bytes());
            if self.unmap_range(region.start, len).is_err() {
                log::warn!(
                    "memory: unmap_all_user_regions fallback cleanup for {:#x}..{:#x}",
                    start,
                    start.saturating_add(len)
                );
                self.teardown_region_for_drop(start, &region);
            }
        }

        let residual_mappings: Vec<EffectiveMapping> = {
            let guard = self.effective_mappings.lock();
            guard.values().copied().collect()
        };
        for mapping in residual_mappings {
            log::warn!(
                "memory: drop cleanup removing orphan effective mapping at {:#x} cap={}",
                mapping.start,
                mapping.cap_id.as_u64()
            );
            self.teardown_effective_mapping_for_drop(mapping);
            crate::silo::release_current_task_memory(mapping.page_size.bytes());
        }
    }

    /// Performs the clone cow operation.
    pub fn clone_cow(&self) -> Result<Arc<AddressSpace>, &'static str> {
        if self.is_kernel {
            return Err("Cannot fork kernel address space");
        }

        let child = Arc::new(AddressSpace::new_user()?);

        let regions: Vec<VirtualMemoryRegion> = {
            let guard = self.regions.lock();
            guard.values().cloned().collect()
        };
        let effective_mappings: Vec<EffectiveMapping> = {
            let guard = self.effective_mappings.lock();
            guard.values().copied().collect()
        };

        let mut tlb_flush_needed = false;
        let mut processed_pages = Vec::new();

        let res: Result<(), &'static str> = (|| {
            let mut parent_mapper = unsafe { self.mapper() };
            let mut child_mapper = unsafe { child.mapper() };
            let mut frame_allocator = BuddyFrameAllocator;

            for region in regions.iter() {
                // Register VMA in child.
                {
                    let mut child_regions = child.regions.lock();
                    child_regions.insert(region.start, region.clone());
                }
            }

            for mapping in effective_mappings.iter().copied() {
                let vaddr = VirtAddr::new(mapping.start);
                let phys_frame_addr = mapping.handle.base;
                let mut new_flags = mapping.flags;
                let is_writable = mapping.flags.contains(PageTableFlags::WRITABLE);
                const COW_BIT: PageTableFlags = PageTableFlags::BIT_9;

                if is_writable {
                    new_flags.remove(PageTableFlags::WRITABLE);
                    new_flags.insert(COW_BIT);

                    unsafe {
                        let res: Result<(), &'static str> = match mapping.page_size {
                            VmaPageSize::Small => parent_mapper
                                .update_flags(
                                    Page::<Size4KiB>::from_start_address(vaddr).unwrap(),
                                    new_flags,
                                )
                                .map(|f| f.ignore())
                                .map_err(|_| "Failed to update parent 4K flags"),
                            VmaPageSize::Huge => parent_mapper
                                .update_flags(
                                    Page::<Size2MiB>::from_start_address(vaddr).unwrap(),
                                    new_flags,
                                )
                                .map(|f| f.ignore())
                                .map_err(|_| "Failed to update parent 2M flags"),
                        };
                        if let Err(e) = res {
                            return Err(e);
                        }
                    }
                    let _ = self.update_effective_mapping_flags(vaddr.as_u64(), new_flags);
                    tlb_flush_needed = true;
                    processed_pages.push((vaddr.as_u64(), mapping.flags, mapping.page_size));
                }

                let handle = mapping.handle;
                crate::memory::cow::handle_inc_ref(handle).map_err(|error| {
                    log::warn!(
                        "clone_cow: failed to pin source handle {:#x}/{} for vaddr={:#x}: {:?}",
                        handle.base.as_u64(),
                        handle.order,
                        vaddr.as_u64(),
                        error
                    );
                    "Failed to pin source COW frame"
                })?;

                // Map in child. We map it as WRITABLE first to ensure intermediate
                // page tables (PDPT, PD) are created with WRITABLE bit set.
                // If we mapped directly as COW (Read-only), some Mapper implementations
                // might create Read-Only intermediate tables, blocking future COW resolution.
                let map_flags = new_flags | PageTableFlags::WRITABLE;

                unsafe {
                    let map_res: Result<(), &'static str> = match mapping.page_size {
                        VmaPageSize::Small => {
                            let page = Page::<Size4KiB>::from_start_address(vaddr).unwrap();
                            let frame = x86_64::structures::paging::PhysFrame::<Size4KiB>::containing_address(phys_frame_addr);
                            child_mapper
                                .map_to(page, frame, map_flags, &mut frame_allocator)
                                .map(|f| f.ignore())
                                .map_err(|_| "Failed to map 4K in child")
                        }
                        VmaPageSize::Huge => {
                            let page = Page::<Size2MiB>::from_start_address(vaddr).unwrap();
                            let frame = x86_64::structures::paging::PhysFrame::<Size2MiB>::containing_address(phys_frame_addr);
                            child_mapper
                                .map_to(page, frame, map_flags, &mut frame_allocator)
                                .map(|f| f.ignore())
                                .map_err(|_| "Failed to map 2M in child")
                        }
                    };

                    if let Err(e) = map_res {
                        crate::memory::cow::handle_dec_ref(handle);
                        return Err(e);
                    }

                    // Now downgrade to the actual COW flags (which may be Read-Only).
                    if !new_flags.contains(PageTableFlags::WRITABLE) {
                        let downgrade_res: Result<(), &'static str> = match mapping.page_size {
                            VmaPageSize::Small => {
                                let page = Page::<Size4KiB>::from_start_address(vaddr).unwrap();
                                child_mapper
                                    .update_flags(page, new_flags)
                                    .map(|f| f.ignore())
                                    .map_err(|_| "Failed to update child 4K flags")
                            }
                            VmaPageSize::Huge => {
                                let page = Page::<Size2MiB>::from_start_address(vaddr).unwrap();
                                child_mapper
                                    .update_flags(page, new_flags)
                                    .map(|f| f.ignore())
                                    .map_err(|_| "Failed to update child 2M flags")
                            }
                        };
                        if let Err(e) = downgrade_res {
                            let unmapped = match mapping.page_size {
                                VmaPageSize::Small => {
                                    let page = Page::<Size4KiB>::from_start_address(vaddr).unwrap();
                                    child_mapper.unmap(page).map(|(_, f)| f.ignore()).is_ok()
                                }
                                VmaPageSize::Huge => {
                                    let page = Page::<Size2MiB>::from_start_address(vaddr).unwrap();
                                    child_mapper.unmap(page).map(|(_, f)| f.ignore()).is_ok()
                                }
                            };
                            if unmapped {
                                crate::memory::cow::handle_dec_ref(handle);
                            }
                            return Err(e);
                        }
                    }
                }

                if child
                    .register_effective_mapping(EffectiveMapping {
                        start: vaddr.as_u64(),
                        cap_id: allocate_mapping_cap_id(),
                        handle,
                        flags: new_flags,
                        page_size: mapping.page_size,
                    })
                    .is_err()
                {
                    match mapping.page_size {
                        VmaPageSize::Small => {
                            let page = Page::<Size4KiB>::from_start_address(vaddr).unwrap();
                            if let Ok((_, flush)) = child_mapper.unmap(page) {
                                flush.ignore();
                            }
                        }
                        VmaPageSize::Huge => {
                            let page = Page::<Size2MiB>::from_start_address(vaddr).unwrap();
                            if let Ok((_, flush)) = child_mapper.unmap(page) {
                                flush.ignore();
                            }
                        }
                    }
                    crate::memory::cow::handle_dec_ref(handle);
                    return Err("Failed to track child COW mapping");
                }

                crate::memory::cow::handle_dec_ref(handle);
            }
            Ok(())
        })();

        let tlb_flush_range = if tlb_flush_needed && !processed_pages.is_empty() {
            let mut range_start = u64::MAX;
            let mut range_end = 0u64;
            for (vaddr, _, page_size) in &processed_pages {
                range_start = range_start.min(*vaddr);
                range_end = range_end.max(vaddr.saturating_add(page_size.bytes()));
            }
            if range_start < range_end {
                Some((range_start, range_end))
            } else {
                None
            }
        } else {
            None
        };

        if let Err(e) = res {
            log::error!("clone_cow error: {}. Rolling back...", e);
            let mut parent_mapper = unsafe { self.mapper() };
            for &(vaddr, original_flags, page_size) in processed_pages.iter().rev() {
                if original_flags.contains(PageTableFlags::WRITABLE) {
                    unsafe {
                        match page_size {
                            VmaPageSize::Small => {
                                let _ = parent_mapper.update_flags(
                                    Page::<Size4KiB>::from_start_address(VirtAddr::new(vaddr))
                                        .unwrap(),
                                    original_flags,
                                );
                            }
                            VmaPageSize::Huge => {
                                let _ = parent_mapper.update_flags(
                                    Page::<Size2MiB>::from_start_address(VirtAddr::new(vaddr))
                                        .unwrap(),
                                    original_flags,
                                );
                            }
                        };
                    }
                    let _ = self.update_effective_mapping_flags(vaddr, original_flags);
                }
            }
            if let Some((range_start, range_end)) = tlb_flush_range {
                crate::arch::x86_64::tlb::shootdown_range(
                    VirtAddr::new(range_start),
                    VirtAddr::new(range_end),
                );
            }
            return Err(e);
        }

        if let Some((range_start, range_end)) = tlb_flush_range {
            crate::arch::x86_64::tlb::shootdown_range(
                VirtAddr::new(range_start),
                VirtAddr::new(range_end),
            );
        }
        Ok(child)
    }

    /// Releases user page tables.
    fn free_user_page_tables(&self) {
        if self.is_kernel {
            return;
        }

        // SAFETY: We have logical ownership of this address space during drop.
        let l4 = unsafe { &mut *self.l4_table_virt.as_mut_ptr::<PageTable>() };

        for i in 0..256 {
            if !l4[i].flags().contains(PageTableFlags::PRESENT) {
                continue;
            }
            let l3_frame = match l4[i].frame() {
                Ok(f) => f,
                Err(_) => {
                    l4[i].set_unused();
                    continue;
                }
            };

            free_l3_table(l3_frame);
            l4[i].set_unused();
        }
    }
}

impl Drop for AddressSpace {
    /// Performs the drop operation.
    fn drop(&mut self) {
        if self.is_kernel {
            return; // Never free the kernel address space.
        }

        log::trace!("AddressSpace::drop begin CR3={:#x}", self.cr3_phys.as_u64());

        // Best-effort cleanup of user mappings.
        self.unmap_all_user_regions();
        #[cfg(not(feature = "selftest"))]
        self.free_user_page_tables();
        #[cfg(feature = "selftest")]
        {
            // Runtime selftests create/destroy many temporary address spaces and
            // currently expose instability in recursive page-table teardown.
            // Keep tests deterministic by skipping deep PT reclaim in this mode.
            log::trace!(
                "AddressSpace::drop selftest mode: skipping deep page-table free for CR3={:#x}",
                self.cr3_phys.as_u64()
            );
        }

        // Free the PML4 frame itself.
        // NOTE: Recursive freeing of intermediate page tables (L3/L2/L1) that
        // belong exclusively to the user half is deferred to P2.
        let phys_frame = crate::memory::PhysFrame {
            start_address: self.cr3_phys,
        };
        crate::sync::with_irqs_disabled(|token| {
            crate::memory::free_frame(token, phys_frame);
        });

        log::trace!("AddressSpace::drop end CR3={:#x}", self.cr3_phys.as_u64());
        log::debug!(
            "User address space dropped: CR3={:#x}",
            self.cr3_phys.as_u64()
        );
    }
}

// ---------------------------------------------------------------------------
// Page table cleanup helpers (user half only)
// ---------------------------------------------------------------------------

/// Releases frame.
fn free_frame(phys: PhysAddr) {
    let phys_frame = crate::memory::PhysFrame {
        start_address: phys,
    };
    crate::sync::with_irqs_disabled(|token| {
        crate::memory::free_frame(token, phys_frame);
    });
}

/// Releases l1 table.
fn free_l1_table(frame: X86PhysFrame<Size4KiB>) {
    let l1_virt = VirtAddr::new(crate::memory::phys_to_virt(frame.start_address().as_u64()));
    // SAFETY: l1_virt points to a valid page table frame in HHDM.
    let l1 = unsafe { &mut *l1_virt.as_mut_ptr::<PageTable>() };
    for entry in l1.iter_mut() {
        if entry.flags().contains(PageTableFlags::PRESENT) {
            // Mapped frames are already freed via unmap_all_user_regions.
            entry.set_unused();
        }
    }
    free_frame(frame.start_address());
}

/// Releases l2 table.
fn free_l2_table(frame: X86PhysFrame<Size4KiB>) {
    let l2_virt = VirtAddr::new(crate::memory::phys_to_virt(frame.start_address().as_u64()));
    let l2 = unsafe { &mut *l2_virt.as_mut_ptr::<PageTable>() };
    for entry in l2.iter_mut() {
        if !entry.flags().contains(PageTableFlags::PRESENT) {
            continue;
        }
        if entry.flags().contains(PageTableFlags::HUGE_PAGE) {
            // 2 MiB pages are not expected in user space today.
            entry.set_unused();
            continue;
        }
        if let Ok(l1_frame) = entry.frame() {
            free_l1_table(l1_frame);
        }
        entry.set_unused();
    }
    free_frame(frame.start_address());
}

/// Releases l3 table.
fn free_l3_table(frame: X86PhysFrame<Size4KiB>) {
    let l3_virt = VirtAddr::new(crate::memory::phys_to_virt(frame.start_address().as_u64()));
    let l3 = unsafe { &mut *l3_virt.as_mut_ptr::<PageTable>() };
    for entry in l3.iter_mut() {
        if !entry.flags().contains(PageTableFlags::PRESENT) {
            continue;
        }
        if entry.flags().contains(PageTableFlags::HUGE_PAGE) {
            // 1 GiB pages are not expected in user space today.
            entry.set_unused();
            continue;
        }
        if let Ok(l2_frame) = entry.frame() {
            free_l2_table(l2_frame);
        }
        entry.set_unused();
    }
    free_frame(frame.start_address());
}

// ---------------------------------------------------------------------------
// Kernel address space singleton
// ---------------------------------------------------------------------------

static KERNEL_ADDRESS_SPACE: Once<Arc<AddressSpace>> = Once::new();

/// Initialize the kernel address space singleton.
///
/// Must be called once during boot, after paging is initialized, before the
/// scheduler creates any tasks.
///
/// # Safety
/// Must be called in single-threaded init context.
pub unsafe fn init_kernel_address_space() {
    KERNEL_ADDRESS_SPACE.call_once(|| {
        // SAFETY: Called once, single-threaded, paging initialized.
        Arc::new(unsafe { AddressSpace::new_kernel() })
    });
}

/// Get a reference to the kernel address space.
///
/// Panics if called before `init_kernel_address_space()`.
pub fn kernel_address_space() -> &'static Arc<AddressSpace> {
    KERNEL_ADDRESS_SPACE
        .get()
        .expect("Kernel address space not initialized")
}