// 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 <stdint.h>
#include <stddef.h>
#include <string.h>
#include <intrin.h>
#include <kernel/mm_file_cache.h>
#include <kernel/io_file.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_list.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_READAHEAD_MAX_REQUESTS 64u
#define MM_FILE_CACHE_READAHEAD_MAX_PAGES 32u
#define MM_FILE_CACHE_READAHEAD_QUEUE_PAGE_BUDGET 128u
#define MM_FILE_CACHE_READAHEAD_MIN_HEADROOM_PAGES 16u
#define MM_FILE_CACHE_BATCH_MAX_PAGES 16u
#define MM_FILE_CACHE_PRESSURE_BATCH_HEADROOM_PAGES 64u
#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;
    uint8_t           Invalidated;
    uint8_t           Dirty;
    uint8_t           Writeback;
    uint64_t          DirtySequence;
} MM_FILE_CACHE_ENTRY;

typedef struct MM_FILE_CACHE_READAHEAD_REQUEST {
    LIST_ENTRY Link;
    FILE_OBJECT* File;
    uint64_t Offset;
    uint32_t Pages;
} MM_FILE_CACHE_READAHEAD_REQUEST;

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 SPINLOCK gReadAheadLock = SPINLOCK_INIT;
static LIST_ENTRY gReadAheadQueue;
static SEMAPHORE* gReadAheadSemaphore;
static uint32_t gReadAheadQueuedRequests;
static uint32_t gReadAheadQueuedPages;
static int gReadAheadQueueReady;
static int gReadAheadStartState;

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 int KiFileCacheIsReclaimCandidateLocked(MM_FILE_CACHE_ENTRY* e, FILE_OBJECT* skipFile, uint64_t skipOffset) {
    if (!e) return 0;
    if (skipFile && e->File == skipFile && e->Offset == skipOffset)
        return 0;
    if (e->Loading || e->Invalidated || !e->Frame || e->Dirty || e->Writeback)
        return 0;

    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 STATUS_SUCCESS;
                }
            }
            KeReleaseSpinlock(&gFileCacheLock, flags);
        }
    }
    return STATUS_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 STATUS_SUCCESS;
        }
    }
    return STATUS_NOT_FOUND;
}

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

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

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


static void KiFileCacheAddReadAheadStat(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);
}

// Updates a maximum counter while the file-cache stats lock is held.
static void KiFileCacheMaxStatLocked(uint64_t* field, uint64_t value) {
    if (*field < value)
        *field = value;
}

static int KiFileCachePageExistsLocked(FILE_OBJECT* file, uint64_t offset) {
    uint32_t bucket = KiFileCacheHash(file, offset);
    return KiFileCacheFindLocked(file, offset, bucket) != (MM_FILE_CACHE_ENTRY*)0;
}

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

static MM_FILE_CACHE_READAHEAD_REQUEST* KiFileCacheReadAheadDequeue(void) {
    uintptr_t flags = KeAcquireSpinlock(&gReadAheadLock);
    LIST_ENTRY* node = ListPopHead(&gReadAheadQueue);
    if (node) {
        MM_FILE_CACHE_READAHEAD_REQUEST* req = CONTAINING_RECORD(node, MM_FILE_CACHE_READAHEAD_REQUEST, Link);
        if (gReadAheadQueuedRequests)
            gReadAheadQueuedRequests--;
        if (gReadAheadQueuedPages >= req->Pages)
            gReadAheadQueuedPages -= req->Pages;
        else
            gReadAheadQueuedPages = 0;
    }
    KeReleaseSpinlock(&gReadAheadLock, flags);

    if (!node)
        return (MM_FILE_CACHE_READAHEAD_REQUEST*)0;
    return CONTAINING_RECORD(node, MM_FILE_CACHE_READAHEAD_REQUEST, Link);
}

