// Coherent file page cache.
//
// The cache owns one resident page per file/offset. Reads, buffered writes, and
// mapped-file faults converge on this page so writeback has a single source of
// truth. Dirty pages are flushed explicitly, by close, or by the background
// writeback worker under pressure.

#include <stdbool.h>
#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <intrin.h>
#include <kernel/mm_file_cache.h>
#include <kernel/io_file.h>
#include <kernel/io_volume.h>
#include <kernel/mm_pfn.h>
#include <kernel/mm_virt.h>
#include <kernel/ob_object.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_irql.h>
#include <kernel/ke_semaphore.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_sched.h>

#define MM_FILE_CACHE_BUCKETS 64u
#define MM_FILE_CACHE_SOFT_MAX_PAGES 512u
#define MM_FILE_CACHE_DIRTY_HIGH_WATER_PAGES 128u
#define MM_FILE_CACHE_DIRTY_HARD_LIMIT_PAGES 192u
#define MM_FILE_CACHE_WRITEBACK_BATCH_PAGES 16u

typedef struct MM_FILE_CACHE_ENTRY {
    struct MM_FILE_CACHE_ENTRY* Next;
    FILE_OBJECT*      File;
    uint64_t          Offset;
    PHYSICAL_ADDRESS  Frame;
    uint32_t          ValidBytes;
    uint8_t           Referenced;
    uint8_t           Standby;
    uint8_t           Loading;
    bool              Invalidated;
    uint8_t           Dirty;
    uint8_t           Writeback;
    uint64_t          DirtySequence;
} MM_FILE_CACHE_ENTRY;

static SPINLOCK gFileCacheLock = SPINLOCK_INIT;
static MM_FILE_CACHE_ENTRY* gFileCacheBuckets[MM_FILE_CACHE_BUCKETS];
static uint64_t gFileCacheResidentPages;
static uint64_t gFileCacheActivePages;
static uint64_t gFileCacheStandbyPages;
static uint64_t gFileCacheDirtyPages;
static uint64_t gFileCacheMaxDirtyPages;
static uint64_t gFileCacheMaxResidentPages;
static MM_FILE_CACHE_STATS gFileCacheStats;


static bool KiFileCacheUnlinkEntryLocked(MM_FILE_CACHE_ENTRY* entry);

static SPINLOCK gWritebackLock = SPINLOCK_INIT;
static SEMAPHORE* gWritebackSemaphore;
static int gWritebackStartState;

static uint32_t KiFileCacheHash(FILE_OBJECT* file, uint64_t offset) {
    uintptr_t f = (uintptr_t)file >> 4;
    uint64_t p = offset >> 12;
    return (uint32_t)((f ^ p ^ (p >> 11)) & (MM_FILE_CACHE_BUCKETS - 1u));
}

static MM_FILE_CACHE_ENTRY* KiFileCacheFindLocked(FILE_OBJECT* file, uint64_t offset, uint32_t bucket) {
    for (MM_FILE_CACHE_ENTRY* e = gFileCacheBuckets[bucket]; e; e = e->Next) {
        if (e->File == file && e->Offset == offset)
            return e;
    }
    return (MM_FILE_CACHE_ENTRY*)0;
}

static void KiFileCacheBumpResidentLocked(int delta, int standby) {
    if (delta > 0) {
        gFileCacheResidentPages += (uint64_t)delta;
        if (standby) gFileCacheStandbyPages += (uint64_t)delta;
        else gFileCacheActivePages += (uint64_t)delta;
        if (gFileCacheResidentPages > gFileCacheMaxResidentPages)
            gFileCacheMaxResidentPages = gFileCacheResidentPages;
    } else if (delta < 0) {
        uint64_t dec = (uint64_t)(-delta);
        gFileCacheResidentPages = gFileCacheResidentPages > dec ? gFileCacheResidentPages - dec : 0;
        if (standby) gFileCacheStandbyPages = gFileCacheStandbyPages > dec ? gFileCacheStandbyPages - dec : 0;
        else gFileCacheActivePages = gFileCacheActivePages > dec ? gFileCacheActivePages - dec : 0;
    }
}

// Updates dirty-page accounting while the cache lock is held.
static void KiFileCacheSetDirtyLocked(MM_FILE_CACHE_ENTRY* e, int dirty) {
    if (!e)
        return;
    if (dirty) {
        if (!e->Dirty) {
            e->Dirty = 1;
            gFileCacheDirtyPages++;
            gFileCacheStats.DirtyPages = gFileCacheDirtyPages;
            if (gFileCacheDirtyPages > gFileCacheMaxDirtyPages) {
                gFileCacheMaxDirtyPages = gFileCacheDirtyPages;
                gFileCacheStats.MaxDirtyPages = gFileCacheMaxDirtyPages;
            }
        }
        e->DirtySequence++;
        gFileCacheStats.DirtyMarks++;
    } else {
        if (e->Dirty) {
            e->Dirty = 0;
            if (gFileCacheDirtyPages)
                gFileCacheDirtyPages--;
            gFileCacheStats.DirtyPages = gFileCacheDirtyPages;
        }
    }
}

static void KiFileCacheMoveEntryToStandbyLocked(MM_FILE_CACHE_ENTRY* e) {
    if (!e || e->Standby)
        return;
    e->Standby = 1;
    if (gFileCacheActivePages) gFileCacheActivePages--;
    gFileCacheStandbyPages++;
    gFileCacheStats.StandbyTransitions++;
}

