dqn: Add head ptr on arena blocks

dqn.h

@@ -1447,10 +1447,6 @@ DQN_API int   Dqn_VMem_Protect (void *ptr, Dqn_usize size, uint32_t page_flags);
 //   @desc Dump the stats of the given arena to the memory log-stream.
 //   @param[in] arena The arena to dump stats for
 
-#if !defined(DQN_ARENA_MIN_BLOCK_SIZE)
-    #define DQN_ARENA_MIN_BLOCK_SIZE DQN_VMEM_RESERVE_GRANULARITY
-#endif
-
 enum Dqn_ArenaBlockFlags
 {
     Dqn_ArenaBlockFlags_Private = 1 << 0, ///< Private blocks can only allocate its memory when used in the 'FromBlock' API variants
@@ -1494,7 +1490,7 @@ struct Dqn_Arena
     bool            use_after_free_guard;
 
     Dqn_String8     label; ///< Optional label to describe the arena
-    Dqn_usize       min_block_size;
+    Dqn_ArenaBlock *head;  ///< Active block the arena is allocating from
     Dqn_ArenaBlock *curr;  ///< Active block the arena is allocating from
     Dqn_ArenaBlock *tail;  ///< Last block in the linked list of blocks
     Dqn_ArenaStat   stats; ///< Current arena stats, reset when reset usage is invoked.
@@ -1502,11 +1498,12 @@ struct Dqn_Arena
 
 struct Dqn_ArenaTempMemory
 {
-    Dqn_Arena      *arena;     ///< The arena the scope is for
+    Dqn_Arena      *arena;     ///< Arena the scope is for
-    Dqn_ArenaBlock *curr;      ///< The current block of the arena at the beginning of the scope
+    Dqn_ArenaBlock *head;      ///< Head block of the arena at the beginning of the scope
-    Dqn_ArenaBlock *tail;      ///< The tail block of the arena at the beginning of the scope
+    Dqn_ArenaBlock *curr;      ///< Current block of the arena at the beginning of the scope
-    Dqn_usize       curr_used; ///< The current used amount of the current block
+    Dqn_ArenaBlock *tail;      ///< Tail block of the arena at the beginning of the scope
-    Dqn_ArenaStat   stats;     ///< The stats of the arena at the beginning of the scope
+    Dqn_usize       curr_used; ///< Current used amount of the current block
+    Dqn_ArenaStat   stats;     ///< Stats of the arena at the beginning of the scope
 };
 
 // Automatically begin and end a temporary memory scope on object construction
@@ -6352,13 +6349,12 @@ DQN_API void Dqn_Arena_Reset(Dqn_Arena *arena, Dqn_ZeroMem zero_mem)
         return;
 
     // NOTE: Zero all the blocks until we reach the first block in the list
-    for (Dqn_ArenaBlock *block = arena->tail; block; block = block->prev) {
+    for (Dqn_ArenaBlock *block = arena->head; block; block = block->next) {
-        if (!block->prev)
-            arena->curr = block;
         Dqn_ArenaBlockResetInfo_ reset_info = {};
         Dqn_Arena_BlockReset_(DQN_LEAK_TRACE block, zero_mem, reset_info);
     }
 
+    arena->curr         = arena->head;
     arena->stats.used   = 0;
     arena->stats.wasted = 0;
 }
@@ -6367,6 +6363,7 @@ DQN_API Dqn_ArenaTempMemory Dqn_Arena_BeginTempMemory(Dqn_Arena *arena)
 {
     Dqn_ArenaTempMemory result = {};
     result.arena               = arena;
+    result.head                = arena->head;
     result.curr                = arena->curr;
     result.tail                = arena->tail;
     result.curr_used           = (arena->curr) ? arena->curr->used : 0;
@@ -6387,6 +6384,7 @@ DQN_API void Dqn_Arena_EndTempMemory_(DQN_LEAK_TRACE_FUNCTION Dqn_ArenaTempMemor
     arena->stats.blocks   = scope.stats.blocks;
 
     // NOTE: Revert the current block to the scope's current block
+    arena->head = scope.head;
     arena->curr = scope.curr;
     if (arena->curr) {
         Dqn_ArenaBlock *curr                = arena->curr;
@@ -6526,6 +6524,7 @@ DQN_API Dqn_ArenaBlock *Dqn_Arena_Grow_(DQN_LEAK_TRACE_FUNCTION Dqn_Arena *arena
         } else {
             DQN_ASSERT(!arena->curr);
             arena->curr = result;
+            arena->head = result;
         }
         arena->tail = result;
 
@@ -6558,15 +6557,18 @@ DQN_API void *Dqn_Arena_Allocate_(DQN_LEAK_TRACE_FUNCTION Dqn_Arena *arena, Dqn_
         while (arena->curr && (arena->curr->flags & Dqn_ArenaBlockFlags_Private))
             arena->curr = arena->curr->next;
 
-        if (!arena->curr) {
+        result = Dqn_Arena_AllocateFromBlock(arena->curr, size, align, zero_mem);
+        if (!result) {
+            if (!arena->curr || arena->curr == arena->tail) {
                 Dqn_usize allocation_size = size + (align - 1);
                 if (!Dqn_Arena_Grow(DQN_LEAK_TRACE_ARG arena, allocation_size, allocation_size /*commit*/, 0 /*flags*/)) {
                     break;
                 }
-            arena->curr = arena->tail;
             }
 
-        result = Dqn_Arena_AllocateFromBlock(arena->curr, size, align, zero_mem);
+            if (arena->curr != arena->tail)
+                arena->curr = arena->curr->next;
+        }
     }
 
     if (result)

dqn_unit_tests.cpp

@@ -31,12 +31,31 @@
 #define DQN_TESTER_IMPLEMENTATION
 #include "dqn_tester.h"
 
+enum Guard {
+    Guard_None,
+    Guard_UseAfterFree,
+    Guard_Count,
+};
+
+static Dqn_String8 ArenaGuardTestSuffix(uint32_t guard)
+{
+    Dqn_String8 result = {};
+    switch (guard) {
+        case Guard_None:         result = DQN_STRING8(" ");      break;
+        case Guard_UseAfterFree: result = DQN_STRING8(" [UAF]"); break;
+    }
+    return result;
+}
+
 Dqn_Tester TestArena()
 {
     Dqn_Tester test = {};
     DQN_TESTER_GROUP(test, "Dqn_Arena") {
-        DQN_TESTER_TEST("Reused memory is zeroed out") {
+        for (Dqn_usize guard = 0; guard < Guard_Count; guard++) {
+            Dqn_String8 test_suffix = ArenaGuardTestSuffix(guard);
+            DQN_TESTER_TEST("Reused memory is zeroed out%.*s", DQN_STRING_FMT(test_suffix)) {
                 Dqn_Arena arena            = {};
+                arena.use_after_free_guard = guard == Guard_UseAfterFree;
 
                 // NOTE: Allocate 128 kilobytes, fill it with garbage, then reset the arena
                 Dqn_usize size              = DQN_KILOBYTES(128);
@@ -60,6 +79,35 @@ Dqn_Tester TestArena()
                     DQN_TESTER_ASSERT(&test, ptr[i] == 0);
                 Dqn_Arena_Free(&arena, Dqn_ZeroMem_No);
             }
+        }
+
+        for (Dqn_usize guard = 0; guard < Guard_Count; guard++) {
+            Dqn_String8 test_suffix = ArenaGuardTestSuffix(guard);
+            DQN_TESTER_TEST("Test arena grows naturally, 1mb + 4mb%.*s", DQN_STRING_FMT(test_suffix)) {
+                Dqn_Arena arena            = {};
+                arena.use_after_free_guard = guard == Guard_UseAfterFree;
+
+                // NOTE: Allocate 1mb, then 4mb, this should force the arena to grow
+                char *ptr_1mb = DQN_CAST(char *)Dqn_Arena_Allocate(&arena, DQN_MEGABYTES(1), 1 /*align*/, Dqn_ZeroMem_Yes);
+                char *ptr_4mb = DQN_CAST(char *)Dqn_Arena_Allocate(&arena, DQN_MEGABYTES(4), 1 /*align*/, Dqn_ZeroMem_Yes);
+                DQN_TESTER_ASSERT(&test, ptr_1mb);
+                DQN_TESTER_ASSERT(&test, ptr_4mb);
+
+                Dqn_ArenaBlock const *block_1mb       = arena.head;
+                char           const *block_1mb_begin = DQN_CAST(char *)block_1mb->memory;
+                char           const *block_1mb_end   = DQN_CAST(char *)block_1mb->memory + block_1mb->size;
+
+                Dqn_ArenaBlock const *block_4mb       = arena.curr;
+                char           const *block_4mb_begin = DQN_CAST(char *)block_4mb->memory;
+                char           const *block_4mb_end   = DQN_CAST(char *)block_4mb->memory + block_4mb->size;
+
+                DQN_TESTER_ASSERTF(&test, block_1mb != block_4mb,                                 "New block should have been allocated and linked");
+                DQN_TESTER_ASSERTF(&test, ptr_1mb >= block_1mb_begin && ptr_1mb <= block_1mb_end, "Pointer was not allocated from correct memory block");
+                DQN_TESTER_ASSERTF(&test, ptr_4mb >= block_4mb_begin && ptr_4mb <= block_4mb_end, "Pointer was not allocated from correct memory block");
+
+                Dqn_Arena_Free(&arena, Dqn_ZeroMem_No);
+            }
+        }
 
         Dqn_usize sizes[] = {DQN_KILOBYTES(1), DQN_KILOBYTES(4), DQN_KILOBYTES(5)};
         for (Dqn_usize size : sizes) {