static void KiFileCacheReadAheadWorker(void* arg) {
    (void)arg;
    KeAssertIrql(PASSIVE_LEVEL);

    for (;;) {
        KSTATUS ws = KeWaitSemaphore(gReadAheadSemaphore);
        if (!KSUCCESS(ws)) {
            KeYield();
            continue;
        }

        MM_FILE_CACHE_READAHEAD_REQUEST* req = KiFileCacheReadAheadDequeue();
        if (!req)
            continue;

        for (uint32_t i = 0; i < req->Pages; ++i) {
            uint64_t offset = req->Offset + ((uint64_t)i * PAGE_SIZE);

            uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
            if (gFileCacheResidentPages + MM_FILE_CACHE_READAHEAD_MIN_HEADROOM_PAGES >=
                MM_FILE_CACHE_SOFT_MAX_PAGES) {
                gFileCacheStats.ReadAheadPagesThrottled += (uint64_t)(req->Pages - i);
                KeReleaseSpinlock(&gFileCacheLock, flags);
                break;
            }
            int exists = KiFileCachePageExistsLocked(req->File, offset);
            if (exists)
                gFileCacheStats.ReadAheadPagesSkipped++;
            KeReleaseSpinlock(&gFileCacheLock, flags);
            if (exists)
                continue;

            MM_FILE_CACHE_PAGE page;
            KSTATUS s = MmFileCacheGetPage(req->File, offset, &page);
            flags = KeAcquireSpinlock(&gFileCacheLock);
            if (KSUCCESS(s))
                gFileCacheStats.ReadAheadPagesLoaded++;
            else
                gFileCacheStats.ReadAheadFailures++;
            KeReleaseSpinlock(&gFileCacheLock, flags);

            if (!KSUCCESS(s))
                break;
            MmFileCacheReleasePage(&page);
        }

        ObDereferenceObject(&req->File->Header);
        MmFreePool(req);
    }
}

static KSTATUS KiFileCacheEnsureReadAheadStarted(void) {
    uintptr_t flags = KeAcquireSpinlock(&gReadAheadLock);
    if (gReadAheadStartState == 2) {
        KeReleaseSpinlock(&gReadAheadLock, flags);
        return STATUS_SUCCESS;
    }
    if (gReadAheadStartState == 1) {
        KeReleaseSpinlock(&gReadAheadLock, flags);
        while (gReadAheadStartState == 1)
            KeYield();
        return gReadAheadStartState == 2 ? STATUS_SUCCESS : STATUS_FAILURE;
    }

    gReadAheadStartState = 1;
    KeReleaseSpinlock(&gReadAheadLock, flags);

    LIST_INIT(&gReadAheadQueue);
    gReadAheadQueuedRequests = 0;
    gReadAheadQueuedPages = 0;
    gReadAheadQueueReady = 1;

    KSTATUS s = KeSemaphoreCreate(0, 0x7fffffff, &gReadAheadSemaphore);
    if (KSUCCESS(s)) {
        THREAD* worker = (THREAD*)0;
        s = KeCreateKernelThread((uintptr_t)KiFileCacheReadAheadWorker,
                                 0,
                                 PriorityNormal,
                                 &worker);
        (void)worker;
    }

    flags = KeAcquireSpinlock(&gReadAheadLock);
    gReadAheadStartState = KSUCCESS(s) ? 2 : 0;
    if (!KSUCCESS(s))
        gReadAheadQueueReady = 0;
    KeReleaseSpinlock(&gReadAheadLock, flags);
    return s;
}

