// Kernel virtual-memory manager. This file owns VAD bookkeeping, faults, and user-copy validation.
// Pool allocation lives in mm_pool.c and page-table construction lives in mm_pte.c.

// Virtual address layout:
//   0000000000001000-00007FFFFFFFF000 is the lower-half user range.
//   0000100000000000-00007FFFFFFFF000 is the current VAD allocation arena.
//   FFFF800000000000-FFFF87FFFFFFFFFF is reserved for kernel images and fixed kernel mappings.
//   FFFF880000000000-FFFF8FFFFFFFFFFF is the physical direct map.
//   FFFF900000000000-FFFF97FFFFFFFFFF is the PCI MMIO aperture owned by pci.sys.
//   FFFF980000000000-FFFF9FFFFFFFFFFF is the generic device-MMIO aperture.
//   FFFFA00000000000-FFFFAFFFFFFFFFFF is the kernel shared-view arena.
//   FFFFC00000000000-FFFFDFFFFFFFFFFF is the kernel pool.
//
// The current VM model uses one shared PML4, so user pointers are globally
// meaningful while protection and lifetime are still tracked per process VAD.

// User mappings:
//   VADs separate reservation from commitment. Committed pages become resident
//   when a user fault or kernel prefault path materializes them.
#include <ntlite/atomic.h>
#include <stdbool.h>
#include <stddef.h>
#include <string.h>
#include <intrin.h>
#include <kernel/mm_virt.h>
#include <kernel/mm_pool.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 bool KiUserRangeValid(uintptr_t addr, size_t size) {
    if (size == 0) return true;
    if (addr < USER_BASE) return false;
    if (addr >= USER_LIMIT) return false;
    if (size > (size_t)(USER_LIMIT - addr)) return false;
    return true;
}

// 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 NTLSTATUS_ACCESS_DENIED;

    if (size == 0)
        return NTLSTATUS_SUCCESS;

    PROCESS* process = KeGetCurrentProcess();
    if (!process)
        return NTLSTATUS_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);
}

// 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;

struct MM_KERNEL_VIEW {
    LIST_ENTRY Link;
    uintptr_t Base;
    size_t PageCount;
};

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;
static LIST_ENTRY gKernelViewList;
static SPINLOCK gKernelViewLock = SPINLOCK_INIT;

#define MM_VAD_ARENA_BASE     0x0000100000000000ULL
#define MM_KERNEL_VIEW_BASE   0xFFFFA00000000000ULL
#define MM_KERNEL_VIEW_LIMIT  0xFFFFB00000000000ULL

// Initializes the VAD list after the shared page table and kernel pool are ready.
static void KiVadInitialize(void) {
    LIST_INIT(&gVadList);
    LIST_INIT(&gKernelViewList);
    gVadLock = (SPINLOCK)SPINLOCK_INIT;
    gKernelViewLock = (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);

    MmPoolInitialize(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);
}

// Inserts a kernel view before one list node while preserving address order.
static void KiKernelViewInsertBefore(LIST_ENTRY* before, MM_KERNEL_VIEW* view) {
    view->Link.Fwd = before;
    view->Link.Bwd = before->Bwd;
    before->Bwd->Fwd = &view->Link;
    before->Bwd = &view->Link;
}

