// Kernel virtual-memory manager. This file owns kernel pool allocation, vad bookkeeping, and user-copy validation.
// Page-table construction lives in mm_pte.c so the higher VMM can grow toward VADs, faults, COW, and swap without mixing those policies with raw PTE edits.

// Kernel pool:
//   A boundary-tag allocator starts at POOL_BASE and commits pages on demand.

// User mappings:
//   VADs separate reservation from commitment. Committed pages become resident
//   when a user fault or kernel prefault path materializes them.
#include <stddef.h>
#include <string.h>
#include <intrin.h>
#include <kernel/mm_virt.h>
#include <kernel/mm_pte.h>
#include <kernel/mm_pfn.h>
#include <kernel/mm_swap.h>
#include <kernel/io_file.h>
#include <kernel/mm_file_cache.h>
#include <kernel/ldr_section.h>
#include <kernel/ob_object.h>
#include <kernel/ke_process.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_irql.h>
#include <kernel/ke_sched.h>

static KSTATUS KiVadEnsureResidentRange(PROCESS* process, uintptr_t start, size_t size, uint64_t accessFlags);
static KSTATUS KiCopyFromUserPaged(void* kernelDst, const void* userSrc, size_t size);
static KSTATUS KiCopyToUserPaged(void* userDst, const void* kernelSrc, size_t size);

static PHYSICAL_ADDRESS gVadZeroPageFrame;

// User/kernel copy boundary helpers.
// Implements the user Range Valid operation.
static int KiUserRangeValid(uintptr_t addr, size_t size) {
    if (size == 0) return 1;
    if (addr < USER_BASE) return 0;
    if (addr >= USER_LIMIT) return 0;
    if (size > (size_t)(USER_LIMIT - addr)) return 0;
    return 1;
}

// Probes user Range and materializes demand-zero pages before the kernel copies through it.
static KSTATUS KiProbeUserRange(uintptr_t addr, size_t size, uint64_t requiredFlags) {
    if (!KiUserRangeValid(addr, size))
        return STATUS_ACCESS_DENIED;

    if (size == 0)
        return STATUS_SUCCESS;

    PROCESS* process = KeGetCurrentProcess();
    if (!process)
        return STATUS_ACCESS_DENIED;

    return KiVadEnsureResidentRange(process, addr, size, requiredFlags);
}

// Probes user Read.
__declspec(dllexport) KSTATUS MmProbeUserRead(const void* ptr, size_t size) {
    return KiProbeUserRange((uintptr_t)ptr, size, PAGE_PRESENT | PAGE_USER);
}

// Probes user Write.
__declspec(dllexport) KSTATUS MmProbeUserWrite(void* ptr, size_t size) {
    return KiProbeUserRange((uintptr_t)ptr, size, PAGE_PRESENT | PAGE_USER | PAGE_WRITE);
}

// Copies from a user buffer through page-sized, pinned segments.
__declspec(dllexport) KSTATUS CopyFromUser(void* kernelDst, const void* userSrc, size_t size) {
    return KiCopyFromUserPaged(kernelDst, userSrc, size);
}

// Copies to a user buffer through page-sized, pinned segments.
__declspec(dllexport) KSTATUS CopyToUser(void* userDst, const void* kernelSrc, size_t size) {
    return KiCopyToUserPaged(userDst, kernelSrc, size);
}

// Boundary-tag allocator (used for both pool and user heap)
// Each allocation has a BLOCK_HDR immediately before the user data, and a
// BLOCK_FTR immediately after. Both carry the block size and a 'free' flag.
// This allows O(1) coalescing of adjacent free blocks.

#define BLOCK_MAGIC  0xB10CB10CUL

#define BLOCK_STATE_USED        0u
#define BLOCK_STATE_FREE        1u
#define BLOCK_STATE_QUARANTINE  2u

#define BLOCK_FLAG_CANARY       0x00000001u
#define BLOCK_FLAG_POISONED     0x00000002u

#define POOL_TAIL_CANARY_BASE   0xC0FFEE1234ABCDEFULL
#define POOL_POISON_FREE        0xDD
#define POOL_QUARANTINE_SLOTS   256u
#define POOL_TAG_DEFAULT        0x6C6F6F70u  /* 'pool' */

typedef struct BLOCK_HDR {
    uint32_t Magic;
    uint32_t Free;      // BLOCK_STATE_*
    size_t   Size;      // usable allocator bytes (excludes header + footer)
    struct BLOCK_HDR* PrevFree;
    struct BLOCK_HDR* NextFree;

    // Debug metadata. Present in all builds so boot.ini bits can enable
    // diagnostics without changing allocator layout or ABI.
    size_t   RequestedSize;
    uint32_t Tag;
    uint32_t Flags;
    uint64_t Sequence;
    uintptr_t AllocCaller;
    uintptr_t FreeCaller;
} BLOCK_HDR;

typedef struct BLOCK_FTR {
    size_t Size;
    uint32_t Free;
    uint32_t Magic;
} BLOCK_FTR;

#define HDR_SIZE  (sizeof(BLOCK_HDR))
#define FTR_SIZE  (sizeof(BLOCK_FTR))
#define OVERHEAD  (HDR_SIZE + FTR_SIZE)
#define MIN_BLOCK (OVERHEAD + sizeof(void*))

// Implements the hdr To Ftr operation.
static inline BLOCK_FTR* KiHdrToFtr(BLOCK_HDR* h) {
    return (BLOCK_FTR *)((char *)h + HDR_SIZE + h->Size);
}
// Implements the ftr To Hdr operation.
static inline BLOCK_HDR* KiFtrToHdr(BLOCK_FTR* f) {
    return (BLOCK_HDR *)((char *)f - HDR_SIZE - f->Size);
}
// Implements the next Block operation.
static inline BLOCK_HDR* KiNextBlock(BLOCK_HDR* h) {
    return (BLOCK_HDR *)((char *)h + HDR_SIZE + h->Size + FTR_SIZE);
}
// Implements the prev Ftr operation.
static inline BLOCK_FTR* KiPrevFtr(BLOCK_HDR* h) {
    return (BLOCK_FTR *)((char *)h - FTR_SIZE);
}

typedef struct HEAP HEAP;

struct HEAP {
    BLOCK_HDR* FreeList;    // doubly-linked list of free blocks
    uintptr_t   Start;       // virtual base of the arena
    uintptr_t   Brk;         // current committed end (page-aligned)
    uintptr_t   Limit;       // hard upper bound
    uint64_t    MapFlags;    // PAGE_* flags for new committed pages
    SPINLOCK    Lock;
    int         DecommitTailOnFree;

    uint64_t    NextSequence;
    BLOCK_HDR*  Quarantine[POOL_QUARANTINE_SLOTS];
    size_t      QuarantineCursor;
    size_t      QuarantineCount;
};

static inline uint32_t KiPoolDebugFlags(void) {
    return KeGetDebugFlags() & (KE_DEBUG_POOL_CANARY |
                                KE_DEBUG_POOL_POISON |
                                KE_DEBUG_POOL_QUARANTINE |
                                KE_DEBUG_POOL_CALLERS);
}

static inline uint64_t KiPoolTailCanaryValue(const BLOCK_HDR* blk) {
    return POOL_TAIL_CANARY_BASE ^
           ((uint64_t)(uintptr_t)blk << 13) ^
           ((uint64_t)blk->RequestedSize << 1) ^
           blk->Sequence;
}

static inline uint8_t* KiBlockUserBase(BLOCK_HDR* blk) {
    return (uint8_t*)blk + HDR_SIZE;
}

static inline uint64_t* KiBlockTailCanary(BLOCK_HDR* blk) {
    return (uint64_t*)(void*)(KiBlockUserBase(blk) + blk->RequestedSize);
}

static const char* KiPoolStateName(uint32_t state) {
    switch (state) {
        case BLOCK_STATE_USED: return "used";
        case BLOCK_STATE_FREE: return "free";
        case BLOCK_STATE_QUARANTINE: return "quarantine";
        default: return "bad-state";
    }
}

static _Noreturn void KiPoolPanic(const char* reason, HEAP* h, BLOCK_HDR* blk, const void* ptr) {
    KePanic(KE_PANIC_CODE_GENERIC,
            "MM: pool corruption: %s ptr=%p hdr=%p heap=[%016llx,%016llx) magic=%08x state=%u(%s) size=%llu req=%llu tag=%08x flags=%08x seq=%llu alloc=%p free=%p\n",
            reason ? reason : "unknown",
            ptr,
            blk,
            (unsigned long long)(h ? h->Start : 0),
            (unsigned long long)(h ? h->Brk : 0),
            blk ? blk->Magic : 0,
            blk ? blk->Free : 0,
            blk ? KiPoolStateName(blk->Free) : "<null>",
            (unsigned long long)(blk ? blk->Size : 0),
            (unsigned long long)(blk ? blk->RequestedSize : 0),
            blk ? blk->Tag : 0,
            blk ? blk->Flags : 0,
            (unsigned long long)(blk ? blk->Sequence : 0),
            blk ? (void*)blk->AllocCaller : (void*)0,
            blk ? (void*)blk->FreeCaller : (void*)0);
}

static void KiPoolValidateFooterLocked(HEAP* h, BLOCK_HDR* blk, const char* where) {
    if (!h || !blk)
        KiPoolPanic(where, h, blk, blk ? KiBlockUserBase(blk) : (void*)0);
    if (blk->Magic != BLOCK_MAGIC)
        KiPoolPanic(where, h, blk, KiBlockUserBase(blk));
    if ((uintptr_t)blk < h->Start || (uintptr_t)blk >= h->Brk)
        KiPoolPanic("block outside heap", h, blk, KiBlockUserBase(blk));
    if (blk->Size > (size_t)(h->Brk - (uintptr_t)blk - OVERHEAD))
        KiPoolPanic("block size outside heap", h, blk, KiBlockUserBase(blk));
    BLOCK_FTR* ftr = KiHdrToFtr(blk);
    if ((uintptr_t)ftr + FTR_SIZE > h->Brk)
        KiPoolPanic("footer outside heap", h, blk, KiBlockUserBase(blk));
    if (ftr->Magic != BLOCK_MAGIC || ftr->Size != blk->Size || ftr->Free != blk->Free)
        KiPoolPanic(where, h, blk, KiBlockUserBase(blk));
}

static void KiPoolWriteTailCanary(BLOCK_HDR* blk) {
    if (!blk || !(blk->Flags & BLOCK_FLAG_CANARY))
        return;
    *KiBlockTailCanary(blk) = KiPoolTailCanaryValue(blk);
}

static void KiPoolValidateTailCanaryLocked(HEAP* h, BLOCK_HDR* blk, const char* where) {
    if (!blk || !(blk->Flags & BLOCK_FLAG_CANARY))
        return;
    if (blk->RequestedSize + sizeof(uint64_t) > blk->Size)
        KiPoolPanic("invalid canary layout", h, blk, KiBlockUserBase(blk));
    uint64_t actual = *KiBlockTailCanary(blk);
    uint64_t expect = KiPoolTailCanaryValue(blk);
    if (actual != expect)
        KiPoolPanic(where, h, blk, KiBlockUserBase(blk));
}

static void KiPoolPoisonFreedPayload(BLOCK_HDR* blk) {
    if (!blk)
        return;
    size_t bytes = blk->RequestedSize;
    if (bytes > blk->Size)
        bytes = blk->Size;
    memset(KiBlockUserBase(blk), POOL_POISON_FREE, bytes);
    blk->Flags |= BLOCK_FLAG_POISONED;
    KiPoolWriteTailCanary(blk);
}

static void KiPoolValidateFreedPoisonLocked(HEAP* h, BLOCK_HDR* blk, const char* where) {
    if (!blk || !(blk->Flags & BLOCK_FLAG_POISONED))
        return;
    size_t bytes = blk->RequestedSize;
    if (bytes > blk->Size)
        bytes = blk->Size;
    const uint8_t* p = KiBlockUserBase(blk);
    for (size_t i = 0; i < bytes; i++) {
        if (p[i] != POOL_POISON_FREE)
            KiPoolPanic(where, h, blk, KiBlockUserBase(blk));
    }
}

// Rolls back a failed heap grow attempt. KiHeapGrow maps pages before the
// new free block is linked into the allocator; if a later frame allocation or
// page-table allocation fails, those earlier mappings must be returned or the
// pool silently leaks physical frames and leaves stale mappings below Brk.
static void KiHeapRollbackGrow(uintptr_t start, uintptr_t end) {
    for (uintptr_t va = start; va < end; va += PAGE_SIZE) {
        PHYSICAL_ADDRESS frame = MmPteTranslate(va);
        if (!frame)
            continue;
        if (KSUCCESS(MmPteUnmapPage(va)))
            (void)MmPfnFree(frame);
    }
}

