From ef8535f0897006d7c88a5e3671dddfb2540ef963 Mon Sep 17 00:00:00 2001
From: "Chloe M." <chloe@mensia.org>
Date: Wed, 24 Jun 2026 05:50:22 +0000
Subject: [PATCH] stos: pool: Add actual allocation groundwork

Still lots to do

Signed-off-by: Chloe M. <chloe@mensia.org>
---
 paw/stos/ex/pool.c      | 130 +++++++++++++++++++++++++++++++++++++++-
 paw/stos/head/ex/pool.h |  31 +++++++++-
 2 files changed, 159 insertions(+), 2 deletions(-)

diff --git a/paw/stos/ex/pool.c b/paw/stos/ex/pool.c
index 2486ce9..4c0a143 100644
--- a/paw/stos/ex/pool.c
+++ b/paw/stos/ex/pool.c
@@ -8,6 +8,7 @@
 
 #include <ex/pool.h>
 #include <ex/trace.h>
+#include <hal/kpcr.h>
 #include <mm/vmm.h>
 #include <string.h>
 
@@ -17,6 +18,11 @@
 /* Globals */
 static USIZE PoolCount = 0;
 
+/*
+ * Initialize a pool block
+ *
+ * @Block: Pool block to initialize
+ */
 static ST_STATUS
 PoolInitBlock(MEMORY_BLOCK *Block)
 {
@@ -35,6 +41,126 @@ PoolInitBlock(MEMORY_BLOCK *Block)
     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)
 {
@@ -46,8 +172,10 @@ ExPoolInit(MEMORY_POOL *Pool)
         return STATUS_INVALID_PARAM;
     }
 
+    Pool->Id = PoolCount;
+    DTRACE("bringing up pool #%d\n", Pool->Id);
+
     /* Initialize the pools */
-    DTRACE("bringing up pool #%d\n", PoolCount);
     for (LevelIdx = 0; LevelIdx < LEVEL_COUNT; ++LevelIdx) {
         Block = &Pool->BlockLevels[LevelIdx];
         Status = PoolInitBlock(Block);
diff --git a/paw/stos/head/ex/pool.h b/paw/stos/head/ex/pool.h
index 57db775..b7cf672 100644
--- a/paw/stos/head/ex/pool.h
+++ b/paw/stos/head/ex/pool.h
@@ -13,9 +13,12 @@
 #include <stapi/status.h>
 #include <stdef.h>
 
+/* Minimum log2 allocation size */
+#define POOL_MIN_LOG2 3
+
 /* Compute the granularity of a specific level */
 #define LEVEL_GRAN(LEVEL) \
-    (1 << (3 + (LEVEL)))
+    (1 << (POOL_MIN_LOG2 + (LEVEL)))
 
 /* Maximum levels per pool */
 #define LEVEL_COUNT 8
@@ -28,6 +31,15 @@
 #define POOL_REAL_PAGESZ \
     (PAGESIZE - sizeof(MEMORY_PAGE))
 
+/*
+ * Valid pool types
+ *
+ * @POOL_NON_PAGED: Not pagable to disk
+ */
+typedef enum {
+    POOL_NON_PAGED
+} POOL_TYPE;
+
 /*
  * Page data structure used internally to link together page
  * lists
@@ -58,9 +70,11 @@ typedef struct {
  * A memory pool holds blocks of varying granularities
  *
  * @BlockLevels: Block levels [granularity : GRAN(level)]
+ * @Id:          Pool ID
  */
 typedef struct {
     MEMORY_BLOCK BlockLevels[LEVEL_COUNT];
+    USIZE Id;
 } MEMORY_POOL;
 
 /*
@@ -70,4 +84,19 @@ typedef struct {
  */
 ST_STATUS ExPoolInit(MEMORY_POOL *Pool);
 
+/*
+ * Allocate a tagged pool
+ *
+ * @Type:       Pool type
+ * @ByteCount:  Number of bytes to allocate
+ * @Tag:        Tag to assign to pool
+ *
+ * Returns the memory base on success, otherwise
+ * NULL on failure.
+ */
+VOID *ExAllocatePoolWithTag(
+    POOL_TYPE Type, USIZE ByteCount,
+    ULONG Tag
+);
+
 #endif  /* !_EX_POOL_H_ */