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;