// Extend heap by at least 'need' bytes (rounded up to page boundary).
static int KiHeapGrow(HEAP* h, size_t need) {
    size_t grow = (need + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (h->Brk + grow > h->Limit) return 0;

    uintptr_t start = h->Brk;
    uintptr_t mappedEnd = start;
    for (uintptr_t va = start; va < start + grow; va += PAGE_SIZE) {
        PHYSICAL_ADDRESS frame = MmPfnAllocate(MmPfnUseKernel);
        if (!frame) {
            KiHeapRollbackGrow(start, mappedEnd);
            return 0;
        }
        if (!KSUCCESS(MmPteMapPage(va, frame, h->MapFlags, MmPteMapFailIfPresent))) {
            MmPfnFree(frame);
            KiHeapRollbackGrow(start, mappedEnd);
            return 0;
        }
        mappedEnd = va + PAGE_SIZE;
        memset((void *)va, 0, PAGE_SIZE);
    }

    // The grow path owns the only transition from reserved heap VA to mapped heap
    // pages, so a successful loop above is the heap's committed-frontier proof.
    // Re-querying every heap page on later allocations only burns page-table locks.
    // Missing pool pages are corruption now, not something the allocator repairs.

    uintptr_t oldBrk = h->Brk;
    uintptr_t newBrk = oldBrk + grow;

    // If the current tail block is already free, extend it in-place instead of
    // leaving adjacent free blocks. This matters once tail decommit is enabled:
    // a later grow should rebuild the same tail block, not fragment the arena
    // into many page-sized free runs.
    if (oldBrk > h->Start) {
        BLOCK_FTR* prevFtr = (BLOCK_FTR*)(oldBrk - FTR_SIZE);
        if (prevFtr->Magic == BLOCK_MAGIC && prevFtr->Free == BLOCK_STATE_FREE) {
            BLOCK_HDR* prev = KiFtrToHdr(prevFtr);
            if ((uintptr_t)prev >= h->Start && prev->Magic == BLOCK_MAGIC && prev->Free == BLOCK_STATE_FREE) {
                prev->Size += grow;
                h->Brk = newBrk;
                KiHdrToFtr(prev)->Magic = BLOCK_MAGIC;
                KiHdrToFtr(prev)->Free  = BLOCK_STATE_FREE;
                KiHdrToFtr(prev)->Size  = prev->Size;
                return 1;
            }
        }
    }

    // Make the new pages a single free block.
    BLOCK_HDR* blk = (BLOCK_HDR *)oldBrk;
    blk->Magic    = BLOCK_MAGIC;
    blk->Free     = BLOCK_STATE_FREE;
    blk->Size     = grow - OVERHEAD;
    blk->PrevFree = (BLOCK_HDR *)0;
    blk->NextFree = h->FreeList;
    blk->RequestedSize = 0;
    blk->Tag = 0;
    blk->Flags = 0;
    blk->Sequence = 0;
    blk->AllocCaller = 0;
    blk->FreeCaller = 0;
    if (h->FreeList) h->FreeList->PrevFree = blk;
    h->FreeList = blk;

    h->Brk = newBrk;
    KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
    KiHdrToFtr(blk)->Size  = blk->Size;

    return 1;
}

// Unlink a free block from the free list.
static void KiFreeListRemove(HEAP* h, BLOCK_HDR* blk) {
    if (blk->PrevFree) blk->PrevFree->NextFree = blk->NextFree;
    else               h->FreeList = blk->NextFree;
    if (blk->NextFree) blk->NextFree->PrevFree = blk->PrevFree;
    blk->PrevFree = blk->NextFree = (BLOCK_HDR *)0;
}

// Implements the heap Try Decommit Tail operation.
static void KiHeapTryDecommitTail(HEAP* h, BLOCK_HDR* blk) {
    if (!h || !blk || blk->Free == BLOCK_STATE_USED) return;

    // Do not eagerly decommit the kernel pool on every free. Large transient
    // syscall bounce buffers, especially SysBlitScreen's frame-sized buffer,
    // are often allocated/freed once per frame. If the tail block is unmapped
    // immediately, the next frame has to map the same pages again, producing
    // massive page-table churn and global TLB shootdowns. Keep the committed
    // pool high-water resident during normal operation; add an explicit idle or
    // pressure pool trim later if needed.
    if (!h->DecommitTailOnFree) {
        KiHdrToFtr(blk)->Size  = blk->Size;
        KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
        KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
        return;
    }

    uintptr_t blockBase = (uintptr_t)blk;
    uintptr_t oldBrk = h->Brk;

    // Only the last free block can return committed heap pages. Interior free
    // blocks must remain mapped because the allocator's boundary tags and free
    // list links are stored inside those blocks.
    if ((uintptr_t)KiNextBlock(blk) != oldBrk) {
        KiHdrToFtr(blk)->Size  = blk->Size;
        KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
        KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
        return;
    }

    // Keep enough mapped space for the tail block's header, footer, and a small
    // usable payload, then decommit whole pages above that. The footer moves to
    // the new Brk so future coalescing can still walk the boundary tag chain.
    uintptr_t minEnd = blockBase + MIN_BLOCK;
    uintptr_t newBrk = (minEnd + PAGE_SIZE - 1) & ~(uintptr_t)(PAGE_SIZE - 1);

    if (newBrk >= oldBrk) {
        KiHdrToFtr(blk)->Size  = blk->Size;
        KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
        KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
        return;
    }

    for (uintptr_t va = newBrk; va < oldBrk; va += PAGE_SIZE) {
        PHYSICAL_ADDRESS frame = MmPteTranslate(va);
        if (!frame)
            continue;
        if (KSUCCESS(MmPteUnmapPage(va)))
            (void)MmPfnFree(frame);
    }

    h->Brk = newBrk;
    blk->Size = newBrk - blockBase - OVERHEAD;
    KiHdrToFtr(blk)->Size  = blk->Size;
    KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
}

// Releases a previously quarantined block into the real free list. Caller holds h->Lock.
static void KiHeapReleaseFreeBlockLocked(HEAP* h, BLOCK_HDR* blk) {
    if (!h || !blk) return;

    KiPoolValidateFooterLocked(h, blk, "release free footer corrupt");
    KiPoolValidateTailCanaryLocked(h, blk, "release free tail canary corrupt");
    KiPoolValidateFreedPoisonLocked(h, blk, "use-after-free payload write detected");

    blk->Free = BLOCK_STATE_FREE;
    blk->RequestedSize = 0;
    blk->Tag = 0;
    blk->Flags = 0;
    blk->AllocCaller = 0;
    // Keep FreeCaller/Sequence for diagnostics until the block is reused.
    KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
    KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(blk)->Size  = blk->Size;

    // Coalesce with next block if it is really free. Quarantined neighbors are
    // intentionally not coalesced; their payload must remain intact for UAF checks.
    BLOCK_HDR* next = KiNextBlock(blk);
    if ((uintptr_t)next < h->Brk) {
        KiPoolValidateFooterLocked(h, next, "next footer corrupt while freeing");
        if (next->Free == BLOCK_STATE_FREE) {
            KiFreeListRemove(h, next);
            blk->Size += OVERHEAD + next->Size;
            KiHdrToFtr(blk)->Size  = blk->Size;
            KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
            KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
        }
    }

    // Coalesce with previous block if it is really free.
    if ((uintptr_t)blk > h->Start) {
        BLOCK_FTR* prevFtr = KiPrevFtr(blk);
        if (prevFtr->Magic == BLOCK_MAGIC) {
            BLOCK_HDR* prev = KiFtrToHdr(prevFtr);
            KiPoolValidateFooterLocked(h, prev, "previous footer corrupt while freeing");
            if (prev->Free == BLOCK_STATE_FREE) {
                KiFreeListRemove(h, prev);
                prev->Size += OVERHEAD + blk->Size;
                KiHdrToFtr(prev)->Size  = prev->Size;
                KiHdrToFtr(prev)->Magic = BLOCK_MAGIC;
                KiHdrToFtr(prev)->Free  = BLOCK_STATE_FREE;
                blk = prev;
            }
        }
    }

    // Insert at head of free list.
    blk->PrevFree  = (BLOCK_HDR *)0;
    blk->NextFree  = h->FreeList;
    if (h->FreeList) h->FreeList->PrevFree = blk;
    h->FreeList = blk;

    KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
    KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(blk)->Size  = blk->Size;

    KiHeapTryDecommitTail(h, blk);
}

static void KiHeapReleaseQuarantineSlotLocked(HEAP* h, size_t index) {
    if (!h || index >= POOL_QUARANTINE_SLOTS)
        return;
    BLOCK_HDR* old = h->Quarantine[index];
    if (!old)
        return;
    h->Quarantine[index] = (BLOCK_HDR*)0;
    if (h->QuarantineCount)
        h->QuarantineCount--;
    if (old->Free != BLOCK_STATE_QUARANTINE)
        KiPoolPanic("quarantine slot state corrupt", h, old, KiBlockUserBase(old));
    KiHeapReleaseFreeBlockLocked(h, old);
}

static void KiHeapDrainQuarantineLocked(HEAP* h) {
    if (!h)
        return;
    for (size_t i = 0; i < POOL_QUARANTINE_SLOTS; i++)
        KiHeapReleaseQuarantineSlotLocked(h, i);
    h->QuarantineCursor = 0;
}

static void KiHeapQuarantineBlockLocked(HEAP* h, BLOCK_HDR* blk) {
    if (!h || !blk)
        return;
    size_t idx = h->QuarantineCursor % POOL_QUARANTINE_SLOTS;
    KiHeapReleaseQuarantineSlotLocked(h, idx);

    blk->Free = BLOCK_STATE_QUARANTINE;
    KiHdrToFtr(blk)->Free  = BLOCK_STATE_QUARANTINE;
    KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(blk)->Size  = blk->Size;

    h->Quarantine[idx] = blk;
    h->QuarantineCursor = (idx + 1u) % POOL_QUARANTINE_SLOTS;
    h->QuarantineCount++;
}

// Implements the heap Alloc operation.
static void* KiHeapAllocEx(HEAP* h, size_t requestedSize, uint32_t tag, uintptr_t caller) {
    if (!requestedSize) return (void *)0;

    uint32_t dbg = KiPoolDebugFlags();
    size_t size = requestedSize;
    if (dbg & KE_DEBUG_POOL_CANARY) {
        if (size > (size_t)-1 - sizeof(uint64_t))
            return (void*)0;
        size += sizeof(uint64_t);
    }
    size = (size + 7) & ~7ULL;   // 8-byte align

    uintptr_t flags = KeAcquireSpinlock(&h->Lock);

    // Best-fit search.
    BLOCK_HDR* best = (BLOCK_HDR *)0;
    for (BLOCK_HDR* cur = h->FreeList; cur; cur = cur->NextFree) {
        KiPoolValidateFooterLocked(h, cur, "free list block corrupt");
        if (cur->Free != BLOCK_STATE_FREE)
            KiPoolPanic("free list contains non-free block", h, cur, KiBlockUserBase(cur));
        if (cur->Size >= size) {
            if (!best || cur->Size < best->Size) best = cur;
            if (best->Size == size) break;
        }
    }

    if (!best && h->QuarantineCount) {
        KiHeapDrainQuarantineLocked(h);
        for (BLOCK_HDR* cur = h->FreeList; cur; cur = cur->NextFree) {
            KiPoolValidateFooterLocked(h, cur, "free list block corrupt after quarantine drain");
            if (cur->Size >= size) {
                if (!best || cur->Size < best->Size) best = cur;
                if (best->Size == size) break;
            }
        }
    }

    if (!best) {
        if (!KiHeapGrow(h, size + OVERHEAD)) {
            KeReleaseSpinlock(&h->Lock, flags);
            return (void *)0;
        }
        // Re-search; KiHeapGrow inserted at the front.
        best = h->FreeList;
        while (best && best->Size < size) best = best->NextFree;
        if (!best) {
            KeReleaseSpinlock(&h->Lock, flags);
            return (void *)0;
        }
        KiPoolValidateFooterLocked(h, best, "grown free block corrupt");
    }

    KiFreeListRemove(h, best);

    // Split the block if the remainder is big enough.
    if (best->Size >= size + OVERHEAD + sizeof(void*)) {
        BLOCK_HDR* rem = (BLOCK_HDR *)((char *)best + HDR_SIZE + size + FTR_SIZE);
        rem->Magic     = BLOCK_MAGIC;
        rem->Free      = BLOCK_STATE_FREE;
        rem->Size      = best->Size - size - OVERHEAD;
        rem->PrevFree  = (BLOCK_HDR *)0;
        rem->NextFree  = h->FreeList;
        rem->RequestedSize = 0;
        rem->Tag = 0;
        rem->Flags = 0;
        rem->Sequence = 0;
        rem->AllocCaller = 0;
        rem->FreeCaller = 0;
        if (h->FreeList) h->FreeList->PrevFree = rem;
        h->FreeList    = rem;
        KiHdrToFtr(rem)->Magic = BLOCK_MAGIC;
        KiHdrToFtr(rem)->Free  = BLOCK_STATE_FREE;
        KiHdrToFtr(rem)->Size  = rem->Size;

        best->Size = size;
    }

    best->Free = BLOCK_STATE_USED;
    best->RequestedSize = requestedSize;
    best->Tag = tag ? tag : POOL_TAG_DEFAULT;
    best->Flags = 0;
    best->Sequence = ++h->NextSequence;
    best->AllocCaller = (dbg & KE_DEBUG_POOL_CALLERS) ? caller : 0;
    best->FreeCaller = 0;
    if (dbg & KE_DEBUG_POOL_CANARY)
        best->Flags |= BLOCK_FLAG_CANARY;

    KiHdrToFtr(best)->Free  = BLOCK_STATE_USED;
    KiHdrToFtr(best)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(best)->Size  = best->Size;

    memset(KiBlockUserBase(best), 0, requestedSize);
    if (best->Size > requestedSize)
        memset(KiBlockUserBase(best) + requestedSize, 0, best->Size - requestedSize);
    KiPoolWriteTailCanary(best);

    KeReleaseSpinlock(&h->Lock, flags);
    return KiBlockUserBase(best);
}

static void* KiHeapAlloc(HEAP* h, size_t size) {
    return KiHeapAllocEx(h, size, POOL_TAG_DEFAULT, (uintptr_t)_ReturnAddress());
}

// Implements the heap Free operation.
static void KiHeapFreeEx(HEAP* h, void* ptr, uintptr_t caller) {
    if (!ptr) return;

    uintptr_t u = (uintptr_t)ptr;
    if (!h || u < h->Start + HDR_SIZE || u >= h->Brk)
        KiPoolPanic("invalid free pointer outside heap", h, (BLOCK_HDR*)0, ptr);

    BLOCK_HDR* blk = (BLOCK_HDR *)(u - HDR_SIZE);

    uintptr_t flags = KeAcquireSpinlock(&h->Lock);

    if (blk->Magic != BLOCK_MAGIC)
        KiPoolPanic("invalid free pointer/bad header magic", h, blk, ptr);
    KiPoolValidateFooterLocked(h, blk, "allocated footer corrupt on free");
    if (blk->Free == BLOCK_STATE_FREE)
        KiPoolPanic("double free", h, blk, ptr);
    if (blk->Free == BLOCK_STATE_QUARANTINE)
        KiPoolPanic("double free/quarantined UAF", h, blk, ptr);
    if (blk->Free != BLOCK_STATE_USED)
        KiPoolPanic("invalid block state on free", h, blk, ptr);

    KiPoolValidateTailCanaryLocked(h, blk, "tail canary corrupt on free");

    uint32_t dbg = KiPoolDebugFlags();
    blk->FreeCaller = (dbg & KE_DEBUG_POOL_CALLERS) ? caller : 0;
    if (dbg & KE_DEBUG_POOL_POISON)
        KiPoolPoisonFreedPayload(blk);

    if (dbg & KE_DEBUG_POOL_QUARANTINE) {
        KiHeapQuarantineBlockLocked(h, blk);
    } else {
        KiHeapReleaseFreeBlockLocked(h, blk);
    }

    KeReleaseSpinlock(&h->Lock, flags);
}

static void KiHeapFree(HEAP* h, void* ptr) {
    KiHeapFreeEx(h, ptr, (uintptr_t)_ReturnAddress());
}

// Implements the heap Realloc operation.
static void* KiHeapReallocEx(HEAP* h, void* ptr, size_t newSize, uintptr_t caller) {
    // If ptr is NULL, realloc acts like malloc.
    if (!ptr) {
        return KiHeapAllocEx(h, newSize, POOL_TAG_DEFAULT, caller);
    }

    // If size is 0, realloc acts like free.
    if (newSize == 0) {
        KiHeapFreeEx(h, ptr, caller);
        return (void*)0;
    }

    uintptr_t u = (uintptr_t)ptr;
    if (!h || u < h->Start + HDR_SIZE || u >= h->Brk)
        KiPoolPanic("invalid realloc pointer outside heap", h, (BLOCK_HDR*)0, ptr);

    BLOCK_HDR* blk = (BLOCK_HDR *)(u - HDR_SIZE);

    uintptr_t flags = KeAcquireSpinlock(&h->Lock);
    if (blk->Magic != BLOCK_MAGIC)
        KiPoolPanic("invalid realloc pointer/bad header magic", h, blk, ptr);
    KiPoolValidateFooterLocked(h, blk, "allocated footer corrupt on realloc");
    if (blk->Free != BLOCK_STATE_USED)
        KiPoolPanic("realloc of non-allocated block", h, blk, ptr);
    KiPoolValidateTailCanaryLocked(h, blk, "tail canary corrupt on realloc");

    size_t oldRequested = blk->RequestedSize;
    size_t oldCopy = oldRequested < newSize ? oldRequested : newSize;
    if (newSize <= oldRequested) {
        blk->RequestedSize = newSize;
        if (blk->Flags & BLOCK_FLAG_CANARY)
            KiPoolWriteTailCanary(blk);
        KeReleaseSpinlock(&h->Lock, flags);
        return ptr;
    }
    KeReleaseSpinlock(&h->Lock, flags);

    void* newPtr = KiHeapAllocEx(h, newSize, POOL_TAG_DEFAULT, caller);
    if (!newPtr)
        return (void*)0;
    memcpy(newPtr, ptr, oldCopy);
    KiHeapFreeEx(h, ptr, caller);
    return newPtr;
}

static void* KiHeapRealloc(HEAP* h, void* ptr, size_t newSize) {
    return KiHeapReallocEx(h, ptr, newSize, (uintptr_t)_ReturnAddress());
}

// Kernel pool.
static HEAP gPoolHeap;

// VAD arena.
typedef enum MM_VAD_PAGE_STATE {
    MmVadPageFree = 0,
    MmVadPageCommitted = 1,
} MM_VAD_PAGE_STATE;

#define MM_VAD_PAGE_COW          0x00000001u
#define MM_VAD_PAGE_ZERO_PAGE    0x00000002u
#define MM_VAD_PAGE_SWAPPED      0x00000004u
#define MM_VAD_PAGE_SHARED_FRAME 0x00000008u
#define MM_VAD_PAGE_PAGING      0x00000010u
#define MM_VAD_PAGE_FILE_BACKED 0x00000020u
#define MM_VAD_PAGE_IMAGE_BACKED 0x00000040u
#define MM_VAD_PAGE_IMAGE_PRIVATE 0x00000080u
#define MM_VAD_PAGE_IMAGE_FIXUP   0x00000100u
#define MM_VAD_PAGE_FILE_SHARED_WRITE 0x00000200u
#define MM_VAD_PAGE_FILE_DIRTY_TRACK  0x00000400u

typedef struct MM_VAD_PAGE {
    uint64_t PageFlags;
    PHYSICAL_ADDRESS Frame;
    uint8_t State;
    uint8_t Resident;
    uint16_t Reserved;
    uint32_t Flags;
    uint32_t SwapSlot;
} MM_VAD_PAGE;

typedef struct MM_VAD {
    LIST_ENTRY ProcessLink;
    LIST_ENTRY GlobalLink;
    PROCESS* Owner;

    uintptr_t ReservationBase;
    size_t    ReservationSize;
    uintptr_t Base;
    size_t    Size;
    size_t    PayloadSize;
    size_t    PageCount;
    size_t    CommittedPages;

    uint64_t PageFlags;
    MM_VAD_TYPE Type;
    uint32_t Flags;
    uint32_t LockCount;
    MM_VAD_PAGE* Pages;

    FILE_OBJECT* FileObject;
    uint64_t FileOffset;
    size_t FileLength;
    uint32_t FileMappingFlags;

    IMAGE_SECTION_OBJECT* ImageSection;
} MM_VAD;

static LIST_ENTRY gVadList;
static SPINLOCK   gVadLock = SPINLOCK_INIT;
static int        gVadListReady = 0;
static MM_VAD_STATS gVadStats;

#define MM_VAD_ARENA_BASE 0x0000100000000000ULL


// Initializes the VAD list after the shared page table and kernel pool are ready.
static void KiVadInitialize(void) {
    LIST_INIT(&gVadList);
    gVadLock = (SPINLOCK)SPINLOCK_INIT;
    memset(&gVadStats, 0, sizeof(gVadStats));
    gVadListReady = 1;
}

// Initializes the shared page tables, kernel pool, and user VAD arena.
void MmInitializeVirtual(PHYSICAL_ADDRESS bootPml4Phys, size_t poolSize) {
    MmPteInitialize(bootPml4Phys);

    gPoolHeap.FreeList = (BLOCK_HDR *)0;
    gPoolHeap.Start    = POOL_BASE;
    gPoolHeap.Brk      = POOL_BASE;
    gPoolHeap.Limit    = POOL_LIMIT;
    gPoolHeap.MapFlags = PAGE_PRESENT | PAGE_WRITE | PAGE_NX;
    gPoolHeap.Lock     = (SPINLOCK)SPINLOCK_INIT;
    gPoolHeap.DecommitTailOnFree = 0;
    gPoolHeap.NextSequence = 0;
    gPoolHeap.QuarantineCursor = 0;
    gPoolHeap.QuarantineCount = 0;
    memset(gPoolHeap.Quarantine, 0, sizeof(gPoolHeap.Quarantine));

    if (poolSize) {
        poolSize = (poolSize + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
        KiHeapGrow(&gPoolHeap, poolSize);
    }

    KiVadInitialize();

    gVadZeroPageFrame = MmPfnAllocate(MmPfnUseAnonymous);
    if (!gVadZeroPageFrame) {
        KePanic(KE_PANIC_CODE_INIT, "MM: failed to allocate global zero page.\n");
    }
    memset(MmPfnToVirtual(gVadZeroPageFrame), 0, PAGE_SIZE);
    (void)MmPfnPin(gVadZeroPageFrame);
}


// Allocates zeroed kernel-pool memory from the non-user heap.
__declspec(dllexport) void* MmAllocatePool(size_t size) {
    return KiHeapAllocEx(&gPoolHeap, size, POOL_TAG_DEFAULT, (uintptr_t)_ReturnAddress());
}

// Allocates zeroed kernel-pool memory with a diagnostic tag.
__declspec(dllexport) void* MmAllocatePoolTag(size_t size, uint32_t tag) {
    return KiHeapAllocEx(&gPoolHeap, size, tag, (uintptr_t)_ReturnAddress());
}

// Releases kernel-pool memory back to the non-user heap.
__declspec(dllexport) void MmFreePool(void* ptr) {
    KiHeapFreeEx(&gPoolHeap, ptr, (uintptr_t)_ReturnAddress());
}

// Reallocates kernel-pool memory while preserving existing contents.
__declspec(dllexport) void* MmReallocatePool(void* ptr, size_t newSize) {
    return KiHeapReallocEx(&gPoolHeap, ptr, newSize, (uintptr_t)_ReturnAddress());
}

// VAD utility helpers.
static inline uintptr_t KiAlignUpPtr(uintptr_t v, size_t align) {
    if (align < PAGE_SIZE) align = PAGE_SIZE;
    if (align & (align - 1)) return 0;
    return (v + (uintptr_t)align - 1) & ~((uintptr_t)align - 1);
}

// Aligns a size to a power-of-two boundary and rejects invalid alignments.
static inline size_t KiAlignUpSize(size_t v, size_t align) {
    if (align < PAGE_SIZE) align = PAGE_SIZE;
    if (align & (align - 1)) return 0;
    return (v + align - 1) & ~(align - 1);
}

// Returns whether two half-open address ranges intersect.
static int KiRangesOverlap(uintptr_t aBase, size_t aSize, uintptr_t bBase, size_t bSize) {
    uintptr_t aEnd = aBase + aSize;
    uintptr_t bEnd = bBase + bSize;
    return aBase < bEnd && bBase < aEnd;
}

// Finds a free range in the global user VA arena.
static KSTATUS KiFindVadRange(size_t reservationSize, size_t alignment, uintptr_t* baseOut) {
    if (!baseOut || !reservationSize) return STATUS_INVALID_PARAMETER;
    if (!gVadListReady) return STATUS_FAILURE;

    uintptr_t candidate = KiAlignUpPtr(MM_VAD_ARENA_BASE, alignment);
    if (!candidate) return STATUS_INVALID_PARAMETER;

    for (;;) {
        if (candidate < MM_VAD_ARENA_BASE || reservationSize > (size_t)(USER_LIMIT - candidate))
            return STATUS_NO_MEMORY;

        uintptr_t next = candidate;
        int collided = 0;

        for (LIST_ENTRY* node = gVadList.Fwd; node != &gVadList; node = node->Fwd) {
            MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, GlobalLink);
            if (KiRangesOverlap(candidate, reservationSize, vad->ReservationBase, vad->ReservationSize)) {
                uintptr_t end = vad->ReservationBase + vad->ReservationSize;
                next = KiAlignUpPtr(end, alignment);
                if (!next || next <= candidate) return STATUS_NO_MEMORY;
                collided = 1;
                break;
            }
        }

        if (!collided) {
            *baseOut = candidate;
            return STATUS_SUCCESS;
        }
        candidate = next;
    }
}

// Returns the page index for an address inside the VAD payload.
static size_t KiVadPageIndex(const MM_VAD* vad, uintptr_t va) {
    return (size_t)((va - vad->Base) / PAGE_SIZE);
}

// Returns the page count covered by a page-aligned half-open range.
static size_t KiVadRangePageCount(uintptr_t alignedStart, uintptr_t alignedEnd) {
    return (size_t)((alignedEnd - alignedStart) / PAGE_SIZE);
}

// Finds a VAD by its public base address while the owning process VAD lock is held.
static MM_VAD* KiFindVadByBaseLocked(PROCESS* process, const void* base) {
    if (!process || !base) return (MM_VAD*)0;
    uintptr_t b = (uintptr_t)base;
    for (LIST_ENTRY* node = process->VadList.Fwd; node != &process->VadList; node = node->Fwd) {
        MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, ProcessLink);
        if (vad->Base == b) return vad;
    }
    return (MM_VAD*)0;
}

