// Kernel page-table manager. This file owns x86-64 PTE manipulation, TLB invalidation, and shared-CR3 activation.

#include <stddef.h>
#include <string.h>
#include <intrin.h>

#include <kernel/mm_pte.h>
#include <kernel/mm_pfn.h>
#include <kernel/ke_cpu.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_spinlock.h>
#include <kernel/hal/hal_apic.h>

// Page-table frames are reached through the physical direct map while the kernel still owns one shared CR3.
static uint64_t gPhysBase = PHYS_MAP_BASE;

// Converts a physical address to the kernel physical-window alias used for page-table editing.
static inline void* KiPhysToVirt(PHYSICAL_ADDRESS p) {
    return (void *)(p + gPhysBase);
}
// Returns the entry pointer for a page-table frame that is already mapped in the physical window.
static inline uint64_t* KiPtePtr(PHYSICAL_ADDRESS tablePhys, unsigned idx) {
    return (uint64_t *)KiPhysToVirt(tablePhys) + idx;
}

// The bootloader builds the initial PML4, and this layer owns all later edits to it.
static PHYSICAL_ADDRESS gPml4Phys = 0;
static SPINLOCK gPageTableLock = SPINLOCK_INIT;

static volatile long gTlbShootdownLock = 0;
static volatile long gTlbShootdownActive = 0;
static volatile long gTlbShootdownAcks = 0;
static uintptr_t gTlbShootdownVirt = 0;
static MM_PTE_STATS gPteStats;

#define TLB_SHOOTDOWN_TIMEOUT_SPINS 20000000u

// These macros extract the hardware page-table index for each level.
#define PML4_IDX(va)  (((va) >> 39) & 0x1FF)
#define PDPT_IDX(va)  (((va) >> 30) & 0x1FF)
#define PDT_IDX(va)   (((va) >> 21) & 0x1FF)
#define PT_IDX(va)    (((va) >> 12) & 0x1FF)

#define PAGE_PS              (1ULL << 7)
#define PAGE_FLAGS_MASK      (0xFFFULL | PAGE_NX)
#define PAGE_FRAME_4K_MASK   0x000FFFFFFFFFF000ULL
#define PAGE_FRAME_2M_MASK   0x000FFFFFFFE00000ULL
#define PAGE_FRAME_1G_MASK   0x000FFFFFC0000000ULL

// Bumps a telemetry counter. These counters are diagnostic only; exactness is
// less important than avoiding extra locks in hot page-table paths.
static inline void KiPteStatInc(uint64_t* counter) {
    (*counter)++;
}

// Clamps query attribution so corrupt callers cannot index outside telemetry arrays.
static MM_PTE_QUERY_PURPOSE KiNormalizeQueryPurpose(MM_PTE_QUERY_PURPOSE purpose) {
    if ((uint32_t)purpose >= (uint32_t)MmPteQueryPurposeCount)
        return MmPteQueryGeneric;
    return purpose;
}

// Ensures that one parent entry points to a child page table.
// New child tables are zeroed before publication so no stale mappings become visible.
static PHYSICAL_ADDRESS KiEnsureTable(PHYSICAL_ADDRESS parentPhys,
                                    unsigned index, uint64_t extraFlags) {
    uint64_t* entry = KiPtePtr(parentPhys, index);
    if (*entry & PAGE_PRESENT) {
        *entry |= extraFlags;
        return *entry & PAGE_FRAME_4K_MASK;
    }
    PHYSICAL_ADDRESS childPhys = MmPfnAllocate(MmPfnUsePageTable);
    if (!childPhys) return 0;
    KiPteStatInc(&gPteStats.PageTablePagesAllocated);
    memset(KiPhysToVirt(childPhys), 0, PAGE_SIZE);
    *entry = childPhys | PAGE_PRESENT | PAGE_WRITE | extraFlags;
    return childPhys;
}

// Initializes the page-table manager around the bootloader-created PML4.
void MmPteInitialize(PHYSICAL_ADDRESS bootPml4Phys) {
    gPml4Phys = bootPml4Phys;
    gPageTableLock = (SPINLOCK)SPINLOCK_INIT;
    gTlbShootdownLock = 0;
    gTlbShootdownActive = 0;
    gTlbShootdownAcks = 0;
    gTlbShootdownVirt = 0;
    memset(&gPteStats, 0, sizeof(gPteStats));
}

