From 05c4c0b0593edeb16ca80dc9718e3f91c3b3b7b5 Mon Sep 17 00:00:00 2001
From: "Chloe M." <chloe@mensia.org>
Date: Wed, 24 Jun 2026 07:31:39 +0000
Subject: [PATCH] stos: pool: Handle page chaining while under pressure

Signed-off-by: Chloe M. <chloe@mensia.org>
---
 paw/stos/ex/pool.c | 58 +++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 55 insertions(+), 3 deletions(-)

diff --git a/paw/stos/ex/pool.c b/paw/stos/ex/pool.c
index 4c0a143..d7ca1ce 100644
--- a/paw/stos/ex/pool.c
+++ b/paw/stos/ex/pool.c
@@ -11,6 +11,7 @@
 #include <hal/kpcr.h>
 #include <mm/vmm.h>
 #include <string.h>
+#include <bitops.h>
 
 #define DTRACE(Fmt, ...) \
     TRACE("[ POOL ]: " Fmt, ##__VA_ARGS__)
@@ -42,9 +43,37 @@ PoolInitBlock(MEMORY_BLOCK *Block)
 }
 
 /*
- * Internally used by AllocatePoolWithTag()
+ * Allocates a new memory page to a block and returns it
  *
- * TODO: Handle running out of pages
+ * @Block: Block to allocate page to
+ */
+static MEMORY_PAGE *
+BlockNewPage(MEMORY_BLOCK *Block)
+{
+    MEMORY_PAGE *Page, *Last;
+
+    if (Block == NULL) {
+        return NULL;
+    }
+
+    Last = Block->LastPage;
+    if (Last == NULL) {
+        return NULL;
+    }
+
+    Page = MmAllocPages(1);
+    if (Page == NULL) {
+        return NULL;
+    }
+
+    RtlMemSet(Page, 0, sizeof(*Page));
+    Last->Next = Page;
+    Block->LastPage = Page;
+    return Page;
+}
+
+/*
+ * Internally used by AllocatePoolWithTag()
  */
 static VOID *
 AllocateFromBlock(UCHAR BlockIndex, MEMORY_BLOCK *Block, USIZE Count)
@@ -52,7 +81,10 @@ AllocateFromBlock(UCHAR BlockIndex, MEMORY_BLOCK *Block, USIZE Count)
     USIZE Gran, Idx;
     USIZE BitsNeeded, BitsFound = 0;
     SSIZE BitBase = -1;
+    USIZE BitsSet, Occupied;
+    USIZE MaxOccupied;
     MEMORY_PAGE *Page;
+    ST_STATUS Status;
 
     if (Block == NULL || Count == 0) {
         return NULL;
@@ -60,7 +92,7 @@ AllocateFromBlock(UCHAR BlockIndex, MEMORY_BLOCK *Block, USIZE Count)
 
     Gran = LEVEL_GRAN(BlockIndex);
     BitsNeeded = Count / Gran;
-    Page = Block->FirstPage;    /* TODO: Handle many */
+    Page = Block->LastPage;
 
     for (Idx = 0; Idx < sizeof(Page->Bitmap) * 8; ++Idx) {
         if (!TESTBIT(&Page->Bitmap, Idx) && BitsFound < BitsNeeded) {
@@ -82,6 +114,26 @@ AllocateFromBlock(UCHAR BlockIndex, MEMORY_BLOCK *Block, USIZE Count)
         SETBIT(&Page->Bitmap, Idx);
     }
 
+    BitsSet = PopCnt(Page->Bitmap);
+    Occupied = BitsSet * Gran;
+    MaxOccupied = PAGESIZE - sizeof(MEMORY_PAGE);
+    MaxOccupied = ALIGN_DOWN(MaxOccupied, Gran);
+
+    /*
+     * If there are no more bytes to grab in this page
+     */
+    if (Occupied >= MaxOccupied || BitsFound < BitsNeeded) {
+        Page = BlockNewPage(Block);
+        if (Page == NULL) {
+            return NULL;
+        }
+
+        /* Need to allocate again? */
+        if (BitsFound < BitsNeeded) {
+            return AllocateFromBlock(BlockIndex, Block, Count);
+        }
+    }
+
     return (BitsFound >= BitsNeeded)
         ? PTR_OFFSET(Page, sizeof(MEMORY_PAGE) + (BitBase * Gran))
         : NULL;