static void KiFileCacheActivateEntryLocked(MM_FILE_CACHE_ENTRY* e) {
    if (!e || !e->Standby)
        return;
    e->Standby = 0;
    if (gFileCacheStandbyPages) gFileCacheStandbyPages--;
    gFileCacheActivePages++;
    gFileCacheStats.StandbyHits++;
}

static void KiFileCacheFreeEntry(MM_FILE_CACHE_ENTRY* e) {
    if (!e) return;
    if (e->File) ObDereferenceObject(&e->File->Header);
    if (e->Frame) (void)MmPfnRelease(e->Frame);
    MmFreePool(e);
}

static void KiFileCacheRemoveEntryLocked(uint32_t bucket, MM_FILE_CACHE_ENTRY* prev, MM_FILE_CACHE_ENTRY* e) {
    if (!e) return;
    if (prev) prev->Next = e->Next;
    else gFileCacheBuckets[bucket] = e->Next;
    if (e->Frame)
        KiFileCacheBumpResidentLocked(-1, e->Standby);
    if (e->Dirty)
        KiFileCacheSetDirtyLocked(e, 0);
}

static bool KiFileCacheIsReclaimCandidateLocked(MM_FILE_CACHE_ENTRY* e, FILE_OBJECT* skipFile, uint64_t skipOffset) {
    if (!e) return false;
    if (skipFile && e->File == skipFile && e->Offset == skipOffset)
        return false;
    if (e->Loading || e->Invalidated || !e->Frame || e->Dirty || e->Writeback)
        return false;

    MM_PFN* pfn = MmPfnGet(e->Frame);
    return pfn && pfn->PinCount == 0 && pfn->MappedCount == 0 && pfn->RefCount == 1;
}

static KSTATUS KiFileCacheDemoteOnePage(FILE_OBJECT* skipFile, uint64_t skipOffset) {
    // A cold clean page first moves to PFN standby so the cache can reactivate it
    // without disk I/O if it becomes hot again before final pressure reclaim.
    for (int pass = 0; pass < 2; ++pass) {
        for (uint32_t b = 0; b < MM_FILE_CACHE_BUCKETS; ++b) {
            MM_FILE_CACHE_ENTRY* demote = (MM_FILE_CACHE_ENTRY*)0;
            uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
            for (MM_FILE_CACHE_ENTRY* e = gFileCacheBuckets[b]; e; e = e->Next) {
                if (e->Standby || !KiFileCacheIsReclaimCandidateLocked(e, skipFile, skipOffset))
                    continue;

                if (e->Referenced && pass == 0) {
                    e->Referenced = 0;
                    continue;
                }

                demote = e;
                break;
            }

            if (demote) {
                KSTATUS s = MmPfnMoveToStandby(demote->Frame);
                if (KSUCCESS(s)) {
                    KiFileCacheMoveEntryToStandbyLocked(demote);
                    KeReleaseSpinlock(&gFileCacheLock, flags);
                    return NTLSTATUS_SUCCESS;
                }
            }
            KeReleaseSpinlock(&gFileCacheLock, flags);
        }
    }
    return NTLSTATUS_NOT_FOUND;
}

static KSTATUS KiFileCacheDropOneStandbyPage(FILE_OBJECT* skipFile, uint64_t skipOffset) {
    for (uint32_t b = 0; b < MM_FILE_CACHE_BUCKETS; ++b) {
        MM_FILE_CACHE_ENTRY* victim = (MM_FILE_CACHE_ENTRY*)0;
        uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
        MM_FILE_CACHE_ENTRY* prev = (MM_FILE_CACHE_ENTRY*)0;
        for (MM_FILE_CACHE_ENTRY* e = gFileCacheBuckets[b]; e; prev = e, e = e->Next) {
            if (!e->Standby || !KiFileCacheIsReclaimCandidateLocked(e, skipFile, skipOffset))
                continue;

            victim = e;
            KiFileCacheRemoveEntryLocked(b, prev, e);
            gFileCacheStats.Drops++;
            break;
        }
        KeReleaseSpinlock(&gFileCacheLock, flags);
        if (victim) {
            KiFileCacheFreeEntry(victim);
            return NTLSTATUS_SUCCESS;
        }
    }
    return NTLSTATUS_NOT_FOUND;
}

static KSTATUS KiFileCacheReclaimOnePage(FILE_OBJECT* skipFile, uint64_t skipOffset) {
    KSTATUS s = KiFileCacheDropOneStandbyPage(skipFile, skipOffset);
    if (KSUCCESS(s))
        return NTLSTATUS_SUCCESS;

    s = KiFileCacheDemoteOnePage(skipFile, skipOffset);
    if (KSUCCESS(s)) {
        s = KiFileCacheDropOneStandbyPage(skipFile, skipOffset);
        if (KSUCCESS(s))
            return NTLSTATUS_SUCCESS;
    }

    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    gFileCacheStats.DropFailures++;
    KeReleaseSpinlock(&gFileCacheLock, flags);
    return NTLSTATUS_NOT_FOUND;
}

// Adds to one file-cache statistic under the cache lock.
static void KiFileCacheAddStat(uint64_t* field, uint64_t value) {
    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    *field += value;
    KeReleaseSpinlock(&gFileCacheLock, flags);
}

