From e0f51bdbb604f1806bdda04e36458071e8e59a18 Mon Sep 17 00:00:00 2001
From: Shane Smiskol <shane@smiskol.com>
Date: Fri, 8 Aug 2025 23:42:54 -0700
Subject: [PATCH] Reapply "LogReader: wrap events to cache which() (#35882)"
 (#35909)

* Reapply "LogReader: wrap events to cache which() (#35882)"

This reverts commit ba2dced54ce2a4fec33565e13f15b610cc8a8b00.

* fix lr

* speed up

* clean up

* more

* should be fast

* clean up

* only supports Event

* rmrmr

* bye

* simple

* gix
---
 tools/lib/logreader.py | 42 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/tools/lib/logreader.py b/tools/lib/logreader.py
index 5915f1dde3..cc84c8b52e 100755
--- a/tools/lib/logreader.py
+++ b/tools/lib/logreader.py
@@ -49,8 +49,46 @@ def decompress_stream(data: bytes):
 
   return decompressed_data
 
+
+class CachedEventReader:
+  __slots__ = ('_evt', '_enum')
+
+  def __init__(self, evt: capnp._DynamicStructReader, _enum: str | None = None):
+    """All capnp attribute accesses are expensive, and which() is often called multiple times"""
+    self._evt = evt
+    self._enum: str | None = _enum
+
+  # fast pickle support
+  def __reduce__(self):
+    return CachedEventReader._reducer, (self._evt.as_builder().to_bytes(), self._enum)
+
+  @staticmethod
+  def _reducer(data: bytes, _enum: str | None = None):
+    with capnp_log.Event.from_bytes(data) as evt:
+      return CachedEventReader(evt, _enum)
+
+  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):
+    if name.startswith("__") and name.endswith("__"):
+      return getattr(self, name)
+    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
 
@@ -75,7 +113,7 @@ class _LogFileReader:
     self._ents = []
     try:
       for e in ents:
-        self._ents.append(e)
+        self._ents.append(CachedEventReader(e))
     except capnp.KjException:
       warnings.warn("Corrupted events detected", RuntimeWarning, stacklevel=1)