add arena

Makefile

@@ -1,7 +1,7 @@
 CXX = clang++
 
 d:
-	$(CXX) -g -O0 -o rntviewer rntviewer.cpp $(shell root-config --cflags --libs) -lROOTNTuple
+	$(CXX) -DDEBUG -g -O0 -o rntviewer rntviewer.cpp $(shell root-config --cflags --libs) -lROOTNTuple -lasan
 
 r:
 	$(CXX) -O2 -o rntviewer rntviewer.cpp $(shell root-config --cflags --libs) -lROOTNTuple

mem.cpp

@@ -0,0 +1,122 @@
+#define align_pow2(x, b) (((x) + (b) - 1) & ( ~((b) - 1)))
+#define SLLStackPush_N(f,n,next) ((n)->next=(f), (f)=(n))
+#define SLLStackPop_N(f,next) ((f)=(f)->next)
+
+#ifdef DEBUG
+#define asan_poison_memory_region(mem, cmt) __asan_poison_memory_region(mem, cmt)
+#define asan_unpoison_memory_region(mem, cmt) __asan_unpoison_memory_region(mem, cmt)
+#else
+#define asan_poison_memory_region(mem, cmt)
+#define asan_unpoison_memory_region(mem, cmt)
+#endif
+
+void *os_reserve(u64 size)
+{
+  return mmap(0, size, PROT_NONE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
+}
+
+void os_release(void *mem, u64 size)
+{
+  munmap(mem, size);
+}
+
+b32x os_commit(void *addr, u64 size)
+{
+  return mprotect(addr, size, PROT_READ|PROT_WRITE) == 0;
+}
+
+// init_res: initial reserved size
+// init_cmt: initial committed size
+Arena *arena_alloc_sized(u64 init_res, u64 init_cmt)
+{
+  assert(ARENA_HEADER_SIZE < init_cmt && init_cmt <= init_res);
+
+  u64 page_size = getpagesize();
+  u64 res = align_pow2(init_res, page_size);
+  u64 cmt = align_pow2(init_cmt, page_size);
+
+  // reserve memory
+  void *mem = os_reserve(res);
+  if (!os_commit(mem, cmt)) {
+    munmap(mem, init_res);
+  }
+
+  Arena *arena = (Arena *)mem;
+  if (arena) {
+    asan_poison_memory_region(mem, cmt);
+    asan_unpoison_memory_region(mem, ARENA_HEADER_SIZE);
+
+    arena->cur = arena;
+    arena->pos = ARENA_HEADER_SIZE;
+    arena->cmt = cmt;
+    arena->res = res;
+    arena->align = 8;
+  }
+
+  return arena;
+}
+
+Arena *arena_alloc()
+{
+  return arena_alloc_sized(ARENA_RESERVE_SIZE, ARENA_COMMIT_SIZE);
+}
+
+void arena_release(Arena *arena)
+{
+  for (Arena *node = arena->cur, *prev = 0; node; node = prev) {
+    prev = node->prev;
+    os_release(node, node->res);
+  }
+}
+
+u64 arena_huge_push_threshold()
+{
+  u64 res = ARENA_RESERVE_SIZE;
+  u64 threshold = (res - ARENA_HEADER_SIZE) / 2 + 1;
+  return threshold;
+}
+
+void *arena_push_impl(Arena *arena, u64 size)
+{
+  Arena *cur = arena->cur;
+  u64 pos_mem = align_pow2(cur->pos, arena->align);
+  u64 pos_new = pos_mem + size;
+
+  if (cur->res < pos_new) {
+    Arena *new_block;
+    if (size < arena_huge_push_threshold()) {
+      new_block = arena_alloc();
+    } else {
+      u64 new_block_size = size + ARENA_HEADER_SIZE;
+      new_block = arena_alloc_sized(new_block_size, new_block_size);
+    }
+
+    if (new_block) {
+      new_block->base_pos = cur->base_pos + cur->res;
+      SLLStackPush_N(arena->cur, new_block, prev);
+
+      cur = new_block;
+      pos_mem = align_pow2(cur->pos, cur->align);
+      pos_new = pos_mem + size;
+    }
+  }
+
+  if (cur->cmt < pos_new) {
+    u64 cmt_new_aligned = align_pow2(pos_new, ARENA_COMMIT_SIZE);
+    u64 cmt_new_clamped = std::min(cmt_new_aligned, cur->res);
+    u64 cmt_new_size = cmt_new_clamped - cur->cmt;
+    b32x is_cmt_ok = os_commit((u8*)cur + cur->cmt, cmt_new_size);
+
+    if (is_cmt_ok)
+      cur->cmt = cmt_new_clamped;
+  }
+
+  void *mem = 0;
+  if (cur->cmt >= pos_new) {
+    mem = (u8*)cur + pos_mem;
+    cur->pos = pos_new;
+    asan_unpoison_memory_region(mem, size);
+  }
+
+  return mem;
+}

mem.h

@@ -0,0 +1,13 @@
+#define ARENA_HEADER_SIZE 128
+#define ARENA_COMMIT_SIZE (64 * 1024)
+#define ARENA_RESERVE_SIZE (64 * 1024 * 1024)
+
+struct Arena {
+  Arena *prev;
+  Arena *cur;
+  u64 base_pos;
+  u64 pos;
+  u64 cmt;
+  u64 res;
+  u64 align;
+};

rntviewer.cpp

@@ -1,10 +1,22 @@
-#include "defer.hpp"
-
 #include <ROOT/RNTupleReader.hxx>
+#include <ROOT/RNTuple.hxx>
+#include <TFile.h>
 
 #include <cstdio>
+#include <cstdint>
 
 #include <sys/mman.h>
+#include <unistd.h>
+
+#ifdef DEBUG
+#include <sanitizer/asan_interface.h>
+#endif
+
+#include "types.h"
+#include "defer.hpp"
+#include "mem.h"
+
+#include "mem.cpp"
 
 using namespace ROOT::Experimental;
 
@@ -18,14 +30,24 @@ static size_t file_size(FILE *f)
 
 int main(int argc, char **argv)
 {
-  if (argc < 2) {
+  if (argc < 3) {
-    fprintf(stderr, "Usage: %s <ntuple_file.root>\n", argv[0]);
+    fprintf(stderr, "Usage: %s <ntuple_name> <ntuple_file.root>\n", argv[0]);
     return 1;
   }
 
-  // Open and map the file
+  // Collect arguments
-  const char *fname = argv[0];
+  const char *ntpl_name = argv[1];
+  const char *fname = argv[2];
 
+  // Allocate memory
+  Arena *arena = arena_alloc();
+  if (!arena) {
+    fprintf(stderr, "Failed to allocate memory\n");
+    return 1;
+  }
+  defer { arena_release(arena); };
+
+  // Open and map the file
   FILE *file = fopen(fname, "rb");
   defer { fclose(file); };
 
@@ -38,7 +60,20 @@ int main(int argc, char **argv)
   }
   defer { munmap(fmem, fsize); };
 
+  // Open the TFile
+  TFile *tfile = TFile::Open(fname, "READ");
+  if (!tfile) {
+    fprintf(stderr, "Failed to open TFile.\n");
+    return 1;
+  }
+  defer { delete tfile; };
+
   // Get the RNTuple information
+  const RNTuple *anchor = tfile->Get<RNTuple>(ntpl_name);
+  if (!anchor) {
+    fprintf(stderr, "RNTuple '%s' not found in %s.\n", ntpl_name, fname);
+    return 1;
+  }
 
   return 0;
 }

types.h

@@ -0,0 +1,13 @@
+using u8 = uint8_t;
+using u16 = uint16_t;
+using u32 = uint32_t;
+using u64 = uint64_t;
+
+using f32 = float;
+using f64 = double;
+
+using b8 = bool;
+using b32x = int;
+
+using i32 = int32_t;
+using i64 = int64_t;