From 0ebee550507c8b577b94da36d313f9fdccacb895 Mon Sep 17 00:00:00 2001
From: Shane Smiskol <shane@smiskol.com>
Date: Fri, 1 Aug 2025 19:07:16 -0700
Subject: [PATCH] LogReader: wrap events to cache which() (#35882)

* speed up lr

* lazy caching

* clean up

* it fast

* stash

* stash

* chatgpt code is bad as usual

* clean up

* clean up

* clean up

* clean up

* clean up

* clean up

* match behavior

* cmt
---
 tools/lib/logreader.py | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/tools/lib/logreader.py b/tools/lib/logreader.py
index 1075bd2f08..90f6f12756 100755
--- a/tools/lib/logreader.py
+++ b/tools/lib/logreader.py
@@ -50,8 +50,35 @@ def decompress_stream(data: bytes):
 
   return decompressed_data
 
+
+class CachedReader:
+  __slots__ = ("_evt", "_enum")
+
+  def __init__(self, evt: capnp._DynamicStructReader):
+    """All capnp attribute accesses are expensive, and which() is often called multiple times"""
+    self._evt = evt
+    self._enum: str | None = None
+
+  def __repr__(self):
+    return self._evt.__repr__()
+
+  def __str__(self):
+    return self._evt.__str__()
+
+  def __dir__(self):
+    return dir(self._evt)
+
+  def which(self) -> str:
+    if self._enum is None:
+      self._enum = self._evt.which()
+    return self._enum
+
+  def __getattr__(self, name: str):
+    return getattr(self._evt, name)
+
+
 class _LogFileReader:
-  def __init__(self, fn, canonicalize=True, only_union_types=False, sort_by_time=False, dat=None):
+  def __init__(self, fn, only_union_types=False, sort_by_time=False, dat=None):
     self.data_version = None
     self._only_union_types = only_union_types
 
@@ -76,7 +103,7 @@ class _LogFileReader:
     self._ents = []
     try:
       for e in ents:
-        self._ents.append(e)
+        self._ents.append(CachedReader(e))
     except capnp.KjException:
       warnings.warn("Corrupted events detected", RuntimeWarning, stacklevel=1)