From 5d73051e052bb6b362e7eaf639d39ab22bb7da92 Mon Sep 17 00:00:00 2001
From: silverweed <silverweed14@proton.me>
Date: Mon, 4 Nov 2024 11:30:08 +0100
Subject: [PATCH] add flag to print rntuple names

---
 src/argparse.cpp  |  4 ++++
 src/hover.cpp     |  2 +-
 src/rntuple.cpp   |  2 +-
 src/rntuple.h     |  9 +++++++--
 src/rntviewer.cpp |  8 +++++++-
 src/str.cpp       | 11 +++++++++++
 src/tfile.cpp     | 19 +++++++++++--------
 7 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/src/argparse.cpp b/src/argparse.cpp
index d7444da..6b0d16e 100644
--- a/src/argparse.cpp
+++ b/src/argparse.cpp
@@ -10,6 +10,7 @@ void print_help(const char *argv0)
     "\n\t-l: display LEN bytes (only in terminal mode)"
     "\n\t-w: display WIDTH bytes per column"
     "\n\t-e: display some extended info (may slow down the startup)"
+    "\n\t-n: list the names of the RNTuples found in the file and exit"
     "\n"
     "\nNOTE: if `ntuple_name' is not passed, rntviewer will look for the first RNTuple in the TFile."
     "\n"
