aaronrene / gabriel-muse public
_query.py python
300 lines 11.2 KB
Raw
sha256:0e3697f29ae6c62f39818a19861d4db6276465de16481fadf866eb93cbf4aae4 test: validate staging push and MP workflow for aaronrene Human 17 days ago
1 """Shared query helpers for the code-domain CLI commands.
2
3 This module provides the low-level primitives that multiple code-domain
4 commands need — symbol extraction from snapshots, commit-graph walking,
5 and language classification — so each command can stay thin.
6
7 None of these functions are part of the public ``CodePlugin`` API. They
8 are internal helpers for the CLI layer and must not be imported by any
9 core module.
10 """
11
12 from __future__ import annotations
13
14 import hashlib
15 import itertools
16 import logging
17 import pathlib
18 from collections import deque
19 from collections.abc import Iterator
20
21 from muse.core._types import Manifest, Metadata
22 from muse.core.object_store import read_object
23 from muse.core.store import CommitRecord, read_commit
24 from muse.core.symbol_cache import SymbolCache, load_symbol_cache
25 from muse.domain import DomainOp
26 from muse.plugins.code.ast_parser import (
27 SymbolTree,
28 parse_symbols,
29 )
30
31 type _SymbolTreeMap = dict[str, SymbolTree]
32 type _CommitIndex = dict[str, CommitRecord]
33
34 logger = logging.getLogger(__name__)
35
36 # ---------------------------------------------------------------------------
37 # Language classification
38 # ---------------------------------------------------------------------------
39
40 _SUFFIX_LANG: Metadata = {
41 ".py": "Python", ".pyi": "Python",
42 ".ts": "TypeScript", ".tsx": "TypeScript",
43 ".js": "JavaScript", ".jsx": "JavaScript",
44 ".mjs": "JavaScript", ".cjs": "JavaScript",
45 ".go": "Go",
46 ".rs": "Rust",
47 ".java": "Java",
48 ".cs": "C#",
49 ".c": "C", ".h": "C",
50 ".cpp": "C++", ".cc": "C++", ".cxx": "C++", ".hpp": "C++", ".hxx": "C++",
51 ".rb": "Ruby",
52 ".kt": "Kotlin", ".kts": "Kotlin",
53 ".swift": "Swift",
54 ".md": "Markdown", ".rst": "reStructuredText", ".txt": "Text",
55 ".toml": "TOML",
56 ".yaml": "YAML", ".yml": "YAML",
57 ".json": "JSON", ".jsonc": "JSON",
58 ".css": "CSS", ".scss": "SCSS",
59 ".html": "HTML", ".htm": "HTML",
60 ".sql": "SQL",
61 ".sh": "Shell", ".bash": "Shell", ".zsh": "Shell",
62 ".proto": "Protobuf",
63 ".tf": "Terraform",
64 }
65
66
67 def language_of(file_path: str) -> str:
68 """Return a display language name for *file_path* based on its suffix."""
69 suffix = pathlib.PurePosixPath(file_path).suffix.lower()
70 return _SUFFIX_LANG.get(suffix, suffix or "(no ext)")
71
72
73 def is_semantic(file_path: str) -> bool:
74 """Return ``True`` if *file_path* has a suffix with AST-level support."""
75 from muse.plugins.code.ast_parser import SEMANTIC_EXTENSIONS
76 suffix = pathlib.PurePosixPath(file_path).suffix.lower()
77 return suffix in SEMANTIC_EXTENSIONS
78
79
80 # ---------------------------------------------------------------------------
81 # Symbol extraction from a snapshot manifest
82 # ---------------------------------------------------------------------------
83
84
85 def _rekey_tree(tree: SymbolTree, file_path: str) -> SymbolTree:
86 """Re-key *tree* so every address starts with *file_path*.
87
88 The ``SymbolCache`` is keyed by content SHA-256. When two files have
89 identical bytes they share one cache entry whose addresses carry the path
90 of whichever file was parsed first. This function patches those addresses
91 to the actual *file_path* being processed — an O(n) pass that is only
92 triggered on a cache hit where the prefix differs.
93 """
94 first = next(iter(tree), None)
95 if first is None:
96 return tree
97 cached_prefix = first.split("::")[0] if "::" in first else None
98 if cached_prefix == file_path:
99 return tree # Addresses already correct — common case.
100 return {
101 f"{file_path}::{addr.split('::', 1)[1]}" if "::" in addr else addr: rec
102 for addr, rec in tree.items()
103 }
104
105
106 def symbols_for_snapshot(
107 root: pathlib.Path,
108 manifest: Manifest,
109 *,
110 kind_filter: str | None = None,
111 file_filter: str | None = None,
112 language_filter: str | None = None,
113 workdir: pathlib.Path | None = None,
114 cache: SymbolCache | None = None,
115 ) -> _SymbolTreeMap:
116 """Extract symbol trees for all semantic files in *manifest*.
117
118 Results are served from the persistent symbol cache when available,
119 cutting a full-snapshot scan from ~1,300 ms to ~22 ms on a warm cache.
120 The cache is keyed by the SHA-256 of the file bytes (``object_id`` for
121 committed files; freshly computed SHA-256 for working-tree reads), so
122 every hit is guaranteed correct.
123
124 Args:
125 root: Repository root (used to locate the object store and
126 load/save the symbol cache).
127 manifest: Snapshot manifest mapping file path → SHA-256.
128 kind_filter: If set, only include symbols with this ``kind``.
129 file_filter: If set, only include symbols from this exact file path.
130 language_filter: If set, only include symbols from files of this language.
131 workdir: When set, file content is read from *workdir* first
132 (working-tree mode), falling back to the object store if
133 the file does not exist on disk. Pass ``root`` to get
134 live, uncommitted changes reflected immediately.
135 cache: Optional pre-loaded ``SymbolCache`` instance. When
136 ``None`` the cache is loaded from disk automatically
137 and saved at the end of the call. Pass an explicit
138 instance when batching multiple calls to share one
139 cache load/save cycle.
140
141 Returns:
142 Dict mapping ``file_path → SymbolTree``; empty trees are omitted.
143 """
144 if cache is None:
145 active_cache: SymbolCache = load_symbol_cache(root)
146 own_cache = True
147 else:
148 active_cache = cache
149 own_cache = False
150
151 result: _SymbolTreeMap = {}
152 for file_path, object_id in sorted(manifest.items()):
153 if not is_semantic(file_path):
154 continue
155 if file_filter and file_path != file_filter:
156 continue
157 if language_filter and language_of(file_path) != language_filter:
158 continue
159
160 # --- Determine cache key and read bytes if needed ----------------
161 # For working-tree files the object_id is the committed version;
162 # the disk content may differ. We compute the SHA-256 of whatever
163 # bytes we actually parse so the cache key is always content-addressed.
164 cache_key = object_id
165 raw: bytes | None = None
166
167 if workdir is not None:
168 disk_path = workdir / file_path
169 if disk_path.is_file():
170 try:
171 raw = disk_path.read_bytes()
172 # Recompute the key for the actual disk content.
173 cache_key = hashlib.sha256(raw).hexdigest()
174 except OSError as exc:
175 logger.debug("Could not read %s from workdir: %s", file_path, exc)
176
177 # --- Cache lookup ------------------------------------------------
178 tree = active_cache.get(cache_key)
179 if tree is None:
180 # Cache miss — fetch bytes and parse.
181 if raw is None:
182 raw = read_object(root, object_id)
183 if raw is None:
184 logger.debug("Object %s missing for %s — skipping", object_id[:8], file_path)
185 continue
186 tree = parse_symbols(raw, file_path)
187 active_cache.put(cache_key, tree)
188 else:
189 # Cache hit — re-key addresses when two files share the same
190 # SHA-256 (identical content). The cache stores the tree built
191 # for whichever file was parsed first; its addresses carry that
192 # file's path prefix. Patch them to the current file_path so
193 # cross-file clone detection and all downstream consumers get
194 # correct symbol addresses.
195 tree = _rekey_tree(tree, file_path)
196
197 # --- Apply filters on top of the (potentially cached) full tree --
198 if kind_filter:
199 tree = {addr: rec for addr, rec in tree.items() if rec["kind"] == kind_filter}
200 if tree:
201 result[file_path] = tree
202
203 if own_cache:
204 active_cache.save()
205
206 return result
207
208
209 # ---------------------------------------------------------------------------
210 # Commit-graph walking
211 # ---------------------------------------------------------------------------
212
213
214 def walk_commits_bfs(
215 root: pathlib.Path,
216 start_commit_id: str,
217 max_commits: int = 500,
218 stop_at_commit_id: str | None = None,
219 ) -> tuple[list[CommitRecord], bool]:
220 """BFS walk of the commit DAG from *start_commit_id*, newest-first.
221
222 Unlike the linear :func:`walk_commits`, this follows ``parent2_commit_id``
223 at merge commits so events on merged feature branches are never missed.
224
225 Args:
226 root: Repository root.
227 start_commit_id: SHA-256 of the commit to start from.
228 max_commits: Safety cap — returns ``truncated=True`` if reached.
229 stop_at_commit_id: When set, stop BFS *before* entering this commit
230 (exclusive lower bound, useful for range queries).
231
232 Returns:
233 ``(commits, truncated)`` — list sorted newest-first by ``committed_at``;
234 ``truncated=True`` when ``max_commits`` was hit before exhausting the DAG.
235 """
236 commits_by_id: _CommitIndex = {}
237 queue: deque[str] = deque([start_commit_id])
238 seen: set[str] = set()
239 truncated = False
240
241 while queue:
242 if len(commits_by_id) >= max_commits:
243 truncated = True
244 break
245 commit_id = queue.popleft()
246 if commit_id in seen:
247 continue
248 if stop_at_commit_id and commit_id == stop_at_commit_id:
249 continue
250 seen.add(commit_id)
251 commit = read_commit(root, commit_id)
252 if commit is None:
253 continue
254 commits_by_id[commit_id] = commit
255 if commit.parent_commit_id:
256 queue.append(commit.parent_commit_id)
257 if commit.parent2_commit_id:
258 queue.append(commit.parent2_commit_id)
259
260 return (
261 sorted(commits_by_id.values(), key=lambda c: c.committed_at, reverse=True),
262 truncated,
263 )
264
265
266 # ---------------------------------------------------------------------------
267 # Op traversal helpers
268 # ---------------------------------------------------------------------------
269
270
271 def flat_symbol_ops(ops: list[DomainOp]) -> Iterator[DomainOp]:
272 """Yield all leaf ops, recursing into PatchOp.child_ops.
273
274 Only yields ops that have a symbol-level address (i.e. contain ``::``).
275 """
276 for op in ops:
277 if op["op"] == "patch":
278 for child in op["child_ops"]:
279 if "::" in child["address"]:
280 yield child
281 elif "::" in op["address"]:
282 yield op
283
284
285 def touched_files(ops: list[DomainOp]) -> frozenset[str]:
286 """Return the set of file paths that appear as PatchOp addresses in *ops*.
287
288 Only counts files that had symbol-level child ops (semantic changes),
289 not coarse file-level replace/insert/delete ops.
290 """
291 files: set[str] = set()
292 for op in ops:
293 if op["op"] == "patch" and op["child_ops"]:
294 files.add(op["address"])
295 return frozenset(files)
296
297
298 def file_pairs(files: frozenset[str]) -> Iterator[tuple[str, str]]:
299 """Yield all ordered pairs ``(a, b)`` with ``a < b`` from *files*."""
300 yield from itertools.combinations(sorted(files), 2)
File History 1 commit
sha256:0e3697f29ae6c62f39818a19861d4db6276465de16481fadf866eb93cbf4aae4 test: validate staging push and MP workflow for aaronrene Human 17 days ago