// Finds a VAD whose reserved payload fully contains the supplied range.
static MM_VAD* KiFindReservedVadContainingLocked(PROCESS* process, uintptr_t base, size_t size) {
    if (!process || !base || !size) return (MM_VAD*)0;
    if (base > (uintptr_t)-1 - (size - 1)) return (MM_VAD*)0;
    uintptr_t end = base + size;

    for (LIST_ENTRY* node = process->VadList.Fwd; node != &process->VadList; node = node->Fwd) {
        MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, ProcessLink);
        uintptr_t vadEnd = vad->Base + vad->PayloadSize;
        if (base >= vad->Base && end <= vadEnd) return vad;
    }
    return (MM_VAD*)0;
}

// Returns whether every page in the aligned range is committed in the VAD metadata.
static int KiVadRangeCommitted(const MM_VAD* vad, uintptr_t alignedStart, uintptr_t alignedEnd) {
    if (!vad || alignedStart < vad->Base || alignedEnd > vad->Base + vad->PayloadSize)
        return 0;
    size_t first = KiVadPageIndex(vad, alignedStart);
    size_t count = KiVadRangePageCount(alignedStart, alignedEnd);
    if (first > vad->PageCount || count > vad->PageCount - first)
        return 0;

    for (size_t i = 0; i < count; ++i) {
        if (vad->Pages[first + i].State != MmVadPageCommitted)
            return 0;
    }
    return 1;
}

static uint64_t KiVadHardwarePageFlags(const MM_VAD_PAGE* page);
static KSTATUS KiVadMapPageTracked(uintptr_t va, PHYSICAL_ADDRESS frame, uint64_t flags, MM_PTE_MAP_MODE mode);
static KSTATUS KiVadUnmapPageTracked(uintptr_t va, PHYSICAL_ADDRESS frame);

static KSTATUS KiVadMarkSharedFilePageDirty(MM_VAD* vad, size_t index, int writeProtect) {
    if (!vad || index >= vad->PageCount || !vad->FileObject)
        return STATUS_INVALID_PARAMETER;

    MM_VAD_PAGE* page = &vad->Pages[index];
    if ((page->Flags & MM_VAD_PAGE_FILE_BACKED) == 0 ||
        (page->Flags & MM_VAD_PAGE_FILE_SHARED_WRITE) == 0)
        return STATUS_NOT_FOUND;

    uint64_t pageDelta = (uint64_t)index * (uint64_t)PAGE_SIZE;
    uint64_t fileOffset = vad->FileOffset + pageDelta;
    if (fileOffset < vad->FileOffset)
        return STATUS_INVALID_PARAMETER;

    uint32_t validBytes = PAGE_SIZE;
    if (pageDelta >= (uint64_t)vad->FileLength) {
        validBytes = PAGE_SIZE;
    } else if ((uint64_t)vad->FileLength - pageDelta < PAGE_SIZE) {
        validBytes = (uint32_t)((uint64_t)vad->FileLength - pageDelta);
        if (validBytes == 0)
            validBytes = PAGE_SIZE;
    }

    MmFileCacheMarkDirty(vad->FileObject, fileOffset, validBytes);

    uint64_t end = fileOffset + validBytes;
    uintptr_t flags = KeAcquireSpinlock(&vad->FileObject->Lock);
    if (end > vad->FileObject->Size)
        vad->FileObject->Size = end;
    KeReleaseSpinlock(&vad->FileObject->Lock, flags);

    if (writeProtect && page->Resident && page->Frame && page->PageFlags) {
        uintptr_t va = vad->Base + index * PAGE_SIZE;
        page->Flags |= MM_VAD_PAGE_FILE_DIRTY_TRACK;
        (void)MmPteProtectRange(va, PAGE_SIZE, KiVadHardwarePageFlags(page));
    }

    gVadStats.FileBackedDirtySyncs++;
    return STATUS_SUCCESS;
}

static void KiVadSyncSharedFilePageBeforeRelease(MM_VAD* vad, size_t index) {
    if (!vad || index >= vad->PageCount)
        return;
    MM_VAD_PAGE* page = &vad->Pages[index];
    if ((page->Flags & MM_VAD_PAGE_FILE_SHARED_WRITE) == 0)
        return;
    uintptr_t va = vad->Base + index * PAGE_SIZE;
    if (page->Resident && page->Frame && MmPteTestAndClearDirty(va))
        (void)KiVadMarkSharedFilePageDirty(vad, index, 0);
}

// Marks a page uncommitted and releases any physical frame it still owns.
static void KiVadReleasePage(MM_VAD* vad, size_t index) {
    if (!vad || index >= vad->PageCount) return;
    MM_VAD_PAGE* page = &vad->Pages[index];
    uintptr_t va = vad->Base + index * PAGE_SIZE;

    KiVadSyncSharedFilePageBeforeRelease(vad, index);

    if (page->Resident && page->PageFlags) {
        (void)KiVadUnmapPageTracked(va, page->Frame);
    }
    if (page->Resident && page->Frame && !(page->Flags & MM_VAD_PAGE_ZERO_PAGE)) {
        (void)MmPfnRelease(page->Frame);
    }
    if ((page->Flags & MM_VAD_PAGE_SWAPPED) && page->SwapSlot != 0xFFFFFFFFu) {
        MmSwapFreeSlot(page->SwapSlot);
    }

    if (page->State == MmVadPageCommitted && vad->CommittedPages)
        vad->CommittedPages--;

    page->PageFlags = 0;
    page->Frame = 0;
    page->State = MmVadPageFree;
    page->Resident = 0;
    page->Flags = 0;
    page->SwapSlot = 0xFFFFFFFFu;
}

// Unmaps all resident pages in a VAD and releases their PFNs.
static void KiUnmapVadPayload(MM_VAD* vad) {
    if (!vad || !vad->Pages) return;
    for (size_t i = 0; i < vad->PageCount; ++i) {
        if (vad->Pages[i].State == MmVadPageCommitted)
            KiVadReleasePage(vad, i);
    }
}

// Returns the hardware PTE flags used for a resident VAD page.
static uint64_t KiVadHardwarePageFlags(const MM_VAD_PAGE* page) {
    if (!page || !page->PageFlags)
        return 0;
    uint64_t flags = page->PageFlags;
    if (page->Flags & MM_VAD_PAGE_COW)
        flags &= ~PAGE_WRITE;
    if (page->Flags & MM_VAD_PAGE_FILE_DIRTY_TRACK)
        flags &= ~PAGE_WRITE;
    return flags;
}

// Maps a user VAD PTE and records the mapping in the PFN database.
static KSTATUS KiVadMapPageTracked(uintptr_t va, PHYSICAL_ADDRESS frame, uint64_t flags, MM_PTE_MAP_MODE mode) {
    KSTATUS s = MmPteMapPage(va, frame, flags, mode);
    if (!KSUCCESS(s))
        return s;

    s = MmPfnAddMapping(frame);
    if (!KSUCCESS(s)) {
        (void)MmPteUnmapPage(va);
        return s;
    }
    return STATUS_SUCCESS;
}

// Unmaps a user VAD PTE and removes the matching PFN mapping reference.
static KSTATUS KiVadUnmapPageTracked(uintptr_t va, PHYSICAL_ADDRESS frame) {
    KSTATUS s = MmPteUnmapPage(va);
    if (KSUCCESS(s))
        (void)MmPfnRemoveMapping(frame);
    return s;
}

// Returns whether the committed page allows the requested fault or probe access.
static KSTATUS KiVadCheckPageAccess(const MM_VAD_PAGE* page, int wantWrite, int wantExecute) {
    if (!page || page->State != MmVadPageCommitted)
        return STATUS_NOT_FOUND;
    if (!page->PageFlags)
        return STATUS_ACCESS_DENIED;
    if ((page->PageFlags & PAGE_USER) == 0)
        return STATUS_ACCESS_DENIED;
    if (wantWrite && (page->PageFlags & PAGE_WRITE) == 0)
        return STATUS_ACCESS_DENIED;
    if (wantExecute && (page->PageFlags & PAGE_NX))
        return STATUS_ACCESS_DENIED;
    return STATUS_SUCCESS;
}

// Waits for another CPU/thread to finish a page-in/page-out transition on a VAD page.
// The pager performs filesystem I/O outside the VAD lock; this transient bit keeps
// concurrent faults from treating half-transitioned metadata as either resident or
// fully swapped.
static void KiVadWaitForPagingComplete(MM_VAD_PAGE* page) {
    while (page && (_InterlockedCompareExchange((volatile long*)&page->Flags, 0, 0) &
                    MM_VAD_PAGE_PAGING))
        KeYield();
}

// Atomically claims the transient paging bit.  This is deliberately not a spinlock:
// a page-in/page-out may allocate frames or do filesystem I/O, so waiters must yield
// instead of holding IRQL-raised locks.  Without this CAS, two CPUs faulting the same
// swapped/demand-zero page can both allocate/read/map it; one map fails and rolls
// metadata back to SWAPPED while the other has already published a valid PTE.  That
// leaves stale or freed swap slots behind and corrupts user data under pressure.
static int KiVadTryBeginPaging(MM_VAD_PAGE* page) {
    if (!page) return 0;

    for (;;) {
        uint32_t oldFlags = (uint32_t)_InterlockedCompareExchange((volatile long*)&page->Flags, 0, 0);
        if (oldFlags & MM_VAD_PAGE_PAGING)
            return 0;
        uint32_t newFlags = oldFlags | MM_VAD_PAGE_PAGING;
        if ((uint32_t)_InterlockedCompareExchange((volatile long*)&page->Flags,
                                                  (long)newFlags,
                                                  (long)oldFlags) == oldFlags)
            return 1;
        _mm_pause();
    }
}

static void KiVadEndPaging(MM_VAD_PAGE* page) {
    if (!page) return;

    for (;;) {
        uint32_t oldFlags = (uint32_t)_InterlockedCompareExchange((volatile long*)&page->Flags, 0, 0);
        uint32_t newFlags = oldFlags & ~MM_VAD_PAGE_PAGING;
        if ((uint32_t)_InterlockedCompareExchange((volatile long*)&page->Flags,
                                                  (long)newFlags,
                                                  (long)oldFlags) == oldFlags)
            return;
        _mm_pause();
    }
}

// Reads one page of a file-backed VAD into a freshly allocated user frame.
// Bytes beyond EOF or beyond the requested mapping length are zero-filled,
// matching the useful mmap/CreateFileMapping contract without materializing the
// whole file up front.
static KSTATUS KiVadFillFileBackedPage(MM_VAD* vad, size_t pageIndex, PHYSICAL_ADDRESS frame) {
    if (!vad || !vad->FileObject || !frame || pageIndex >= vad->PageCount)
        return STATUS_INVALID_PARAMETER;

    void* pageVa = MmPfnToVirtual(frame);
    if (!pageVa)
        return STATUS_FAILURE;
    memset(pageVa, 0, PAGE_SIZE);

    uint64_t pageDelta = (uint64_t)pageIndex * (uint64_t)PAGE_SIZE;
    if (pageDelta >= (uint64_t)vad->FileLength)
        return STATUS_SUCCESS;

    size_t remaining = vad->FileLength - (size_t)pageDelta;
    size_t toRead = remaining < PAGE_SIZE ? remaining : PAGE_SIZE;

    uint64_t fileOffset = vad->FileOffset + pageDelta;
    if (fileOffset < vad->FileOffset)
        return STATUS_INVALID_PARAMETER;

    size_t got = 0;
    KSTATUS s = IoReadFileAt(vad->FileObject, fileOffset, pageVa, toRead, &got);
    if (!KSUCCESS(s)) {
        gVadStats.FileBackedReadFailures++;
        return s;
    }
    gVadStats.FileBackedReads++;
    if (got < toRead)
        memset((uint8_t*)pageVa + got, 0, toRead - got);
    return STATUS_SUCCESS;
}