// Maps existing frames into the kernel shared-view arena.
KSTATUS MmCreateKernelViewFromFrames(const PHYSICAL_ADDRESS* frames,
                                     size_t pageCount,
                                     MM_KERNEL_VIEW** outView) {
    if (!outView)
        return NTLSTATUS_INVALID_PARAMETER;
    *outView = (MM_KERNEL_VIEW*)0;
    if (!frames || pageCount == 0 || pageCount > (SIZE_MAX / PAGE_SIZE))
        return NTLSTATUS_INVALID_PARAMETER;

    size_t viewBytes = pageCount * PAGE_SIZE;
    MM_KERNEL_VIEW* view = (MM_KERNEL_VIEW*)MmAllocatePoolTag(sizeof(MM_KERNEL_VIEW),
                                                               MM_POOL_TAG('k', 'v', 'i', 'w'));
    if (!view)
        return NTLSTATUS_NO_MEMORY;
    memset(view, 0, sizeof(*view));
    LIST_INIT(&view->Link);

    uintptr_t candidate = MM_KERNEL_VIEW_BASE;
    LIST_ENTRY* before = &gKernelViewList;
    uintptr_t lockFlags = KeAcquireSpinlock(&gKernelViewLock);
    for (LIST_ENTRY* it = gKernelViewList.Fwd; it != &gKernelViewList; it = it->Fwd) {
        MM_KERNEL_VIEW* existing = CONTAINING_RECORD(it, MM_KERNEL_VIEW, Link);
        if (candidate <= existing->Base &&
            viewBytes <= (size_t)(existing->Base - candidate)) {
            before = it;
            break;
        }

        uintptr_t existingEnd = existing->Base + existing->PageCount * PAGE_SIZE;
        if (existingEnd > candidate)
            candidate = existingEnd;
    }

    if (candidate > MM_KERNEL_VIEW_LIMIT || viewBytes > MM_KERNEL_VIEW_LIMIT - candidate) {
        KeReleaseSpinlock(&gKernelViewLock, lockFlags);
        MmFreePool(view);
        return NTLSTATUS_NO_MEMORY;
    }

    view->Base = candidate;
    view->PageCount = pageCount;
    KiKernelViewInsertBefore(before, view);
    KeReleaseSpinlock(&gKernelViewLock, lockFlags);

    const uint64_t hardwareFlags = PAGE_PRESENT | PAGE_WRITE | PAGE_NX;
    size_t mappedPages = 0;
    KSTATUS s = NTLSTATUS_SUCCESS;
    for (; mappedPages < pageCount; mappedPages++) {
        if (!frames[mappedPages]) {
            s = NTLSTATUS_INVALID_PARAMETER;
            break;
        }
        s = MmPteMapPage(candidate + mappedPages * PAGE_SIZE, frames[mappedPages],
                         hardwareFlags, MmPteMapFailIfPresent);
        if (!KSUCCESS(s))
            break;
    }

    if (!KSUCCESS(s)) {
        while (mappedPages != 0) {
            mappedPages--;
            (void)MmPteUnmapPage(candidate + mappedPages * PAGE_SIZE);
        }
        lockFlags = KeAcquireSpinlock(&gKernelViewLock);
        ListRemove(&view->Link);
        KeReleaseSpinlock(&gKernelViewLock, lockFlags);
        MmFreePool(view);
        return s;
    }

    *outView = view;
    return NTLSTATUS_SUCCESS;
}

// Returns the base address of one kernel shared view.
void* MmKernelViewBase(const MM_KERNEL_VIEW* view) {
    return view ? (void*)view->Base : (void*)0;
}

// Unmaps and releases one kernel shared view without releasing its frames.
void MmDestroyKernelView(MM_KERNEL_VIEW* view) {
    if (!view)
        return;

    for (size_t i = 0; i < view->PageCount; i++)
        (void)MmPteUnmapPage(view->Base + i * PAGE_SIZE);

    uintptr_t lockFlags = KeAcquireSpinlock(&gKernelViewLock);
    ListRemove(&view->Link);
    KeReleaseSpinlock(&gKernelViewLock, lockFlags);
    MmFreePool(view);
}