// Installs the shared PML4 into CR3.
void MmPteActivate(void) {
    __writecr3(gPml4Phys);
}

// Flushes one virtual address from the local TLB.
void MmPteFlushLocal(uintptr_t virt) {
    KiPteStatInc(&gPteStats.LocalFlushes);
    __invlpg((void*)virt);
}

// Acquires the shootdown state lock without disabling interrupts.
static void KiAcquireTlbShootdownLock(void) {
    // Do not use the normal spinlock here: it masks interrupts while spinning.
    // A CPU spinning with IF=0 on the shootdown lock could not acknowledge a
    // concurrent fixed-vector TLB IPI, creating the exact deadlock this lock is
    // meant to avoid.
    while (_InterlockedCompareExchange(&gTlbShootdownLock, 1, 0) != 0)
        _mm_pause();
    _ReadWriteBarrier();
}

// Releases the shootdown state lock after all shared fields are stable.
static void KiReleaseTlbShootdownLock(void) {
    _ReadWriteBarrier();
    _InterlockedExchange(&gTlbShootdownLock, 0);
}

// Flushes one kernel mapping on every online CPU that shares this PML4.
void MmPteFlushKernelAddress(uintptr_t virt) {
    KiPteStatInc(&gPteStats.GlobalFlushes);
    MmPteFlushLocal(virt);

    if (gCpuCount <= 1) {
        KiPteStatInc(&gPteStats.ShootdownSkippedNoTargets);
        return;
    }

    KCPU* self = KeGetCurrentCpu();
    uint32_t selfIndex = self ? self->CpuIndex : KE_INVALID_CPU;

    KiAcquireTlbShootdownLock();

    long targets = 0;
    gTlbShootdownVirt = virt;
    _InterlockedExchange(&gTlbShootdownAcks, 0);
    _ReadWriteBarrier();
    _InterlockedExchange(&gTlbShootdownActive, 1);
    _ReadWriteBarrier();

    for (uint32_t i = 0; i < gCpuCount; i++) {
        KCPU* cpu = &gCpus[i];
        if (i == selfIndex)
            continue;
        if (!cpu->Present || !cpu->Started || !cpu->Online)
            continue;
        targets++;
        KiPteStatInc(&gPteStats.ShootdownTargetCpus);
        HalApicSendTlbShootdownIpi(cpu->LapicId);
    }

    if (!targets) {
        KiPteStatInc(&gPteStats.ShootdownSkippedNoTargets);
        _InterlockedExchange(&gTlbShootdownActive, 0);
        _ReadWriteBarrier();
        KiReleaseTlbShootdownLock();
        return;
    }

    KiPteStatInc(&gPteStats.ShootdownRequests);

    for (uint32_t spin = 0; targets > 0 && spin < TLB_SHOOTDOWN_TIMEOUT_SPINS; spin++) {
        if (_InterlockedCompareExchange(&gTlbShootdownAcks, 0, 0) >= targets)
            break;
        _mm_pause();
    }

    long acks = _InterlockedCompareExchange(&gTlbShootdownAcks, 0, 0);
    _InterlockedExchange(&gTlbShootdownActive, 0);
    _ReadWriteBarrier();
    KiReleaseTlbShootdownLock();

    gPteStats.ShootdownAcks += (uint64_t)acks;
    if (acks < targets) {
        KiPteStatInc(&gPteStats.ShootdownTimeouts);
        KePanic(KE_PANIC_CODE_GENERIC,
                "TLB shootdown timeout: va=0x%016llX targets=%ld acks=%ld selfCpu=%u\n",
                (uint64_t)virt, targets, acks, selfIndex);
    }
}

// Handles an incoming shootdown request on the current CPU.
static int KiHandleTlbShootdown(void) {
    if (_InterlockedCompareExchange(&gTlbShootdownActive, 0, 0) == 0)
        return 0;

    uintptr_t va = gTlbShootdownVirt;
    _ReadWriteBarrier();
    MmPteFlushLocal(va);
    _ReadWriteBarrier();
    _InterlockedIncrement(&gTlbShootdownAcks);
    return 1;
}

// Handles the NMI vector reserved for TLB shootdowns.
int MmPteHandleShootdownNmi(void) {
    return KiHandleTlbShootdown();
}