// Materializes one image-backed prototype page or zero-fill page for an image VAD.
static KSTATUS KiVadFillImageBackedPage(MM_VAD* vad, size_t pageIndex, int wantWrite) {
    if (!vad || !vad->ImageSection || pageIndex >= vad->PageCount)
        return STATUS_INVALID_PARAMETER;

    MM_VAD_PAGE* page = &vad->Pages[pageIndex];
    uintptr_t va = vad->Base + pageIndex * PAGE_SIZE;

    PHYSICAL_ADDRESS frame = 0;
    KSTATUS s = LdrRetainImageSectionPage(vad->ImageSection, pageIndex, &frame);
    if (!KSUCCESS(s))
        return s;

    if (frame) {
        page->Frame = frame;
        page->Resident = 1;
        page->Flags &= ~(MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_SWAPPED);
        page->Flags |= MM_VAD_PAGE_IMAGE_BACKED | MM_VAD_PAGE_SHARED_FRAME | MM_VAD_PAGE_COW;
        page->SwapSlot = 0xFFFFFFFFu;

        s = KiVadMapPageTracked(va, frame, KiVadHardwarePageFlags(page), MmPteMapFailIfPresent);
        if (!KSUCCESS(s)) {
            page->Frame = 0;
            page->Resident = 0;
            page->Flags &= ~(MM_VAD_PAGE_SHARED_FRAME | MM_VAD_PAGE_COW);
            (void)MmPfnRelease(frame);
            return s;
        }

        gVadStats.ImageBackedFaults++;
        return STATUS_SUCCESS;
    }

    if (!wantWrite && gVadZeroPageFrame) {
        page->Frame = gVadZeroPageFrame;
        page->Resident = 1;
        page->Flags |= MM_VAD_PAGE_IMAGE_BACKED | MM_VAD_PAGE_ZERO_PAGE;
        if (page->PageFlags & PAGE_WRITE)
            page->Flags |= MM_VAD_PAGE_COW;
        page->SwapSlot = 0xFFFFFFFFu;

        s = KiVadMapPageTracked(va, page->Frame, KiVadHardwarePageFlags(page), MmPteMapFailIfPresent);
        if (!KSUCCESS(s)) {
            page->Frame = 0;
            page->Resident = 0;
            page->Flags &= ~(MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_COW);
            return s;
        }

        gVadStats.ImageBackedZeroFaults++;
        return STATUS_SUCCESS;
    }

    frame = MmPfnAllocate(MmPfnUseImage);
    if (!frame)
        return STATUS_NO_MEMORY;

    void* pageVa = MmPfnToVirtual(frame);
    if (!pageVa) {
        MmPfnFree(frame);
        return STATUS_FAILURE;
    }
    memset(pageVa, 0, PAGE_SIZE);

    page->Frame = frame;
    page->Resident = 1;
    page->Flags &= ~(MM_VAD_PAGE_COW | MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_SHARED_FRAME | MM_VAD_PAGE_SWAPPED);
    page->Flags |= MM_VAD_PAGE_IMAGE_BACKED;
    page->SwapSlot = 0xFFFFFFFFu;

    s = KiVadMapPageTracked(va, frame, KiVadHardwarePageFlags(page), MmPteMapFailIfPresent);
    if (!KSUCCESS(s)) {
        page->Frame = 0;
        page->Resident = 0;
        MmPfnFree(frame);
        return s;
    }

    gVadStats.ImageBackedPrivateFaults++;
    return STATUS_SUCCESS;
}

// Materializes one committed VAD page as either the shared zero page or a private frame.
static KSTATUS KiVadEnsurePageResident(MM_VAD* vad, size_t pageIndex, int wantWrite, int wantExecute) {
    if (!vad || pageIndex >= vad->PageCount)
        return STATUS_NOT_FOUND;

    MM_VAD_PAGE* page = &vad->Pages[pageIndex];
retry_after_paging:
    KiVadWaitForPagingComplete(page);

    KSTATUS s = KiVadCheckPageAccess(page, wantWrite, wantExecute);
    if (!KSUCCESS(s))
        return s;

    uintptr_t va = vad->Base + pageIndex * PAGE_SIZE;
    if (page->Resident) {
        gVadStats.AlreadyResidentPages++;
        return STATUS_SUCCESS;
    }

    if (!KiVadTryBeginPaging(page))
        goto retry_after_paging;

    // Re-check after claiming the transition bit.  Another CPU may have resolved
    // the page between our initial wait/access check and the CAS above.
    s = KiVadCheckPageAccess(page, wantWrite, wantExecute);
    if (!KSUCCESS(s)) {
        KiVadEndPaging(page);
        return s;
    }
    if (page->Resident) {
        KiVadEndPaging(page);
        gVadStats.AlreadyResidentPages++;
        return STATUS_SUCCESS;
    }

    if ((page->Flags & MM_VAD_PAGE_SWAPPED) && page->SwapSlot != 0xFFFFFFFFu) {
        uint32_t oldSlot = page->SwapSlot;

        PHYSICAL_ADDRESS frame = MmPfnAllocate(MmPfnUseAnonymous);
        if (!frame) {
            KiVadEndPaging(page);
            return STATUS_NO_MEMORY;
        }

        KSTATUS swapStatus = MmSwapPageIn(oldSlot, frame);
        if (!KSUCCESS(swapStatus)) {
            MmPfnFree(frame);
            KiVadEndPaging(page);
            return swapStatus;
        }

        page->Frame = frame;
        page->Resident = 1;
        page->Flags &= ~MM_VAD_PAGE_SWAPPED;
        page->SwapSlot = 0xFFFFFFFFu;

        KSTATUS mapStatus = KiVadMapPageTracked(va, frame, KiVadHardwarePageFlags(page), MmPteMapFailIfPresent);
        if (!KSUCCESS(mapStatus)) {
            page->Frame = 0;
            page->Resident = 0;
            page->Flags |= MM_VAD_PAGE_SWAPPED;
            page->SwapSlot = oldSlot;
            KiVadEndPaging(page);
            MmPfnFree(frame);
            return mapStatus;
        }

        KiVadEndPaging(page);
        MmSwapFreeSlot(oldSlot);
        gVadStats.SwapPageIns++;
        return STATUS_SUCCESS;
    }

    if (page->Flags & MM_VAD_PAGE_IMAGE_BACKED) {
        s = KiVadFillImageBackedPage(vad, pageIndex, wantWrite);
        KiVadEndPaging(page);
        return s;
    }

    if (page->Flags & MM_VAD_PAGE_FILE_BACKED) {
        PHYSICAL_ADDRESS frame = 0;
        MM_FILE_CACHE_PAGE cachePage;
        memset(&cachePage, 0, sizeof(cachePage));

        uint64_t fileOffset = vad->FileOffset + (uint64_t)pageIndex * (uint64_t)PAGE_SIZE;
        if (fileOffset < vad->FileOffset) {
            KiVadEndPaging(page);
            return STATUS_INVALID_PARAMETER;
        }

        if (page->Flags & MM_VAD_PAGE_FILE_SHARED_WRITE) {
            uintptr_t f = KeAcquireSpinlock(&vad->FileObject->Lock);
            uint64_t oldSize = vad->FileObject->Size;
            KeReleaseSpinlock(&vad->FileObject->Lock, f);
            int needExisting = fileOffset < oldSize;
            s = MmFileCacheGetWritablePage(vad->FileObject, fileOffset, oldSize, needExisting, &cachePage);
            if (!KSUCCESS(s)) {
                gVadStats.FileBackedReadFailures++;
                KiVadEndPaging(page);
                return s;
            }
            frame = cachePage.Frame;
            page->Flags |= MM_VAD_PAGE_FILE_DIRTY_TRACK;
            gVadStats.FileBackedSharedWritableFaults++;
        } else if ((page->PageFlags & PAGE_WRITE) == 0) {
            s = MmFileCacheGetPage(vad->FileObject, fileOffset, &cachePage);
            if (!KSUCCESS(s)) {
                gVadStats.FileBackedReadFailures++;
                KiVadEndPaging(page);
                return s;
            }
            frame = cachePage.Frame;
        } else {
            frame = MmPfnAllocate(MmPfnUseMappedFile);
            if (!frame) {
                KiVadEndPaging(page);
                return STATUS_NO_MEMORY;
            }

            s = KiVadFillFileBackedPage(vad, pageIndex, frame);
            if (!KSUCCESS(s)) {
                KiVadEndPaging(page);
                MmPfnFree(frame);
                return s;
            }
        }

        page->Frame = frame;
        page->Resident = 1;
        page->Flags &= ~(MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_SHARED_FRAME | MM_VAD_PAGE_SWAPPED);
        page->SwapSlot = 0xFFFFFFFFu;

        s = KiVadMapPageTracked(va, frame, KiVadHardwarePageFlags(page), MmPteMapFailIfPresent);
        if (!KSUCCESS(s)) {
            page->Frame = 0;
            page->Resident = 0;
            KiVadEndPaging(page);
            if (cachePage.Frame) MmFileCacheReleasePage(&cachePage);
            else MmPfnFree(frame);
            return s;
        }

        KiVadEndPaging(page);
        gVadStats.FileBackedFaults++;
        return STATUS_SUCCESS;
    }

    if (!wantWrite && gVadZeroPageFrame) {
        uint32_t savedFlags = page->Flags;
        page->Frame = gVadZeroPageFrame;
        page->Resident = 1;
        page->Flags |= MM_VAD_PAGE_ZERO_PAGE;
        if (page->PageFlags & PAGE_WRITE)
            page->Flags |= MM_VAD_PAGE_COW;

        s = KiVadMapPageTracked(va, page->Frame, KiVadHardwarePageFlags(page), MmPteMapFailIfPresent);
        if (!KSUCCESS(s)) {
            page->Frame = 0;
            page->Resident = 0;
            page->Flags = savedFlags;
            KiVadEndPaging(page);
            return s;
        }

        KiVadEndPaging(page);
        gVadStats.ZeroPageMaps++;
        return STATUS_SUCCESS;
    }

    PHYSICAL_ADDRESS frame = MmPfnAllocate(MmPfnUseAnonymous);
    if (!frame) {
        KiVadEndPaging(page);
        return STATUS_NO_MEMORY;
    }

    void* pageVa = MmPfnToVirtual(frame);
    if (!pageVa) {
        KiVadEndPaging(page);
        MmPfnFree(frame);
        return STATUS_FAILURE;
    }
    memset(pageVa, 0, PAGE_SIZE);

    page->Frame = frame;
    page->Resident = 1;
    page->Flags &= ~(MM_VAD_PAGE_COW | MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_SHARED_FRAME | MM_VAD_PAGE_SWAPPED);
    page->SwapSlot = 0xFFFFFFFFu;

    s = KiVadMapPageTracked(va, frame, KiVadHardwarePageFlags(page), MmPteMapFailIfPresent);
    if (!KSUCCESS(s)) {
        page->Frame = 0;
        page->Resident = 0;
        KiVadEndPaging(page);
        MmPfnFree(frame);
        return s;
    }

    KiVadEndPaging(page);
    gVadStats.DemandZeroPages++;
    return STATUS_SUCCESS;
}

// Splits a copy-on-write page so this VAD receives a private writable frame.
static KSTATUS KiVadBreakCopyOnWrite(MM_VAD* vad, size_t pageIndex) {
    if (!vad || pageIndex >= vad->PageCount)
        return STATUS_NOT_FOUND;

    MM_VAD_PAGE* page = &vad->Pages[pageIndex];
    if (page->State != MmVadPageCommitted || !(page->Flags & MM_VAD_PAGE_COW))
        return STATUS_ACCESS_DENIED;
    if (!page->Resident || !page->Frame)
        return STATUS_NOT_FOUND;
    if ((page->PageFlags & PAGE_WRITE) == 0)
        return STATUS_ACCESS_DENIED;

    PHYSICAL_ADDRESS newFrame = MmPfnAllocate(MmPfnUseCow);
    if (!newFrame)
        return STATUS_NO_MEMORY;

    void* newPage = MmPfnToVirtual(newFrame);
    if (!newPage) {
        MmPfnFree(newFrame);
        return STATUS_FAILURE;
    }

    if (page->Flags & MM_VAD_PAGE_ZERO_PAGE) {
        memset(newPage, 0, PAGE_SIZE);
        gVadStats.ZeroPageCowSplits++;
    } else {
        void* oldPage = MmPfnToVirtual(page->Frame);
        if (!oldPage) {
            MmPfnFree(newFrame);
            return STATUS_FAILURE;
        }
        memcpy(newPage, oldPage, PAGE_SIZE);
    }

    uintptr_t va = vad->Base + pageIndex * PAGE_SIZE;
    KSTATUS s = KiVadUnmapPageTracked(va, page->Frame);
    if (!KSUCCESS(s)) {
        MmPfnFree(newFrame);
        return s;
    }

    s = KiVadMapPageTracked(va, newFrame, page->PageFlags, MmPteMapFailIfPresent);
    if (!KSUCCESS(s)) {
        (void)KiVadMapPageTracked(va, page->Frame, KiVadHardwarePageFlags(page), MmPteMapFailIfPresent);
        MmPfnFree(newFrame);
        return s;
    }

    if (!(page->Flags & MM_VAD_PAGE_ZERO_PAGE))
        (void)MmPfnRelease(page->Frame);
    page->Frame = newFrame;
    page->Flags &= ~(MM_VAD_PAGE_COW | MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_SHARED_FRAME | MM_VAD_PAGE_IMAGE_BACKED);
    page->Resident = 1;
    gVadStats.CowPagesCopied++;
    return STATUS_SUCCESS;
}

// Ensures that every committed page in a range is resident and accessible.
static KSTATUS KiVadEnsureResidentRange(PROCESS* process, uintptr_t start, size_t size, uint64_t accessFlags) {
    if (!process || !size)
        return STATUS_INVALID_PARAMETER;
    if (start > (uintptr_t)-1 - (size - 1))
        return STATUS_NO_MEMORY;

    uintptr_t alignedStart = start & PAGE_MASK;
    uintptr_t alignedEnd = (start + size + PAGE_SIZE - 1) & PAGE_MASK;
    if (alignedEnd <= alignedStart)
        return STATUS_NO_MEMORY;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindReservedVadContainingLocked(process, alignedStart, alignedEnd - alignedStart);
    KeReleaseSpinlock(&process->VadLock, procFlags);
    if (!vad)
        return STATUS_NOT_FOUND;

    int wantWrite = (accessFlags & PAGE_WRITE) ? 1 : 0;
    int wantExecute = (accessFlags & PAGE_NX) ? 1 : 0;
    size_t first = KiVadPageIndex(vad, alignedStart);
    size_t count = KiVadRangePageCount(alignedStart, alignedEnd);
    if (first > vad->PageCount || count > vad->PageCount - first)
        return STATUS_ACCESS_DENIED;

    gVadStats.EnsureResidentCalls++;
    for (size_t i = 0; i < count; ++i) {
        KSTATUS s = KiVadEnsurePageResident(vad, first + i, wantWrite, wantExecute);
        if (!KSUCCESS(s))
            return s;
        if (wantWrite && (vad->Pages[first + i].Flags & MM_VAD_PAGE_COW)) {
            s = KiVadBreakCopyOnWrite(vad, first + i);
            if (!KSUCCESS(s))
                return s;
        }
        gVadStats.EnsureResidentPages++;
    }

    return STATUS_SUCCESS;
}

// Makes an image VAD range private before the loader patches relocations or imports.
KSTATUS MmVadMakeImageRangePrivate(PROCESS* process, void* base, size_t size) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!process || !base || !size)
        return STATUS_INVALID_PARAMETER;

    uintptr_t start = (uintptr_t)base;
    if (start > (uintptr_t)-1 - (size - 1))
        return STATUS_NO_MEMORY;

    uintptr_t alignedStart = start & PAGE_MASK;
    uintptr_t alignedEnd = (start + size + PAGE_SIZE - 1) & PAGE_MASK;
    if (alignedEnd <= alignedStart)
        return STATUS_NO_MEMORY;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindReservedVadContainingLocked(process, alignedStart, alignedEnd - alignedStart);
    KeReleaseSpinlock(&process->VadLock, procFlags);
    if (!vad)
        return STATUS_NOT_FOUND;
    if (vad->Type != MmVadImage)
        return STATUS_INVALID_PARAMETER;

    size_t first = KiVadPageIndex(vad, alignedStart);
    size_t count = KiVadRangePageCount(alignedStart, alignedEnd);
    if (first > vad->PageCount || count > vad->PageCount - first)
        return STATUS_ACCESS_DENIED;

    gVadStats.ImageFixupPrivateCalls++;
    for (size_t i = 0; i < count; ++i) {
        MM_VAD_PAGE* page = &vad->Pages[first + i];

        KSTATUS s = KiVadEnsurePageResident(vad, first + i, 1, 0);
        if (!KSUCCESS(s))
            return s;

        if (page->Flags & MM_VAD_PAGE_COW) {
            s = KiVadBreakCopyOnWrite(vad, first + i);
            if (!KSUCCESS(s))
                return s;
        }

        if ((page->Flags & MM_VAD_PAGE_IMAGE_PRIVATE) == 0) {
            page->Flags |= MM_VAD_PAGE_IMAGE_PRIVATE | MM_VAD_PAGE_IMAGE_FIXUP;
            page->Flags &= ~(MM_VAD_PAGE_IMAGE_BACKED | MM_VAD_PAGE_SHARED_FRAME | MM_VAD_PAGE_ZERO_PAGE);
            gVadStats.ImageFixupPrivatePages++;
        } else {
            page->Flags |= MM_VAD_PAGE_IMAGE_FIXUP;
        }
    }

    return STATUS_SUCCESS;
}

