SystemPaw3/paw/stos/ex/pool.c

/*
 * Copyright (c) 2026, Chloe M.
 * Provided under the BSD-3 clause.
 *
 * Description: Pool allocator
 * Author:      Chloe M.
 */

#include <ex/pool.h>
#include <ex/trace.h>
#include <hal/kpcr.h>
#include <mm/vmm.h>
#include <string.h>

#define DTRACE(Fmt, ...) \
    TRACE("[ POOL ]: " Fmt, ##__VA_ARGS__)

/* Globals */
static USIZE PoolCount = 0;

/*
 * Initialize a pool block
 *
 * @Block: Pool block to initialize
 */
static ST_STATUS
PoolInitBlock(MEMORY_BLOCK *Block)
{
    if (Block == NULL) {
        return STATUS_INVALID_PARAM;
    }

    Block->FirstPage = MmAllocPages(1);
    if (Block->FirstPage == NULL) {
        return STATUS_NO_MEMORY;
    }

    Block->LastPage = Block->FirstPage;
    Block->PageCount = 1;
    RtlMemSet(Block->FirstPage, 0, PAGESIZE);
    return STATUS_SUCCESS;
}

/*
 * Internally used by AllocatePoolWithTag()
 *
 * TODO: Handle running out of pages
 */
static VOID *
AllocateFromBlock(UCHAR BlockIndex, MEMORY_BLOCK *Block, USIZE Count)
{
    USIZE Gran, Idx;
    USIZE BitsNeeded, BitsFound = 0;
    SSIZE BitBase = -1;
    MEMORY_PAGE *Page;

    if (Block == NULL || Count == 0) {
        return NULL;
    }

    Gran = LEVEL_GRAN(BlockIndex);
    BitsNeeded = Count / Gran;
    Page = Block->FirstPage;    /* TODO: Handle many */

    for (Idx = 0; Idx < sizeof(Page->Bitmap) * 8; ++Idx) {
        if (!TESTBIT(&Page->Bitmap, Idx) && BitsFound < BitsNeeded) {
            if (BitBase < 0)
                BitBase = Idx;

            ++BitsFound;
        } else if (TESTBIT(&Page->Bitmap, Idx)) {
            BitBase = -1;
            BitsFound = 0;
        }

        if (BitsFound >= BitsNeeded) {
            break;
        }
    }

    for (Idx = BitBase; Idx < BitBase + BitsNeeded; ++Idx) {
        SETBIT(&Page->Bitmap, Idx);
    }

    return (BitsFound >= BitsNeeded)
        ? PTR_OFFSET(Page, sizeof(MEMORY_PAGE) + (BitBase * Gran))
        : NULL;
}

/*
 * Internally used by ExAllocatePoolWithTag()
 */
static VOID *
AllocatePoolWithTag(MEMORY_POOL *Pool, POOL_TYPE Type, USIZE ByteCount, ULONG Tag)
{
    USIZE TmpCount;
    UCHAR BlockIndex = 0;
    MEMORY_BLOCK *Block;

    if (Pool == NULL || ByteCount == 0) {
        return NULL;
    }

    /*
     * Obtain the block index with the formula:
     *
     * i = (log2(count) - POOL_MIN_LOG2) - 1
     */
    ByteCount = ALIGN_UP(ByteCount, LEVEL_GRAN(0));
    TmpCount = ByteCount;
    while (TmpCount != 0) {
        ++BlockIndex;
        TmpCount >>= 1;
    }

    BlockIndex -= (POOL_MIN_LOG2 + 1);
    BlockIndex %= LEVEL_COUNT;
    Block = &Pool->BlockLevels[BlockIndex];
    return AllocateFromBlock(BlockIndex, Block, ByteCount);
}

/*
 * TODO: Handle tag assignment and prepend a pool header
 */
VOID *
ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag)
{
    KPCR *ThisCore;
    MEMORY_POOL *AllocPool;

    if (ByteCount == 0) {
        return NULL;
    }

    /* Verify that the pool type is valid */
    switch (Type) {
    case POOL_NON_PAGED:
        break;
    default:
        DTRACE("got bad pool type in allocation\n");
        return NULL;
    }

    ThisCore = HalKpcrCurrent();
    AllocPool = &ThisCore->AllocPool;

    /*
     * Ensure that the byte count is aligned by the minimum granularity
     * that we can allocate. Furthermore, ensure that it does not exceed
     * the size of a single page.
     */
    ByteCount = ALIGN_UP(ByteCount, LEVEL_GRAN(0));
    if (ByteCount >= PAGESIZE) {
        DTRACE(
            "allocation in pool #%d exceeds page size\n",
            AllocPool->Id
        );
        return NULL;
    }

    return AllocatePoolWithTag(AllocPool, Type, ByteCount, Tag);
}

ST_STATUS
ExPoolInit(MEMORY_POOL *Pool)
{
    UCHAR LevelIdx;
    MEMORY_BLOCK *Block;
    ST_STATUS Status;

    if (Pool == NULL) {
        return STATUS_INVALID_PARAM;
    }

    Pool->Id = PoolCount;
    DTRACE("bringing up pool #%d\n", Pool->Id);

    /* Initialize the pools */
    for (LevelIdx = 0; LevelIdx < LEVEL_COUNT; ++LevelIdx) {
        Block = &Pool->BlockLevels[LevelIdx];
        Status = PoolInitBlock(Block);

        /* XXX: We may wanna clean up our mess we made */
        if (Status != STATUS_SUCCESS) {
            DTRACE("critical: failed to initialize pool\n");
            return STATUS_NO_MEMORY;
        }
    }

    DTRACE("brought up %d levels for pool #%d : ok\n", LEVEL_COUNT, PoolCount);
    ++PoolCount;
    return STATUS_SUCCESS;
}