From ac8ee66af93b5aeb152846696349654cbd6f689f Mon Sep 17 00:00:00 2001
From: silverweed <silverweed14@proton.me>
Date: Tue, 21 Jan 2025 17:33:38 +0100
Subject: [PATCH] wip

---
 src/rntuple.cpp   | 63 +++++++++++++++++++++++++++++++++++++++++++++--
 src/rntuple.h     |  1 +
 src/rntviewer.cpp |  2 +-
 3 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/src/rntuple.cpp b/src/rntuple.cpp
index 9bf5656..9aa0f14 100644
--- a/src/rntuple.cpp
+++ b/src/rntuple.cpp
@@ -162,11 +162,11 @@ u64 calc_page_uncomp_size(const u8 *fmem, const Page_Info_Node *page_head)
 }
 
 internal
-RNTuple_Data get_rntuple_data(Arena *arena, const Inspected_File &file, const TFile_Data &tfile_data, b8 extended_info)
+RNTuple_Data get_rntuple_data(Arena *arena, const Inspected_File &file, const TFile_Data &tfile_data, const ROOT::RNTuple &anchor, b8 extended_info)
 {
   RNTuple_Data rndata {};
 
-  if (!tfile_data.anchors_head->anchor.GetNBytesHeader()) // FIXME :multiple_anchors:
+  if (!anchor.GetNBytesHeader())
     return rndata;
   
   using ROOT::Experimental::RNTupleDescriptor;
@@ -436,6 +436,65 @@ RNTuple_Data get_rntuple_data(Arena *arena, const Inspected_File &file, const TF
   return rndata;
 }
 
+template <typename T>
+internal
+T *merge_arrays(Arena *arena, const T *a, u64 n_a, const T *b, u64 n_b, u64 &n_merged)
+{
+  T *merged = nullptr;
+  n_merged = n_a + n_b;
+  if (n_merged) {
+    merged = arena_push_array<T>(arena, n_merged);
+    if (n_a)
+      memcpy(merged, a, n_a * sizeof(T));
+    if (n_b)
+      memcpy(merged + n_a * sizeof(T), b, n_b * sizeof(T));
+  } 
+  return merged;
+}
+
+internal
+RNTuple_Data merge_rntuple_data(Arena *arena, const RNTuple_Data &a, const RNTuple_Data &b)
+{
+  // XXX: this is totally wrong, as we cannot simply memcpy things and expect them to work! (they are linked lists
+  // in almost all cases, or they have pointers to other stuff).
+  // Either construct a list of RNTuple_Data and keep them all separate, or properly concatenate their content
+  // where appropriate (this means that we might leak some memory because we cannot reset the individual RNTuple_Datas)
+  // Hm. Gotta think about it.
+  RNTuple_Data merged {};
+  merged.pages = merge_arrays(arena, a.pages, a.n_pages, b.pages, b.n_pages, merged.n_pages);
+  merged.n_elems = a.n_elems + b.n_elems;
+  merged.tot_page_comp_size = a.tot_page_comp_size + b.tot_page_comp_size;
+  merged.tot_page_uncomp_size = a.tot_page_uncomp_size + b.tot_page_uncomp_size;
+  // TODO: checksums
+  merged.cluster_groups = merge_arrays(arena, a.cluster_groups, a.n_cluster_groups, b.cluster_groups, b.n_cluster_groups, merged.n_cluster_groups);
+  merged.tot_page_list_size = a.tot_page_list_size + b.tot_page_list_size;
+  merged.clusters = merge_arrays(arena, a.clusters, a.n_clusters, b.clusters, b.n_clusters, merged.n_clusters);
+  merged.n_entries = a.n_entries + b.n_entries;
+  merged.page_groups = merge_arrays(arena, a.page_groups, a.n_page_groups, b.page_groups, b.n_page_groups, merged.n_page_groups);
+  merged.page_chunks = merge_arrays(arena, a.page_chunks, a.n_page_chunks, b.page_chunks, b.n_page_chunks, merged.n_page_chunks);
+
+  return merged;
+}
+
+internal
+RNTuple_Data get_all_rntuple_data(Arena *arena, const Inspected_File &file, const TFile_Data &tfile_data, b8 extended_info)
+{
+  Temp scratch = scratch_begin(&arena, 1);
+  defer { scratch_end(scratch); };
+
+  RNTuple_Data all_data {};
+  
+  for (RNTuple_Anchor_Node *node = tfile_data.anchors_head; node; node = node->next) {
+    RNTuple_Data rndata = get_rntuple_data(scratch.arena, file, tfile_data, node->anchor, extended_info);
+    // Only copy data to permanent storage when we merge the last anchor data
+    Arena *a = node == tfile_data.anchors_tail ? arena : scratch.arena;
+    printf("merging to arena: %p (scratch: %p)\n", a, scratch.arena);
+    all_data = merge_rntuple_data(a, all_data, rndata);
+  }
+
+  return all_data;
+}
+
 internal
 Byte_Range get_section_range(const App_State &app, Section_Id sec)
 {
diff --git a/src/rntuple.h b/src/rntuple.h
index 00c5913..1b95030 100644
--- a/src/rntuple.h
+++ b/src/rntuple.h
@@ -157,6 +157,7 @@ struct TFile_Data {
   RNTuple_Anchor_Node *anchors_head, *anchors_tail;
 };
 
+// @Volatile: when changed, remember to update merge_rntuple_data()
 struct RNTuple_Data {
   Page_Info_Node *pages;
   u64 n_pages;
diff --git a/src/rntviewer.cpp b/src/rntviewer.cpp
index 885f434..3520381 100644
--- a/src/rntviewer.cpp
+++ b/src/rntviewer.cpp
@@ -174,7 +174,7 @@ int main(int argc, char **argv)
   if (!app.ntpl_name.str)
     fprintf(stderr, "Warning: found no RNTuples in %s\n", args.file_name.c());
   else if (success)
-    app.rndata = get_rntuple_data(arena, app.inspected_file, app.tfile_data, args.extended_info);
+    app.rndata = get_all_rntuple_data(arena, app.inspected_file, app.tfile_data, args.extended_info);
 
   if (args.print_to_terminal) {
     u64 nbytes_displayed = args.nbytes_displayed;