fix page grouping

src/render.cpp

@@ -40,16 +40,25 @@ u32 mem_edit_bg_color_fn(const u8 *, u64 off, void *user_data)
   off += app->vsettings.base_display_addr;
 
 #define COL(c) (ImColor((c)[0], (c)[1], (c)[2]))
-  // Handle pages
+  // TFile start
-  // fast case: offset from same page info as previous
+  if (off <= rdata.root_file_header_size) return COL(app->vsettings.col_tfile);
-  if (off < app->last_pinfo->range.end()) return COL(app->vsettings.col_page);
 
-  // still fast case: offset from next page info
+  // Handle pages
-  if (app->last_pinfo->next) app->last_pinfo = app->last_pinfo->next;
+  // fast case: `off` is in the same page info as previous `off`.
-  if (app->last_pinfo && app->last_pinfo->range.start <= off && off <= app->last_pinfo->range.end())
+  if (app->last_pinfo->range.start < off && off <= app->last_pinfo->range.end())
     return COL(app->vsettings.col_page);
 
-  if (off <= rdata.root_file_header_size) return COL(app->vsettings.col_tfile);
+  // still fast case: `off is in the next page info as the previous.
+  if (app->last_pinfo->next)
+    app->last_pinfo = app->last_pinfo->next;
+
+  if (app->last_pinfo && app->last_pinfo->range.start <= off && off <= app->last_pinfo->range.end()) {
+    if (off == app->last_pinfo->range.start)
+      return COL(app->vsettings.col_page_start);
+    else
+      return COL(app->vsettings.col_page);
+  }
+
   if (rdata.rng_anchor_key.start <= off && off <= rdata.rng_anchor_key.end()) return COL(app->vsettings.col_key);
   if (rdata.rng_header.start - rblob_sz <= off && off <= rdata.rng_header.start) return COL(app->vsettings.col_key);
   if (rdata.rng_footer.start - rblob_sz <= off && off <= rdata.rng_footer.start) return COL(app->vsettings.col_key);
@@ -57,11 +66,9 @@ u32 mem_edit_bg_color_fn(const u8 *, u64 off, void *user_data)
   if (rdata.rng_header.start <= off && off <= rdata.rng_header.end()) return COL(app->vsettings.col_header);
   if (rdata.rng_footer.start <= off && off <= rdata.rng_footer.end()) return COL(app->vsettings.col_footer);
 