// Wakes the writeback worker after dirty pressure crosses the high watermark.
static void KiFileCacheSignalWriteback(void) {
    if (gWritebackSemaphore)
        KeSignalSemaphore(gWritebackSemaphore);
}

// Finds the lowest dirty page for a file, or globally when file is NULL.
static MM_FILE_CACHE_ENTRY* KiFileCacheFindLowestDirtyLocked(FILE_OBJECT* file) {
    MM_FILE_CACHE_ENTRY* best = (MM_FILE_CACHE_ENTRY*)0;
    for (uint32_t b = 0; b < MM_FILE_CACHE_BUCKETS; ++b) {
        for (MM_FILE_CACHE_ENTRY* e = gFileCacheBuckets[b]; e; e = e->Next) {
            if (!e->Dirty || e->Writeback || e->Loading || e->Invalidated || !e->Frame)
                continue;
            if (file && e->File != file)
                continue;
            if (!best || e->Offset < best->Offset)
                best = e;
        }
    }
    return best;
}

// Flushes the next dirty cache page for file, or globally when file is NULL.
static KSTATUS KiFileCacheFlushOneDirtyPage(FILE_OBJECT* file, int* didWork) {
    if (didWork) *didWork = 0;
    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    MM_FILE_CACHE_ENTRY* e = KiFileCacheFindLowestDirtyLocked(file);
    if (!e) {
        KeReleaseSpinlock(&gFileCacheLock, flags);
        return NTLSTATUS_NOT_FOUND;
    }

    KSTATUS rs = MmPfnRetain(e->Frame);
    if (!KSUCCESS(rs)) {
        gFileCacheStats.WritebackFailures++;
        KeReleaseSpinlock(&gFileCacheLock, flags);
        return rs;
    }

    FILE_OBJECT* target = e->File;
    ObReferenceObject(&target->Header);
    PHYSICAL_ADDRESS frame = e->Frame;
    uint64_t offset = e->Offset;
    uint32_t validBytes = e->ValidBytes ? e->ValidBytes : PAGE_SIZE;
    uint64_t sequence = e->DirtySequence;
    e->Writeback = 1;
    gFileCacheStats.WritebackRequests++;
    KeReleaseSpinlock(&gFileCacheLock, flags);

    void* page = MmPfnToVirtual(frame);
    size_t written = 0;
    KSTATUS s = NTLSTATUS_FAILURE;
    if (page)
        s = IoVolumeWrite(target->Volume, target->FileCtx, offset, page, validBytes, &written);
    if (KSUCCESS(s) && written != validBytes)
        s = NTLSTATUS_FAILURE;

    MM_FILE_CACHE_ENTRY* invalidatedVictim = (MM_FILE_CACHE_ENTRY*)0;
    flags = KeAcquireSpinlock(&gFileCacheLock);
    uint32_t bucket = KiFileCacheHash(target, offset);
    MM_FILE_CACHE_ENTRY* cur = KiFileCacheFindLocked(target, offset, bucket);
    if (cur && cur->Frame == frame) {
        cur->Writeback = 0;
        if (cur->Invalidated) {
            if (cur->Dirty)
                KiFileCacheSetDirtyLocked(cur, 0);
            if (KiFileCacheUnlinkEntryLocked(cur))
                invalidatedVictim = cur;
        } else if (KSUCCESS(s) && cur->Dirty && cur->DirtySequence == sequence) {
            KiFileCacheSetDirtyLocked(cur, 0);
            gFileCacheStats.WritebackPages++;
            gFileCacheStats.WritebackBytes += validBytes;
        } else if (!KSUCCESS(s)) {
            gFileCacheStats.WritebackFailures++;
        }
    }
    KeReleaseSpinlock(&gFileCacheLock, flags);

    (void)MmPfnRelease(frame);
    if (invalidatedVictim)
        KiFileCacheFreeEntry(invalidatedVictim);
    ObDereferenceObject(&target->Header);
    if (didWork) *didWork = 1;
    return s;
}

// Background writeback worker drains dirty cache pressure best-effort.
static void KiFileCacheWritebackWorker(void* arg) {
    (void)arg;
    KeAssertIrql(PASSIVE_LEVEL);

    while (true) {
        KSTATUS ws = KeWaitSemaphore(gWritebackSemaphore, KE_WAIT_INFINITE);
        if (!KSUCCESS(ws)) {
            KeYield();
            continue;
        }

        for (uint32_t i = 0; i < MM_FILE_CACHE_WRITEBACK_BATCH_PAGES; ++i) {
            uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
            int overHigh = gFileCacheDirtyPages >= MM_FILE_CACHE_DIRTY_HIGH_WATER_PAGES;
            KeReleaseSpinlock(&gFileCacheLock, flags);
            if (!overHigh)
                break;

            int didWork = 0;
            KSTATUS s = KiFileCacheFlushOneDirtyPage((FILE_OBJECT*)0, &didWork);
            if (!didWork || !KSUCCESS(s))
                break;
        }
    }
}

