Added arena checkpointing

src/arena.c

@@ -13,6 +13,12 @@ static void memory_release(void *mem)
     free(mem);
 }
 
+static INLINE U32 arena_absolute_pos(const Arena *arena)
+{
+    const Arena *current = arena->current;
+    return current->base_pos + current->pos;
+}
+
 Arena *arena_alloc_(const ArenaParams *params)
 {
     Assert(params->size > ARENA_HEADER_SZ);
@@ -27,6 +33,7 @@ Arena *arena_alloc_(const ArenaParams *params)
     arena->current  = arena;
     arena->size     = params->size;
     arena->pos      = ARENA_HEADER_SZ;
+    arena->base_pos = 0;
     arena->prev     = NULL;
 
     return arena;
@@ -35,10 +42,9 @@ Arena *arena_alloc_(const ArenaParams *params)
 void arena_free(Arena *arena)
 {
     Arena *it = arena->current;
-    Arena *prev = NULL;
 
     while (it != NULL) {
-        prev = it->prev;
+        Arena *prev = it->prev;
         memory_release(it);
         it = prev;
     }
@@ -49,11 +55,11 @@ void *arena_push(Arena *arena, U32 size, U32 align)
     Assert(size < (U32_Max - align - ARENA_HEADER_SZ - 1));
     Assert(IsPow2(align));
     Arena *current = arena->current;
-    U32 start_pos = AlignUpPow2(arena->pos, align);
+    U32 start_pos = AlignUpPow2(current->pos, align);
     U32 available = arena_available(current);
 
     if (available < size) {
-        U32 new_size = current->size;
+        U32 new_size = arena->size; // original user requeted size
         if (new_size < size) {
             new_size = AlignUpPow2(size + ARENA_HEADER_SZ,align);
         }
@@ -62,18 +68,57 @@ void *arena_push(Arena *arena, U32 size, U32 align)
         new_params.flags        = current->flags;
         new_params.size         = new_size;
 
-        Arena *new_arena = arena_alloc_(&new_params);
-        new_arena->prev  = current;
-        current->current = new_arena;
-        current          = new_arena;
-        start_pos        = AlignUpPow2(current->pos, align);
+        Arena *new_arena    = arena_alloc_(&new_params);
+        new_arena->base_pos = current->base_pos + current->size;
+        new_arena->prev     = current;
+        arena->current      = new_arena;
+        current             = new_arena;
+        start_pos           = AlignUpPow2(current->pos, align);
     }
 
     U32 end_pos = start_pos + size;
     Assert(end_pos <= current->size);
 
     U8 *mem     = &current->base[start_pos];
-    arena->pos  = start_pos + size;
+    current->pos  = end_pos;
 
     return mem;
 }
+
+void arena_pop_to(Arena *arena, U32 abs_pos)
+{
+    Arena *current = arena->current;
+    abs_pos = Max(ARENA_HEADER_SZ, abs_pos);
+    while (current) {
+        Arena *prev = current->prev;
+
+        if (current->base_pos >= abs_pos) {
+            memory_release(current->base);
+            arena->current = prev;
+        } else {
+            U32 new_pos = abs_pos - current->base_pos;
+            // Note, when doing commit/release, new_pos <= current->pos
+            // is the right thing . shouldn't be able to pop to larger than what's already commited
+            Assert(new_pos <= current->size);
+            current->pos = new_pos;
+            break;
+        }
+
+        current = prev;
+    }
+}
+
+ArenaCheckPoint arena_temp_begin(Arena *arena)
+{
+    ArenaCheckPoint checkpoint = {};
+    checkpoint.arena = arena;
+    checkpoint.absolute_pos = arena_absolute_pos(arena);
+
+    return checkpoint;
+}
+
+void arena_temp_end(ArenaCheckPoint *checkpoint)
+{
+    Arena *arena = checkpoint->arena;
+    arena_pop_to(arena, checkpoint->absolute_pos);
+}