From b7a0481ed1659529231a2f93f335827f05101a8a Mon Sep 17 00:00:00 2001
From: silverweed <silverweed14@proton.me>
Date: Tue, 16 Jul 2024 17:50:30 +0200
Subject: [PATCH] introduce proper thread-local scratch arenas

---
 src/mem.cpp       | 43 +++++++++++++++++++++++++++++++++++++++++++
 src/mem.h         | 15 +++++++++++++++
 src/render.cpp    |  9 ++++++---
 src/render.h      |  3 +++
 src/rntuple.cpp   |  4 ++--
 src/rntviewer.cpp |  4 ++++
 6 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/src/mem.cpp b/src/mem.cpp
index 73d5935..d896316 100644
--- a/src/mem.cpp
+++ b/src/mem.cpp
@@ -11,6 +11,49 @@
 #define asan_unpoison_memory_region(mem, cmt)
 #endif
 
+thread_local Thread_Ctx *thread_local_ctx = nullptr;
+
+internal
+void tctx_init(Thread_Ctx &tctx) {
+  memset(&tctx, 0, sizeof(Thread_Ctx));
+  Arena **arena_ptr = tctx.arenas;
+  for (u64 i = 0; i < countof(tctx.arenas); ++i, ++arena_ptr) {
+    *arena_ptr = arena_alloc();
+  }
+  thread_local_ctx = &tctx;
+}
+
+internal
+void tctx_release()
+{
+  for (u64 i = 0; i < countof(thread_local_ctx->arenas); ++i)
+    arena_release(thread_local_ctx->arenas[i]);
+}
+
+internal
+Arena *tctx_get_scratch(Arena **conflicts, u64 count)
+{
+  Thread_Ctx *tctx = thread_local_ctx;
+  Arena *res = nullptr;
+  Arena **arena_ptr = tctx->arenas;
+  for (u64 i = 0; i < countof(tctx->arenas); ++i, ++arena_ptr) {
+    Arena **conflict_ptr = conflicts;
+    b8 has_conflict = false;
+    for (u64 j = 0; j < count; ++j, ++conflict_ptr) {
+      if (*arena_ptr == *conflict_ptr) {
+        has_conflict = true;
+        break;
+      }
+    }
+    if (!has_conflict) {
+      res = *arena_ptr;
+      break;
+    }
+  }
+
+  return res;
+}
+
 // init_res: initial reserved size
 // init_cmt: initial committed size
 internal 
diff --git a/src/mem.h b/src/mem.h
index 1bad0c6..c2b8b6e 100644
--- a/src/mem.h
+++ b/src/mem.h
@@ -19,3 +19,18 @@ struct Temp {
   Arena *arena;
   u64 pos;
 };
+
+struct Thread_Ctx {
+  Arena *arenas[2];
+};
+
+#define countof(a) (sizeof(a) / sizeof((a)[0]))
+
+
+internal Arena *arena_alloc();
+internal void arena_release(Arena *arena);
+
+internal Arena *tctx_get_scratch(Arena **conflicts, u64 count);
+
+#define scratch_begin(conflicts, count) temp_begin(tctx_get_scratch((conflicts), (count)))
+#define scratch_end(scratch)            temp_end(scratch)
diff --git a/src/render.cpp b/src/render.cpp
index e27d714..dbbcc1a 100644
--- a/src/render.cpp
+++ b/src/render.cpp
@@ -179,8 +179,8 @@ void update_and_render(Arena *arena, App_State &app, f32 delta_time_ms)
   ImGui::SetNextWindowPos({ 0, 0 });
   ImGui::SetNextWindowSize({ (f32)app.win_data.width, (f32)app.win_data.height });
 
-  Temp scratch = temp_begin(arena);
-  defer { temp_end(scratch); };
+  Temp scratch = scratch_begin(&arena, 1);
+  defer { scratch_end(scratch); };
 
   const auto main_win_flags = ImGuiWindowFlags_NoResize | ImGuiWindowFlags_NoTitleBar | ImGuiWindowFlags_NoDecoration;
   if (ImGui::Begin("main", nullptr, main_win_flags)) {
@@ -219,7 +219,10 @@ void update_and_render(Arena *arena, App_State &app, f32 delta_time_ms)
       void *content = app.inspected_file.mem + app.viewer.base_display_addr;
       u64 content_size = app.inspected_file.size - app.viewer.base_display_addr;
       app.last_pinfo = &invalid_pinfo;
-      app.viewer.mem_edit.DrawContents(content, content_size, app.viewer.base_display_addr);
+      {
+        TIMED_SCOPE();
+        app.viewer.mem_edit.DrawContents(content, content_size, app.viewer.base_display_addr);
+      }
 
       ImGui::TableNextColumn();
       ImGuiColorEditFlags flags =  ImGuiColorEditFlags_NoInputs|ImGuiColorEditFlags_NoLabel;
diff --git a/src/render.h b/src/render.h
index 664e2e6..665f676 100644
--- a/src/render.h
+++ b/src/render.h
@@ -13,6 +13,9 @@ struct Viewer {
 
   u64 base_display_addr;
   u64 latest_page_gone_to;
+
+  // has size equal to app.inspected_file.size
+  u32 *byte_colors;
 };
 
 struct Edit_Bg_Color_Data {
diff --git a/src/rntuple.cpp b/src/rntuple.cpp
index 2197340..f0ab40c 100644
--- a/src/rntuple.cpp
+++ b/src/rntuple.cpp
@@ -15,8 +15,8 @@ ROOT::Experimental::RNTupleDescriptor create_descriptor(Arena *arena, RMicroFile
   using namespace ROOT::Experimental;
   using namespace ROOT::Experimental::Internal;
 
-  Temp scratch = temp_begin(arena);
-  defer { temp_end(scratch); };
+  Temp scratch = scratch_begin(&arena, 1);
+  defer { scratch_end(scratch); };
 
   const RNTuple_Anchor &anchor = info.anchor;
   
diff --git a/src/rntviewer.cpp b/src/rntviewer.cpp
index 37aae02..bdce714 100644
--- a/src/rntviewer.cpp
+++ b/src/rntviewer.cpp
@@ -77,6 +77,10 @@ void app_cleanup(App_State &app)
 
 int main(int argc, char **argv)
 {
+  Thread_Ctx tctx;
+  tctx_init(tctx);
+  defer { tctx_release(); };
+  
   if (argc > 1 && argv[1][0] == '-') {
     fprintf(stderr, "Usage: %s <ntuple_name> <ntuple_file.root>\n", argv[0]);
     return 1;