// Starts the writeback worker once.
static KSTATUS KiFileCacheEnsureWritebackStarted(void) {
    uintptr_t flags = KeAcquireSpinlock(&gWritebackLock);
    if (gWritebackStartState == 2) {
        KeReleaseSpinlock(&gWritebackLock, flags);
        return NTLSTATUS_SUCCESS;
    }
    if (gWritebackStartState == 1) {
        KeReleaseSpinlock(&gWritebackLock, flags);
        while (gWritebackStartState == 1)
            KeYield();
        return gWritebackStartState == 2 ? NTLSTATUS_SUCCESS : NTLSTATUS_FAILURE;
    }

    gWritebackStartState = 1;
    KeReleaseSpinlock(&gWritebackLock, flags);

    KSTATUS s = KeSemaphoreCreate(0, 0x7fffffff, &gWritebackSemaphore);
    if (KSUCCESS(s)) {
        THREAD* worker = (THREAD*)0;
        const KE_THREAD_CREATE_PARAMETERS parameters =
            KE_THREAD_CREATE_NAMED_NORMAL(PriorityNormal, "mm.writeback");
        s = KeCreateKernelThread((uintptr_t)KiFileCacheWritebackWorker,
                                 0,
                                 &parameters,
                                 &worker);
        (void)worker;
    }

    flags = KeAcquireSpinlock(&gWritebackLock);
    gWritebackStartState = KSUCCESS(s) ? 2 : 0;
    KeReleaseSpinlock(&gWritebackLock, flags);
    return s;
}

static KSTATUS KiFileCacheTryHit(FILE_OBJECT* file, uint64_t aligned, MM_FILE_CACHE_PAGE* outPage, int waitForLoading) {
    if (outPage)
        memset(outPage, 0, sizeof(*outPage));

    uint32_t bucket = KiFileCacheHash(file, aligned);
    int counted = 0;
    while (true) {
        uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
        if (!counted) {
            gFileCacheStats.Lookups++;
            counted = 1;
        }

        MM_FILE_CACHE_ENTRY* hit = KiFileCacheFindLocked(file, aligned, bucket);
        if (!hit) {
            gFileCacheStats.Misses++;
            KeReleaseSpinlock(&gFileCacheLock, flags);
            return NTLSTATUS_NOT_FOUND;
        }

        if (hit->Loading) {
            if (!waitForLoading) {
                gFileCacheStats.Misses++;
                KeReleaseSpinlock(&gFileCacheLock, flags);
                return NTLSTATUS_PENDING;
            }

            gFileCacheStats.InflightWaits++;
            KeReleaseSpinlock(&gFileCacheLock, flags);
            KeYield();
            continue;
        }

        if (hit->Invalidated || !hit->Frame) {
            gFileCacheStats.Misses++;
            KeReleaseSpinlock(&gFileCacheLock, flags);
            return NTLSTATUS_NOT_FOUND;
        }

        if (hit->Standby) {
            KSTATUS as = MmPfnActivate(hit->Frame);
            if (KSUCCESS(as))
                KiFileCacheActivateEntryLocked(hit);
        }

        KSTATUS rs = NTLSTATUS_SUCCESS;
        if (outPage)
            rs = MmPfnRetain(hit->Frame);
        if (KSUCCESS(rs)) {
            hit->Referenced = 1;
            gFileCacheStats.Hits++;
            if (outPage) {
                outPage->File = file;
                outPage->Offset = aligned;
                outPage->Frame = hit->Frame;
                outPage->ValidBytes = hit->ValidBytes;
            }
            KeReleaseSpinlock(&gFileCacheLock, flags);
            return NTLSTATUS_SUCCESS;
        }

        // A retained cache entry should never point at an invalid frame, but do not
        // return corrupt data if the PFN layer rejects it.
        gFileCacheStats.Misses++;
        KeReleaseSpinlock(&gFileCacheLock, flags);
        return rs;
    }
}

static void KiFileCacheKeepBounded(FILE_OBJECT* skipFile, uint64_t skipOffset) {
    while (true) {
        uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
        int over = gFileCacheResidentPages > MM_FILE_CACHE_SOFT_MAX_PAGES;
        KeReleaseSpinlock(&gFileCacheLock, flags);
        if (!over) break;
        if (!KSUCCESS(KiFileCacheReclaimOnePage(skipFile, skipOffset))) break;
    }
}

static MM_FILE_CACHE_ENTRY* KiFileCacheAllocateLoadingEntry(FILE_OBJECT* file, uint64_t aligned) {
    MM_FILE_CACHE_ENTRY* entry = (MM_FILE_CACHE_ENTRY*)MmAllocatePoolTag(sizeof(MM_FILE_CACHE_ENTRY), MM_POOL_TAG('p', 'c', 'a', 'c'));
    if (!entry)
        return (MM_FILE_CACHE_ENTRY*)0;

    memset(entry, 0, sizeof(*entry));
    entry->File = file;
    entry->Offset = aligned;
    entry->Referenced = 1;
    entry->Loading = 1;
    ObReferenceObject(&file->Header);
    return entry;
}

static void KiFileCacheFreeLoadingEntry(MM_FILE_CACHE_ENTRY* entry) {
    if (!entry)
        return;
    entry->Frame = 0;
    entry->Loading = 0;
    KiFileCacheFreeEntry(entry);
}