// Pins the user page that contains the supplied address after materializing it.
// The returned KernelAddress is a direct-map alias for the physical frame, so
// callers do not continue touching the original user VA after validation. This
// closes the unmap-after-pin race and is the primitive future zero-copy-ish I/O
// paths should build on.
static KSTATUS KiVadPinUserPage(PROCESS* process, uintptr_t address,
                                uint64_t accessFlags, MM_PINNED_USER_PAGE* out) {
    if (!process || !out)
        return STATUS_INVALID_PARAMETER;
    memset(out, 0, sizeof(*out));

    if (!KiUserRangeValid(address, 1))
        return STATUS_ACCESS_DENIED;

    uintptr_t pageBase = address & PAGE_MASK;
    KSTATUS s = KiVadEnsureResidentRange(process, pageBase, PAGE_SIZE, accessFlags | PAGE_USER);
    if (!KSUCCESS(s))
        return s;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindReservedVadContainingLocked(process, pageBase, PAGE_SIZE);
    PHYSICAL_ADDRESS frame = 0;
    uint64_t pageFlags = 0;
    if (vad) {
        size_t index = KiVadPageIndex(vad, pageBase);
        if (index < vad->PageCount) {
            MM_VAD_PAGE* page = &vad->Pages[index];
            if (page->State == MmVadPageCommitted && page->Resident && page->Frame && page->PageFlags) {
                frame = page->Frame;
                pageFlags = page->PageFlags;
            }
        }
    }
    KeReleaseSpinlock(&process->VadLock, procFlags);

    if (!frame)
        return STATUS_NOT_FOUND;

    s = MmPfnPin(frame);
    if (!KSUCCESS(s))
        return s;

    procFlags = KeAcquireSpinlock(&process->VadLock);
    vad = KiFindReservedVadContainingLocked(process, pageBase, PAGE_SIZE);
    int stillValid = 0;
    if (vad) {
        size_t index = KiVadPageIndex(vad, pageBase);
        if (index < vad->PageCount) {
            MM_VAD_PAGE* page = &vad->Pages[index];
            stillValid = page->State == MmVadPageCommitted &&
                         page->Resident &&
                         page->Frame == frame &&
                         page->PageFlags != 0 &&
                         ((accessFlags & PAGE_WRITE) == 0 || (page->PageFlags & PAGE_WRITE));
            if (stillValid)
                pageFlags = page->PageFlags;
        }
    }
    KeReleaseSpinlock(&process->VadLock, procFlags);

    if (!stillValid) {
        (void)MmPfnUnpin(frame);
        return STATUS_NOT_FOUND;
    }

    out->UserPage = pageBase;
    out->PageOffset = (size_t)(address & (PAGE_SIZE - 1));
    out->Frame = frame;
    out->KernelAddress = MmPfnToVirtual(frame);
    out->PageFlags = pageFlags;
    if (!out->KernelAddress) {
        MmUnpinUserPage(out);
        return STATUS_FAILURE;
    }
    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS MmPinUserPageForRead(PROCESS* process, const void* userAddress, MM_PINNED_USER_PAGE* out) {
    return KiVadPinUserPage(process, (uintptr_t)userAddress, PAGE_PRESENT | PAGE_USER, out);
}

__declspec(dllexport) KSTATUS MmPinUserPageForWrite(PROCESS* process, void* userAddress, MM_PINNED_USER_PAGE* out) {
    return KiVadPinUserPage(process, (uintptr_t)userAddress, PAGE_PRESENT | PAGE_USER | PAGE_WRITE, out);
}

__declspec(dllexport) void MmUnpinUserPage(MM_PINNED_USER_PAGE* page) {
    if (!page || !page->Frame)
        return;
    (void)MmPfnUnpin(page->Frame);
    memset(page, 0, sizeof(*page));
}

typedef struct MM_LOCKED_USER_BUFFER {
    PROCESS* Process;
    void* UserAddress;
    size_t Size;
    size_t PageCount;
    int WriteAccess;
    MM_PINNED_USER_PAGE Pages[1];
} MM_LOCKED_USER_BUFFER;

// Returns the per-page VAD lock array that follows the pinned-page array.
static MM_VAD** KiLockedUserBufferVads(MM_LOCKED_USER_BUFFER* buffer) {
    return buffer ? (MM_VAD**)(buffer->Pages + buffer->PageCount) : (MM_VAD**)0;
}

// Adds one outstanding user-buffer lock to the VAD that owns a page.
static KSTATUS KiVadAddUserBufferPageLock(PROCESS* process, uintptr_t pageBase, MM_VAD** vadOut) {
    if (vadOut) *vadOut = (MM_VAD*)0;
    if (!process || !vadOut)
        return STATUS_INVALID_PARAMETER;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindReservedVadContainingLocked(process, pageBase, PAGE_SIZE);
    if (!vad) {
        KeReleaseSpinlock(&process->VadLock, procFlags);
        return STATUS_NOT_FOUND;
    }
    if (vad->LockCount == UINT32_MAX) {
        KeReleaseSpinlock(&process->VadLock, procFlags);
        return STATUS_FAILURE;
    }
    vad->LockCount++;
    *vadOut = vad;
    KeReleaseSpinlock(&process->VadLock, procFlags);
    return STATUS_SUCCESS;
}

// Removes one outstanding user-buffer lock from a VAD.
static void KiVadRemoveUserBufferPageLock(PROCESS* process, MM_VAD* vad) {
    if (!process || !vad)
        return;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    if (vad->LockCount)
        vad->LockCount--;
    KeReleaseSpinlock(&process->VadLock, procFlags);
}

// Locks a user buffer for asynchronous kernel access.
__declspec(dllexport) KSTATUS MmLockUserBuffer(PROCESS* process,
                                               void* userAddress,
                                               size_t size,
                                               int writeAccess,
                                               MM_LOCKED_USER_BUFFER** out) {
    if (!out)
        return STATUS_INVALID_PARAMETER;
    *out = (MM_LOCKED_USER_BUFFER*)0;

    if (size == 0)
        return STATUS_SUCCESS;
    if (!process || !userAddress)
        return STATUS_INVALID_PARAMETER;

    uintptr_t start = (uintptr_t)userAddress;
    if (!KiUserRangeValid(start, size))
        return STATUS_ACCESS_DENIED;

    uintptr_t alignedStart = start & PAGE_MASK;
    uintptr_t alignedEnd = (start + size + PAGE_SIZE - 1) & PAGE_MASK;
    if (alignedEnd <= alignedStart)
        return STATUS_NO_MEMORY;

    size_t pageCount = KiVadRangePageCount(alignedStart, alignedEnd);
    if (!pageCount)
        return STATUS_INVALID_PARAMETER;
    if (pageCount > (((size_t)-1) - sizeof(MM_LOCKED_USER_BUFFER)) /
                    (sizeof(MM_PINNED_USER_PAGE) + sizeof(MM_VAD*)))
        return STATUS_NO_MEMORY;

    size_t allocSize = sizeof(MM_LOCKED_USER_BUFFER) +
                       (pageCount - 1) * sizeof(MM_PINNED_USER_PAGE) +
                       pageCount * sizeof(MM_VAD*);
    MM_LOCKED_USER_BUFFER* buffer = (MM_LOCKED_USER_BUFFER*)MmAllocatePoolTag(allocSize, MM_POOL_TAG('u', 'l', 'o', 'k'));
    if (!buffer)
        return STATUS_NO_MEMORY;

    memset(buffer, 0, allocSize);
    buffer->Process = process;
    buffer->UserAddress = userAddress;
    buffer->Size = size;
    buffer->PageCount = pageCount;
    buffer->WriteAccess = writeAccess ? 1 : 0;
    ObReferenceObject(&process->Header);

    MM_VAD** vads = KiLockedUserBufferVads(buffer);
    KSTATUS s = STATUS_SUCCESS;
    size_t lockedPages = 0;
    for (size_t i = 0; i < pageCount; ++i) {
        uintptr_t pageBase = alignedStart + i * PAGE_SIZE;
        s = KiVadAddUserBufferPageLock(process, pageBase, &vads[i]);
        if (!KSUCCESS(s))
            break;

        s = writeAccess
            ? MmPinUserPageForWrite(process, (void*)pageBase, &buffer->Pages[i])
            : MmPinUserPageForRead(process, (const void*)pageBase, &buffer->Pages[i]);
        if (!KSUCCESS(s)) {
            KiVadRemoveUserBufferPageLock(process, vads[i]);
            vads[i] = (MM_VAD*)0;
            break;
        }
        lockedPages++;
    }

    if (!KSUCCESS(s)) {
        for (size_t i = 0; i < lockedPages; ++i)
            MmUnpinUserPage(&buffer->Pages[i]);
        for (size_t i = 0; i < pageCount; ++i)
            KiVadRemoveUserBufferPageLock(process, vads[i]);
        ObDereferenceObject(&process->Header);
        MmFreePool(buffer);
        return s;
    }

    *out = buffer;
    return STATUS_SUCCESS;
}

// Unlocks a user buffer locked for asynchronous kernel access.
__declspec(dllexport) void MmUnlockUserBuffer(MM_LOCKED_USER_BUFFER* buffer) {
    if (!buffer)
        return;

    MM_VAD** vads = KiLockedUserBufferVads(buffer);
    for (size_t i = 0; i < buffer->PageCount; ++i)
        MmUnpinUserPage(&buffer->Pages[i]);
    for (size_t i = 0; i < buffer->PageCount; ++i)
        KiVadRemoveUserBufferPageLock(buffer->Process, vads[i]);

    if (buffer->Process) {
        ObDereferenceObject(&buffer->Process->Header);
        buffer->Process = (PROCESS*)0;
    }
    MmFreePool(buffer);
}

// Copies user memory one page at a time while the current page is pinned.
// The pin is short-lived and exists only to keep pressure reclaim from removing
// the PTE while the kernel is executing the memcpy.
static KSTATUS KiCopyUserPaged(void* dst, const void* src, size_t size, int toUser) {
    if (size == 0)
        return STATUS_SUCCESS;
    if (!dst || !src)
        return STATUS_INVALID_PARAMETER;

    uintptr_t user = toUser ? (uintptr_t)dst : (uintptr_t)src;
    if (!KiUserRangeValid(user, size))
        return STATUS_ACCESS_DENIED;

    PROCESS* process = KeGetCurrentProcess();
    if (!process)
        return STATUS_ACCESS_DENIED;

    uint8_t* d = (uint8_t*)dst;
    const uint8_t* s = (const uint8_t*)src;
    size_t done = 0;
    uint64_t access = PAGE_PRESENT | PAGE_USER | (toUser ? PAGE_WRITE : 0);

    while (done < size) {
        uintptr_t userAddress = user + done;
        size_t pageOffset = (size_t)(userAddress & (PAGE_SIZE - 1));
        size_t chunk = PAGE_SIZE - pageOffset;
        if (chunk > size - done)
            chunk = size - done;

        MM_PINNED_USER_PAGE pinned;
        KSTATUS st = toUser
            ? MmPinUserPageForWrite(process, (void*)userAddress, &pinned)
            : MmPinUserPageForRead(process, (const void*)userAddress, &pinned);
        if (!KSUCCESS(st))
            return st;

        uint8_t* kernelPage = (uint8_t*)pinned.KernelAddress;
        if (toUser)
            memcpy(kernelPage + pinned.PageOffset, s + done, chunk);
        else
            memcpy(d + done, kernelPage + pinned.PageOffset, chunk);

        MmUnpinUserPage(&pinned);
        done += chunk;
    }

    return STATUS_SUCCESS;
}

// Copies from user memory after pinning each source page for the duration of the copy.
static KSTATUS KiCopyFromUserPaged(void* kernelDst, const void* userSrc, size_t size) {
    return KiCopyUserPaged(kernelDst, userSrc, size, 0);
}

// Copies to user memory after pinning each destination page for the duration of the copy.
static KSTATUS KiCopyToUserPaged(void* userDst, const void* kernelSrc, size_t size) {
    return KiCopyUserPaged(userDst, kernelSrc, size, 1);
}


// Reserves a VAD in the global user VA arena without committing physical pages.
KSTATUS MmVadReserve(PROCESS* process, size_t size, size_t alignment,
                     uint64_t pageFlags, uint32_t vadFlags,
                     MM_VAD_TYPE type, void** outBase) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!process || !outBase || !size) return STATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    size_t payloadSize = (size + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (payloadSize < size) return STATUS_NO_MEMORY;

    size_t guardLow = (vadFlags & MM_VAD_GUARD_LOW) ? PAGE_SIZE : 0;
    size_t guardHigh = (vadFlags & MM_VAD_GUARD_HIGH) ? PAGE_SIZE : 0;
    size_t reservationSize = guardLow + payloadSize + guardHigh;
    if (reservationSize < payloadSize) return STATUS_NO_MEMORY;

    if (alignment < PAGE_SIZE) alignment = PAGE_SIZE;
    if (alignment & (alignment - 1)) return STATUS_INVALID_PARAMETER;

    size_t pageCount = payloadSize / PAGE_SIZE;
    if (!pageCount) return STATUS_INVALID_PARAMETER;

    if (pageFlags && !(pageFlags & PAGE_USER)) pageFlags |= PAGE_USER;

    MM_VAD* vad = (MM_VAD*)MmAllocatePoolTag(sizeof(MM_VAD), MM_POOL_TAG('v', 'a', 'd', ' '));
    if (!vad) return STATUS_NO_MEMORY;
    memset(vad, 0, sizeof(*vad));

    MM_VAD_PAGE* pages = (MM_VAD_PAGE*)MmAllocatePoolTag(pageCount * sizeof(MM_VAD_PAGE), MM_POOL_TAG('v', 'p', 'g', 's'));
    if (!pages) {
        MmFreePool(vad);
        return STATUS_NO_MEMORY;
    }
    memset(pages, 0, pageCount * sizeof(MM_VAD_PAGE));

    LIST_INIT(&vad->ProcessLink);
    LIST_INIT(&vad->GlobalLink);

    uintptr_t reservationBase = 0;
    uintptr_t globalFlags = KeAcquireSpinlock(&gVadLock);
    KSTATUS s = KiFindVadRange(reservationSize, alignment, &reservationBase);
    if (KSUCCESS(s)) {
        vad->Owner = process;
        vad->ReservationBase = reservationBase;
        vad->ReservationSize = reservationSize;
        vad->Base = reservationBase + guardLow;
        vad->Size = size;
        vad->PayloadSize = payloadSize;
        vad->PageCount = pageCount;
        vad->CommittedPages = 0;
        vad->PageFlags = pageFlags;
        vad->Type = type;
        vad->Flags = vadFlags;
        vad->Pages = pages;
        ListInsertTail(&gVadList, &vad->GlobalLink);
    }
    KeReleaseSpinlock(&gVadLock, globalFlags);

    if (!KSUCCESS(s)) {
        MmFreePool(pages);
        MmFreePool(vad);
        return s;
    }

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    ListInsertTail(&process->VadList, &vad->ProcessLink);
    KeReleaseSpinlock(&process->VadLock, procFlags);

    *outBase = (void*)vad->Base;
    return STATUS_SUCCESS;
}

// Commits pages in an existing VAD without allocating PFNs until access.
KSTATUS MmVadCommit(PROCESS* process, void* base, size_t size,
                    uint64_t pageFlags, uint32_t commitFlags) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!process || !base || !size) return STATUS_INVALID_PARAMETER;

    uintptr_t start = (uintptr_t)base;
    if (start > (uintptr_t)-1 - (size - 1)) return STATUS_NO_MEMORY;

    uintptr_t alignedStart = start & PAGE_MASK;
    uintptr_t alignedEnd = (start + size + PAGE_SIZE - 1) & PAGE_MASK;
    if (alignedEnd <= alignedStart) return STATUS_NO_MEMORY;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindReservedVadContainingLocked(process, alignedStart, alignedEnd - alignedStart);
    KeReleaseSpinlock(&process->VadLock, procFlags);

    if (!vad) {
        DbgPrint("MM: commit failed: pid=%u range=%016llx-%016llx not inside reserved VAD\n",
                 process->ProcessId,
                 (unsigned long long)alignedStart,
                 (unsigned long long)alignedEnd);
        return STATUS_NOT_FOUND;
    }

    uint64_t flags = pageFlags ? pageFlags : vad->PageFlags;
    if (flags) flags |= PAGE_PRESENT | PAGE_USER;

    size_t first = KiVadPageIndex(vad, alignedStart);
    size_t count = KiVadRangePageCount(alignedStart, alignedEnd);
    if (first > vad->PageCount || count > vad->PageCount - first)
        return STATUS_ACCESS_DENIED;

    uint8_t* newlyCommitted = (uint8_t*)MmAllocatePool(count ? count : 1);
    if (!newlyCommitted) return STATUS_NO_MEMORY;
    memset(newlyCommitted, 0, count ? count : 1);

    KSTATUS s = STATUS_SUCCESS;
    for (size_t i = 0; i < count; ++i) {
        size_t pageIndex = first + i;
        MM_VAD_PAGE* page = &vad->Pages[pageIndex];

        if (page->State == MmVadPageCommitted) {
            continue;
        }

        page->State = MmVadPageCommitted;
        page->PageFlags = flags;
        page->Frame = 0;
        page->Resident = 0;
        page->Flags = 0;
        page->SwapSlot = 0xFFFFFFFFu;
        vad->CommittedPages++;
        newlyCommitted[i] = 1;
    }

    if ((commitFlags & MM_VAD_PREFETCH) && flags) {
        s = KiVadEnsureResidentRange(process, alignedStart, alignedEnd - alignedStart, flags);
    }

    if (!KSUCCESS(s)) {
        for (size_t i = 0; i < count; ++i) {
            if (newlyCommitted[i])
                KiVadReleasePage(vad, first + i);
        }
    }

    MmFreePool(newlyCommitted);
    return s;
}

// Reserves and commits a VAD in one step.
KSTATUS MmVadCreateCommitted(PROCESS* process, size_t size, size_t alignment,
                             uint64_t pageFlags, uint32_t vadFlags,
                             MM_VAD_TYPE type, void** outBase) {
    if (!outBase) return STATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    void* base = (void*)0;
    KSTATUS s = MmVadReserve(process, size, alignment, pageFlags, vadFlags, type, &base);
    if (!KSUCCESS(s)) return s;

    s = MmVadCommit(process, base, size, pageFlags, vadFlags);
    if (!KSUCCESS(s)) {
        MmVadDestroy(process, base);
        return s;
    }

    *outBase = base;
    return STATUS_SUCCESS;
}

