Deduplicate frames for 95% size reduction on large traces

Stack traces now use frame-level deduplication: a shared FRAMES array
of unique strings, with STACKS storing integer indices. For a 600K
allocation trace this reduces the stacks section from 1.3GB to ~62MB.

Also: regex search toggle, title from file stem, tab reset fix.

Author: drisspg

test/test_memory_viz.py

@@ -11,6 +11,11 @@
     process_snapshot,
 )
 
+def _unpack(snapshot, device=0):
+    timeline, allocs, frames, stacks, categories, max_ts = process_snapshot(snapshot, device)
+    resolved_stacks = [[frames[fi] for fi in s] for s in stacks]
+    return timeline, allocs, resolved_stacks, categories, max_ts
+
 DATA_DIR = Path(__file__).parent / "data"
 SNAPSHOT_PATH = DATA_DIR / "mini_snapshot.pickle"
 
@@ -78,13 +83,20 @@ def test_is_cpython_c_frame(self):
 
 class TestProcessSnapshot:
     def test_returns_correct_tuple_shape(self, snapshot):
-        timeline, allocs, stacks, categories, max_ts = process_snapshot(snapshot)
+        timeline, allocs, frames, stacks, categories, max_ts = process_snapshot(snapshot)
         assert len(timeline) > 0
         assert len(allocs) > 0
+        assert len(frames) > 0
         assert len(stacks) > 0
         assert len(categories) == len(stacks)
         assert max_ts > 0
 
+    def test_stacks_reference_valid_frame_indices(self, snapshot):
+        _, _, frames, stacks, *_ = process_snapshot(snapshot)
+        for stack in stacks:
+            for fi in stack:
+                assert 0 <= fi < len(frames)
+
     def test_timeline_fields(self, snapshot):
         timeline, *_ = process_snapshot(snapshot)
         entry = timeline[0]
@@ -111,15 +123,13 @@ def test_allocated_never_negative(self, snapshot):
         assert all(e["a"] >= 0 for e in timeline)
 
     def test_stack_indices_valid(self, snapshot):
-        timeline, allocs, stacks, *_ = process_snapshot(snapshot)
-        for e in timeline:
-            assert 0 <= e["si"] < len(stacks)
+        _, allocs, _, stacks, *_ = process_snapshot(snapshot)
         for a in allocs:
             assert 0 <= a["si"] < len(stacks)
 
     def test_empty_device_returns_empty(self, snapshot):
         result = process_snapshot(snapshot, device=99)
-        assert result == ([], [], [], [], 0)
+        assert result == ([], [], [], [], [], 0)
 
     def test_polygon_offsets_non_negative(self, snapshot):
         _, allocs, *_ = process_snapshot(snapshot)

transformer_nuggets/utils/memory_viz.py

@@ -59,20 +59,29 @@ def _categorize_stack(frames: list[str]) -> str:
 
 def process_snapshot(
     snapshot: dict, device: int = 0
-) -> tuple[list[dict], list[dict], list[list[str]], list[str], int]:
+) -> tuple[list[dict], list[dict], list[str], list[list[int]], list[str], int]:
     traces = snapshot.get("device_traces", [])
     if device >= len(traces):
-        return [], [], [], [], 0
-
-    stack_to_idx: dict[tuple[str, ...], int] = {}
-    stacks: list[list[str]] = []
-
-    def get_stack_idx(frames: list[dict]) -> int:
-        extracted = _extract_frames(frames)
-        key = tuple(extracted)
+        return [], [], [], [], [], 0
+
+    frame_to_idx: dict[str, int] = {}
+    frames: list[str] = []
+    stack_to_idx: dict[tuple[int, ...], int] = {}
+    stacks: list[list[int]] = []
+
+    def intern_frame(f: str) -> int:
+        if f not in frame_to_idx:
+            frame_to_idx[f] = len(frames)
+            frames.append(f)
+        return frame_to_idx[f]
+
+    def get_stack_idx(raw_frames: list[dict]) -> int:
+        extracted = _extract_frames(raw_frames)
+        frame_indices = [intern_frame(f) for f in extracted]
+        key = tuple(frame_indices)
         if key not in stack_to_idx:
             stack_to_idx[key] = len(stacks)
-            stacks.append(extracted)
+            stacks.append(frame_indices)
         return stack_to_idx[key]
 
     allocated = 0
@@ -159,16 +168,18 @@ def get_stack_idx(frames: list[dict]) -> int:
         poly["ts"].append(timestep)
         poly["offsets"].append(poly["offsets"][-1])
 
-    categories = [_categorize_stack(stack) for stack in stacks]
-    return timeline, alloc_polys, stacks, categories, timestep
+    categories = [_categorize_stack([frames[fi] for fi in stack]) for stack in stacks]
+    return timeline, alloc_polys, frames, stacks, categories, timestep
 
 
 def generate_memory_html(
     snapshot: dict,
     device: int = 0,
     title: str = "Memory Timeline",
 ) -> str:
-    timeline, alloc_polys, stacks, categories, max_ts = process_snapshot(snapshot, device)
+    timeline, alloc_polys, frames, stacks, categories, max_ts = process_snapshot(
+        snapshot, device
+    )
     hwm = max((p["h"] for p in timeline), default=0)
     hwm_timestep = next((i for i, p in enumerate(timeline) if p["a"] == hwm), 0)
 
@@ -193,6 +204,7 @@ def generate_memory_html(
         _MEMORY_VIZ_TEMPLATE.replace("__TITLE__", title)
         .replace("__TIMELINE__", json.dumps(timeline))
         .replace("__ALLOCS__", json.dumps(alloc_polys))
+        .replace("__FRAMES__", json.dumps(frames))
         .replace("__STACKS__", json.dumps(stacks))
         .replace("__CATEGORIES__", json.dumps(cat_indices))
         .replace("__META__", json.dumps(meta))
@@ -683,10 +695,16 @@ def generate_memory_html(
 <script>
 const TIMELINE = __TIMELINE__;
 const ALLOCS = __ALLOCS__;
+const FRAMES = __FRAMES__;
 const STACKS = __STACKS__;
 const CATEGORIES = __CATEGORIES__;
 const META = __META__;
 
+function resolveStack(stackIdx) {
+  const indices = STACKS[stackIdx] || [];
+  return indices.map(i => FRAMES[i]);
+}
+
 function formatBytes(b) {
   if (Math.abs(b) >= 1024**3) return (b / 1024**3).toFixed(2) + ' GiB';
   if (Math.abs(b) >= 1024**2) return (b / 1024**2).toFixed(1) + ' MiB';
@@ -743,7 +761,7 @@ def generate_memory_html(
 }
 
 function bestFrame(stackIdx) {
-  const stack = STACKS[stackIdx] || [];
+  const stack = resolveStack(stackIdx);
   for (const f of stack) {
     if (classifyFrame(f) === 'user') return f;
   }
@@ -782,7 +800,7 @@ def generate_memory_html(
 function renderStack(stackIdx, label) {
   lastStackIdx = stackIdx;
   lastStackLabel = label;
-  const stack = STACKS[stackIdx] || [];
+  const stack = resolveStack(stackIdx);
   detailStats.textContent = label;
   if (!stack.length) {
     detailBody.innerHTML = '<div class="empty-detail">No frames recorded</div>';
@@ -1025,7 +1043,7 @@ def generate_memory_html(
       el.classed('dimmed', false).classed('highlighted', false);
       return;
     }
-    const stack = STACKS[d.si] || [];
+    const stack = resolveStack(d.si) || [];
     const match = stack.some(f => matcher.test(f));
     el.classed('dimmed', !match).classed('highlighted', match);
   });