Revert "Simplify Arena allocator"

This reverts commit d4313fb8e4.

mem/include/arena/mem_arena.h

@@ -2,6 +2,7 @@
 #define MEM_ARENA_H
 
 #include "aliases.h"
+#include "mem_allocator.h"
 #include <stdbool.h>
 
 #ifdef __cplusplus
@@ -10,9 +11,13 @@ extern "C" {
 
 typedef struct growing_arena Arena;
 
+Allocator wapp_mem_arena_allocator(const Arena *arena);
 bool wapp_mem_arena_init(Arena **arena, u64 base_capacity);
 void *wapp_mem_arena_alloc(Arena *arena, u64 size);
 void *wapp_mem_arena_alloc_aligned(Arena *arena, u64 size, u64 alignment);
+void *wapp_mem_arena_realloc(Arena *arena, void *ptr, u64 size);
+void *wapp_mem_arena_realloc_aligned(Arena *arena, void *ptr, u64 size,
+                                     u64 alignment);
 void wapp_mem_arena_clear(Arena *arena);
 void wapp_mem_arena_free(Arena **arena);
 

mem/src/arena/mem_arena.c

@@ -1,6 +1,8 @@
 #include "mem_arena.h"
 #include "aliases.h"
+#include "mem_allocator.h"
 #include "mem_utils.h"
+#include <math.h>
 #include <stdbool.h>
 #include <stdlib.h>
 #include <string.h>
@@ -11,12 +13,24 @@
 #define DEFAULT_ALIGNMENT (2 * sizeof(void *))
 #endif /* ifndef DEFAULT_ALIGNMENT */
 
+#define HDR_MAGIC_BYTE_COUNT 8
+#define HDR_MAGIC                                                              \
+  { 0x57, 0x41, 0x41, 0x52, 0x4e, 0x48, 0x44, 0x52 }
+#define MAX_HDR_SEARCH_LENGTH 256
+
 #define ARENA_MINIMUM_CAPACITY 1024
 
+typedef struct arena_alloc_hdr ArenaAllocHDR;
+struct arena_alloc_hdr {
+  u8 magic[HDR_MAGIC_BYTE_COUNT];
+  u64 alloc_size;
+  u64 alignment;
+  u8 *alloc_start;
+};
+
 typedef struct base_arena BaseArena;
 struct base_arena {
   u8 *buf;
-  u8 *last;
   u8 *offset;
   u64 capacity;
   BaseArena *prev;
@@ -29,14 +43,35 @@ struct growing_arena {
   u64 initial_capacity;
 };
 
+internal BaseArena *find_arena_from_pointer(const Arena *arena, void *ptr);
+
 internal bool base_arena_init(BaseArena *arena, u64 capacity);
 internal void *base_arena_alloc_aligned(BaseArena *arena, u64 size,
                                         u64 alignment);
 internal void base_arena_clear(BaseArena *arena);
 internal void base_arena_free(BaseArena *arena);
+internal ArenaAllocHDR *find_alloc_header(BaseArena *arena, void *alloc_ptr);
+
+internal void *mem_arena_alloc(u64 size, void *alloc_obj);
+internal void *mem_arena_alloc_aligned(u64 size, u64 alignment,
+                                       void *alloc_obj);
+internal void *mem_arena_realloc(void *ptr, u64 size, void *alloc_obj);
+internal void *mem_arena_realloc_aligned(void *ptr, u64 size, u64 alignment,
+                                         void *alloc_obj);
 
 // PUBLIC API
 
+Allocator wapp_mem_arena_allocator(const Arena *arena) {
+  return (Allocator){
+      .obj = (void *)arena,
+      .alloc = mem_arena_alloc,
+      .alloc_aligned = mem_arena_alloc_aligned,
+      .realloc = mem_arena_realloc,
+      .realloc_aligned = mem_arena_realloc_aligned,
+      .free = NULL,
+  };
+}
+
 bool wapp_mem_arena_init(Arena **arena, u64 base_capacity) {
   if (!arena || *arena || base_capacity == 0) {
     return false;
@@ -107,6 +142,54 @@ void *wapp_mem_arena_alloc_aligned(Arena *arena, u64 size, u64 alignment) {
   return output;
 }
 
+void *wapp_mem_arena_realloc(Arena *arena, void *ptr, u64 size) {
+  return wapp_mem_arena_realloc_aligned(arena, ptr, size, DEFAULT_ALIGNMENT);
+}
+
+void *wapp_mem_arena_realloc_aligned(Arena *arena, void *ptr, u64 size,
+                                     u64 alignment) {
+  if (!arena) {
+    return NULL;
+  }
+
+  BaseArena *base_arena = find_arena_from_pointer(arena, ptr);
+  if (!base_arena) {
+    return NULL;
+  }
+
+  ArenaAllocHDR *header = find_alloc_header(base_arena, ptr);
+  if (!header) {
+    return NULL;
+  }
+
+  if (header->alloc_start + header->alloc_size == base_arena->offset) {
+    // Avoid allocating new pointer and copying memory if pointer is at the end
+    // of the arena
+    i64 diff = size - header->alloc_size;
+
+    u8 *new_offset = base_arena->offset + diff;
+
+    u8 *clear_start = diff < 0 ? new_offset : base_arena->offset;
+    memset(clear_start, 0, llabs(diff));
+
+    header->alloc_size = size;
+    base_arena->offset = new_offset;
+
+    return header->alloc_start;
+  }
+
+  void *new_alloc = wapp_mem_arena_alloc_aligned(arena, size, alignment);
+  if (!new_alloc) {
+    return NULL;
+  }
+
+  u64 to_copy = size < header->alloc_size ? size : header->alloc_size;
+
+  memcpy(new_alloc, ptr, to_copy);
+
+  return new_alloc;
+}
+
 void wapp_mem_arena_clear(Arena *arena) {
   if (!arena) {
     return;
@@ -172,6 +255,30 @@ void wapp_mem_arena_free(Arena **arena) {
 
 // INTERNAL FUNCTIONS
 
+internal BaseArena *find_arena_from_pointer(const Arena *arena, void *ptr) {
+  if (!arena || !ptr) {
+    return NULL;
+  }
+
+  // Ensure pointer is not out of bounds
+  u8 *alloc = (u8 *)ptr;
+  BaseArena *active = arena->active_arena;
+
+  if (alloc > active->buf + arena->initial_capacity) {
+    return NULL;
+  }
+
+  for (u64 i = 0; i < arena->count; ++i) {
+    if (alloc >= active->buf && alloc < active->buf + arena->initial_capacity) {
+      return (BaseArena *)active;
+    }
+
+    active = active->prev;
+  }
+
+  return NULL;
+}
+
 internal bool base_arena_init(BaseArena *arena, u64 capacity) {
   if (!arena || arena->buf || capacity == 0) {
     return false;
@@ -198,15 +305,23 @@ internal void *base_arena_alloc_aligned(BaseArena *arena, u64 size,
     return NULL;
   }
 
-  u8 *alloc_start = arena->offset;
+  u8 *start_offset = arena->offset;
+  u8 *alloc_start = arena->offset + sizeof(ArenaAllocHDR);
 
   u8 *output = wapp_mem_util_align_forward((void *)alloc_start, alignment);
   if (output + size >= arena->buf + arena->capacity) {
     return NULL;
   }
 
+  ArenaAllocHDR *header = (ArenaAllocHDR *)start_offset;
+  *header = (ArenaAllocHDR){
+      .magic = HDR_MAGIC,
+      .alloc_size = size,
+      .alignment = alignment,
+      .alloc_start = output,
+  };
+
   arena->offset = output + size;
-  arena->last = output;
 
   return (void *)output;
 }
@@ -233,3 +348,59 @@ internal void base_arena_free(BaseArena *arena) {
   arena->capacity = 0;
   arena->prev = arena->next = NULL;
 }
+
+internal ArenaAllocHDR *find_alloc_header(BaseArena *arena, void *alloc_ptr) {
+  persistent const u8 magic[HDR_MAGIC_BYTE_COUNT] = HDR_MAGIC;
+
+  u8 *current = (u8 *)alloc_ptr;
+
+  u8 *max_search_end = current - MAX_HDR_SEARCH_LENGTH;
+  u8 *arena_buf_start = arena->buf;
+
+  u8 *search_end =
+      max_search_end > arena_buf_start ? max_search_end : arena_buf_start;
+
+  bool match;
+  for (; current >= search_end; --current) {
+    match = true;
+
+    for (u64 i = 0; i < HDR_MAGIC_BYTE_COUNT; ++i) {
+      if (current[i] != magic[i]) {
+        match = false;
+        break;
+      }
+    }
+
+    if (match) {
+      return (ArenaAllocHDR *)current;
+    }
+  }
+
+  return NULL;
+}
+
+internal void *mem_arena_alloc(u64 size, void *alloc_obj) {
+  Arena *arena = (Arena *)alloc_obj;
+
+  return wapp_mem_arena_alloc(arena, size);
+}
+
+internal void *mem_arena_alloc_aligned(u64 size, u64 alignment,
+                                       void *alloc_obj) {
+  Arena *arena = (Arena *)alloc_obj;
+
+  return wapp_mem_arena_alloc_aligned(arena, size, alignment);
+}
+
+internal void *mem_arena_realloc(void *ptr, u64 size, void *alloc_obj) {
+  Arena *arena = (Arena *)alloc_obj;
+
+  return wapp_mem_arena_realloc(arena, ptr, size);
+}
+
+internal void *mem_arena_realloc_aligned(void *ptr, u64 size, u64 alignment,
+                                         void *alloc_obj) {
+  Arena *arena = (Arena *)alloc_obj;
+
+  return wapp_mem_arena_realloc_aligned(arena, ptr, size, alignment);
+}