// Creates a demand-paged file-backed user mapping.  The VAD owns a reference to
// the FILE_OBJECT so the caller may close the original file handle immediately
// after a view is established.  Pages are read from the file on first fault and
// then behave like ordinary private user pages: if memory pressure trims them
// after they are resident, the swap path preserves later private modifications.
KSTATUS MmVadCreateFileMapping(PROCESS* process, FILE_OBJECT* file,
                               uint64_t fileOffset, size_t size,
                               uint64_t pageFlags, uint32_t vadFlags,
                               void** outBase) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!process || !file || !size || !outBase)
        return STATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    if (fileOffset & (PAGE_SIZE - 1))
        return STATUS_INVALID_PARAMETER;

    size_t alignedSize = (size + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (!alignedSize || alignedSize < size)
        return STATUS_NO_MEMORY;

    if (pageFlags && !(pageFlags & PAGE_USER))
        pageFlags |= PAGE_USER;
    pageFlags |= PAGE_PRESENT | PAGE_USER;

    void* base = (void*)0;
    KSTATUS s = MmVadReserve(process, alignedSize, PAGE_SIZE, pageFlags,
                             vadFlags, MmVadMappedFile, &base);
    if (!KSUCCESS(s))
        return s;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindVadByBaseLocked(process, base);
    KeReleaseSpinlock(&process->VadLock, procFlags);
    if (!vad) {
        (void)MmVadDestroy(process, base);
        return STATUS_NOT_FOUND;
    }

    ObReferenceObject(&file->Header);
    vad->FileObject = file;
    vad->FileOffset = fileOffset;
    vad->FileLength = size;
    vad->FileMappingFlags = vadFlags;

    for (size_t i = 0; i < vad->PageCount; ++i) {
        MM_VAD_PAGE* page = &vad->Pages[i];
        page->State = MmVadPageCommitted;
        page->PageFlags = pageFlags;
        page->Frame = 0;
        page->Resident = 0;
        page->Flags = MM_VAD_PAGE_FILE_BACKED;
        if (vadFlags & MM_VAD_SHARED_WRITE)
            page->Flags |= MM_VAD_PAGE_FILE_SHARED_WRITE | MM_VAD_PAGE_FILE_DIRTY_TRACK;
        page->SwapSlot = 0xFFFFFFFFu;
        vad->CommittedPages++;
    }

    *outBase = base;
    return STATUS_SUCCESS;
}


// Creates an image-backed VAD whose pages fault from a shared image section object.
KSTATUS MmVadCreateImageSectionMapping(PROCESS* process, IMAGE_SECTION_OBJECT* section,
                                       uint64_t pageFlags, uint32_t vadFlags,
                                       void** outBase) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!process || !section || !outBase)
        return STATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    const IMAGE_SECTION_METADATA* metadata = LdrImageSectionMetadata(section);
    if (!metadata || !metadata->ImageSize || !metadata->PageCount)
        return STATUS_INVALID_PARAMETER;

    pageFlags |= PAGE_PRESENT | PAGE_USER;

    void* base = (void*)0;
    KSTATUS s = MmVadReserve(process, metadata->ImageSize, PAGE_SIZE, pageFlags,
                             vadFlags, MmVadImage, &base);
    if (!KSUCCESS(s))
        return s;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindVadByBaseLocked(process, base);
    KeReleaseSpinlock(&process->VadLock, procFlags);
    if (!vad) {
        (void)MmVadDestroy(process, base);
        return STATUS_NOT_FOUND;
    }

    LdrReferenceImageSection(section);
    vad->ImageSection = section;

    for (size_t i = 0; i < vad->PageCount; ++i) {
        MM_VAD_PAGE* page = &vad->Pages[i];
        page->State = MmVadPageCommitted;
        page->PageFlags = pageFlags;
        page->Frame = 0;
        page->Resident = 0;
        page->Flags = MM_VAD_PAGE_IMAGE_BACKED;
        page->SwapSlot = 0xFFFFFFFFu;
        vad->CommittedPages++;
    }

    *outBase = base;
    return STATUS_SUCCESS;
}

// Materializes committed VAD pages before the kernel or a recovered user fault touches them.
KSTATUS MmVadEnsureResident(PROCESS* process, void* base, size_t size, uint64_t accessFlags) {
    if (!process || !base || !size)
        return STATUS_INVALID_PARAMETER;
    return KiVadEnsureResidentRange(process, (uintptr_t)base, size, accessFlags | PAGE_USER);
}

// Creates a private copy-on-write view of a resident committed VAD range.
KSTATUS MmVadCreateCowView(PROCESS* sourceProcess, void* sourceBase, size_t size,
                           PROCESS* targetProcess, void** outBase) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!sourceProcess || !sourceBase || !targetProcess || !outBase || !size)
        return STATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    uintptr_t start = (uintptr_t)sourceBase;
    if (start & (PAGE_SIZE - 1))
        return STATUS_INVALID_PARAMETER;
    size_t alignedSize = (size + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (!alignedSize || alignedSize < size)
        return STATUS_NO_MEMORY;

    KSTATUS s = KiVadEnsureResidentRange(sourceProcess, start, alignedSize, PAGE_USER);
    if (!KSUCCESS(s))
        return s;

    uintptr_t srcFlags = KeAcquireSpinlock(&sourceProcess->VadLock);
    MM_VAD* sourceVad = KiFindReservedVadContainingLocked(sourceProcess, start, alignedSize);
    KeReleaseSpinlock(&sourceProcess->VadLock, srcFlags);
    if (!sourceVad || !KiVadRangeCommitted(sourceVad, start, start + alignedSize))
        return STATUS_NOT_FOUND;

    void* targetBase = (void*)0;
    s = MmVadReserve(targetProcess, alignedSize, PAGE_SIZE, sourceVad->PageFlags, sourceVad->Flags, sourceVad->Type, &targetBase);
    if (!KSUCCESS(s))
        return s;

    uintptr_t dstFlags = KeAcquireSpinlock(&targetProcess->VadLock);
    MM_VAD* targetVad = KiFindVadByBaseLocked(targetProcess, targetBase);
    KeReleaseSpinlock(&targetProcess->VadLock, dstFlags);
    if (!targetVad) {
        (void)MmVadDestroy(targetProcess, targetBase);
        return STATUS_NOT_FOUND;
    }

    uintptr_t dst = (uintptr_t)targetBase;
    size_t first = KiVadPageIndex(sourceVad, start);
    size_t count = alignedSize / PAGE_SIZE;

    for (size_t i = 0; i < count; ++i) {
        MM_VAD_PAGE* srcPage = &sourceVad->Pages[first + i];
        if (srcPage->State != MmVadPageCommitted || !srcPage->Resident || !srcPage->Frame) {
            s = STATUS_NOT_FOUND;
            break;
        }

        s = MmPfnRetain(srcPage->Frame);
        if (!KSUCCESS(s))
            break;

        MM_VAD_PAGE* dstPage = &targetVad->Pages[i];
        dstPage->State = MmVadPageCommitted;
        dstPage->PageFlags = srcPage->PageFlags;
        dstPage->Frame = srcPage->Frame;
        dstPage->Resident = 1;
        dstPage->Flags = MM_VAD_PAGE_SHARED_FRAME | ((srcPage->PageFlags & PAGE_WRITE) ? MM_VAD_PAGE_COW : 0);
        dstPage->SwapSlot = 0xFFFFFFFFu;
        targetVad->CommittedPages++;

        if (dstPage->Flags) {
            srcPage->Flags |= MM_VAD_PAGE_COW | MM_VAD_PAGE_SHARED_FRAME;
            if (srcPage->Resident && srcPage->PageFlags)
                (void)MmPteProtectRange(sourceVad->Base + (first + i) * PAGE_SIZE, PAGE_SIZE, KiVadHardwarePageFlags(srcPage));
        }

        s = KiVadMapPageTracked(dst + i * PAGE_SIZE, dstPage->Frame, KiVadHardwarePageFlags(dstPage), MmPteMapFailIfPresent);
        if (!KSUCCESS(s))
            break;

        gVadStats.CowPagesShared++;
    }

    if (!KSUCCESS(s)) {
        (void)MmVadDestroy(targetProcess, targetBase);
        return s;
    }

    gVadStats.CowViewsCreated++;
    *outBase = targetBase;
    return STATUS_SUCCESS;
}

// Creates a COW view from loader-owned template frames.
KSTATUS MmVadCreateCowViewFromFrames(PROCESS* targetProcess,
                                     const PHYSICAL_ADDRESS* frames,
                                     size_t pageCount,
                                     size_t payloadSize,
                                     uint64_t pageFlags,
                                     uint32_t vadFlags,
                                     MM_VAD_TYPE type,
                                     void** outBase) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!targetProcess || !frames || !pageCount || !payloadSize || !outBase)
        return STATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    size_t alignedSize = (payloadSize + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (!alignedSize || alignedSize / PAGE_SIZE > pageCount)
        return STATUS_INVALID_PARAMETER;

    void* targetBase = (void*)0;
    KSTATUS s = MmVadReserve(targetProcess, alignedSize, PAGE_SIZE, pageFlags, vadFlags, type, &targetBase);
    if (!KSUCCESS(s))
        return s;

    uintptr_t dstFlags = KeAcquireSpinlock(&targetProcess->VadLock);
    MM_VAD* targetVad = KiFindVadByBaseLocked(targetProcess, targetBase);
    KeReleaseSpinlock(&targetProcess->VadLock, dstFlags);
    if (!targetVad) {
        (void)MmVadDestroy(targetProcess, targetBase);
        return STATUS_NOT_FOUND;
    }

    uintptr_t dst = (uintptr_t)targetBase;
    size_t count = alignedSize / PAGE_SIZE;
    for (size_t i = 0; i < count; ++i) {
        MM_VAD_PAGE* page = &targetVad->Pages[i];
        page->State = MmVadPageCommitted;
        page->PageFlags = pageFlags | PAGE_PRESENT | PAGE_USER;
        page->Frame = frames[i];
        page->Resident = frames[i] ? 1 : 0;
        page->Flags = (frames[i] ? MM_VAD_PAGE_SHARED_FRAME : 0) | ((page->PageFlags & PAGE_WRITE) ? MM_VAD_PAGE_COW : 0);
        page->SwapSlot = 0xFFFFFFFFu;
        targetVad->CommittedPages++;

        if (!frames[i])
            continue;

        s = MmPfnRetain(frames[i]);
        if (!KSUCCESS(s))
            break;

        s = KiVadMapPageTracked(dst + i * PAGE_SIZE, page->Frame, KiVadHardwarePageFlags(page), MmPteMapFailIfPresent);
        if (!KSUCCESS(s)) {
            (void)MmPfnRelease(frames[i]);
            break;
        }

        gVadStats.CowPagesShared++;
    }

    if (!KSUCCESS(s)) {
        (void)MmVadDestroy(targetProcess, targetBase);
        return s;
    }

    gVadStats.CowViewsCreated++;
    *outBase = targetBase;
    return STATUS_SUCCESS;
}

KSTATUS MmVadCreateSharedViewFromFrames(PROCESS* targetProcess,
                                        const PHYSICAL_ADDRESS* frames,
                                        size_t pageCount,
                                        size_t payloadSize,
                                        uint64_t pageFlags,
                                        uint32_t vadFlags,
                                        MM_VAD_TYPE type,
                                        void** outBase) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!targetProcess || !frames || !pageCount || !payloadSize || !outBase)
        return STATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    size_t alignedSize = (payloadSize + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (!alignedSize || alignedSize / PAGE_SIZE > pageCount)
        return STATUS_INVALID_PARAMETER;

    pageFlags |= PAGE_PRESENT | PAGE_USER;

    void* targetBase = (void*)0;
    KSTATUS s = MmVadReserve(targetProcess, alignedSize, PAGE_SIZE, pageFlags,
                             vadFlags, type, &targetBase);
    if (!KSUCCESS(s))
        return s;

    uintptr_t dstFlags = KeAcquireSpinlock(&targetProcess->VadLock);
    MM_VAD* targetVad = KiFindVadByBaseLocked(targetProcess, targetBase);
    KeReleaseSpinlock(&targetProcess->VadLock, dstFlags);
    if (!targetVad) {
        (void)MmVadDestroy(targetProcess, targetBase);
        return STATUS_NOT_FOUND;
    }

    uintptr_t dst = (uintptr_t)targetBase;
    size_t count = alignedSize / PAGE_SIZE;
    for (size_t i = 0; i < count; ++i) {
        if (!frames[i]) {
            s = STATUS_INVALID_PARAMETER;
            break;
        }

        s = MmPfnRetain(frames[i]);
        if (!KSUCCESS(s))
            break;

        MM_VAD_PAGE* page = &targetVad->Pages[i];
        page->State = MmVadPageCommitted;
        page->PageFlags = pageFlags;
        page->Frame = frames[i];
        page->Resident = 1;
        page->Flags = MM_VAD_PAGE_SHARED_FRAME;
        page->SwapSlot = 0xFFFFFFFFu;
        targetVad->CommittedPages++;

        s = KiVadMapPageTracked(dst + i * PAGE_SIZE, page->Frame,
                         KiVadHardwarePageFlags(page), MmPteMapFailIfPresent);
        if (!KSUCCESS(s)) {
            page->PageFlags = 0;
            page->Frame = 0;
            page->Resident = 0;
            page->Flags = 0;
            page->State = MmVadPageFree;
            page->SwapSlot = 0xFFFFFFFFu;
            if (targetVad->CommittedPages)
                targetVad->CommittedPages--;
            (void)MmPfnRelease(frames[i]);
            break;
        }

        gVadStats.CowPagesShared++;
    }

    if (!KSUCCESS(s)) {
        (void)MmVadDestroy(targetProcess, targetBase);
        return s;
    }

    gVadStats.CowViewsCreated++;
    *outBase = targetBase;
    return STATUS_SUCCESS;
}

// Returns whether a resident file-backed page can be discarded instead of
// swapped.  This is safe only while the page is still clean by construction:
// read-only mappings cannot have been modified, and the backing file can
// repopulate the page on the next fault.  Writable/private file mappings still
// fall back to the swap path because this VMM does not yet track dirty bits.
static int KiVadPageCanDiscardCleanFile(const MM_VAD* vad, const MM_VAD_PAGE* page) {
    if (!vad || !page || vad->Type != MmVadMappedFile)
        return 0;
    if (!vad->FileObject || page->State != MmVadPageCommitted)
        return 0;
    if (!page->Resident || !page->Frame || !page->PageFlags)
        return 0;
    if ((page->Flags & MM_VAD_PAGE_FILE_BACKED) == 0)
        return 0;
    if (page->PageFlags & PAGE_WRITE)
        return 0;
    if (page->Flags & (MM_VAD_PAGE_COW | MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_SHARED_FRAME |
                       MM_VAD_PAGE_SWAPPED | MM_VAD_PAGE_PAGING))
        return 0;
    return 1;
}

static int KiVadPageCanDiscardCleanFileAfterPagingClaim(const MM_VAD* vad, const MM_VAD_PAGE* page) {
    if (!vad || !page || vad->Type != MmVadMappedFile)
        return 0;
    if (!vad->FileObject || page->State != MmVadPageCommitted)
        return 0;
    if (!page->Resident || !page->Frame || !page->PageFlags)
        return 0;
    if ((page->Flags & MM_VAD_PAGE_FILE_BACKED) == 0)
        return 0;
    if (page->PageFlags & PAGE_WRITE)
        return 0;
    if (page->Flags & (MM_VAD_PAGE_COW | MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_SHARED_FRAME |
                       MM_VAD_PAGE_SWAPPED))
        return 0;
    return (page->Flags & MM_VAD_PAGE_PAGING) != 0;
}

// Returns whether a resident VAD page is safe to move to swap.  Clean read-only
// file-backed pages are intentionally excluded; the pressure trimmer discards
// them and reloads them from the file instead of wasting swap slots.
static int KiVadPageCanSwap(const MM_VAD_PAGE* page) {
    if (!page || page->State != MmVadPageCommitted)
        return 0;
    if (!page->Resident || !page->Frame || !page->PageFlags)
        return 0;
    if ((page->Flags & MM_VAD_PAGE_FILE_BACKED) && ((page->PageFlags & PAGE_WRITE) == 0))
        return 0;
    if (page->Flags & MM_VAD_PAGE_FILE_SHARED_WRITE)
        return 0;
    if (page->Flags & (MM_VAD_PAGE_COW | MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_SHARED_FRAME |
                       MM_VAD_PAGE_SWAPPED | MM_VAD_PAGE_PAGING))
        return 0;
    return 1;
}

static int KiVadPageCanSwapAfterPagingClaim(const MM_VAD_PAGE* page) {
    if (!page || page->State != MmVadPageCommitted)
        return 0;
    if (!page->Resident || !page->Frame || !page->PageFlags)
        return 0;
    if ((page->Flags & MM_VAD_PAGE_FILE_BACKED) && ((page->PageFlags & PAGE_WRITE) == 0))
        return 0;
    if (page->Flags & MM_VAD_PAGE_FILE_SHARED_WRITE)
        return 0;
    if (page->Flags & (MM_VAD_PAGE_COW | MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_SHARED_FRAME |
                       MM_VAD_PAGE_SWAPPED))
        return 0;
    if ((page->Flags & MM_VAD_PAGE_PAGING) == 0)
        return 0;
    return 1;
}

