From a4cae4211429ac34dff6cf9e92129a0169cbd94d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nils=20O=2E=20Sel=C3=A5sdal?= <nos@utel.tech>
Date: Tue, 21 Apr 2026 01:00:49 +0200
Subject: [PATCH] small_allocator improvements

Remove unnneded align_up. last_alloc could point to the middle
of a merged block. Fix by resetting the last_alloc on free
---
 src/small_allocator.c | 40 +++++++++++++++++++++++-----------------
 1 file changed, 23 insertions(+), 17 deletions(-)

diff --git a/src/small_allocator.c b/src/small_allocator.c
index 0d63943..4618d99 100644
--- a/src/small_allocator.c
+++ b/src/small_allocator.c
@@ -35,7 +35,7 @@ typedef struct {
     uint8_t* base;            // backing store
     uint32_t size;            // total size
     BlockHeader* heap_start;  // first block
-    BlockHeader* last_alloc;  // next-fit pointer
+    BlockHeader* next_fit;  // next-fit pointer
 } SmallAllocator;
 
 /* --- Header helpers --- */
@@ -95,14 +95,16 @@ void small_allocator_init(SmallAllocator* a, void* backing_store, uint32_t size)
     a->size = size;
 
     BlockHeader* first = backing_store;
-    uint32_t usable = ALIGN_DOWN(size - sizeof(BlockHeader));
+    // allocator logic requires the first block to be properly aligned
+    // as well as size being eligned.
+    uint32_t usable = ALIGN_DOWN(size) - sizeof(BlockHeader);
     block_set_free(first);
     block_set_size(first, usable);
     first->next = NULL_BLOCK;
     first->prev = NULL_BLOCK;
 
     a->heap_start = first;
-    a->last_alloc = first;
+    a->next_fit = first;
 }
 
 static void maybe_split_block(SmallAllocator* a, BlockHeader* block, uint32_t size)
@@ -110,12 +112,14 @@ static void maybe_split_block(SmallAllocator* a, BlockHeader* block, uint32_t si
     // assumes size is already aligned
     uint32_t bsize = block_size(block);
 
-    if (bsize < size + sizeof(BlockHeader) + ALIGNMENT) return;
+    if (bsize < size + sizeof(BlockHeader) + ALIGNMENT) {
+        return;
+    }
 
     uint8_t* base = (uint8_t*)block;
     BlockHeader* new_block = (BlockHeader*)(base + sizeof(BlockHeader) + size);
 
-    uint32_t new_size = ALIGN_UP(bsize - size - sizeof(BlockHeader));
+    uint32_t new_size = bsize - size - sizeof(BlockHeader);
 
     block_set_size(block, size);
     block_set_used(block);
@@ -137,7 +141,7 @@ static void merge_with_next(SmallAllocator* a, BlockHeader* block)
     BlockHeader* next = block_from_offset(a, block->next);
     if (next == NULL || !block_is_free(next)) return;
 
-    uint32_t new_size = ALIGN_UP(block_size(block) + sizeof(BlockHeader) + block_size(next));
+    uint32_t new_size = block_size(block) + sizeof(BlockHeader) + block_size(next);
 
     block_set_size(block, new_size);
     block->next = next->next;
@@ -147,19 +151,24 @@ static void merge_with_next(SmallAllocator* a, BlockHeader* block)
 
 void* small_alloc(SmallAllocator* a, uint32_t size)
 {
+    if (size > UINT32_MAX - (ALIGNMENT - 1)) {
+        return NULL;
+    }
     size = ALIGN_UP(size);
 
-    BlockHeader* start = a->last_alloc ? a->last_alloc : a->heap_start;
+    BlockHeader* start = a->next_fit ? a->next_fit : a->heap_start;
     BlockHeader* curr = start;
 
     do {
         if (block_is_free(curr) && block_size(curr) >= size) {
             maybe_split_block(a, curr, size);
             block_set_used(curr);
-            a->last_alloc = curr;
+            BlockHeader *next = block_from_offset(a, curr->next);
+            a->next_fit = next ? next : a->heap_start;;
             return (uint8_t*)curr + sizeof(BlockHeader);
         }
-        curr = block_from_offset(a, curr->next) ? block_from_offset(a, curr->next) : a->heap_start;
+        BlockHeader *next = block_from_offset(a, curr->next);
+        curr = next ? next : a->heap_start;
     } while (curr != start);
 
     return NULL;
@@ -175,7 +184,7 @@ void small_free(SmallAllocator* a, void* ptr)
 
     BlockHeader* block = (BlockHeader*)((uint8_t*)ptr - sizeof(BlockHeader));
 
-    assert((uint8_t*)ptr >= a->base && (uint8_t*)ptr < a->base + a->size);
+    assert((uint8_t*)ptr - sizeof(BlockHeader) >= a->base && (uint8_t*)ptr < a->base + a->size);
     // double free detection
     assert(!block_is_free(block));
 
@@ -185,14 +194,12 @@ void small_free(SmallAllocator* a, void* ptr)
     BlockHeader* prev = block_from_offset(a, block->prev);
     if (prev != NULL && block_is_free(prev)) {
         merge_with_next(a, prev);
-
-        if (a->last_alloc == block) {
-            a->last_alloc = prev;
-        }
+        a->next_fit = prev;
+    } else {
+        a->next_fit = block;
     }
 }
 
-#if 0
 #include <string.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -214,7 +221,7 @@ int main(void)
     SmallAllocator a;
     char* ptrs[SZ];
 
-    small_allocator_init(&a, buff, sizeof buff);
+    small_allocator_init(&a, buff, sizeof(buff));
     debug_alloc(&a);
     ptrs[0] = small_alloc(&a, 32);
     ptrs[1] = small_alloc(&a, 80);
@@ -241,4 +248,3 @@ int main(void)
     debug_alloc(&a);
 }
 
-#endif