From 5a50fece2fdb4b88d4daaa003bf6911977f55db6 Mon Sep 17 00:00:00 2001
From: silverweed <silverweed14@proton.me>
Date: Tue, 12 Nov 2024 12:02:43 +0100
Subject: [PATCH] add flag to print tkeys info

---
 src/argparse.cpp  |  4 ++++
 src/rntuple.cpp   |  4 ++--
 src/rntviewer.cpp |  2 +-
 src/tfile.cpp     | 14 ++++++++++++--
 4 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/argparse.cpp b/src/argparse.cpp
index beb12ab..5c5ebaa 100644
--- a/src/argparse.cpp
+++ b/src/argparse.cpp
@@ -12,6 +12,7 @@ void print_help(const char *argv0)
     "\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\t-v: print version and copyright info"
+    "\n\t-k: print information about the TKeys in the file"
     "\n"
     "\nNOTES:"
     "\n- if `ntuple_name' is not passed, rntviewer will look for the first RNTuple in the TFile."
@@ -58,6 +59,7 @@ struct Cmdline_Args {
   u16 n_cols;
   b8 extended_info;
   b8 only_print_rntuple_names;
+  b8 print_keys_info;
 
   String8 ntpl_name;
   String8 file_name;
@@ -148,6 +150,8 @@ Cmdline_Args parse_args(i32 argc, char **argv)
           args.only_print_rntuple_names = true;
         else if (arg[1] == 'v')
           args.show_version_and_exit = true;
+        else if (arg[1] == 'k')
+          args.print_keys_info = true;
         else if (arg[1] == 'w') {
           u64 n_cols = 0;
           parse_int_arg(i, argc, argv, n_cols);
diff --git a/src/rntuple.cpp b/src/rntuple.cpp
index d288199..2ee6d00 100644
--- a/src/rntuple.cpp
+++ b/src/rntuple.cpp
@@ -148,9 +148,9 @@ const char *get_column_type_name_from_ondisk_type(u16 type)
 }  
 
 internal
-b8 get_tfile_data(Arena *arena, const Inspected_File &file, String8 &ntpl_name, TFile_Data &tfile_data)
+b8 get_tfile_data(Arena *arena, const Inspected_File &file, b8 print_keys_info, String8 &ntpl_name, TFile_Data &tfile_data)
 {  
-  b8 success = walk_tkeys(arena, file.mem, file.size, tfile_data);
+  b8 success = walk_tkeys(arena, file.mem, file.size, print_keys_info, tfile_data);
   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)
diff --git a/src/rntviewer.cpp b/src/rntviewer.cpp
index c6eca00..50e8d35 100644
--- a/src/rntviewer.cpp
+++ b/src/rntviewer.cpp
@@ -157,7 +157,7 @@ int main(int argc, char **argv)
 
   app.ntpl_name = args.ntpl_name; // may be null
   app.base_display_addr = args.start_addr;
-  b8 success = get_tfile_data(arena, app.inspected_file, app.ntpl_name, app.tfile_data);
+  b8 success = get_tfile_data(arena, app.inspected_file, args.print_keys_info, app.ntpl_name, app.tfile_data);
   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);
diff --git a/src/tfile.cpp b/src/tfile.cpp
index 55a8f4d..9d9a0f3 100644
--- a/src/tfile.cpp
+++ b/src/tfile.cpp
@@ -50,7 +50,7 @@ struct RNTuple_Anchor_On_Disk {
 // Additionally, it returns the names of all the RNTuples stored in it.
 // `data` should point to the beginning of a ROOT file.
 internal
-b8 walk_tkeys(Arena *arena, const u8 *data, u64 data_len, TFile_Data &tfile_data)
+b8 walk_tkeys(Arena *arena, const u8 *data, u64 data_len, b8 print_keys_info, TFile_Data &tfile_data)
 { 
   u64 cur = 4; // offset of header version
   if (data_len < cur) {
@@ -60,6 +60,9 @@ b8 walk_tkeys(Arena *arena, const u8 *data, u64 data_len, TFile_Data &tfile_data
 
   TKeys_Data &tkeys_data = tfile_data.tkeys_data;
 
+  Temp scratch = scratch_begin(&arena, 1);
+  defer { scratch_end(scratch); };
+
   // read tfile header
   TFile_Header header;
   memcpy(&header, data, sizeof(header));
@@ -161,6 +164,7 @@ b8 walk_tkeys(Arena *arena, const u8 *data, u64 data_len, TFile_Data &tfile_data
 
   RNTuple_Anchor_Info *rntuple_info_tail = nullptr;
 
+  u32 n_keys = 0;
   // 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.
@@ -194,7 +198,10 @@ b8 walk_tkeys(Arena *arena, const u8 *data, u64 data_len, TFile_Data &tfile_data
     char cname[256];
     cname[cname_len] = 0;
     memcpy(cname, data + cname_off + 1, cname_len);
-    printf("key 0x%lX, len: %d: %s (name len: %u)\n", cur, n_bytes, cname, cname_len);
+    if (print_keys_info)
+      printf("TKey '%s' at 0x%lX, len: %d (%s)\n", cname, cur, n_bytes, to_pretty_size(scratch.arena, n_bytes).c());
+
+    ++n_keys;
 
     if (cur == tkeys_data.sections[Sec_TKey_List].range.start) {
       if (keylen != tkeys_data.sections[Sec_TKey_List].pre_size ||
@@ -269,6 +276,9 @@ b8 walk_tkeys(Arena *arena, const u8 *data, u64 data_len, TFile_Data &tfile_data
     cur += n_bytes;
   }
 
+  if (print_keys_info)
+    printf("Found %u TKeys.\n", n_keys);
+
   for (u32 id = 1; id < Sec_COUNT; ++id)
     tkeys_data.sections[id].id = (Section_Id)id;