// Maps or replaces one 4 KiB leaf PTE in the shared page tables.
__declspec(dllexport) KSTATUS MmPteMapPage(uintptr_t virt, PHYSICAL_ADDRESS phys, uint64_t flags, MM_PTE_MAP_MODE mode) {
    KiPteStatInc(&gPteStats.MapCalls);
    unsigned userFlag = (flags & PAGE_USER) ? PAGE_USER : 0;

    uintptr_t lockFlags = KeAcquireSpinlock(&gPageTableLock);

    PHYSICAL_ADDRESS pdpt = KiEnsureTable(gPml4Phys, PML4_IDX(virt), userFlag);
    if (!pdpt) {
        KiPteStatInc(&gPteStats.MapFailures);
        KeReleaseSpinlock(&gPageTableLock, lockFlags);
        return STATUS_NO_MEMORY;
    }
    PHYSICAL_ADDRESS pdt  = KiEnsureTable(pdpt, PDPT_IDX(virt), userFlag);
    if (!pdt) {
        KiPteStatInc(&gPteStats.MapFailures);
        KeReleaseSpinlock(&gPageTableLock, lockFlags);
        return STATUS_NO_MEMORY;
    }
    PHYSICAL_ADDRESS pt   = KiEnsureTable(pdt,  PDT_IDX(virt),  userFlag);
    if (!pt) {
        KiPteStatInc(&gPteStats.MapFailures);
        KeReleaseSpinlock(&gPageTableLock, lockFlags);
        return STATUS_NO_MEMORY;
    }

    uint64_t* pte = KiPtePtr(pt, PT_IDX(virt));
    if ((*pte & PAGE_PRESENT) && mode == MmPteMapFailIfPresent) {
        KiPteStatInc(&gPteStats.MapFailures);
        KeReleaseSpinlock(&gPageTableLock, lockFlags);
        return STATUS_ALREADY_EXISTS;
    }
    int wasPresent = ((*pte & PAGE_PRESENT) != 0);
    *pte = phys | flags;
    if (!wasPresent) {
        KiPteStatInc(&gPteStats.CurrentLeafMappings);
        KiPteStatInc(&gPteStats.NewMappingNoFlushes);
    }

    KeReleaseSpinlock(&gPageTableLock, lockFlags);

    if (wasPresent) {
        KiPteStatInc(&gPteStats.ReplacedMappingFlushes);
        MmPteFlushKernelAddress(virt);
    }
    return STATUS_SUCCESS;
}

// Removes one 4 KiB leaf PTE from the shared page tables.
__declspec(dllexport) KSTATUS MmPteUnmapPage(uintptr_t virt) {
    KiPteStatInc(&gPteStats.UnmapCalls);
    uintptr_t lockFlags = KeAcquireSpinlock(&gPageTableLock);

    uint64_t* e4 = KiPtePtr(gPml4Phys, PML4_IDX(virt));
    if (!(*e4 & PAGE_PRESENT)) {
        KiPteStatInc(&gPteStats.UnmapFailures);
        KeReleaseSpinlock(&gPageTableLock, lockFlags);
        return STATUS_NOT_FOUND;
    }
    uint64_t* e3 = KiPtePtr(*e4 & PAGE_FRAME_4K_MASK, PDPT_IDX(virt));
    if (!(*e3 & PAGE_PRESENT)) {
        KiPteStatInc(&gPteStats.UnmapFailures);
        KeReleaseSpinlock(&gPageTableLock, lockFlags);
        return STATUS_NOT_FOUND;
    }
    uint64_t* e2 = KiPtePtr(*e3 & PAGE_FRAME_4K_MASK, PDT_IDX(virt));
    if (!(*e2 & PAGE_PRESENT)) {
        KiPteStatInc(&gPteStats.UnmapFailures);
        KeReleaseSpinlock(&gPageTableLock, lockFlags);
        return STATUS_NOT_FOUND;
    }
    uint64_t* pte = KiPtePtr(*e2 & PAGE_FRAME_4K_MASK, PT_IDX(virt));
    if (!(*pte & PAGE_PRESENT)) {
        KiPteStatInc(&gPteStats.UnmapFailures);
        KeReleaseSpinlock(&gPageTableLock, lockFlags);
        return STATUS_NOT_FOUND;
    }
    *pte = 0;
    if (gPteStats.CurrentLeafMappings)
        gPteStats.CurrentLeafMappings--;

    KeReleaseSpinlock(&gPageTableLock, lockFlags);
    MmPteFlushKernelAddress(virt);
    return STATUS_SUCCESS;
}

