From 3d074ce0ac8ab05b7b95609ab7bdf0f8110dc6f0 Mon Sep 17 00:00:00 2001
From: Krzosa Karol <krzosa.karol@gmail.com>
Date: Sun, 31 Dec 2023 19:41:49 +0100
Subject: [PATCH] Add bootstrap exclusive arena

---
 arena.c             | 16 ++++++++++++----
 arena.h             |  4 +++-
 test/test_arena.cpp | 24 ++++++++++++++++++++++++
 3 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/arena.c b/arena.c
index e82c208..ef1231c 100644
--- a/arena.c
+++ b/arena.c
@@ -156,6 +156,13 @@ MA_API MA_Arena *MA_Bootstrap(void) {
     return arena;
 }
 
+MA_API M_Allocator MA_BootstrapExclusive(void) {
+    MA_Arena bootstrap_arena = {0};
+    MA_Arena *arena = MA_PushStruct(&bootstrap_arena, MA_Arena);
+    *arena = bootstrap_arena;
+    return MA_GetExclusiveAllocator(arena);
+}
+
 MA_API void MA_InitFromBuffer(MA_Arena *arena, void *buffer, size_t size) {
     arena->memory.data = (uint8_t *)buffer;
     arena->memory.commit = size;
@@ -285,10 +292,11 @@ MA_API void *MA_AllocatorProc(void *allocator, M_AllocatorOp kind, void *p, size
 MA_API void *MA_ExclusiveAllocatorProc(void *allocator, M_AllocatorOp kind, void *p, size_t size, size_t old_size) {
     MA_Arena *arena = (MA_Arena *)allocator;
     if (kind == M_AllocatorOp_Reallocate) {
-        if (size > arena->len) {
-            size_t size_to_push = size - arena->len;
-            MA_PushSizeNonZeroed(arena, size_to_push);
-            return arena->memory.data;
+        if (size > old_size) {
+            size_t size_to_push = size - old_size;
+            void *result = MA_PushSizeNonZeroed(arena, size_to_push);
+            if (!p) p = result;
+            return p;
         }
     }
 
diff --git a/arena.h b/arena.h
index 490871a..28b1de4 100644
--- a/arena.h
+++ b/arena.h
@@ -116,9 +116,10 @@ MA_API void            MA_InitEx(MA_Arena *a, size_t reserve);
 MA_API void            MA_Init(MA_Arena *a);
 MA_API MA_Arena        MA_Create();
 MA_API void            MA_MakeSureInitialized(MA_Arena *a);
-MA_API MA_Arena *      MA_Bootstrap(void);
 MA_API void            MA_InitFromBuffer(MA_Arena *arena, void *buffer, size_t size);
 MA_API MA_Arena        MA_MakeFromBuffer(void *buffer, size_t size);
+MA_API MA_Arena *      MA_Bootstrap(void);
+MA_API M_Allocator     MA_BootstrapExclusive(void);
 MA_API MA_Arena        MA_PushArena(MA_Arena *arena, size_t size);
 
 MA_API void *          MA_PushSizeNonZeroed(MA_Arena *a, size_t size);
@@ -154,6 +155,7 @@ MA_API void *          M_Alloc(M_Allocator allocator, size_t size);
 MA_API void *          M_AllocCopy(M_Allocator allocator, void *p, size_t size);
 MA_API void *          M_ReallocNonZeroed(M_Allocator allocator, void *p, size_t size, size_t old_size);
 MA_API void            M_Dealloc(M_Allocator allocator, void *p);
+
 MA_API M_Allocator     M_GetSystemAllocator(void);
 MA_API M_Allocator     MA_GetExclusiveAllocator(MA_Arena *arena);
 MA_API M_Allocator     MA_GetAllocator(MA_Arena *arena);
diff --git a/test/test_arena.cpp b/test/test_arena.cpp
index f67535d..b061d13 100644
--- a/test/test_arena.cpp
+++ b/test/test_arena.cpp
@@ -48,8 +48,32 @@ void TestCreateAllocate() {
     MA_DeallocateArena(&created_arena);
 }
 
+void TestBootstrap() {
+    MA_Arena *arena = MA_Bootstrap();
+    void *r0 = MA_PushSize(arena, 1024);
+    void *r1 = MA_PushSize(arena, 1024);
+    IO_Assert(arena->len >= 2048);
+    IO_Assert(r0 != r1);
+    MA_DeallocateArena(arena);
+}
+
+void TestBootstrapExclusive() {
+    Array<int> v = {MA_BootstrapExclusive()};
+    v.add(0);
+    int *ptr = &v[0];
+    for (int i = 1; i < 1000; i += 1) {
+        IO_Assert(&v[0] == ptr);
+        v.add(i);
+    }
+    int i = 0;
+    For(v) IO_Assert(it == i++);
+    v.dealloc();
+}
+
 int main() {
     TestScratch();
     TestBuffer();
     TestCreateAllocate();
+    TestBootstrap();
+    TestBootstrapExclusive();
 }
\ No newline at end of file