// Discards one resident clean file-backed page.  No swap slot is allocated; the
// VAD keeps its FILE_BACKED state and refaults the page from the file later.
static KSTATUS KiVadDiscardCleanFilePage(MM_VAD* vad, size_t pageIndex, void** baseOut) {
    if (baseOut) *baseOut = (void*)0;
    if (!vad || pageIndex >= vad->PageCount)
        return STATUS_INVALID_PARAMETER;

    MM_VAD_PAGE* page = &vad->Pages[pageIndex];
    if (!KiVadPageCanDiscardCleanFile(vad, page))
        return STATUS_NOT_FOUND;

    PHYSICAL_ADDRESS frame = page->Frame;
    MM_PFN* pfn = MmPfnGet(frame);
    if (!pfn || pfn->PinCount)
        return STATUS_NOT_FOUND;

    uintptr_t va = vad->Base + pageIndex * PAGE_SIZE;

    // Second chance for hot mapped-file pages.  Dropping the active video/WAD
    // working set is much worse than swapping cold anonymous pages: it causes
    // repeated file reads, page-table churn, and global TLB shootdowns.  Clear
    // the hardware accessed bit and skip the page once; if it stays cold until
    // a later pressure scan, it may be discarded then.
    if (MmPteTestAndClearAccessed(va))
        return STATUS_NOT_FOUND;

    if (!KiVadTryBeginPaging(page))
        return STATUS_NOT_FOUND;

    if (!KiVadPageCanDiscardCleanFileAfterPagingClaim(vad, page) || page->Frame != frame) {
        KiVadEndPaging(page);
        return STATUS_NOT_FOUND;
    }

    KSTATUS s = MmPfnPin(frame);
    if (!KSUCCESS(s)) {
        KiVadEndPaging(page);
        return s;
    }

    s = KiVadUnmapPageTracked(va, frame);
    if (!KSUCCESS(s)) {
        KiVadEndPaging(page);
        (void)MmPfnUnpin(frame);
        return s;
    }

    page->Frame = 0;
    page->Resident = 0;
    page->Flags &= ~(MM_VAD_PAGE_SWAPPED | MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_SHARED_FRAME | MM_VAD_PAGE_COW);
    page->Flags |= MM_VAD_PAGE_FILE_BACKED;
    page->SwapSlot = 0xFFFFFFFFu;
    KiVadEndPaging(page);

    (void)MmPfnUnpin(frame);
    (void)MmPfnRelease(frame);
    if (vad->FileObject) {
        uint64_t fileOffset = vad->FileOffset + (uint64_t)pageIndex * (uint64_t)PAGE_SIZE;
        MmFileCacheInvalidate(vad->FileObject, fileOffset, PAGE_SIZE);
    }

    gVadStats.FileBackedDiscards++;
    if (baseOut) *baseOut = (void*)va;
    return STATUS_SUCCESS;
}

// Moves one private resident VAD page to the swap backend.
static KSTATUS KiVadTrimPageToSwap(MM_VAD* vad, size_t pageIndex, void** baseOut) {
    if (baseOut) *baseOut = (void*)0;
    if (!vad || pageIndex >= vad->PageCount)
        return STATUS_INVALID_PARAMETER;

    MM_VAD_PAGE* page = &vad->Pages[pageIndex];
    if (!KiVadPageCanSwap(page))
        return STATUS_NOT_FOUND;

    PHYSICAL_ADDRESS frame = page->Frame;
    MM_PFN* pfn = MmPfnGet(frame);
    if (!pfn || pfn->RefCount != 1 || pfn->PinCount)
        return STATUS_NOT_FOUND;

    // Claim the transient paging bit atomically before pinning/removing the
    // user PTE.  This serializes page-out against another page-out and against
    // a concurrent fault on the same VAD page without holding a spinlock across
    // pagefile I/O.
    if (!KiVadTryBeginPaging(page))
        return STATUS_NOT_FOUND;

    if (!KiVadPageCanSwapAfterPagingClaim(page) || page->Frame != frame) {
        KiVadEndPaging(page);
        return STATUS_NOT_FOUND;
    }

    KSTATUS s = MmPfnPin(frame);
    if (!KSUCCESS(s)) {
        KiVadEndPaging(page);
        return s;
    }

    uintptr_t va = vad->Base + pageIndex * PAGE_SIZE;
    s = KiVadUnmapPageTracked(va, frame);
    if (!KSUCCESS(s)) {
        KiVadEndPaging(page);
        (void)MmPfnUnpin(frame);
        return s;
    }

    page->Resident = 0;

    uint32_t slot = 0;
    s = MmSwapPageOut(frame, &slot);
    if (!KSUCCESS(s)) {
        page->Resident = 1;
        KSTATUS mapStatus = KiVadMapPageTracked(va, frame, KiVadHardwarePageFlags(page), MmPteMapFailIfPresent);
        KiVadEndPaging(page);
        (void)MmPfnUnpin(frame);
        return KSUCCESS(mapStatus) ? s : mapStatus;
    }

    page->Frame = 0;
    page->Flags |= MM_VAD_PAGE_SWAPPED;
    page->SwapSlot = slot;
    KiVadEndPaging(page);

    (void)MmPfnUnpin(frame);
    (void)MmPfnRelease(frame);

    gVadStats.SwapPageOuts++;
    if (baseOut) *baseOut = (void*)va;
    return STATUS_SUCCESS;
}

// Finds one private resident page in a process and swaps it out.
static KSTATUS KiVadTrimOnePageInProcess(PROCESS* process, void** baseOut) {
    if (baseOut) *baseOut = (void*)0;
    if (!process) return STATUS_INVALID_PARAMETER;

    for (LIST_ENTRY* node = process->VadList.Fwd; node != &process->VadList; node = node->Fwd) {
        MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, ProcessLink);
        for (size_t i = 0; i < vad->PageCount; ++i) {
            KSTATUS s = KiVadTrimPageToSwap(vad, i, baseOut);
            if (KSUCCESS(s))
                return STATUS_SUCCESS;
        }
    }
    return STATUS_NOT_FOUND;
}

// Moves one private resident user page to the swap backend for self-tests.
KSTATUS MmVadTrimOnePageForTest(PROCESS* process, void** baseOut) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!MmSwapIsEnabled()) return STATUS_NOT_SUPPORTED;
    return KiVadTrimOnePageInProcess(process, baseOut);
}

// Reclaims one resident user page from any process during allocation pressure.
// Prefer clean file-backed pages because they can be dropped and re-read without
// consuming swap.  Fall back to anonymous/private swap only when no such page is
// available.
KSTATUS MmVadTrimOnePageForPressure(void) {
    if (!gVadListReady) return STATUS_NOT_SUPPORTED;

    KSTATUS cacheDrop = MmFileCacheDropOnePageForPressure();
    if (KSUCCESS(cacheDrop))
        return STATUS_SUCCESS;

    for (LIST_ENTRY* node = gVadList.Fwd; node != &gVadList; node = node->Fwd) {
        MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, GlobalLink);
        for (size_t i = 0; i < vad->PageCount; ++i) {
            KSTATUS s = KiVadDiscardCleanFilePage(vad, i, (void**)0);
            if (KSUCCESS(s))
                return STATUS_SUCCESS;
        }
    }

    if (!MmSwapIsEnabled()) return STATUS_NOT_SUPPORTED;

    for (LIST_ENTRY* node = gVadList.Fwd; node != &gVadList; node = node->Fwd) {
        MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, GlobalLink);
        for (size_t i = 0; i < vad->PageCount; ++i) {
            KSTATUS s = KiVadTrimPageToSwap(vad, i, (void**)0);
            if (KSUCCESS(s))
                return STATUS_SUCCESS;
        }
    }
    return STATUS_NOT_FOUND;
}

KSTATUS MmVadSyncFileMappings(FILE_OBJECT* file) {
    if (!file)
        return STATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gVadLock);
    for (LIST_ENTRY* node = gVadList.Fwd; node != &gVadList; node = node->Fwd) {
        MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, GlobalLink);
        if (vad->Type != MmVadMappedFile || vad->FileObject != file || !vad->Pages)
            continue;

        for (size_t i = 0; i < vad->PageCount; ++i) {
            MM_VAD_PAGE* page = &vad->Pages[i];
            if ((page->Flags & MM_VAD_PAGE_FILE_SHARED_WRITE) == 0)
                continue;
            if (!page->Resident || !page->Frame || !page->PageFlags)
                continue;

            uintptr_t va = vad->Base + i * PAGE_SIZE;
            int wasDirty = MmPteTestAndClearDirty(va);
            if (wasDirty)
                (void)KiVadMarkSharedFilePageDirty(vad, i, 0);

            if ((page->Flags & MM_VAD_PAGE_FILE_DIRTY_TRACK) == 0) {
                page->Flags |= MM_VAD_PAGE_FILE_DIRTY_TRACK;
                (void)MmPteProtectRange(va, PAGE_SIZE, KiVadHardwarePageFlags(page));
            }
        }
    }
    KeReleaseSpinlock(&gVadLock, flags);
    return STATUS_SUCCESS;
}

// Resolves a user-mode page fault against committed VAD metadata.
// Resolves a user-mode page fault against committed VAD metadata.
KSTATUS MmVadResolveUserPageFault(PROCESS* process, uintptr_t faultAddress, uint64_t errorCode) {
    gVadStats.ResolveFaultCalls++;

    if (!process)
        return STATUS_ACCESS_DENIED;
    if (!KiUserRangeValid(faultAddress, 1)) {
        gVadStats.NoVadFaults++;
        gVadStats.FailedFaults++;
        return STATUS_ACCESS_DENIED;
    }

    uintptr_t va = faultAddress & PAGE_MASK;
    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindReservedVadContainingLocked(process, va, PAGE_SIZE);
    KeReleaseSpinlock(&process->VadLock, procFlags);
    if (!vad) {
        gVadStats.NoVadFaults++;
        gVadStats.FailedFaults++;
        return STATUS_NOT_FOUND;
    }

    size_t index = KiVadPageIndex(vad, va);
    if (index >= vad->PageCount || vad->Pages[index].State != MmVadPageCommitted) {
        gVadStats.UncommittedFaults++;
        gVadStats.FailedFaults++;
        return STATUS_NOT_FOUND;
    }
    if (!vad->Pages[index].PageFlags) {
        gVadStats.NoAccessFaults++;
        gVadStats.FailedFaults++;
        return STATUS_ACCESS_DENIED;
    }

    int presentFault = (errorCode & 0x1) ? 1 : 0;
    int wantWrite = (errorCode & 0x2) ? 1 : 0;
    int wantExecute = (errorCode & 0x10) ? 1 : 0;
    if (presentFault) {
        if (wantWrite &&
            (vad->Pages[index].Flags & MM_VAD_PAGE_FILE_SHARED_WRITE) &&
            (vad->Pages[index].Flags & MM_VAD_PAGE_FILE_DIRTY_TRACK)) {
            KSTATUS dirty = KiVadMarkSharedFilePageDirty(vad, index, 0);
            if (KSUCCESS(dirty)) {
                vad->Pages[index].Flags &= ~MM_VAD_PAGE_FILE_DIRTY_TRACK;
                KSTATUS ps = MmPteProtectRange(va, PAGE_SIZE, KiVadHardwarePageFlags(&vad->Pages[index]));
                if (KSUCCESS(ps)) {
                    gVadStats.FileBackedDirtyFaults++;
                    gVadStats.ResolvedFaults++;
                    return STATUS_SUCCESS;
                }
                dirty = ps;
            }
            gVadStats.ProtectionFaults++;
            gVadStats.FailedFaults++;
            return dirty;
        }
        if (wantWrite && (vad->Pages[index].Flags & MM_VAD_PAGE_COW)) {
            KSTATUS cow = KiVadBreakCopyOnWrite(vad, index);
            if (KSUCCESS(cow)) {
                gVadStats.CowFaults++;
                gVadStats.ResolvedFaults++;
                return STATUS_SUCCESS;
            }
        }
        gVadStats.ProtectionFaults++;
        gVadStats.FailedFaults++;
        return STATUS_ACCESS_DENIED;
    }

    KSTATUS s = KiVadEnsurePageResident(vad, index, wantWrite, wantExecute);
    if (!KSUCCESS(s)) {
        gVadStats.FailedFaults++;
        return s;
    }

    gVadStats.ResolvedFaults++;
    return STATUS_SUCCESS;
}


// Resolves a page fault by looking up the owning VAD in the global user arena.
// The single address space intentionally allows the kernel and other processes
// to touch any lower-half mapping, so fault ownership is derived from the VAD
// that owns the virtual address rather than from the currently running process.
KSTATUS MmVadResolvePageFaultByAddress(uintptr_t faultAddress, uint64_t errorCode) {
    if (!KiUserRangeValid(faultAddress, 1))
        return STATUS_ACCESS_DENIED;
    if (!gVadListReady)
        return STATUS_NOT_SUPPORTED;

    uintptr_t va = faultAddress & PAGE_MASK;
    PROCESS* owner = (PROCESS*)0;

    uintptr_t flags = KeAcquireSpinlock(&gVadLock);
    for (LIST_ENTRY* node = gVadList.Fwd; node != &gVadList; node = node->Fwd) {
        MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, GlobalLink);
        uintptr_t end = vad->Base + vad->PayloadSize;
        if (va >= vad->Base && va < end) {
            owner = vad->Owner;
            break;
        }
    }
    KeReleaseSpinlock(&gVadLock, flags);

    if (!owner)
        return STATUS_NOT_FOUND;

    return MmVadResolveUserPageFault(owner, faultAddress, errorCode);
}

// Changes page protections on a committed VAD range.
KSTATUS MmVadProtect(PROCESS* process, void* base, size_t size, uint64_t pageFlags) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!process || !base || !size) return STATUS_INVALID_PARAMETER;

    uintptr_t start = (uintptr_t)base;
    if (start > (uintptr_t)-1 - (size - 1)) return STATUS_NO_MEMORY;

    uintptr_t alignedStart = start & PAGE_MASK;
    uintptr_t alignedEnd = (start + size + PAGE_SIZE - 1) & PAGE_MASK;
    if (alignedEnd <= alignedStart) return STATUS_NO_MEMORY;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindReservedVadContainingLocked(process, alignedStart, alignedEnd - alignedStart);
    KeReleaseSpinlock(&process->VadLock, procFlags);

    if (!vad || !KiVadRangeCommitted(vad, alignedStart, alignedEnd)) {
        DbgPrint("MM: protect VAD failed: pid=%u range=%016llx-%016llx not committed\n",
                 process->ProcessId,
                 (unsigned long long)alignedStart,
                 (unsigned long long)alignedEnd);
        return STATUS_NOT_FOUND;
    }
    if (vad->LockCount)
        return STATUS_ACCESS_DENIED;

    uint64_t finalFlags = pageFlags ? (pageFlags | PAGE_PRESENT | PAGE_USER) : 0;
    size_t first = KiVadPageIndex(vad, alignedStart);
    size_t count = KiVadRangePageCount(alignedStart, alignedEnd);

    for (size_t i = 0; i < count; ++i) {
        size_t pageIndex = first + i;
        MM_VAD_PAGE* page = &vad->Pages[pageIndex];
        uintptr_t va = vad->Base + pageIndex * PAGE_SIZE;

        if (!finalFlags) {
            if (page->Resident && page->PageFlags)
                (void)KiVadUnmapPageTracked(va, page->Frame);
            page->PageFlags = 0;
            continue;
        }

        uint64_t oldFlags = page->PageFlags;
        page->PageFlags = finalFlags;
        if ((finalFlags & PAGE_WRITE) == 0 &&
            (page->Flags & (MM_VAD_PAGE_SHARED_FRAME | MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_IMAGE_BACKED)) == 0) {
            page->Flags &= ~MM_VAD_PAGE_COW;
        }

        if (page->Resident) {
            uint64_t hardwareFlags = KiVadHardwarePageFlags(page);
            KSTATUS s = oldFlags
                        ? MmPteProtectRange(va, PAGE_SIZE, hardwareFlags)
                        : KiVadMapPageTracked(va, page->Frame, hardwareFlags, MmPteMapFailIfPresent);
            if (!KSUCCESS(s)) return s;
        }
    }

    return STATUS_SUCCESS;
}

// Destroys a VAD and releases all resident payload pages.
KSTATUS MmVadDestroy(PROCESS* process, void* base) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!process || !base) return STATUS_INVALID_PARAMETER;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindVadByBaseLocked(process, base);
    if (!vad) {
        KeReleaseSpinlock(&process->VadLock, procFlags);
        return STATUS_NOT_FOUND;
    }
    if (vad->LockCount) {
        KeReleaseSpinlock(&process->VadLock, procFlags);
        return STATUS_ACCESS_DENIED;
    }
    ListRemove(&vad->ProcessLink);
    KeReleaseSpinlock(&process->VadLock, procFlags);

    uintptr_t globalFlags = KeAcquireSpinlock(&gVadLock);
    ListRemove(&vad->GlobalLink);
    KeReleaseSpinlock(&gVadLock, globalFlags);

    KiUnmapVadPayload(vad);
    if (vad->FileObject) {
        ObDereferenceObject(&vad->FileObject->Header);
        vad->FileObject = (FILE_OBJECT*)0;
    }
    if (vad->ImageSection) {
        LdrDereferenceImageSection(vad->ImageSection);
        vad->ImageSection = (IMAGE_SECTION_OBJECT*)0;
    }
    MmFreePool(vad->Pages);
    MmFreePool(vad);
    return STATUS_SUCCESS;
}