// Replaces leaf flags for a 4 KiB range and refuses huge pages instead of silently corrupting them.
__declspec(dllexport) KSTATUS MmPteProtectRange(uintptr_t virt, size_t size, uint64_t flags) {
    KiPteStatInc(&gPteStats.ProtectCalls);
    if (!size) return STATUS_SUCCESS;

    // The protection operation is page-granular even when callers pass an interior address.
    uintptr_t base = virt & PAGE_MASK;
    uintptr_t end  = (virt + size + PAGE_SIZE - 1) & PAGE_MASK;

    uintptr_t lockFlags = KeAcquireSpinlock(&gPageTableLock);
    KSTATUS status = STATUS_SUCCESS;
    uintptr_t flushEnd = base;

    for (uintptr_t va = base; va < end; va += PAGE_SIZE) {
        uint64_t* e4 = KiPtePtr(gPml4Phys,          PML4_IDX(va));
        if (!(*e4 & PAGE_PRESENT)) {
            status = STATUS_NOT_FOUND;
            break;
        }
        uint64_t* e3 = KiPtePtr(*e4 & PAGE_FRAME_4K_MASK,    PDPT_IDX(va));
        if (!(*e3 & PAGE_PRESENT)) {
            status = STATUS_NOT_FOUND;
            break;
        }
        if (*e3 & PAGE_PS) {
            status = STATUS_NOT_SUPPORTED;
            break;
        }
        uint64_t* e2 = KiPtePtr(*e3 & PAGE_FRAME_4K_MASK,    PDT_IDX(va));
        if (!(*e2 & PAGE_PRESENT)) {
            status = STATUS_NOT_FOUND;
            break;
        }
        if (*e2 & PAGE_PS) {
            status = STATUS_NOT_SUPPORTED;
            break;
        }
        uint64_t* pte = KiPtePtr(*e2 & PAGE_FRAME_4K_MASK,   PT_IDX(va));
        if (!(*pte & PAGE_PRESENT)) {
            status = STATUS_NOT_FOUND;
            break;
        }

        // The physical frame stays attached while access and cache attributes change.
        PHYSICAL_ADDRESS phys = *pte & PAGE_FRAME_4K_MASK;
        *pte = phys | flags;
        flushEnd = va + PAGE_SIZE;
    }

    if (!KSUCCESS(status))
        KiPteStatInc(&gPteStats.ProtectFailures);

    KeReleaseSpinlock(&gPageTableLock, lockFlags);

    for (uintptr_t va = base; va < flushEnd; va += PAGE_SIZE)
        MmPteFlushKernelAddress(va);

    return status;
}

// Tests and clears the hardware accessed bit on a mapped 4 KiB leaf. This is
// used by the pressure trimmer as a second-chance test for clean file-backed
// pages. If the page was accessed recently, reclaim skips it once and flushes
// the TLB entry so future accesses can set PAGE_ACCESSED again.
__declspec(dllexport) int MmPteTestAndClearAccessed(uintptr_t virt) {
    if (!gPml4Phys)
        return 0;

    uintptr_t va = virt & PAGE_MASK;
    int wasAccessed = 0;

    uintptr_t lockFlags = KeAcquireSpinlock(&gPageTableLock);

    uint64_t* e4 = KiPtePtr(gPml4Phys, PML4_IDX(va));
    if (!(*e4 & PAGE_PRESENT))
        goto out;
    uint64_t* e3 = KiPtePtr(*e4 & PAGE_FRAME_4K_MASK, PDPT_IDX(va));
    if (!(*e3 & PAGE_PRESENT) || (*e3 & PAGE_PS))
        goto out;
    uint64_t* e2 = KiPtePtr(*e3 & PAGE_FRAME_4K_MASK, PDT_IDX(va));
    if (!(*e2 & PAGE_PRESENT) || (*e2 & PAGE_PS))
        goto out;
    uint64_t* pte = KiPtePtr(*e2 & PAGE_FRAME_4K_MASK, PT_IDX(va));
    if (!(*pte & PAGE_PRESENT))
        goto out;

    if (*pte & PAGE_ACCESSED) {
        *pte &= ~PAGE_ACCESSED;
        wasAccessed = 1;
    }

out:
    KeReleaseSpinlock(&gPageTableLock, lockFlags);

    if (wasAccessed)
        MmPteFlushKernelAddress(va);
    return wasAccessed;
}