// 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 STATUS_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 = STATUS_FAILURE;
    if (page && target->Fs && target->Fs->Write)
        s = target->Fs->Write(target->Fs, target->FileCtx, offset, page, validBytes, &written);
    if (KSUCCESS(s) && written != validBytes)
        s = STATUS_FAILURE;

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

    (void)MmPfnRelease(frame);
    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);

    for (;;) {
        KSTATUS ws = KeWaitSemaphore(gWritebackSemaphore);
        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 STATUS_SUCCESS;
    }
    if (gWritebackStartState == 1) {
        KeReleaseSpinlock(&gWritebackLock, flags);
        while (gWritebackStartState == 1)
            KeYield();
        return gWritebackStartState == 2 ? STATUS_SUCCESS : STATUS_FAILURE;
    }

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

    KSTATUS s = KeSemaphoreCreate(0, 0x7fffffff, &gWritebackSemaphore);
    if (KSUCCESS(s)) {
        THREAD* worker = (THREAD*)0;
        s = KeCreateKernelThread((uintptr_t)KiFileCacheWritebackWorker,
                                 0,
                                 PriorityNormal,
                                 &worker);
        (void)worker;
    }

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

static size_t KiMinSize(size_t a, size_t b) { return a < b ? a : b; }

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;
    for (;;) {
        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 STATUS_NOT_FOUND;
        }

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

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

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

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

        KSTATUS rs = STATUS_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 STATUS_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) {
    for (;;) {
        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 int KiFileCacheUnlinkEntryLocked(MM_FILE_CACHE_ENTRY* entry) {
    if (!entry)
        return 0;

    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 1;
        }
    }
    return 0;
}

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

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

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

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

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

    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    int 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 STATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    if (entry->Invalidated) {
        (void)KiFileCacheUnlinkEntryLocked(entry);
        KeReleaseSpinlock(&gFileCacheLock, flags);
        (void)MmPfnRelease(frame);
        KiFileCacheFreeLoadingEntry(entry);
        return STATUS_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 STATUS_SUCCESS;
}


static uint32_t KiFileCacheChooseForegroundBatchPages(void) {
    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    uint64_t resident = gFileCacheResidentPages;
    KeReleaseSpinlock(&gFileCacheLock, flags);

    if (resident >= MM_FILE_CACHE_SOFT_MAX_PAGES)
        return 1;

    uint64_t headroom = MM_FILE_CACHE_SOFT_MAX_PAGES - resident;
    if (headroom <= MM_FILE_CACHE_READAHEAD_MIN_HEADROOM_PAGES)
        return 1;
    if (headroom <= MM_FILE_CACHE_PRESSURE_BATCH_HEADROOM_PAGES)
        return 2;
    if (headroom <= (MM_FILE_CACHE_PRESSURE_BATCH_HEADROOM_PAGES * 2u))
        return 4;
    if (headroom <= (MM_FILE_CACHE_PRESSURE_BATCH_HEADROOM_PAGES * 4u))
        return 8;
    return MM_FILE_CACHE_BATCH_MAX_PAGES;
}