// 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 bool 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 NTLSTATUS_INVALID_PARAMETER;
    if (!gVadListReady) return NTLSTATUS_FAILURE;

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

    while (true) {
        if (candidate < MM_VAD_ARENA_BASE || reservationSize > (size_t)(USER_LIMIT - candidate))
            return NTLSTATUS_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 NTLSTATUS_NO_MEMORY;
                collided = 1;
                break;
            }
        }

        if (!collided) {
            *baseOut = candidate;
            return NTLSTATUS_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 bool KiVadRangeCommitted(const MM_VAD* vad, uintptr_t alignedStart, uintptr_t alignedEnd) {
    if (!vad || alignedStart < vad->Base || alignedEnd > vad->Base + vad->PayloadSize)
        return false;
    size_t first = KiVadPageIndex(vad, alignedStart);
    size_t count = KiVadRangePageCount(alignedStart, alignedEnd);
    if (first > vad->PageCount || count > vad->PageCount - first)
        return false;

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

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 NTLSTATUS_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 NTLSTATUS_NOT_FOUND;

    uint64_t pageDelta = (uint64_t)index * (uint64_t)PAGE_SIZE;
    uint64_t fileOffset = vad->FileOffset + pageDelta;
    if (fileOffset < vad->FileOffset)
        return NTLSTATUS_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 NTLSTATUS_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)
        return;

    int wasDirty = MmPteTestAndClearDirty(va);
    bool wasWriteEnabled = (page->Flags & MM_VAD_PAGE_FILE_DIRTY_TRACK) == 0;
    if (wasDirty || wasWriteEnabled)
        (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 NTLSTATUS_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 NTLSTATUS_NOT_FOUND;
    if (!page->PageFlags)
        return NTLSTATUS_ACCESS_DENIED;
    if ((page->PageFlags & PAGE_USER) == 0)
        return NTLSTATUS_ACCESS_DENIED;
    if (wantWrite && (page->PageFlags & PAGE_WRITE) == 0)
        return NTLSTATUS_ACCESS_DENIED;
    if (wantExecute && (page->PageFlags & PAGE_NX))
        return NTLSTATUS_ACCESS_DENIED;
    return NTLSTATUS_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 && (NtlAtomicLoadU32(&page->Flags) &
                    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 bool KiVadTryBeginPaging(MM_VAD_PAGE* page) {
    if (!page) return false;

    while (true) {
        uint32_t oldFlags = (uint32_t)NtlAtomicLoadU32(&page->Flags);
        if (oldFlags & MM_VAD_PAGE_PAGING)
            return false;
        uint32_t newFlags = oldFlags | MM_VAD_PAGE_PAGING;
        if ((uint32_t)NtlAtomicCompareExchangeU32(&page->Flags,
                                                   newFlags,
                                                   oldFlags) == oldFlags)
            return true;
        _mm_pause();
    }
}

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

    while (true) {
        uint32_t oldFlags = (uint32_t)NtlAtomicLoadU32(&page->Flags);
        uint32_t newFlags = oldFlags & ~MM_VAD_PAGE_PAGING;
        if ((uint32_t)NtlAtomicCompareExchangeU32(&page->Flags,
                                                   newFlags,
                                                   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 NTLSTATUS_INVALID_PARAMETER;

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

    uint64_t pageDelta = (uint64_t)pageIndex * (uint64_t)PAGE_SIZE;
    if (pageDelta >= (uint64_t)vad->FileLength)
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_SUCCESS;
    }

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

    void* pageVa = MmPfnToVirtual(frame);
    if (!pageVa) {
        MmPfnFree(frame);
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_SUCCESS;
    }

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

    void* pageVa = MmPfnToVirtual(frame);
    if (!pageVa) {
        KiVadEndPaging(page);
        MmPfnFree(frame);
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_NOT_FOUND;

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

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

    void* newPage = MmPfnToVirtual(newFrame);
    if (!newPage) {
        MmPfnFree(newFrame);
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    if (start > (uintptr_t)-1 - (size - 1))
        return NTLSTATUS_NO_MEMORY;

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

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindReservedVadContainingLocked(process, alignedStart, alignedEnd - alignedStart);
    KeReleaseSpinlock(&process->VadLock, procFlags);
    if (!vad)
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;

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

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

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

    size_t first = KiVadPageIndex(vad, alignedStart);
    size_t count = KiVadRangePageCount(alignedStart, alignedEnd);
    if (first > vad->PageCount || count > vad->PageCount - first)
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    memset(out, 0, sizeof(*out));

    if (!KiUserRangeValid(address, 1))
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_FAILURE;
    }
    return NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindReservedVadContainingLocked(process, pageBase, PAGE_SIZE);
    if (!vad) {
        KeReleaseSpinlock(&process->VadLock, procFlags);
        return NTLSTATUS_NOT_FOUND;
    }
    if (vad->LockCount == UINT32_MAX) {
        KeReleaseSpinlock(&process->VadLock, procFlags);
        return NTLSTATUS_FAILURE;
    }
    vad->LockCount++;
    *vadOut = vad;
    KeReleaseSpinlock(&process->VadLock, procFlags);
    return NTLSTATUS_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,
                                               bool writeAccess,
                                               MM_LOCKED_USER_BUFFER** out) {
    if (!out)
        return NTLSTATUS_INVALID_PARAMETER;
    *out = (MM_LOCKED_USER_BUFFER*)0;

    if (size == 0)
        return NTLSTATUS_SUCCESS;
    if (!process || !userAddress)
        return NTLSTATUS_INVALID_PARAMETER;

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

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

    size_t pageCount = KiVadRangePageCount(alignedStart, alignedEnd);
    if (!pageCount)
        return NTLSTATUS_INVALID_PARAMETER;
    if (pageCount > (((size_t)-1) - sizeof(MM_LOCKED_USER_BUFFER)) /
                    (sizeof(MM_PINNED_USER_PAGE) + sizeof(MM_VAD*)))
        return NTLSTATUS_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 NTLSTATUS_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 = NTLSTATUS_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 NTLSTATUS_SUCCESS;
}

// Returns the page count held by a locked user buffer.
size_t MmLockedUserBufferPageCount(const MM_LOCKED_USER_BUFFER* buffer) {
    return buffer ? buffer->PageCount : 0;
}

// Returns one pinned frame held by a locked user buffer.
PHYSICAL_ADDRESS MmLockedUserBufferPageFrame(const MM_LOCKED_USER_BUFFER* buffer, size_t index) {
    if (!buffer || index >= buffer->PageCount)
        return (PHYSICAL_ADDRESS)0;
    return buffer->Pages[index].Frame;
}

// 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 NTLSTATUS_SUCCESS;
    if (!dst || !src)
        return NTLSTATUS_INVALID_PARAMETER;

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

    PROCESS* process = KeGetCurrentProcess();
    if (!process)
        return NTLSTATUS_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 NTLSTATUS_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);
}

// Copies between process-owned user ranges without bouncing through a kernel payload buffer.
__declspec(dllexport) KSTATUS MmCopyUserToUser(PROCESS* sourceProcess, const void* source,
                                               PROCESS* targetProcess, void* target, size_t size) {
    if (size == 0)
        return NTLSTATUS_SUCCESS;
    if (!sourceProcess || !targetProcess || !source || !target)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t sourceAddress = (uintptr_t)source;
    uintptr_t targetAddress = (uintptr_t)target;
    if (!KiUserRangeValid(sourceAddress, size) || !KiUserRangeValid(targetAddress, size))
        return NTLSTATUS_ACCESS_DENIED;

    size_t done = 0;
    while (done < size) {
        MM_PINNED_USER_PAGE sourcePage;
        MM_PINNED_USER_PAGE targetPage;
        memset(&sourcePage, 0, sizeof(sourcePage));
        memset(&targetPage, 0, sizeof(targetPage));

        KSTATUS status = MmPinUserPageForRead(sourceProcess,
                                              (const void*)(sourceAddress + done),
                                              &sourcePage);
        if (!KSUCCESS(status))
            return status;

        status = MmPinUserPageForWrite(targetProcess,
                                       (void*)(targetAddress + done),
                                       &targetPage);
        if (!KSUCCESS(status)) {
            MmUnpinUserPage(&sourcePage);
            return status;
        }

        size_t sourceRemaining = PAGE_SIZE - sourcePage.PageOffset;
        size_t targetRemaining = PAGE_SIZE - targetPage.PageOffset;
        size_t chunk = size - done;
        if (chunk > sourceRemaining) chunk = sourceRemaining;
        if (chunk > targetRemaining) chunk = targetRemaining;

        memcpy((uint8_t*)targetPage.KernelAddress + targetPage.PageOffset,
               (const uint8_t*)sourcePage.KernelAddress + sourcePage.PageOffset,
               chunk);

        MmUnpinUserPage(&targetPage);
        MmUnpinUserPage(&sourcePage);
        done += chunk;
    }

    return NTLSTATUS_SUCCESS;
}

// 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 NTLSTATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    size_t payloadSize = (size + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (payloadSize < size) return NTLSTATUS_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 NTLSTATUS_NO_MEMORY;

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

    size_t pageCount = payloadSize / PAGE_SIZE;
    if (!pageCount) return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;

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

    uintptr_t alignedStart = start & PAGE_MASK;
    uintptr_t alignedEnd = (start + size + PAGE_SIZE - 1) & PAGE_MASK;
    if (alignedEnd <= alignedStart) return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_ACCESS_DENIED;

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

    KSTATUS s = NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

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

    size_t alignedSize = (size + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (!alignedSize || alignedSize < size)
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    const IMAGE_SECTION_METADATA* metadata = LdrImageSectionMetadata(section);
    if (!metadata || !metadata->ImageSize || !metadata->PageCount)
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    uintptr_t start = (uintptr_t)sourceBase;
    if (start & (PAGE_SIZE - 1))
        return NTLSTATUS_INVALID_PARAMETER;
    size_t alignedSize = (size + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (!alignedSize || alignedSize < size)
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 = NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    size_t alignedSize = (payloadSize + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (!alignedSize || alignedSize / PAGE_SIZE > pageCount)
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    *outBase = (void*)0;

    size_t alignedSize = (payloadSize + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (!alignedSize || alignedSize / PAGE_SIZE > pageCount)
        return NTLSTATUS_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 NTLSTATUS_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 = NTLSTATUS_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 NTLSTATUS_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 bool KiVadPageCanDiscardCleanFile(const MM_VAD* vad, const MM_VAD_PAGE* page) {
    if (!vad || !page || vad->Type != MmVadMappedFile)
        return false;
    if (!vad->FileObject || page->State != MmVadPageCommitted)
        return false;
    if (!page->Resident || !page->Frame || !page->PageFlags)
        return false;
    if ((page->Flags & MM_VAD_PAGE_FILE_BACKED) == 0)
        return false;
    if (page->PageFlags & PAGE_WRITE)
        return false;
    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 false;
    return true;
}

static bool KiVadPageCanDiscardCleanFileAfterPagingClaim(const MM_VAD* vad, const MM_VAD_PAGE* page) {
    if (!vad || !page || vad->Type != MmVadMappedFile)
        return false;
    if (!vad->FileObject || page->State != MmVadPageCommitted)
        return false;
    if (!page->Resident || !page->Frame || !page->PageFlags)
        return false;
    if ((page->Flags & MM_VAD_PAGE_FILE_BACKED) == 0)
        return false;
    if (page->PageFlags & PAGE_WRITE)
        return false;
    if (page->Flags & (MM_VAD_PAGE_COW | MM_VAD_PAGE_ZERO_PAGE | MM_VAD_PAGE_SHARED_FRAME |
                       MM_VAD_PAGE_SWAPPED))
        return false;
    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 bool KiVadPageCanSwap(const MM_VAD_PAGE* page) {
    if (!page || page->State != MmVadPageCommitted)
        return false;
    if (!page->Resident || !page->Frame || !page->PageFlags)
        return false;
    if ((page->Flags & MM_VAD_PAGE_FILE_BACKED) && ((page->PageFlags & PAGE_WRITE) == 0))
        return false;
    if (page->Flags & MM_VAD_PAGE_FILE_SHARED_WRITE)
        return false;
    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 false;
    return true;
}

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

// 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 NTLSTATUS_INVALID_PARAMETER;

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

    PHYSICAL_ADDRESS frame = page->Frame;
    MM_PFN* pfn = MmPfnGet(frame);
    if (!pfn || pfn->PinCount)
        return NTLSTATUS_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 NTLSTATUS_NOT_FOUND;

    if (!KiVadTryBeginPaging(page))
        return NTLSTATUS_NOT_FOUND;

    if (!KiVadPageCanDiscardCleanFileAfterPagingClaim(vad, page) || page->Frame != frame) {
        KiVadEndPaging(page);
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;

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

    PHYSICAL_ADDRESS frame = page->Frame;
    MM_PFN* pfn = MmPfnGet(frame);
    if (!pfn || pfn->RefCount != 1 || pfn->PinCount)
        return NTLSTATUS_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 NTLSTATUS_NOT_FOUND;

    if (!KiVadPageCanSwapAfterPagingClaim(page) || page->Frame != frame) {
        KiVadEndPaging(page);
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_SUCCESS;
        }
    }
    return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_NOT_SUPPORTED;

    KSTATUS cacheDrop = MmFileCacheDropOnePageForPressure();
    if (KSUCCESS(cacheDrop))
        return NTLSTATUS_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 NTLSTATUS_SUCCESS;
        }
    }

    if (!MmSwapIsEnabled()) return NTLSTATUS_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 NTLSTATUS_SUCCESS;
        }
    }
    return NTLSTATUS_NOT_FOUND;
}

KSTATUS MmVadSyncFileMappings(FILE_OBJECT* file) {
    if (!file)
        return NTLSTATUS_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);
            bool wasWriteEnabled = (page->Flags & MM_VAD_PAGE_FILE_DIRTY_TRACK) == 0;
            if (wasDirty || wasWriteEnabled)
                (void)KiVadMarkSharedFilePageDirty(vad, i, 0);

            if (wasWriteEnabled) {
                page->Flags |= MM_VAD_PAGE_FILE_DIRTY_TRACK;
                (void)MmPteProtectRange(va, PAGE_SIZE, KiVadHardwarePageFlags(page));
            }
        }
    }
    KeReleaseSpinlock(&gVadLock, flags);
    return NTLSTATUS_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 NTLSTATUS_ACCESS_DENIED;
    if (!KiUserRangeValid(faultAddress, 1)) {
        gVadStats.NoVadFaults++;
        gVadStats.FailedFaults++;
        return NTLSTATUS_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 NTLSTATUS_NOT_FOUND;
    }

    size_t index = KiVadPageIndex(vad, va);
    if (index >= vad->PageCount || vad->Pages[index].State != MmVadPageCommitted) {
        gVadStats.UncommittedFaults++;
        gVadStats.FailedFaults++;
        return NTLSTATUS_NOT_FOUND;
    }
    if (!vad->Pages[index].PageFlags) {
        gVadStats.NoAccessFaults++;
        gVadStats.FailedFaults++;
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_SUCCESS;
            }
        }
        gVadStats.ProtectionFaults++;
        gVadStats.FailedFaults++;
        return NTLSTATUS_ACCESS_DENIED;
    }

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

    gVadStats.ResolvedFaults++;
    return NTLSTATUS_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 NTLSTATUS_ACCESS_DENIED;
    if (!gVadListReady)
        return NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;

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

    uintptr_t alignedStart = start & PAGE_MASK;
    uintptr_t alignedEnd = (start + size + PAGE_SIZE - 1) & PAGE_MASK;
    if (alignedEnd <= alignedStart) return NTLSTATUS_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 NTLSTATUS_NOT_FOUND;
    }
    if (vad->LockCount)
        return NTLSTATUS_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 NTLSTATUS_SUCCESS;
}

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

    uintptr_t procFlags = KeAcquireSpinlock(&process->VadLock);
    MM_VAD* vad = KiFindVadByBaseLocked(process, base);
    if (!vad) {
        KeReleaseSpinlock(&process->VadLock, procFlags);
        return NTLSTATUS_NOT_FOUND;
    }
    if (vad->LockCount) {
        KeReleaseSpinlock(&process->VadLock, procFlags);
        return NTLSTATUS_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 NTLSTATUS_SUCCESS;
}

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

    while (true) {
        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 NTLSTATUS_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 NTLSTATUS_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 NTLSTATUS_SUCCESS;
    }
    KeReleaseSpinlock(&process->VadLock, procFlags);
    return NTLSTATUS_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 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) {
        MmPoolDebugDescribeAddress(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 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;
}