// Tests and clears the hardware dirty bit on a mapped 4 KiB leaf.
__declspec(dllexport) int MmPteTestAndClearDirty(uintptr_t virt) {
    if (!gPml4Phys)
        return 0;

    uintptr_t va = virt & PAGE_MASK;
    int wasDirty = 0;

    uintptr_t lockFlags = KeAcquireSpinlock(&gPageTableLock);

    uint64_t* e4 = KiPtePtr(gPml4Phys, PML4_IDX(va));
    if (!(*e4 & PAGE_PRESENT))
        goto out;
    uint64_t* e3 = KiPtePtr(*e4 & PAGE_FRAME_4K_MASK, PDPT_IDX(va));
    if (!(*e3 & PAGE_PRESENT) || (*e3 & PAGE_PS))
        goto out;
    uint64_t* e2 = KiPtePtr(*e3 & PAGE_FRAME_4K_MASK, PDT_IDX(va));
    if (!(*e2 & PAGE_PRESENT) || (*e2 & PAGE_PS))
        goto out;
    uint64_t* pte = KiPtePtr(*e2 & PAGE_FRAME_4K_MASK, PT_IDX(va));
    if (!(*pte & PAGE_PRESENT))
        goto out;

    if (*pte & PAGE_DIRTY) {
        *pte &= ~PAGE_DIRTY;
        wasDirty = 1;
    }

out:
    KeReleaseSpinlock(&gPageTableLock, lockFlags);

    if (wasDirty)
        MmPteFlushKernelAddress(va);
    return wasDirty;
}

// Translates a mapped virtual address to a physical address.
__declspec(dllexport) PHYSICAL_ADDRESS MmPteTranslate(uintptr_t virt) {
    KiPteStatInc(&gPteStats.TranslateCalls);
    if (!gPml4Phys) return 0;

    PHYSICAL_ADDRESS result = 0;
    uintptr_t lockFlags = KeAcquireSpinlock(&gPageTableLock);

    uint64_t pml4_idx = (virt >> 39) & 0x1FF;
    uint64_t pdpt_idx = (virt >> 30) & 0x1FF;
    uint64_t pd_idx   = (virt >> 21) & 0x1FF;
    uint64_t pt_idx   = (virt >> 12) & 0x1FF;

    uint64_t* pml4 = KiPhysToVirt(gPml4Phys);
    if (!(pml4[pml4_idx] & PAGE_PRESENT)) goto out;

    uint64_t* pdpt = KiPhysToVirt(pml4[pml4_idx] & PAGE_FRAME_4K_MASK);
    if (!(pdpt[pdpt_idx] & PAGE_PRESENT)) goto out;

    // A present PDPT leaf maps a 1 GiB huge page.
    if (pdpt[pdpt_idx] & PAGE_PS) {
        result = (pdpt[pdpt_idx] & PAGE_FRAME_1G_MASK) + (virt & 0x3FFFFFFF);
        goto out;
    }

    uint64_t* pd = KiPhysToVirt(pdpt[pdpt_idx] & PAGE_FRAME_4K_MASK);
    if (!(pd[pd_idx] & PAGE_PRESENT)) goto out;
    
    // A present PDT leaf maps a 2 MiB huge page.
    if (pd[pd_idx] & PAGE_PS) {
        result = (pd[pd_idx] & PAGE_FRAME_2M_MASK) + (virt & 0x1FFFFF);
        goto out;
    }

    uint64_t* pt = KiPhysToVirt(pd[pd_idx] & PAGE_FRAME_4K_MASK);
    if (!(pt[pt_idx] & PAGE_PRESENT)) goto out;

    // A normal PT leaf maps a 4 KiB page.
    result = (pt[pt_idx] & PAGE_FRAME_4K_MASK) + (virt & 0xFFF);

out:
    KeReleaseSpinlock(&gPageTableLock, lockFlags);
    return result;
}