static KSTATUS KiFileCacheLoadSpan(FILE_OBJECT* file,
                                   uint64_t aligned,
                                   uint32_t maxPages,
                                   MM_FILE_CACHE_PAGE* requestedPage,
                                   uint32_t* loadedPagesOut) {
    if (loadedPagesOut)
        *loadedPagesOut = 0;
    if (!file || maxPages == 0)
        return STATUS_INVALID_PARAMETER;
    if (maxPages > MM_FILE_CACHE_BATCH_MAX_PAGES)
        maxPages = MM_FILE_CACHE_BATCH_MAX_PAGES;

    uintptr_t statFlags = KeAcquireSpinlock(&gFileCacheLock);
    gFileCacheStats.BatchPagesRequested += maxPages;
    KiFileCacheMaxStatLocked(&gFileCacheStats.MaxBatchPagesRequested, maxPages);
    KeReleaseSpinlock(&gFileCacheLock, statFlags);

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

    MM_FILE_CACHE_ENTRY* entries[MM_FILE_CACHE_BATCH_MAX_PAGES];
    PHYSICAL_ADDRESS frames[MM_FILE_CACHE_BATCH_MAX_PAGES];
    memset(entries, 0, sizeof(entries));
    memset(frames, 0, sizeof(frames));

    uint32_t pages = 0;
    KSTATUS reserveStatus = STATUS_SUCCESS;
    while (pages < maxPages) {
        uint64_t pageOffset = aligned + ((uint64_t)pages * PAGE_SIZE);
        if (pageOffset >= fileSize)
            break;

        reserveStatus = KiFileCacheReserveLoadingPage(file, pageOffset, &entries[pages]);
        if (!KSUCCESS(reserveStatus))
            break;
        pages++;
    }

    if (pages == 0) {
        if (requestedPage)
            return KiFileCacheTryHit(file, aligned, requestedPage, 1);
        return reserveStatus == STATUS_ALREADY_EXISTS ? STATUS_ALREADY_EXISTS : reserveStatus;
    }

    uint8_t* readBuffer = (uint8_t*)MmAllocatePoolTag((size_t)pages * PAGE_SIZE, MM_POOL_TAG('f', 'c', 'b', 'l'));
    if (!readBuffer) {
        for (uint32_t i = 0; i < pages; ++i)
            KiFileCacheCancelLoadingPage(entries[i]);
        KiFileCacheAddReadAheadStat(&gFileCacheStats.LoadFailures, 1);
        return STATUS_NO_MEMORY;
    }
    memset(readBuffer, 0, (size_t)pages * PAGE_SIZE);

    for (uint32_t i = 0; i < pages; ++i) {
        frames[i] = MmPfnAllocate(MmPfnUseCleanFileCache);
        if (!frames[i]) {
            for (uint32_t j = 0; j < i; ++j)
                (void)MmPfnRelease(frames[j]);
            for (uint32_t j = 0; j < pages; ++j)
                KiFileCacheCancelLoadingPage(entries[j]);
            MmFreePool(readBuffer);
            KiFileCacheAddReadAheadStat(&gFileCacheStats.LoadFailures, 1);
            return STATUS_NO_MEMORY;
        }
    }

    size_t toRead = KiMinSize((size_t)pages * PAGE_SIZE, (size_t)(fileSize - aligned));
    size_t got = 0;
    KSTATUS s = STATUS_NOT_SUPPORTED;
    if (file->Fs && file->Fs->Read)
        s = file->Fs->Read(file->Fs, file->FileCtx, aligned, readBuffer, toRead, &got);
    if (!KSUCCESS(s)) {
        for (uint32_t i = 0; i < pages; ++i) {
            if (frames[i])
                (void)MmPfnRelease(frames[i]);
            KiFileCacheCancelLoadingPage(entries[i]);
        }
        MmFreePool(readBuffer);
        KiFileCacheAddReadAheadStat(&gFileCacheStats.LoadFailures, 1);
        return s;
    }
    if (got > toRead)
        got = toRead;

    uint64_t published = 0;
    KSTATUS requestedStatus = STATUS_NOT_FOUND;
    for (uint32_t i = 0; i < pages; ++i) {
        uint64_t pageOffset = aligned + ((uint64_t)i * PAGE_SIZE);
        size_t pageBase = (size_t)i * PAGE_SIZE;
        if (pageBase >= got && pageOffset >= fileSize) {
            (void)MmPfnRelease(frames[i]);
            frames[i] = 0;
            KiFileCacheCancelLoadingPage(entries[i]);
            continue;
        }

        uint32_t validBytes = 0;
        if (pageBase < got) {
            size_t avail = got - pageBase;
            if (avail > PAGE_SIZE) avail = PAGE_SIZE;
            validBytes = (uint32_t)avail;
        }

        void* page = MmPfnToVirtual(frames[i]);
        if (!page) {
            (void)MmPfnRelease(frames[i]);
            frames[i] = 0;
            KiFileCacheCancelLoadingPage(entries[i]);
            continue;
        }
        memcpy(page, readBuffer + pageBase, validBytes);
        if (validBytes < PAGE_SIZE)
            memset((uint8_t*)page + validBytes, 0, PAGE_SIZE - validBytes);

        MM_FILE_CACHE_PAGE* out = (pageOffset == aligned) ? requestedPage : (MM_FILE_CACHE_PAGE*)0;
        KSTATUS ps = KiFileCacheCompleteLoadingPage(entries[i], frames[i], validBytes, out, &published);
        if (KSUCCESS(ps)) {
            if (pageOffset == aligned)
                requestedStatus = STATUS_SUCCESS;
            frames[i] = 0;
        } else {
            if (pageOffset == aligned)
                requestedStatus = ps;
            frames[i] = 0;
        }
    }

    MmFreePool(readBuffer);

    uintptr_t flags = KeAcquireSpinlock(&gFileCacheLock);
    if (published) {
        gFileCacheStats.BatchReads++;
        gFileCacheStats.BatchPagesLoaded += published;
        KiFileCacheMaxStatLocked(&gFileCacheStats.MaxBatchPagesLoaded, published);
    }
    KeReleaseSpinlock(&gFileCacheLock, flags);

    if (loadedPagesOut)
        *loadedPagesOut = (uint32_t)published;

    KiFileCacheKeepBounded(file, aligned);

    if (requestedPage) {
        if (KSUCCESS(requestedStatus))
            return STATUS_SUCCESS;
        return KiFileCacheTryHit(file, aligned, requestedPage, 1);
    }
    return published ? STATUS_SUCCESS : STATUS_NOT_FOUND;
}

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

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

    if (needExistingBytes && aligned < oldFileSize) {
        s = KiFileCacheLoadSpan(file, aligned, 1, outPage, (uint32_t*)0);
        if (KSUCCESS(s))
            return STATUS_SUCCESS;
        if (s != STATUS_NOT_FOUND)
            return s;
    }

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

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

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

    uint32_t validBytes = PAGE_SIZE;
    if (oldFileSize > aligned && oldFileSize - aligned < PAGE_SIZE)
        validBytes = (uint32_t)(oldFileSize - aligned);

    s = KiFileCacheCompleteLoadingPage(entry, frame, validBytes, outPage, (uint64_t*)0);
    if (!KSUCCESS(s))
        return s;
    return STATUS_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, int needExistingBytes,
                                   MM_FILE_CACHE_PAGE* outPage) {
    if (!outPage)
        return STATUS_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 STATUS_INVALID_PARAMETER;
    if (!KeCanWait())
        return STATUS_INVALID_IRQL;
    if (!file->Fs || !file->Fs->Write)
        return STATUS_NOT_SUPPORTED;
    if (size == 0)
        return STATUS_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 ? STATUS_SUCCESS : s;

        void* dst = MmPfnToVirtual(page.Frame);
        if (!dst) {
            MmFileCacheReleasePage(&page);
            return total ? STATUS_SUCCESS : STATUS_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 STATUS_SUCCESS;
}

KSTATUS MmFileCacheFlushFile(FILE_OBJECT* file) {
    if (!file)
        return STATUS_INVALID_PARAMETER;
    if (!KeCanWait())
        return STATUS_INVALID_IRQL;

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

    KSTATUS firstError = STATUS_SUCCESS;
    for (;;) {
        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 STATUS_INVALID_PARAMETER;
    memset(outPage, 0, sizeof(*outPage));
    if (!file) return STATUS_INVALID_PARAMETER;
    if (!KeCanWait()) return STATUS_INVALID_IRQL;

    uint64_t aligned = offset & ~(uint64_t)(PAGE_SIZE - 1);
    KSTATUS s = KiFileCacheTryHit(file, aligned, outPage, 1);
    if (KSUCCESS(s))
        return STATUS_SUCCESS;
    return KiFileCacheLoadSpan(file, aligned, KiFileCacheChooseForegroundBatchPages(), outPage, (uint32_t*)0);
}

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

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


KSTATUS MmFileCacheScheduleReadAhead(FILE_OBJECT* file, uint64_t offset, uint32_t pages) {
    if (!file || pages == 0)
        return STATUS_INVALID_PARAMETER;
    if (!KeCanWait())
        return STATUS_INVALID_IRQL;

    uint32_t requestedPages = pages;
    if (pages > MM_FILE_CACHE_READAHEAD_MAX_PAGES)
        pages = MM_FILE_CACHE_READAHEAD_MAX_PAGES;

    uintptr_t cacheFlags = KeAcquireSpinlock(&gFileCacheLock);
    gFileCacheStats.ReadAheadPagesRequested += requestedPages;
    KiFileCacheMaxStatLocked(&gFileCacheStats.MaxReadAheadPagesPerRequest, requestedPages);
    uint64_t resident = gFileCacheResidentPages;
    uint64_t headroom = resident < MM_FILE_CACHE_SOFT_MAX_PAGES
        ? MM_FILE_CACHE_SOFT_MAX_PAGES - resident
        : 0;
    if (headroom <= MM_FILE_CACHE_READAHEAD_MIN_HEADROOM_PAGES) {
        gFileCacheStats.ReadAheadPagesThrottled += pages;
        KeReleaseSpinlock(&gFileCacheLock, cacheFlags);
        return STATUS_QUEUE_FULL;
    }
    headroom -= MM_FILE_CACHE_READAHEAD_MIN_HEADROOM_PAGES;
    if ((uint64_t)pages > headroom)
        pages = (uint32_t)headroom;
    KeReleaseSpinlock(&gFileCacheLock, cacheFlags);
    if (pages == 0)
        return STATUS_QUEUE_FULL;

    KSTATUS s = KiFileCacheEnsureReadAheadStarted();
    if (!KSUCCESS(s)) {
        KiFileCacheAddReadAheadStat(&gFileCacheStats.ReadAheadFailures, 1);
        return s;
    }

    MM_FILE_CACHE_READAHEAD_REQUEST* req =
        (MM_FILE_CACHE_READAHEAD_REQUEST*)MmAllocatePoolTag(sizeof(*req), MM_POOL_TAG('r', 'd', 'a', 'h'));
    if (!req) {
        KiFileCacheAddReadAheadStat(&gFileCacheStats.ReadAheadFailures, 1);
        return STATUS_NO_MEMORY;
    }

    LIST_INIT(&req->Link);
    req->File = file;
    req->Offset = offset & ~(uint64_t)(PAGE_SIZE - 1);
    req->Pages = pages;
    ObReferenceObject(&file->Header);

    uintptr_t flags = KeAcquireSpinlock(&gReadAheadLock);
    if (!gReadAheadQueueReady ||
        gReadAheadQueuedRequests >= MM_FILE_CACHE_READAHEAD_MAX_REQUESTS ||
        gReadAheadQueuedPages >= MM_FILE_CACHE_READAHEAD_QUEUE_PAGE_BUDGET) {
        KeReleaseSpinlock(&gReadAheadLock, flags);
        ObDereferenceObject(&file->Header);
        MmFreePool(req);
        KiFileCacheAddReadAheadStat(&gFileCacheStats.ReadAheadPagesThrottled, pages);
        return STATUS_QUEUE_FULL;
    }
    uint32_t queueHeadroom = MM_FILE_CACHE_READAHEAD_QUEUE_PAGE_BUDGET - gReadAheadQueuedPages;
    if (pages > queueHeadroom)
        pages = queueHeadroom;
    req->Pages = pages;

    ListInsertTail(&gReadAheadQueue, &req->Link);
    gReadAheadQueuedRequests++;
    gReadAheadQueuedPages += pages;
    KeReleaseSpinlock(&gReadAheadLock, flags);

    flags = KeAcquireSpinlock(&gFileCacheLock);
    gFileCacheStats.ReadAheadRequests++;
    gFileCacheStats.ReadAheadPagesQueued += pages;
    KeReleaseSpinlock(&gFileCacheLock, flags);

    KeSignalSemaphore(gReadAheadSemaphore);
    return STATUS_SUCCESS;
}

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