// Destroys every VAD owned by a process during teardown.
void MmVadDestroyAllForProcess(PROCESS* process) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!process) return;

    for (;;) {
        uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
        LIST_ENTRY* node = (LIST_ENTRY*)0;
        int sawLocked = 0;
        for (LIST_ENTRY* it = process->VadList.Fwd; it != &process->VadList; it = it->Fwd) {
            MM_VAD* candidate = CONTAINING_RECORD(it, MM_VAD, ProcessLink);
            if (candidate->LockCount) {
                sawLocked = 1;
                continue;
            }
            node = it;
            ListRemove(node);
            break;
        }
        KeReleaseSpinlock(&process->VadLock, procFlags);
        if (!node) {
            if (sawLocked) {
                KeYield();
                continue;
            }
            break;
        }

        MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, ProcessLink);
        uintptr_t globalFlags = KeAcquireSpinlock(&gVadLock);
        ListRemove(&vad->GlobalLink);
        KeReleaseSpinlock(&gVadLock, globalFlags);
        KiUnmapVadPayload(vad);
        if (vad->FileObject) {
            ObDereferenceObject(&vad->FileObject->Header);
            vad->FileObject = (FILE_OBJECT*)0;
        }
        if (vad->ImageSection) {
            LdrDereferenceImageSection(vad->ImageSection);
            vad->ImageSection = (IMAGE_SECTION_OBJECT*)0;
        }
        MmFreePool(vad->Pages);
        MmFreePool(vad);
    }
}

// Queries the committed VAD range containing an address.
KSTATUS MmVadQuery(PROCESS* process, const void* address, void** baseOut, size_t* sizeOut,
                   uint64_t* pageFlagsOut, MM_VAD_TYPE* typeOut, uint32_t* vadFlagsOut) {
    if (!process || !address) return STATUS_INVALID_PARAMETER;
    uintptr_t addr = (uintptr_t)address;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    for (LIST_ENTRY* node = process->VadList.Fwd; node != &process->VadList; node = node->Fwd) {
        MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, ProcessLink);
        if (addr < vad->Base || addr >= vad->Base + vad->PayloadSize)
            continue;

        size_t index = KiVadPageIndex(vad, addr & PAGE_MASK);
        if (index >= vad->PageCount || vad->Pages[index].State != MmVadPageCommitted) {
            KeReleaseSpinlock(&process->VadLock, procFlags);
            return STATUS_NOT_FOUND;
        }

        uint64_t flags = vad->Pages[index].PageFlags;
        size_t runStart = index;
        while (runStart > 0 &&
               vad->Pages[runStart - 1].State == MmVadPageCommitted &&
               vad->Pages[runStart - 1].PageFlags == flags) {
            runStart--;
        }

        size_t runEnd = index + 1;
        while (runEnd < vad->PageCount &&
               vad->Pages[runEnd].State == MmVadPageCommitted &&
               vad->Pages[runEnd].PageFlags == flags) {
            runEnd++;
        }

        if (baseOut) *baseOut = (void*)(vad->Base + runStart * PAGE_SIZE);
        if (sizeOut) *sizeOut = (runEnd - runStart) * PAGE_SIZE;
        if (pageFlagsOut) *pageFlagsOut = flags;
        if (typeOut) *typeOut = vad->Type;
        if (vadFlagsOut) *vadFlagsOut = vad->Flags;
        KeReleaseSpinlock(&process->VadLock, procFlags);
        return STATUS_SUCCESS;
    }
    KeReleaseSpinlock(&process->VadLock, procFlags);
    return STATUS_NOT_FOUND;
}

// Returns per-process VAD accounting for ps and memory telemetry.
void MmVadQueryProcessMemory(PROCESS* process, size_t* vadCountOut, size_t* reservedBytesOut, size_t* committedBytesOut) {
    size_t vads = 0;
    size_t reserved = 0;
    size_t committed = 0;

    if (process) {
        uintptr_t flags = KeAcquireSpinlock(&process->VadLock);
        for (LIST_ENTRY* node = process->VadList.Fwd; node != &process->VadList; node = node->Fwd) {
            MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, ProcessLink);
            vads++;
            reserved += vad->ReservationSize;
            committed += vad->CommittedPages * PAGE_SIZE;
        }
        KeReleaseSpinlock(&process->VadLock, flags);
    }

    if (vadCountOut) *vadCountOut = vads;
    if (reservedBytesOut) *reservedBytesOut = reserved;
    if (committedBytesOut) *committedBytesOut = committed;
}

// Returns VAD and demand-fault telemetry.
void MmVadQueryStats(MM_VAD_STATS* out) {
    if (!out) return;
    *out = gVadStats;
}

// Returns the requested payload size for a VAD base.
size_t MmVadGetSize(PROCESS* process, const void* base) {
    if (!process || !base) return 0;
    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindVadByBaseLocked(process, base);
    size_t size = vad ? vad->Size : 0;
    KeReleaseSpinlock(&process->VadLock, procFlags);
    return size;
}

// Returns the requested payload size for a VAD base without knowing the owner.
size_t MmVadGetSizeByBase(const void* base) {
    if (!base) return 0;
    uintptr_t b = (uintptr_t)base;
    uintptr_t flags = KeAcquireSpinlock(&gVadLock);
    for (LIST_ENTRY* node = gVadList.Fwd; node != &gVadList; node = node->Fwd) {
        MM_VAD* vad = CONTAINING_RECORD(node, MM_VAD, GlobalLink);
        if (vad->Base == b) {
            size_t size = vad->Size;
            KeReleaseSpinlock(&gVadLock, flags);
            return size;
        }
    }
    KeReleaseSpinlock(&gVadLock, flags);
    return 0;
}

// Implements the VAD Type Name operation.
static const char* KiVadTypeName(MM_VAD_TYPE type) {
    switch (type) {
        case MmVadGeneric:     return "generic";
        case MmVadHeap:        return "heap";
        case MmVadImage:       return "image";
        case MmVadStack:       return "stack";
        case MmVadTeb:         return "teb";
        case MmVadPeb:         return "peb";
        case MmVadEnvironment: return "env";
        case MmVadMappedFile:  return "mapped";
        case MmVadShared:      return "shared";
        default:               return "?";
    }
}

// Implements the debug Dump VADs operation.
void MmVadDebugDumpProcess(PROCESS* process) {
    if (!process) {
        DbgPrint("MM: VADs: <no process>\n");
        return;
    }

    DbgPrint("MM: VADs for pid=%u name=%s:\n", process->ProcessId, process->Name);
    uintptr_t flags = KeAcquireSpinlock(&process->VadLock);
    uint32_t count = 0;
    for (LIST_ENTRY* node = process->VadList.Fwd; node != &process->VadList; node = node->Fwd) {
        MM_VAD* r = CONTAINING_RECORD(node, MM_VAD, ProcessLink);
        DbgPrint("  #%u %-7s res=%016llx-%016llx base=%016llx size=%llx committed=%llx/%llx flags=%x pte=%llx\n",
                 count,
                 KiVadTypeName(r->Type),
                 (unsigned long long)r->ReservationBase,
                 (unsigned long long)(r->ReservationBase + r->ReservationSize),
                 (unsigned long long)r->Base,
                 (unsigned long long)r->Size,
                 (unsigned long long)(r->CommittedPages * PAGE_SIZE),
                 (unsigned long long)r->PayloadSize,
                 r->Flags,
                 (unsigned long long)r->PageFlags);
        count++;
    }
    KeReleaseSpinlock(&process->VadLock, flags);
    if (!count) DbgPrint("  <none>\n");
}

// Implements the debug Describe Pool Address operation.
static void KiDebugDescribePoolAddress(uintptr_t addr) {
    DbgPrint("MM: kernel pool start=%016llx brk=%016llx limit=%016llx usage=%llu KiB\n",
             (unsigned long long)gPoolHeap.Start,
             (unsigned long long)gPoolHeap.Brk,
             (unsigned long long)gPoolHeap.Limit,
             (unsigned long long)((gPoolHeap.Brk - gPoolHeap.Start) / 1024));

    if (addr < gPoolHeap.Start || addr >= gPoolHeap.Limit) return;
    if (addr >= gPoolHeap.Brk) {
        DbgPrint("MM: address %016llx is inside pool VA range but above committed brk by +0x%llx\n",
                 (unsigned long long)addr,
                 (unsigned long long)(addr - gPoolHeap.Brk));
        return;
    }

    uintptr_t flags = KeAcquireSpinlock(&gPoolHeap.Lock);
    for (BLOCK_HDR* cur = (BLOCK_HDR*)gPoolHeap.Start;
         (uintptr_t)cur < gPoolHeap.Brk;
         cur = KiNextBlock(cur)) {
        if (cur->Magic != BLOCK_MAGIC) {
            DbgPrint("MM: pool walk corrupt at hdr=%p magic=%08x\n", cur, cur->Magic);
            break;
        }
        uintptr_t hdr = (uintptr_t)cur;
        uintptr_t data = hdr + HDR_SIZE;
        uintptr_t end = data + cur->Size;
        uintptr_t ftr = end;
        BLOCK_FTR* footer = KiHdrToFtr(cur);
        if ((uintptr_t)footer + FTR_SIZE > gPoolHeap.Brk ||
            footer->Magic != BLOCK_MAGIC || footer->Size != cur->Size || footer->Free != cur->Free) {
            DbgPrint("MM: pool walk corrupt footer hdr=%p state=%u size=%llx footer=%p fmagic=%08x fstate=%u fsize=%llx\n",
                     cur,
                     cur->Free,
                     (unsigned long long)cur->Size,
                     footer,
                     ((uintptr_t)footer + FTR_SIZE <= gPoolHeap.Brk) ? footer->Magic : 0,
                     ((uintptr_t)footer + FTR_SIZE <= gPoolHeap.Brk) ? footer->Free : 0,
                     ((uintptr_t)footer + FTR_SIZE <= gPoolHeap.Brk) ? (unsigned long long)footer->Size : 0ULL);
            break;
        }
        if (addr >= hdr && addr < ftr + FTR_SIZE) {
            DbgPrint("MM: address %016llx in pool block hdr=%016llx data=%016llx-%016llx size=%llx req=%llx state=%u(%s) tag=%08x flags=%08x seq=%llu alloc=%p free=%p\n",
                     (unsigned long long)addr,
                     (unsigned long long)hdr,
                     (unsigned long long)data,
                     (unsigned long long)end,
                     (unsigned long long)cur->Size,
                     (unsigned long long)cur->RequestedSize,
                     cur->Free,
                     KiPoolStateName(cur->Free),
                     cur->Tag,
                     cur->Flags,
                     (unsigned long long)cur->Sequence,
                     (void*)cur->AllocCaller,
                     (void*)cur->FreeCaller);
            break;
        }
        if (KiNextBlock(cur) <= cur) {
            DbgPrint("MM: pool walk stopped: non-advancing block hdr=%p size=%llx\n", cur, (unsigned long long)cur->Size);
            break;
        }
    }
    KeReleaseSpinlock(&gPoolHeap.Lock, flags);
}

// Implements the debug Describe Address operation.
void MmDebugDescribeAddress(uintptr_t addr) {
    PHYSICAL_ADDRESS phys = 0;
    uint64_t flags = 0;
    KSTATUS qs = MmPteQueryPageFor(MmPteQueryDebug, addr, &phys, &flags);
    DbgPrint("MM: describe va=%016llx query=%d phys=%016llx flags=%016llx\n",
             (unsigned long long)addr,
             qs,
             (unsigned long long)phys,
             (unsigned long long)flags);

    if (addr >= POOL_BASE && addr < POOL_LIMIT) {
        KiDebugDescribePoolAddress(addr);
        return;
    }

    if (addr >= MM_VAD_ARENA_BASE && addr < USER_LIMIT) {
        uintptr_t gl = KeAcquireSpinlock(&gVadLock);
        for (LIST_ENTRY* node = gVadList.Fwd; node != &gVadList; node = node->Fwd) {
            MM_VAD* r = CONTAINING_RECORD(node, MM_VAD, GlobalLink);
            if (addr >= r->ReservationBase && addr < r->ReservationBase + r->ReservationSize) {
                const char* pageState = "guard";
                if (addr >= r->Base && addr < r->Base + r->PayloadSize) {
                    size_t pageIndex = KiVadPageIndex(r, addr & PAGE_MASK);
                    if (pageIndex < r->PageCount) {
                        MM_VAD_PAGE* page = &r->Pages[pageIndex];
                        if (page->State == MmVadPageCommitted)
                            pageState = page->PageFlags ? "mapped" : "noaccess";
                        else
                            pageState = "reserved";
                    }
                }

                DbgPrint("MM: address belongs to pid=%u name=%s vad=%s page=%s res=%016llx-%016llx base=%016llx size=%llx flags=%x\n",
                         r->Owner ? r->Owner->ProcessId : 0,
                         r->Owner ? r->Owner->Name : "?",
                         KiVadTypeName(r->Type),
                         pageState,
                         (unsigned long long)r->ReservationBase,
                         (unsigned long long)(r->ReservationBase + r->ReservationSize),
                         (unsigned long long)r->Base,
                         (unsigned long long)r->Size,
                         r->Flags);
                KeReleaseSpinlock(&gVadLock, gl);
                return;
            }
        }
        KeReleaseSpinlock(&gVadLock, gl);
        DbgPrint("MM: address is in user arena but not reserved by any MM_VAD\n");
    }
}

// Creates user Memory For Process.
__declspec(dllexport) void* MmAllocateUserMemoryForProcess(PROCESS* process, size_t size) {
    void* base = (void*)0;
    if (!process || !size) return (void*)0;
    KSTATUS s = MmVadCreateCommitted(process, size, PAGE_SIZE,
                                   PAGE_PRESENT | PAGE_WRITE | PAGE_USER | PAGE_NX,
                                   MM_VAD_ZERO,
                                   MmVadHeap,
                                   &base);
    if (!KSUCCESS(s)) return (void*)0;
    s = MmVadEnsureResident(process, base, size, PAGE_WRITE | PAGE_USER);
    if (!KSUCCESS(s)) {
        (void)MmVadDestroy(process, base);
        return (void*)0;
    }
    return base;
}

// Releases user Memory For Process.
__declspec(dllexport) void MmFreeUserMemoryForProcess(PROCESS* process, void* ptr) {
    if (!process || !ptr) return;
    (void)MmVadDestroy(process, ptr);
}

// Implements the reallocate User Memory For Process operation.
__declspec(dllexport) void* MmReallocateUserMemoryForProcess(PROCESS* process, void* ptr, size_t newSize) {
    if (!process) return (void*)0;
    if (!ptr) return MmAllocateUserMemoryForProcess(process, newSize);
    if (!newSize) {
        MmFreeUserMemoryForProcess(process, ptr);
        return (void*)0;
    }

    size_t oldSize = MmVadGetSize(process, ptr);
    if (!oldSize) return (void*)0;
    if (newSize <= oldSize) return ptr;

    void* newPtr = MmAllocateUserMemoryForProcess(process, newSize);
    if (!newPtr) return (void*)0;
    memcpy(newPtr, ptr, oldSize);
    MmFreeUserMemoryForProcess(process, ptr);
    return newPtr;
}

// Creates user Memory.
__declspec(dllexport) void* MmAllocateUserMemory(size_t size) {
    PROCESS* process = KeGetCurrentProcess();
    return MmAllocateUserMemoryForProcess(process, size);
}

// Releases user Memory.
__declspec(dllexport) void MmFreeUserMemory(void* ptr) {
    PROCESS* process = KeGetCurrentProcess();
    MmFreeUserMemoryForProcess(process, ptr);
}

// Implements the reallocate User Memory operation.
__declspec(dllexport) void* MmReallocateUserMemory(void* ptr, size_t newSize) {
    PROCESS* process = KeGetCurrentProcess();
    return MmReallocateUserMemoryForProcess(process, ptr, newSize);
}

// Creates user Image For Process.
__declspec(dllexport) void* MmAllocateUserImageForProcess(PROCESS* process, size_t size) {
    void* base = (void*)0;
    if (!process || !size) return (void*)0;
    KSTATUS s = MmVadCreateCommitted(process, size, PAGE_SIZE,
                                   PAGE_PRESENT | PAGE_WRITE | PAGE_USER | PAGE_NX,
                                   MM_VAD_ZERO | MM_VAD_GUARD_LOW | MM_VAD_GUARD_HIGH,
                                   MmVadImage,
                                   &base);
    if (!KSUCCESS(s)) return (void*)0;
    return base;
}

// Releases user Image For Process.
__declspec(dllexport) void MmFreeUserImageForProcess(PROCESS* process, void* ptr, size_t size) {
    (void)size;
    if (!process || !ptr) return;
    (void)MmVadDestroy(process, ptr);
}

// Creates user Image.
__declspec(dllexport) void* MmAllocateUserImage(size_t size) {
    PROCESS* process = KeGetCurrentProcess();
    return MmAllocateUserImageForProcess(process, size);
}

// Releases user Image.
__declspec(dllexport) void MmFreeUserImage(void* ptr, size_t size) {
    PROCESS* process = KeGetCurrentProcess();
    MmFreeUserImageForProcess(process, ptr, size);
}

// Returns kernel Pool Usage.
size_t MmGetKernelPoolUsage(void) {
    // Returns bytes committed to the kernel pool.
    return (size_t)(gPoolHeap.Brk - gPoolHeap.Start);
}

// Returns user Heap Usage.
size_t MmGetUserHeapUsage(void) {
    size_t total = 0;
    if (!gVadListReady) return 0;
    uintptr_t flags = KeAcquireSpinlock(&gVadLock);
    for (LIST_ENTRY* node = gVadList.Fwd; node != &gVadList; node = node->Fwd) {
        MM_VAD* r = CONTAINING_RECORD(node, MM_VAD, GlobalLink);
        total += r->CommittedPages * PAGE_SIZE;
    }
    KeReleaseSpinlock(&gVadLock, flags);
    return total;
}