// Queries one virtual address through the shared page-table hierarchy.
KSTATUS MmPteQueryPageFor(MM_PTE_QUERY_PURPOSE purpose,
                           uintptr_t virt,
                           PHYSICAL_ADDRESS* physOut,
                           uint64_t* flagsOut) {
    purpose = KiNormalizeQueryPurpose(purpose);
    KiPteStatInc(&gPteStats.QueryCalls);
    KiPteStatInc(&gPteStats.QueryByPurpose[purpose]);
    if (!gPml4Phys) {
        KiPteStatInc(&gPteStats.QueryFailures);
        KiPteStatInc(&gPteStats.QueryFailuresByPurpose[purpose]);
        return STATUS_NOT_FOUND;
    }

    KSTATUS status = STATUS_NOT_FOUND;
    uintptr_t lockFlags = KeAcquireSpinlock(&gPageTableLock);

    uint64_t* pml4 = (uint64_t*)KiPhysToVirt(gPml4Phys);
    uint64_t e4 = pml4[PML4_IDX(virt)];
    if (!(e4 & PAGE_PRESENT)) goto out;

    uint64_t effectiveWrite = e4 & PAGE_WRITE;
    uint64_t effectiveUser  = e4 & PAGE_USER;
    uint64_t effectiveNx    = e4 & PAGE_NX;

    uint64_t* pdpt = (uint64_t*)KiPhysToVirt(e4 & PAGE_FRAME_4K_MASK);
    uint64_t e3 = pdpt[PDPT_IDX(virt)];
    if (!(e3 & PAGE_PRESENT)) goto out;
    effectiveWrite &= e3 & PAGE_WRITE;
    effectiveUser  &= e3 & PAGE_USER;
    effectiveNx    |= e3 & PAGE_NX;

    if (e3 & PAGE_PS) {
        if (physOut)
            *physOut = (e3 & PAGE_FRAME_1G_MASK) + (virt & 0x3FFFFFFFULL);
        if (flagsOut) {
            uint64_t flags = (e3 & PAGE_FLAGS_MASK) | effectiveNx;
            if (!effectiveWrite) flags &= ~PAGE_WRITE;
            if (!effectiveUser)  flags &= ~PAGE_USER;
            *flagsOut = flags;
        }
        status = STATUS_SUCCESS;
        goto out;
    }

    uint64_t* pd = (uint64_t*)KiPhysToVirt(e3 & PAGE_FRAME_4K_MASK);
    uint64_t e2 = pd[PDT_IDX(virt)];
    if (!(e2 & PAGE_PRESENT)) goto out;
    effectiveWrite &= e2 & PAGE_WRITE;
    effectiveUser  &= e2 & PAGE_USER;
    effectiveNx    |= e2 & PAGE_NX;

    if (e2 & PAGE_PS) {
        if (physOut)
            *physOut = (e2 & PAGE_FRAME_2M_MASK) + (virt & 0x1FFFFFULL);
        if (flagsOut) {
            uint64_t flags = (e2 & PAGE_FLAGS_MASK) | effectiveNx;
            if (!effectiveWrite) flags &= ~PAGE_WRITE;
            if (!effectiveUser)  flags &= ~PAGE_USER;
            *flagsOut = flags;
        }
        status = STATUS_SUCCESS;
        goto out;
    }

    uint64_t* pt = (uint64_t*)KiPhysToVirt(e2 & PAGE_FRAME_4K_MASK);
    uint64_t e1 = pt[PT_IDX(virt)];
    if (!(e1 & PAGE_PRESENT)) goto out;
    effectiveWrite &= e1 & PAGE_WRITE;
    effectiveUser  &= e1 & PAGE_USER;
    effectiveNx    |= e1 & PAGE_NX;

    if (physOut)
        *physOut = (e1 & PAGE_FRAME_4K_MASK) + (virt & 0xFFFULL);
    if (flagsOut) {
        uint64_t flags = (e1 & PAGE_FLAGS_MASK) | effectiveNx;
        if (!effectiveWrite) flags &= ~PAGE_WRITE;
        if (!effectiveUser)  flags &= ~PAGE_USER;
        *flagsOut = flags;
    }
    status = STATUS_SUCCESS;

out:
    if (!KSUCCESS(status)) {
        KiPteStatInc(&gPteStats.QueryFailures);
        KiPteStatInc(&gPteStats.QueryFailuresByPurpose[purpose]);
    }
    KeReleaseSpinlock(&gPageTableLock, lockFlags);
    return status;
}

// Queries one virtual address without caller-specific attribution.
__declspec(dllexport) KSTATUS MmPteQueryPage(uintptr_t virt,
                                          PHYSICAL_ADDRESS* physOut,
                                          uint64_t* flagsOut) {
    return MmPteQueryPageFor(MmPteQueryGeneric, virt, physOut, flagsOut);
}

// Returns page-table telemetry for diagnostics.
void MmPteQueryStats(MM_PTE_STATS* out) {
    if (!out)
        return;
    *out = gPteStats;
}