aaronrene / gabriel-muse public
symbol_cache.py python
252 lines 9.2 KB
Raw
sha256:660fcac1df3ab28f61862e961890bd2ca8b754fa0242079d93ca1e25037ec8a6 chore(tests): add docstring to tests/__init__.py so rc14 tr… Human 32 days ago
1 """Persistent symbol-tree cache — eliminates re-parsing on every CLI invocation.
2
3 Architecture
4 ------------
5 Every call to ``symbols_for_snapshot`` currently re-parses all semantic files
6 from their raw bytes: ~1,300 ms for a 400-file Python codebase. The parse
7 result is fully determined by the file content — the same bytes always produce
8 the same ``SymbolTree``.
9
10 ``SymbolCache`` exploits this by persisting the parse results keyed by the
11 SHA-256 of the file bytes (``object_id`` from the content-addressed object
12 store, or a freshly computed SHA-256 for working-tree reads). On a warm cache
13 hit the parse step is skipped entirely; the full snapshot loads in ~22 ms
14 instead of ~1,300 ms — a 60× speedup.
15
16 Cache key
17 ---------
18 The key is the 64-character lowercase SHA-256 hex digest of the **raw file
19 bytes**, which is:
20
21 * For committed files — the ``object_id`` from the snapshot manifest
22 (content-addressed, so it is already the SHA-256 of the bytes).
23 * For working-tree files — ``hashlib.sha256(raw_bytes).hexdigest()``,
24 computed from the bytes already read from disk.
25
26 Because the key is content-addressed, every hit is guaranteed correct.
27 A file edit produces a new SHA-256 → cache miss → fresh parse → new entry.
28 Old entries for changed files become dead weight but never produce wrong
29 results.
30
31 Storage
32 -------
33 ``.muse/symbol_cache.msgpack`` — a single msgpack document::
34
35 {
36 "version": 1,
37 "entries": {
38 "<sha256-hex>": {
39 "<file::address>": {
40 "kind": "function",
41 "name": "run",
42 ...
43 }
44 }
45 }
46 }
47
48 msgpack is used instead of JSON because:
49
50 * 2.6× faster serialisation, 1.4× faster deserialisation.
51 * Binary-safe (no base64 for blob fields).
52 * Already a first-class dependency (used by the MWP wire protocol).
53
54 Writes are atomic: data is flushed to a ``.tmp`` sibling then renamed over
55 the target file so a crash mid-write never corrupts the cache.
56
57 Pruning
58 -------
59 Entries accumulate as new file versions are committed. The cache is naturally
60 bounded by the object store: once ``muse gc`` collects an object, its symbol
61 cache entry is dead weight. A future ``muse gc`` integration can prune the
62 cache via :meth:`SymbolCache.prune`. Until then, the cache grows at the rate
63 of unique file-content changes — typically kilobytes per commit.
64 """
65
66 from __future__ import annotations
67
68 import hashlib
69 import logging
70 import pathlib
71 from typing import TypeGuard, get_args
72
73 import msgpack
74
75 from muse.core.store import MsgpackValue, read_msgpack_file
76 from muse.plugins.code.ast_parser import SymbolKind, SymbolRecord, SymbolTree
77
78
79 type _SymbolTreeMap = dict[str, "SymbolTree"]
80 logger = logging.getLogger(__name__)
81
82 _CACHE_VERSION = 1
83 _CACHE_FILENAME = "symbol_cache.msgpack"
84
85 # All string fields in SymbolRecord, used for structural validation on load.
86 _STR_FIELDS: frozenset[str] = frozenset({
87 "kind", "name", "qualified_name", "content_id",
88 "body_hash", "signature_id", "metadata_id", "canonical_key",
89 })
90 _INT_FIELDS: frozenset[str] = frozenset({"lineno", "end_lineno"})
91 _VALID_KINDS: frozenset[str] = frozenset(get_args(SymbolKind))
92
93
94 def _object_id_of(raw: bytes) -> str:
95 """Return the SHA-256 hex digest of *raw* bytes — the content-addressed key."""
96 return hashlib.sha256(raw).hexdigest()
97
98
99 def _is_symbol_record(obj: MsgpackValue) -> TypeGuard[SymbolRecord]:
100 """Return ``True`` if *obj* is structurally a valid ``SymbolRecord``.
101
102 Used during cache deserialization to validate entries from msgpack without
103 constructing a new dict — the validated dict IS the SymbolRecord at runtime.
104 """
105 if not isinstance(obj, dict):
106 return False
107 for field in _STR_FIELDS:
108 if not isinstance(obj.get(field), str):
109 return False
110 for field in _INT_FIELDS:
111 if not isinstance(obj.get(field), int):
112 return False
113 return obj["kind"] in _VALID_KINDS
114
115
116 class SymbolCache:
117 """Persistent msgpack cache mapping object_id → SymbolTree.
118
119 Typical lifecycle inside ``symbols_for_snapshot``::
120
121 cache = SymbolCache.load(root / ".muse")
122 for file_path, object_id in manifest.items():
123 tree = cache.get(object_id)
124 if tree is None:
125 raw = read_object(root, object_id)
126 tree = parse_symbols(raw, file_path)
127 cache.put(object_id, tree)
128 cache.save()
129 """
130
131 def __init__(
132 self,
133 muse_dir: pathlib.Path | None,
134 entries: _SymbolTreeMap,
135 ) -> None:
136 self._muse_dir = muse_dir
137 self._entries = entries
138 self._dirty = False
139
140 # ------------------------------------------------------------------
141 # Construction
142 # ------------------------------------------------------------------
143
144 @classmethod
145 def load(cls, muse_dir: pathlib.Path) -> SymbolCache:
146 """Load the cache from *muse_dir*/symbol_cache.msgpack.
147
148 Returns a fresh empty cache if the file is absent, unreadable, or
149 version-mismatched — never raises. Invalid entries are skipped so
150 a partially corrupt file does not poison the entire cache.
151 """
152 cache_file = muse_dir / _CACHE_FILENAME
153 if not cache_file.is_file():
154 return cls(muse_dir, {})
155 try:
156 doc = read_msgpack_file(cache_file)
157 if not isinstance(doc, dict) or doc.get("version") != _CACHE_VERSION:
158 logger.debug("⚠️ symbol_cache version mismatch — starting fresh")
159 return cls(muse_dir, {})
160 raw_entries = doc.get("entries")
161 if not isinstance(raw_entries, dict):
162 return cls(muse_dir, {})
163 entries: _SymbolTreeMap = {}
164 for obj_id, tree_raw in raw_entries.items():
165 if not isinstance(obj_id, str) or not isinstance(tree_raw, dict):
166 continue
167 tree: SymbolTree = {}
168 valid = True
169 for addr, rec_raw in tree_raw.items():
170 # _is_symbol_record validates structure and narrows to SymbolRecord.
171 if not _is_symbol_record(rec_raw):
172 valid = False
173 break
174 tree[addr] = rec_raw
175 if valid:
176 entries[obj_id] = tree
177 return cls(muse_dir, entries)
178 except Exception as exc: # noqa: BLE001
179 logger.debug("⚠️ symbol_cache unreadable (%s) — starting fresh", exc)
180 return cls(muse_dir, {})
181
182 @classmethod
183 def empty(cls) -> SymbolCache:
184 """Return a no-op cache for contexts without a ``.muse`` directory."""
185 return cls(None, {})
186
187 # ------------------------------------------------------------------
188 # Data access
189 # ------------------------------------------------------------------
190
191 def get(self, object_id: str) -> SymbolTree | None:
192 """Return the cached ``SymbolTree`` for *object_id*, or ``None`` on miss."""
193 return self._entries.get(object_id)
194
195 def put(self, object_id: str, tree: SymbolTree) -> None:
196 """Store *tree* under *object_id* and mark the cache dirty."""
197 self._entries[object_id] = tree
198 self._dirty = True
199
200 def prune(self, live_ids: set[str]) -> None:
201 """Remove entries whose object IDs are not in *live_ids*.
202
203 Call this from ``muse gc`` after identifying all reachable object IDs
204 to keep the cache from growing unboundedly.
205 """
206 stale = set(self._entries) - live_ids
207 if stale:
208 for k in stale:
209 del self._entries[k]
210 self._dirty = True
211 logger.debug("🗑️ symbol_cache pruned %d stale entries", len(stale))
212
213 @property
214 def size(self) -> int:
215 """Number of cached object IDs."""
216 return len(self._entries)
217
218 # ------------------------------------------------------------------
219 # Persistence
220 # ------------------------------------------------------------------
221
222 def save(self) -> None:
223 """Atomically persist the cache to disk if it has changed.
224
225 Uses a temp-file-then-rename pattern so a crash mid-write never
226 leaves a corrupt cache file. Silently skips when there is no
227 ``.muse`` directory (e.g. in-memory unit tests).
228 """
229 if not self._dirty or self._muse_dir is None:
230 return
231 doc = {"version": _CACHE_VERSION, "entries": self._entries}
232 cache_file = self._muse_dir / _CACHE_FILENAME
233 tmp = self._muse_dir / (_CACHE_FILENAME + ".tmp")
234 try:
235 tmp.write_bytes(msgpack.packb(doc, use_bin_type=True))
236 tmp.replace(cache_file)
237 self._dirty = False
238 logger.debug("✅ symbol_cache saved (%d entries)", len(self._entries))
239 except OSError as exc:
240 logger.warning("⚠️ symbol_cache save failed: %s", exc)
241
242
243 def load_symbol_cache(root: pathlib.Path) -> SymbolCache:
244 """Convenience loader: return a ``SymbolCache`` for a repository root.
245
246 Returns ``SymbolCache.empty()`` when *root* has no ``.muse`` directory
247 so callers never need to guard against a missing repo.
248 """
249 muse_dir = root / ".muse"
250 if muse_dir.is_dir():
251 return SymbolCache.load(muse_dir)
252 return SymbolCache.empty()
File History 2 commits
sha256:660fcac1df3ab28f61862e961890bd2ca8b754fa0242079d93ca1e25037ec8a6 chore(tests): add docstring to tests/__init__.py so rc14 tr… Human 32 days ago
sha256:d8316ffae901be06347e16ab55be11868eb519dd16ade3e8aa16a99e662f7e62 baseline: rc14 re-baseline after rc3 store corruption recovery Human patch 32 days ago