-  // Slow page group lookup, hopefully rare (ideally only done once)
+  // Slow page group lookup, ideally only done once per render when last_pinfo is invalid.
-  // printf("%lu vs %lu\n", off, rdata.page_groups[rdata.n_page_groups - 1].range.end());
   for (Page_Info_Chunk *chunk = rdata.page_chunks; chunk; chunk = chunk->next) {
     if (chunk->range.start <= off && off <= chunk->range.end()) {
-      printf("slow path for 0x%lX (last pinfo: 0x%lX - 0x%lX)\n", off, app->last_pinfo->range.start, app->last_pinfo->range.end());
       for (u64 group_idx = chunk->first_group; group_idx < rdata.n_page_groups; ++group_idx) {
         const Page_Info_Group &group = rdata.page_groups[group_idx];
         if (off < group.range.start || off > group.range.end()) 
@@ -195,7 +202,15 @@ void update_and_render(Arena *arena, App_State &app, f32 delta_time_ms)
 
       ImGui::ColorEdit3("_TKey Header",    app.vsettings.col_key,    flags);
       ImGui::SameLine(); 
-      if (ImGui::Button("TKey Header"))    {} // TODO app.vsettings.base_display_addr = app.rndatarng_footer.start;
+      if (ImGui::Button("TKey Header"))    {} // TODO jump to next key
+
+      ImGui::ColorEdit3("_Page Start", app.vsettings.col_page_start, flags);
+      ImGui::SameLine(); 
+      if (ImGui::Button("Page Start"))     {} // TODO: jump to next page
+
+      ImGui::ColorEdit3("_Page", app.vsettings.col_page, flags);
+      ImGui::SameLine(); 
+      if (ImGui::Button("Page"))           {} // TODO: jump to next page
 
       ImGui::Separator();
       ImGui::Text("Num pages: %lu", app.rndata.n_pages);

src/rntuple.cpp

@@ -72,13 +72,17 @@ void gather_metadata(Arena *arena, RMicroFileReader &reader, const RNTuple_File_
   u64 n_pages = 0;
   u64 n_elems = 0;
   Page_Info_Node *pinfo_head = nullptr, *pinfo_tail = nullptr;
+  Page_Info_Node *last_inserted_pinfo = nullptr;
   
   fprintf(stderr, "Loading pages...\n");
   
+  chr::time_point start_t = chr::high_resolution_clock::now();
+  
   // for all clusters, gather page metadata
   for (const RClusterDescriptor &cluster_desc : descriptor.GetClusterIterable()) {
     for (const RClusterDescriptor::RColumnRange &col_range : cluster_desc.GetColumnRangeIterable()) {
       // insert page infos sorted by byte range
+      // @Speed: this is slow! speed it up!
       const auto &page_range = cluster_desc.GetPageRange(col_range.fPhysicalColumnId);
       for (const auto &page_info : page_range.fPageInfos) {
         Page_Info_Node *pinfo = arena_push<Page_Info_Node>(arena);
@@ -98,20 +102,30 @@ void gather_metadata(Arena *arena, RMicroFileReader &reader, const RNTuple_File_
           // before head
           pinfo->next = pinfo_head;
           pinfo_head = pinfo;
+        } else if (last_inserted_pinfo && pinfo->range.start == last_inserted_pinfo->range.end()) {
+          // common case: insert after previous
+          pinfo->next = last_inserted_pinfo->next;
+          last_inserted_pinfo->next = pinfo;
         } else for (Page_Info_Node *node = pinfo_head->next, *prev = pinfo_head; node; prev = node, node = node->next) {
-          if (pinfo->range.end() < node->range.start) {
+          if (pinfo->range.end() <= node->range.start) {
             prev->next = pinfo;
             pinfo->next = node;
+            break;
           }
         }
 
+        last_inserted_pinfo = pinfo;
+
         ++n_pages;
         n_elems += page_info.fNElements;
       }
     }
   }
 
-  fprintf(stderr, "Loaded %lu pages.\nGenerating groups...", n_pages);
+  chr::time_point end_t = chr::high_resolution_clock::now();
+  u64 time_spent_ms = chr::duration_cast<chr::milliseconds>(end_t - start_t).count();
+
+  fprintf(stderr, "Loaded %lu pages in %lu ms.\nGenerating groups...\n", n_pages, time_spent_ms);
 
   // Create page groups and chunks.
   // Each page group is a grouping of GROUP_SIZE page infos whose range is equal to the combined ranges
@@ -121,13 +135,16 @@ void gather_metadata(Arena *arena, RMicroFileReader &reader, const RNTuple_File_
   // of a page or not.
   const u64 GROUP_SIZE = 100;
   Page_Info_Group *groups = arena_push_array_nozero<Page_Info_Group>(arena, n_pages / GROUP_SIZE + 1);
-  Page_Info_Node *last = nullptr;
+  static const Page_Info_Node invalid_last = { nullptr, { 0, 0 }, 0 };
+  const Page_Info_Node *last = &invalid_last;
   Page_Info_Chunk *chunks_head = nullptr, *chunks_tail = nullptr;
   u64 idx = 0;
   u64 n_groups = 0;
   u64 n_chunks = 0;
   for (Page_Info_Node *pinfo = pinfo_head; pinfo; pinfo = pinfo->next) {
+    assert(last->range.end() <= pinfo->range.start);    
     last = pinfo;
+
     if (idx++ % GROUP_SIZE != 0)
       continue;
     
@@ -136,7 +153,7 @@ void gather_metadata(Arena *arena, RMicroFileReader &reader, const RNTuple_File_
     Page_Info_Group &cur_group = groups[n_groups];
     cur_group.first = pinfo;
     cur_group.range.start = pinfo->range.start;
-    if (n_groups > 1) {
+    if (n_groups > 0) {
       Page_Info_Group &prev_group = groups[n_groups - 1];
       prev_group.range.len = cur_group.range.start - prev_group.range.start;
     }
@@ -160,17 +177,22 @@ void gather_metadata(Arena *arena, RMicroFileReader &reader, const RNTuple_File_
       chunks_tail->next = chunk;
       chunks_tail = chunk;
       ++n_chunks;
-    } else {
-      chunks_tail->range.len += cur_group.range.len;
     }
 
     ++n_groups;
   }
-  if (last)
+
-    groups[n_groups - 1].range.len = last->range.end() - groups[n_groups - 1].range.start;
+  if (n_groups) {
+    Page_Info_Group &last_group = groups[n_groups - 1];
+    last_group.range.len = last->range.end() - last_group.range.start;
+    chunks_tail->range.len = last->range.end() - chunks_tail->range.start;
+  }
 
   fprintf(stderr, "Generated %lu groups and %lu chunks.\n", n_groups, n_chunks);
 
+  assert(!chunks_tail->next);
+  assert(!pinfo_tail->next);
+
   rndata.pages = pinfo_head;
   rndata.page_groups = groups;
   rndata.n_page_groups = n_groups;