static bool KiFileCacheUnlinkEntryLocked(MM_FILE_CACHE_ENTRY* entry) {
    if (!entry)
        return false;

    uint32_t bucket = KiFileCacheHash(entry->File, entry->Offset);
    MM_FILE_CACHE_ENTRY* prev = (MM_FILE_CACHE_ENTRY*)0;
    for (MM_FILE_CACHE_ENTRY* e = gFileCacheBuckets[bucket]; e; prev = e, e = e->Next) {
        if (e == entry) {
            KiFileCacheRemoveEntryLocked(bucket, prev, e);
            return true;
        }
    }
    return false;
}

static KSTATUS KiFileCacheReserveLoadingPage(FILE_OBJECT* file, uint64_t aligned, MM_FILE_CACHE_ENTRY** outEntry) {
    if (outEntry)
        *outEntry = (MM_FILE_CACHE_ENTRY*)0;

    MM_FILE_CACHE_ENTRY* entry = KiFileCacheAllocateLoadingEntry(file, aligned);
    if (!entry)
        return NTLSTATUS_NO_MEMORY;

    uint32_t bucket = KiFileCacheHash(file, aligned);
    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    if (KiFileCacheFindLocked(file, aligned, bucket)) {
        KeReleaseSpinlock(&gFileCacheLock, flags);
        KiFileCacheFreeLoadingEntry(entry);
        return NTLSTATUS_ALREADY_EXISTS;
    }

    entry->Next = gFileCacheBuckets[bucket];
    gFileCacheBuckets[bucket] = entry;
    KeReleaseSpinlock(&gFileCacheLock, flags);

    if (outEntry)
        *outEntry = entry;
    return NTLSTATUS_SUCCESS;
}

static void KiFileCacheCancelLoadingPage(MM_FILE_CACHE_ENTRY* entry) {
    if (!entry)
        return;

    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    bool unlinked = KiFileCacheUnlinkEntryLocked(entry);
    KeReleaseSpinlock(&gFileCacheLock, flags);

    if (unlinked)
        KiFileCacheFreeLoadingEntry(entry);
}

