Implement realloc function for Arena

mem/include/arena/mem_arena.h

@@ -13,6 +13,9 @@ typedef struct growing_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,7 @@
 #include "mem_arena.h"
 #include "aliases.h"
 #include "mem_utils.h"
+#include <math.h>
 #include <stdbool.h>
 #include <stdlib.h>
 #include <string.h>
@@ -14,6 +15,7 @@
 #define HDR_MAGIC_BYTE_COUNT 8
 #define HDR_MAGIC                                                              \
   { 0x57, 0x41, 0x41, 0x52, 0x4e, 0x48, 0x44, 0x52 }
+#define MAX_HDR_SEARCH_LENGTH 256
 
 typedef struct arena_alloc_hdr ArenaAllocHDR;
 struct arena_alloc_hdr {
@@ -38,6 +40,8 @@ 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);
@@ -117,6 +121,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;
@@ -182,6 +234,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;
@@ -249,4 +325,34 @@ internal void base_arena_free(BaseArena *arena) {
   arena->prev = arena->next = NULL;
 }
 
-internal ArenaAllocHDR *find_alloc_header(BaseArena *arena, void *alloc_ptr) {}
+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;
+  while (current > search_end) {
+    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;
+    }
+
+    current -= HDR_MAGIC_BYTE_COUNT;
+  }
+
+  return NULL;
+}