From 0c07437ef2d0bf4435339515e6637420a12472ff Mon Sep 17 00:00:00 2001
From: Abdelrahman Said <said.abdelrahman@flawlessai.com>
Date: Sun, 24 Mar 2024 07:59:46 +0000
Subject: [PATCH] Implement memory context with main and temp arenas

---
 mem/include/ctx/mem_ctx.h | 27 +++++++++++++
 mem/src/ctx/mem_ctx.c     | 82 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+)
 create mode 100644 mem/include/ctx/mem_ctx.h
 create mode 100644 mem/src/ctx/mem_ctx.c

diff --git a/mem/include/ctx/mem_ctx.h b/mem/include/ctx/mem_ctx.h
new file mode 100644
index 0000000..e9cd343
--- /dev/null
+++ b/mem/include/ctx/mem_ctx.h
@@ -0,0 +1,27 @@
+#ifndef MEM_CTX_H
+#define MEM_CTX_H
+
+#include "mem_arena.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif // __cplusplus
+
+typedef enum {
+  CTX_DEST_BUFFER_MAIN,
+  CTX_DEST_BUFFER_TEMP,
+
+  COUNT_CTX_DEST_BUFFER,
+} CTXDestBuffer;
+
+void wapp_mem_ctx_init(u64 main_buf_capacity, u64 temp_buf_capacity);
+void *wapp_mem_ctx_alloc(CTXDestBuffer buffer, u64 size);
+void *wapp_mem_ctx_alloc_aligned(CTXDestBuffer buffer, u64 size, u64 alignment);
+void wapp_mem_ctx_clear(CTXDestBuffer buffer);
+void wapp_mem_ctx_free(void);
+
+#ifdef __cplusplus
+}
+#endif // __cplusplus
+
+#endif // !MEM_CTX_H
diff --git a/mem/src/ctx/mem_ctx.c b/mem/src/ctx/mem_ctx.c
new file mode 100644
index 0000000..b27d029
--- /dev/null
+++ b/mem/src/ctx/mem_ctx.c
@@ -0,0 +1,82 @@
+#include "mem_ctx.h"
+#include "aliases.h"
+#include "mem_arena.h"
+#include <assert.h>
+#include <stdlib.h>
+
+typedef struct mem_ctx MemCTX;
+struct mem_ctx {
+  Arena *main;
+  Arena *temp;
+  bool main_initialised;
+  bool temp_initialised;
+};
+
+internal MemCTX g_context = {0};
+
+internal Arena *get_arena(CTXDestBuffer buffer);
+
+void wapp_mem_ctx_init(u64 main_buf_capacity, u64 temp_buf_capacity) {
+  g_context.main_initialised =
+      wapp_mem_arena_init(&g_context.main, main_buf_capacity);
+  g_context.temp_initialised =
+      wapp_mem_arena_init(&g_context.temp, temp_buf_capacity);
+}
+
+void *wapp_mem_ctx_alloc(CTXDestBuffer buffer, u64 size) {
+  Arena *arena = get_arena(buffer);
+  if (!arena) {
+    return NULL;
+  }
+
+  return wapp_mem_arena_alloc(arena, size);
+}
+
+void *wapp_mem_ctx_alloc_aligned(CTXDestBuffer buffer, u64 size,
+                                 u64 alignment) {
+  Arena *arena = get_arena(buffer);
+  if (!arena) {
+    return NULL;
+  }
+
+  return wapp_mem_arena_alloc_aligned(arena, size, alignment);
+}
+
+void wapp_mem_ctx_clear(CTXDestBuffer buffer) {
+  Arena *arena = get_arena(buffer);
+  if (!arena) {
+    return;
+  }
+
+  wapp_mem_arena_clear(arena);
+}
+
+void wapp_mem_ctx_free(void) {
+  wapp_mem_arena_free(&(g_context.main));
+  g_context.main_initialised = false;
+
+  wapp_mem_arena_free(&(g_context.temp));
+  g_context.temp_initialised = false;
+}
+
+internal Arena *get_arena(CTXDestBuffer buffer) {
+  Arena *output = NULL;
+
+  switch (buffer) {
+  case CTX_DEST_BUFFER_MAIN:
+    if (g_context.main_initialised) {
+      output = g_context.main;
+    }
+    break;
+  case CTX_DEST_BUFFER_TEMP:
+    if (g_context.temp_initialised) {
+      output = g_context.temp;
+    }
+    break;
+  default:
+    assert(false && "Not all context destination buffers are handled");
+    break;
+  }
+
+  return output;
+}