@@ -23,6 +24,7 @@ struct Cmdline_Args {
   u64 nbytes_displayed;
   u16 n_cols;
   b8 extended_info;
+  b8 only_print_rntuple_names;
 
   String8 ntpl_name;
   String8 file_name;
@@ -109,6 +111,8 @@ Cmdline_Args parse_args(i32 argc, char **argv)
           parse_int_arg(i, argc, argv, args.start_addr);
         else if (arg[1] == 'l')
           parse_int_arg(i, argc, argv, args.nbytes_displayed);
+        else if (arg[1] == 'n')
+          args.only_print_rntuple_names = true;
         else if (arg[1] == 'w') {
           u64 n_cols = 0;
           parse_int_arg(i, argc, argv, n_cols);
diff --git a/src/hover.cpp b/src/hover.cpp
index 8f3afcb..f1fcd4e 100644
--- a/src/hover.cpp
+++ b/src/hover.cpp
@@ -106,7 +106,7 @@ String8_Node *display_val_rootzip(Arena *arena, String8_Node *prev, const char *
   sn = push_str8_node_child(arena, sn, "Zip method: %s", zip_method.c());
   sn = push_str8_node(arena, sn, "Compressed size: %s", to_pretty_size(arena, comp_size).c());
   sn = push_str8_node(arena, sn, "Uncompressed size: %s", to_pretty_size(arena, uncomp_size).c());
-  sn = push_str8_node(arena, sn, "Comp. ratio: %.2f", (f32)comp_size / uncomp_size);
+  sn = push_str8_node(arena, sn, "Comp. ratio: %f", (f32)comp_size / uncomp_size);
 
   return sn;
 }
diff --git a/src/rntuple.cpp b/src/rntuple.cpp
index 3080c56..7f1bfc5 100644
--- a/src/rntuple.cpp
+++ b/src/rntuple.cpp
@@ -146,7 +146,7 @@ b8 get_tfile_data(Arena *arena, const Inspected_File &file, String8 &ntpl_name,
   if (success) {
     // If we weren't given a rntuple name, use the first one in the file (if any)
     if (!ntpl_name.size && tfile_data.tkeys_data.rntuples)
-      ntpl_name = tfile_data.tkeys_data.rntuples->head->str;
+      ntpl_name = tfile_data.tkeys_data.rntuples->name;
 
     if (ntpl_name.size) {
       { // DEBUG
diff --git a/src/rntuple.h b/src/rntuple.h
index 5f52e28..c686744 100644
--- a/src/rntuple.h
+++ b/src/rntuple.h
@@ -100,11 +100,16 @@ struct Byte_Range_Node {
   Byte_Range rng;
 };
 
+struct RNTuple_Anchor_Info {
+  RNTuple_Anchor_Info *next;
+  String8 name;
+  u64 offset_in_file;
+};
+
 struct TKeys_Data {
   Section sections[Sec_COUNT];
 
-  // XXX: currently we're only using the first RNTuple of the file.
-  String8_Node *rntuples;
+  RNTuple_Anchor_Info *rntuples;
   Byte_Range_Node *rblob_keys;
 };
 
diff --git a/src/rntviewer.cpp b/src/rntviewer.cpp
index f8ae208..30c429e 100644
--- a/src/rntviewer.cpp
+++ b/src/rntviewer.cpp
@@ -156,8 +156,14 @@ int main(int argc, char **argv)
   b8 success = get_tfile_data(arena, app.inspected_file, app.ntpl_name, app.tfile_data);
   if (!app.ntpl_name.str)
     fprintf(stderr, "Warning: found no RNTuples in %s\n", args.file_name.c());
-  else if (success)
+  else if (success) {
+    if (args.only_print_rntuple_names) {
+      for (RNTuple_Anchor_Info *info = app.tfile_data.tkeys_data.rntuples; info; info = info->next)
+        printf("%s at 0x%" PRIX64 "\n", info->name.c(), info->offset_in_file);
+      return 0;
+    }
     app.rndata = get_rntuple_data(arena, app.inspected_file, app.tfile_data, args.extended_info);
+  }
 
   if (args.print_to_terminal) {
     u64 nbytes_displayed = args.nbytes_displayed;
diff --git a/src/str.cpp b/src/str.cpp
index 9bf82d0..8898c7b 100644
--- a/src/str.cpp
+++ b/src/str.cpp
@@ -29,6 +29,17 @@ String8 str8_from_c(const char *str)
   return String8 { (u8*)str, size };
 }
 
+internal
+String8 str8_from_buf(Arena *arena, const u8 *buf, u64 size)
+{
+  String8 s;
+  s.str = arena_push_array_nozero<u8>(arena, size + 1);
+  s.size = size;
+  s.str[size] = 0;
+  memcpy(s.str, buf, size);
+  return s;
+}
+
 internal
 String8 to_pretty_size(Arena *arena, u64 bytes)
 {
diff --git a/src/tfile.cpp b/src/tfile.cpp
index 654aa0e..2267139 100644
--- a/src/tfile.cpp
+++ b/src/tfile.cpp
@@ -144,7 +144,6 @@ b8 walk_tkeys(Arena *arena, const u8 *data, u64 data_len, TFile_Data &tfile_data
     }
 
     tkeys_data.sections[Sec_TKey_List].range.start = seek_keys;
-    printf("seek keys: 0x%lX\n", seek_keys);
 
     // Read right away the tkey list information, so we have it even if the key walk fails to find it
     u32 seek_keys_nbytes_be;
@@ -160,6 +159,8 @@ b8 walk_tkeys(Arena *arena, const u8 *data, u64 data_len, TFile_Data &tfile_data
 
   cur = pre + tfile_obj_nbytes;
 
+  RNTuple_Anchor_Info *rntuple_info_tail = nullptr;
+
   // Walk through all the TKeys in the file and do two things:
   // 1. assign proper start, len and pre_size to all sections as we go
   // 2. collect all RBlob keys and save them for later.
@@ -218,15 +219,17 @@ b8 walk_tkeys(Arena *arena, const u8 *data, u64 data_len, TFile_Data &tfile_data
       if (key_has_class_name("ROOT::Experimental::RNTuple")) {
         u64 name_off = cname_off + cname_len + 1;
         u8 name_len = data[name_off];
+        RNTuple_Anchor_Info *rntuple_info = arena_push<RNTuple_Anchor_Info>(arena);
+        rntuple_info->offset_in_file = cur;
         if (name_len) {
-          char name[256];
-          name[name_len] = 0;
-          memcpy(name, data + name_off + 1, name_len);
-          tkeys_data.rntuples = push_str8_node(arena, tkeys_data.rntuples, "%s", name);
-        } else {
-          tkeys_data.rntuples = push_str8_node(arena, tkeys_data.rntuples, "");
+          rntuple_info->name = str8_from_buf(arena, data + name_off + 1, name_len);
         }
-        printf("Found RNTuple at 0x%lX: %s\n", cur, tkeys_data.rntuples->str.c());
+        if (rntuple_info_tail)
+          rntuple_info_tail->next = rntuple_info;
+        else
+          tkeys_data.rntuples = rntuple_info;
+        rntuple_info_tail = rntuple_info;
+        // printf("Found RNTuple at 0x%lX: %s\n", cur, tkeys_data.rntuples->str.c());
 
         if (!tkeys_data.sections[Sec_RNTuple_Anchor].range.start) {
           u64 anchor_seek = cur + keylen;