Allocator memes

src/basic/basic_alloc.cpp

@@ -148,13 +148,6 @@ API void *TrackingAllocatorProc(void *object, int kind, void *p, size_t size) {
     return result;
 }
 
-API void TrackingAllocatorCheck() {
-    // For (MemoryTrackingRecord) {
-    //     ReportConsolef("%s(%d): error: memory leak");
-    // }
-    Assert(MemoryTrackingRecord.len == 0);
-}
-
 API Allocator GetTrackingAllocator() {
     Allocator result = {TrackingAllocatorProc};
     return result;
@@ -267,20 +260,28 @@ API void *ArenaAllocatorProc(void *object, int kind, void *p, size_t size) {
 ///////////////////////////////
 // Block Arena
 
+void AddBlock(BlockArena *arena, size_t size = MiB(1)) {
+    size_t block_size = MiB(1);
+    if (size > block_size) {
+        block_size = size;
+    }
+    if (arena->allocator.proc == NULL) {
+        arena->allocator = GetSystemAllocator();
+    }
+    BlockArenaNode *new_block = (BlockArenaNode *)AllocSize(arena->allocator, block_size + sizeof(BlockArenaNode) + 16);
+    Assert(GetAlignOffset((size_t)new_block->start, DEFAULT_ALIGNMENT) == 0);
+    arena->start = new_block->start;
+    new_block->end = arena->end = new_block->start + block_size;
+    SLL_STACK_ADD(arena->blocks, new_block);
+}
+
+API void Init(BlockArena *arena) {
+    AddBlock(arena);
+}
+
 API void *PushSize(BlockArena *arena, size_t size) {
     if (size > (size_t)(arena->end - arena->start)) {
-        size_t block_size = MiB(1);
-        if (size > block_size) {
-            block_size = size;
-        }
-        if (arena->allocator.proc == NULL) {
-            arena->allocator = GetSystemAllocator();
-        }
-        BlockArenaNode *new_block = (BlockArenaNode *)AllocSize(arena->allocator, block_size + sizeof(BlockArenaNode));
-        Assert(GetAlignOffset((size_t)new_block->start, DEFAULT_ALIGNMENT) == 0);
-        arena->start = new_block->start;
-        new_block->end = arena->end = new_block->start + block_size;
-        SLL_STACK_ADD(arena->blocks, new_block);
+        AddBlock(arena, size);
     }
     U8 *result = arena->start;
     Assert(GetAlignOffset((size_t)result, DEFAULT_ALIGNMENT) == 0);
@@ -293,23 +294,31 @@ API void Release(BlockArena *arena) {
         next = it->next;
         Dealloc(arena->allocator, it);
     }
+    bool safe = arena->when_unwinding_make_sure_to_keep_last_block_alive;
     MemoryZero(arena, sizeof(BlockArena));
+    arena->when_unwinding_make_sure_to_keep_last_block_alive = safe;
 }
 
 API void Unwind(BlockArena *arena, U8 *pos) {
     bool contains = false;
     for (BlockArenaNode *it = arena->blocks, *next = NULL; it; it = next) {
         next = it->next;
-        if ((pos >= it->start) && (pos < it->end)) {
+        if ((pos >= it->start) && (pos <= it->end)) {
             contains = true;
             break;
         } else {
-            arena->blocks = arena->blocks->next;
-            Dealloc(arena->allocator, it);
+            bool keep = arena->when_unwinding_make_sure_to_keep_last_block_alive && arena->blocks->next == NULL;
+            if (keep) {
+                pos = arena->blocks->start;
+            } else {
+                arena->blocks = arena->blocks->next;
+                Dealloc(arena->allocator, it);
+            }
         }
     }
-    Assert(contains || pos == NULL);
+    Assert(contains || pos == NULL || (arena->when_unwinding_make_sure_to_keep_last_block_alive && pos == arena->blocks->start));
     arena->start = pos;
+    arena->end = arena->blocks ? arena->blocks->end : NULL;
 }
 
 API void *BlockArenaAllocatorProc(void *object, int kind, void *p, size_t size) {
@@ -323,6 +332,37 @@ API void *BlockArenaAllocatorProc(void *object, int kind, void *p, size_t size)
     return NULL;
 }
 
+thread_local VirtualArena ScratchArenas[4];
+
+API void InitScratch() {
+    for (int i = 0; i < 4; i += 1) {
+        InitArena(&ScratchArenas[i]);
+    }
+}
+
+API VirtualArena *GetScratchEx(VirtualArena **conflicts, int conflict_count) {
+    VirtualArena *unoccupied = 0;
+    for (int i = 0; i < Lengthof(ScratchArenas); i += 1) {
+        VirtualArena *from_pool = &ScratchArenas[i];
+        unoccupied = from_pool;
+        for (int conflict_i = 0; conflict_i < conflict_count; conflict_i += 1) {
+            VirtualArena *from_conflict = conflicts[conflict_i];
+            if (from_pool == from_conflict) {
+                unoccupied = 0;
+                break;
+            }
+        }
+
+        if (unoccupied) {
+            break;
+        }
+    }
+
+    // Failed to get free scratch memory, this is a fatal error, this shouldnt happen
+    Assert(unoccupied);
+    return unoccupied;
+}
+
 void RunArenaTest() {
     Allocator memory_tracking_allocator = GetTrackingAllocator();
     {
@@ -335,7 +375,7 @@ void RunArenaTest() {
             }
         }
         Release(&arena);
-        TrackingAllocatorCheck();
+        Assert(MemoryTrackingRecord.len == 0);
     }
 
     {
@@ -351,7 +391,7 @@ void RunArenaTest() {
         Assert(arena.blocks[0].start == (U8 *)vals);
         Unwind(&arena, NULL);
         Dealloc(arena.allocator, arena.blocks);
-        TrackingAllocatorCheck();
+        Assert(MemoryTrackingRecord.len == 0);
     }
 
     {
@@ -372,7 +412,7 @@ void RunArenaTest() {
         Assert(arena.start == p);
 
         Release(&arena);
-        TrackingAllocatorCheck();
+        Assert(MemoryTrackingRecord.len == 0);
     }
 
     {
@@ -381,6 +421,22 @@ void RunArenaTest() {
         U8 *a = (U8 *)PushSize(&arena, KiB(2000));
         Assert((size_t)(arena.blocks[0].end - arena.blocks[0].start) == KiB(2000));
         Release(&arena);
-        TrackingAllocatorCheck();
+        Assert(MemoryTrackingRecord.len == 0);
+    }
+
+    {
+        BlockArena arena = {};
+        arena.allocator = memory_tracking_allocator;
+        arena.when_unwinding_make_sure_to_keep_last_block_alive = true;
+        PushSize(&arena, KiB(2000));
+        Unwind(&arena, 0);
+        Assert(MemoryTrackingRecord.len == 1);
+        PushSize(&arena, MiB(1));
+        PushSize(&arena, MiB(3));
+        Assert(MemoryTrackingRecord.len == 2);
+        Unwind(&arena, 0);
+        Assert(MemoryTrackingRecord.len == 1);
+        Release(&arena);
+        Assert(MemoryTrackingRecord.len == 0);
     }
 }