static KSTATUS KiFileCacheCompleteLoadingPage(MM_FILE_CACHE_ENTRY* entry,
                                              PHYSICAL_ADDRESS frame,
                                              uint32_t validBytes,
                                              MM_FILE_CACHE_PAGE* requestedPage,
                                              uint64_t* publishedCount) {
    if (!entry || !frame)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    if (entry->Invalidated) {
        (void)KiFileCacheUnlinkEntryLocked(entry);
        KeReleaseSpinlock(&gFileCacheLock, flags);
        (void)MmPfnRelease(frame);
        KiFileCacheFreeLoadingEntry(entry);
        return NTLSTATUS_NOT_FOUND;
    }

    entry->Frame = frame;
    entry->ValidBytes = validBytes;
    entry->Referenced = 1;
    entry->Standby = 0;
    entry->Loading = 0;
    gFileCacheStats.Loads++;
    KiFileCacheBumpResidentLocked(1, 0);
    if (publishedCount)
        (*publishedCount)++;

    if (requestedPage) {
        KSTATUS rs = MmPfnRetain(frame);
        if (!KSUCCESS(rs)) {
            (void)KiFileCacheUnlinkEntryLocked(entry);
            gFileCacheStats.LoadFailures++;
            KeReleaseSpinlock(&gFileCacheLock, flags);
            KiFileCacheFreeEntry(entry);
            return rs;
        }
        requestedPage->File = entry->File;
        requestedPage->Offset = entry->Offset;
        requestedPage->Frame = frame;
        requestedPage->ValidBytes = validBytes;
    }

    KeReleaseSpinlock(&gFileCacheLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Loads exactly one demanded file page into the coherent cache.
static KSTATUS KiFileCacheLoadPage(FILE_OBJECT* file,
                                   uint64_t aligned,
                                   MM_FILE_CACHE_PAGE* outPage) {
    if (!file || !outPage)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t fileFlags = KeAcquireSpinlock(&file->Lock);
    uint64_t fileSize = file->Size;
    KeReleaseSpinlock(&file->Lock, fileFlags);
    if (aligned >= fileSize)
        return NTLSTATUS_SUCCESS;

    MM_FILE_CACHE_ENTRY* entry = (MM_FILE_CACHE_ENTRY*)0;
    KSTATUS s = KiFileCacheReserveLoadingPage(file, aligned, &entry);
    if (s == NTLSTATUS_ALREADY_EXISTS)
        return KiFileCacheTryHit(file, aligned, outPage, 1);
    if (!KSUCCESS(s))
        return s;

    PHYSICAL_ADDRESS frame = MmPfnAllocate(MmPfnUseCleanFileCache);
    if (!frame) {
        KiFileCacheCancelLoadingPage(entry);
        KiFileCacheAddStat(&gFileCacheStats.LoadFailures, 1);
        return NTLSTATUS_NO_MEMORY;
    }

    void* page = MmPfnToVirtual(frame);
    if (!page) {
        (void)MmPfnRelease(frame);
        KiFileCacheCancelLoadingPage(entry);
        KiFileCacheAddStat(&gFileCacheStats.LoadFailures, 1);
        return NTLSTATUS_FAILURE;
    }
    memset(page, 0, PAGE_SIZE);

    size_t toRead = PAGE_SIZE;
    uint64_t remaining = fileSize - aligned;
    if (remaining < PAGE_SIZE)
        toRead = (size_t)remaining;

    size_t got = 0;
    s = IoVolumeRead(file->Volume, file->FileCtx, aligned, page, toRead, &got);
    if (!KSUCCESS(s)) {
        (void)MmPfnRelease(frame);
        KiFileCacheCancelLoadingPage(entry);
        KiFileCacheAddStat(&gFileCacheStats.LoadFailures, 1);
        return s;
    }
    if (got > toRead)
        got = toRead;
    if (got < PAGE_SIZE)
        memset((uint8_t*)page + got, 0, PAGE_SIZE - got);

    s = KiFileCacheCompleteLoadingPage(entry, frame, (uint32_t)got, outPage, (uint64_t*)0);
    if (!KSUCCESS(s))
        return s;

    KiFileCacheKeepBounded(file, aligned);
    return NTLSTATUS_SUCCESS;
}

// Returns a retained cache page suitable for in-place buffered writes.
static KSTATUS KiFileCacheGetWritablePage(FILE_OBJECT* file,
                                          uint64_t aligned,
                                          uint64_t oldFileSize,
                                          bool needExistingBytes,
                                          MM_FILE_CACHE_PAGE* outPage) {
    if (!file || !outPage)
        return NTLSTATUS_INVALID_PARAMETER;
    memset(outPage, 0, sizeof(*outPage));

    KSTATUS s = KiFileCacheTryHit(file, aligned, outPage, 1);
    if (KSUCCESS(s))
        return NTLSTATUS_SUCCESS;

    if (needExistingBytes && aligned < oldFileSize) {
        s = KiFileCacheLoadPage(file, aligned, outPage);
        if (KSUCCESS(s))
            return NTLSTATUS_SUCCESS;
        if (s != NTLSTATUS_NOT_FOUND)
            return s;
    }

    MM_FILE_CACHE_ENTRY* entry = (MM_FILE_CACHE_ENTRY*)0;
    s = KiFileCacheReserveLoadingPage(file, aligned, &entry);
    if (s == NTLSTATUS_ALREADY_EXISTS)
        return KiFileCacheTryHit(file, aligned, outPage, 1);
    if (!KSUCCESS(s))
        return s;

    PHYSICAL_ADDRESS frame = MmPfnAllocate(MmPfnUseCleanFileCache);
    if (!frame) {
        KiFileCacheCancelLoadingPage(entry);
        return NTLSTATUS_NO_MEMORY;
    }

    void* page = MmPfnToVirtual(frame);
    if (!page) {
        (void)MmPfnRelease(frame);
        KiFileCacheCancelLoadingPage(entry);
        return NTLSTATUS_FAILURE;
    }
    memset(page, 0, PAGE_SIZE);

    uint32_t validBytes = 0;
    if (oldFileSize > aligned) {
        uint64_t oldBytes = oldFileSize - aligned;
        validBytes = oldBytes < PAGE_SIZE ? (uint32_t)oldBytes : PAGE_SIZE;
    }

    s = KiFileCacheCompleteLoadingPage(entry, frame, validBytes, outPage, (uint64_t*)0);
    if (!KSUCCESS(s))
        return s;
    return NTLSTATUS_SUCCESS;
}

// Marks one cache page dirty and updates its valid byte range.
static void KiFileCacheMarkPageDirty(FILE_OBJECT* file, uint64_t aligned, uint32_t validBytes) {
    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    uint32_t bucket = KiFileCacheHash(file, aligned);
    MM_FILE_CACHE_ENTRY* e = KiFileCacheFindLocked(file, aligned, bucket);
    if (e) {
        if (validBytes > e->ValidBytes)
            e->ValidBytes = validBytes;
        if (e->Standby)
            KiFileCacheActivateEntryLocked(e);
        KiFileCacheSetDirtyLocked(e, 1);
    }
    uint64_t dirty = gFileCacheDirtyPages;
    KeReleaseSpinlock(&gFileCacheLock, flags);

    if (dirty >= MM_FILE_CACHE_DIRTY_HIGH_WATER_PAGES) {
        if (KSUCCESS(KiFileCacheEnsureWritebackStarted()))
            KiFileCacheSignalWriteback();
    }
}

KSTATUS MmFileCacheGetWritablePage(FILE_OBJECT* file, uint64_t offset,
                                   uint64_t fileSizeBefore, bool needExistingBytes,
                                   MM_FILE_CACHE_PAGE* outPage) {
    if (!outPage)
        return NTLSTATUS_INVALID_PARAMETER;
    uint64_t aligned = offset & ~(uint64_t)(PAGE_SIZE - 1);
    return KiFileCacheGetWritablePage(file, aligned, fileSizeBefore, needExistingBytes, outPage);
}

void MmFileCacheMarkDirty(FILE_OBJECT* file, uint64_t offset, uint32_t validBytes) {
    uint64_t aligned = offset & ~(uint64_t)(PAGE_SIZE - 1);
    if (validBytes > PAGE_SIZE)
        validBytes = PAGE_SIZE;
    if (validBytes == 0)
        validBytes = PAGE_SIZE;
    KiFileCacheMarkPageDirty(file, aligned, validBytes);

    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    gFileCacheStats.MappedDirtySyncs++;
    gFileCacheStats.MappedDirtyPages++;
    KeReleaseSpinlock(&gFileCacheLock, flags);
}

KSTATUS MmFileCacheWrite(FILE_OBJECT* file, uint64_t offset, const void* buf,
                         size_t size, uint64_t fileSizeBefore, size_t* writtenOut) {
    if (writtenOut) *writtenOut = 0;
    if (!file || (!buf && size) || !writtenOut)
        return NTLSTATUS_INVALID_PARAMETER;
    if (!KeCanWait())
        return NTLSTATUS_INVALID_IRQL;
    if (!IoVolumeIsMounted(file->Volume))
        return NTLSTATUS_OBJECT_DEAD;
    if (size == 0)
        return NTLSTATUS_SUCCESS;

    size_t total = 0;
    const uint8_t* src = (const uint8_t*)buf;
    while (total < size) {
        uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
        uint64_t dirty = gFileCacheDirtyPages;
        KeReleaseSpinlock(&gFileCacheLock, flags);
        if (dirty >= MM_FILE_CACHE_DIRTY_HARD_LIMIT_PAGES) {
            KSTATUS fs = MmFileCacheFlushFile(file);
            if (!KSUCCESS(fs) && total == 0)
                return fs;
        }

        uint64_t absolute = offset + total;
        if (absolute < offset)
            break;
        uint64_t aligned = absolute & ~(uint64_t)(PAGE_SIZE - 1);
        size_t inPage = (size_t)(absolute - aligned);
        size_t chunk = PAGE_SIZE - inPage;
        if (chunk > size - total)
            chunk = size - total;

        int coversFullPage = (inPage == 0 && chunk == PAGE_SIZE);
        int needExisting = !coversFullPage;
        MM_FILE_CACHE_PAGE page;
        KSTATUS s = KiFileCacheGetWritablePage(file, aligned, fileSizeBefore, needExisting, &page);
        if (!KSUCCESS(s))
            return total ? NTLSTATUS_SUCCESS : s;

        void* dst = MmPfnToVirtual(page.Frame);
        if (!dst) {
            MmFileCacheReleasePage(&page);
            return total ? NTLSTATUS_SUCCESS : NTLSTATUS_FAILURE;
        }

        memcpy((uint8_t*)dst + inPage, src + total, chunk);
        uint32_t valid = (uint32_t)(inPage + chunk);
        uint64_t pageEnd = aligned + valid;
        uint64_t newEnd = offset + total + chunk;
        if (pageEnd > newEnd)
            valid = (uint32_t)(newEnd - aligned);
        KiFileCacheMarkPageDirty(file, aligned, valid);
        MmFileCacheReleasePage(&page);

        total += chunk;
    }

    *writtenOut = total;
    return NTLSTATUS_SUCCESS;
}

KSTATUS MmFileCacheFlushAll(void) {
    if (!KeCanWait()) return NTLSTATUS_INVALID_IRQL;
    for (uint32_t i = 0; i < 65536u; ++i) {
        int didWork = 0;
        KSTATUS st = KiFileCacheFlushOneDirtyPage((FILE_OBJECT*)0, &didWork);
        if (st == NTLSTATUS_NOT_FOUND && !didWork) {
            uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
            bool clean = gFileCacheDirtyPages == 0;
            KeReleaseSpinlock(&gFileCacheLock, flags);
            if (clean) return NTLSTATUS_SUCCESS;
            KeYield();
            continue;
        }
        if (!KSUCCESS(st)) return st;
        if (!didWork) return NTLSTATUS_SUCCESS;
    }
    return NTLSTATUS_LIMIT_REACHED;
}

KSTATUS MmFileCacheFlushFile(FILE_OBJECT* file) {
    if (!file)
        return NTLSTATUS_INVALID_PARAMETER;
    if (!KeCanWait())
        return NTLSTATUS_INVALID_IRQL;

    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    gFileCacheStats.FlushFileCalls++;
    KeReleaseSpinlock(&gFileCacheLock, flags);

    KSTATUS firstError = NTLSTATUS_SUCCESS;
    while (true) {
        int didWork = 0;
        KSTATUS s = KiFileCacheFlushOneDirtyPage(file, &didWork);
        if (!didWork)
            break;
        if (!KSUCCESS(s) && KSUCCESS(firstError))
            firstError = s;
    }

    if (!KSUCCESS(firstError)) {
        flags = KeAcquireSpinlock(&gFileCacheLock);
        gFileCacheStats.FlushFileFailures++;
        KeReleaseSpinlock(&gFileCacheLock, flags);
    }
    return firstError;
}

KSTATUS MmFileCacheGetPage(FILE_OBJECT* file, uint64_t offset, MM_FILE_CACHE_PAGE* outPage) {
    if (!outPage) return NTLSTATUS_INVALID_PARAMETER;
    memset(outPage, 0, sizeof(*outPage));
    if (!file) return NTLSTATUS_INVALID_PARAMETER;
    if (!KeCanWait()) return NTLSTATUS_INVALID_IRQL;

    uint64_t aligned = offset & ~(uint64_t)(PAGE_SIZE - 1);
    KSTATUS s = KiFileCacheTryHit(file, aligned, outPage, 1);
    if (KSUCCESS(s))
        return NTLSTATUS_SUCCESS;
    return KiFileCacheLoadPage(file, aligned, outPage);
}

void MmFileCacheReleasePage(MM_FILE_CACHE_PAGE* page) {
    if (!page || !page->Frame) return;
    (void)MmPfnRelease(page->Frame);
    memset(page, 0, sizeof(*page));
}

void MmFileCacheInvalidate(FILE_OBJECT* file, uint64_t offset, size_t size) {
    if (!file) return;
    if (size == 0) return;
    uint64_t start = offset & ~(uint64_t)(PAGE_SIZE - 1);
    uint64_t end = offset + size;
    if (end < offset || size == (size_t)-1) {
        end = UINT64_MAX;
    } else {
        end = (end + PAGE_SIZE - 1) & ~(uint64_t)(PAGE_SIZE - 1);
    }

    for (uint32_t b = 0; b < MM_FILE_CACHE_BUCKETS; ++b) {
        while (true) {
            MM_FILE_CACHE_ENTRY* victim = (MM_FILE_CACHE_ENTRY*)0;
            uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
            MM_FILE_CACHE_ENTRY* prev = (MM_FILE_CACHE_ENTRY*)0;
            for (MM_FILE_CACHE_ENTRY* e = gFileCacheBuckets[b]; e; prev = e, e = e->Next) {
                uint64_t pageEnd = e->Offset + PAGE_SIZE;
                if (e->File == file && e->Offset < end && pageEnd > start) {
                    if (e->Loading) {
                        if (!e->Invalidated) {
                            e->Invalidated = 1;
                            gFileCacheStats.Invalidations++;
                        }
                        continue;
                    }
                    if (e->Dirty || e->Writeback) {
                        e->Invalidated = 1;
                        continue;
                    }
                    victim = e;
                    KiFileCacheRemoveEntryLocked(b, prev, e);
                    gFileCacheStats.Invalidations++;
                    break;
                }
            }
            KeReleaseSpinlock(&gFileCacheLock, flags);
            if (!victim) break;
            KiFileCacheFreeEntry(victim);
        }
    }
}

void MmFileCacheInvalidateAll(FILE_OBJECT* file) {
    if (!file) return;
    for (uint32_t b = 0; b < MM_FILE_CACHE_BUCKETS; ++b) {
        while (true) {
            MM_FILE_CACHE_ENTRY* victim = (MM_FILE_CACHE_ENTRY*)0;
            uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
            MM_FILE_CACHE_ENTRY* prev = (MM_FILE_CACHE_ENTRY*)0;
            for (MM_FILE_CACHE_ENTRY* e = gFileCacheBuckets[b]; e; prev = e, e = e->Next) {
                if (e->File == file) {
                    if (e->Loading) {
                        if (!e->Invalidated) {
                            e->Invalidated = 1;
                            gFileCacheStats.Invalidations++;
                        }
                        continue;
                    }
                    if (e->Dirty || e->Writeback) {
                        e->Invalidated = 1;
                        continue;
                    }
                    victim = e;
                    KiFileCacheRemoveEntryLocked(b, prev, e);
                    gFileCacheStats.Invalidations++;
                    break;
                }
            }
            KeReleaseSpinlock(&gFileCacheLock, flags);
            if (!victim) break;
            KiFileCacheFreeEntry(victim);
        }
    }
}

// Discards all cache state owned by one dead volume without attempting writeback.
void MmFileCacheDiscardVolume(IO_VOLUME* volume) {
    if (!volume) return;

    for (uint32_t b = 0; b < MM_FILE_CACHE_BUCKETS; ++b) {
        while (true) {
            MM_FILE_CACHE_ENTRY* victim = (MM_FILE_CACHE_ENTRY*)0;
            bool pending = false;
            uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
            MM_FILE_CACHE_ENTRY* prev = (MM_FILE_CACHE_ENTRY*)0;
            for (MM_FILE_CACHE_ENTRY* e = gFileCacheBuckets[b]; e; prev = e, e = e->Next) {
                if (!e->File || e->File->Volume != volume)
                    continue;

                if (e->Loading || e->Writeback) {
                    e->Invalidated = 1;
                    if (e->Dirty)
                        KiFileCacheSetDirtyLocked(e, 0);
                    pending = true;
                    continue;
                }

                victim = e;
                KiFileCacheRemoveEntryLocked(b, prev, e);
                gFileCacheStats.Invalidations++;
                break;
            }
            KeReleaseSpinlock(&gFileCacheLock, flags);
            if (victim) {
                KiFileCacheFreeEntry(victim);
                continue;
            }
            if (!pending)
                break;
            // Loading/writeback entries unlink themselves when their in-flight operation completes.
            break;
        }
    }
}

KSTATUS MmFileCacheDropOnePageForPressure(void) {
    return KiFileCacheReclaimOnePage((FILE_OBJECT*)0, 0);
}

void MmFileCacheQueryStats(MM_FILE_CACHE_STATS* out) {
    if (!out) return;
    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    MM_FILE_CACHE_STATS st = gFileCacheStats;
    st.ResidentPages = gFileCacheResidentPages;
    st.ActivePages = gFileCacheActivePages;
    st.StandbyPages = gFileCacheStandbyPages;
    st.DirtyPages = gFileCacheDirtyPages;
    st.MaxDirtyPages = gFileCacheMaxDirtyPages;
    st.MaxResidentPages = gFileCacheMaxResidentPages;
    KeReleaseSpinlock(&gFileCacheLock, flags);
    *out = st;
}