"""Shared query helpers for the code-domain CLI commands. This module provides the low-level primitives that multiple code-domain commands need — symbol extraction from snapshots, commit-graph walking, and language classification — so each command can stay thin. None of these functions are part of the public ``CodePlugin`` API. They are internal helpers for the CLI layer and must not be imported by any core module. """ from __future__ import annotations import hashlib import itertools import logging import pathlib from collections import deque from collections.abc import Iterator from muse.core._types import Manifest, Metadata from muse.core.object_store import read_object from muse.core.store import CommitRecord, read_commit from muse.core.symbol_cache import SymbolCache, load_symbol_cache from muse.domain import DomainOp from muse.plugins.code.ast_parser import ( SymbolTree, parse_symbols, ) type _SymbolTreeMap = dict[str, SymbolTree] type _CommitIndex = dict[str, CommitRecord] logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Language classification # --------------------------------------------------------------------------- _SUFFIX_LANG: Metadata = { ".py": "Python", ".pyi": "Python", ".ts": "TypeScript", ".tsx": "TypeScript", ".js": "JavaScript", ".jsx": "JavaScript", ".mjs": "JavaScript", ".cjs": "JavaScript", ".go": "Go", ".rs": "Rust", ".java": "Java", ".cs": "C#", ".c": "C", ".h": "C", ".cpp": "C++", ".cc": "C++", ".cxx": "C++", ".hpp": "C++", ".hxx": "C++", ".rb": "Ruby", ".kt": "Kotlin", ".kts": "Kotlin", ".swift": "Swift", ".md": "Markdown", ".rst": "reStructuredText", ".txt": "Text", ".toml": "TOML", ".yaml": "YAML", ".yml": "YAML", ".json": "JSON", ".jsonc": "JSON", ".css": "CSS", ".scss": "SCSS", ".html": "HTML", ".htm": "HTML", ".sql": "SQL", ".sh": "Shell", ".bash": "Shell", ".zsh": "Shell", ".proto": "Protobuf", ".tf": "Terraform", } def language_of(file_path: str) -> str: """Return a display language name for *file_path* based on its suffix.""" suffix = pathlib.PurePosixPath(file_path).suffix.lower() return _SUFFIX_LANG.get(suffix, suffix or "(no ext)") def is_semantic(file_path: str) -> bool: """Return ``True`` if *file_path* has a suffix with AST-level support.""" from muse.plugins.code.ast_parser import SEMANTIC_EXTENSIONS suffix = pathlib.PurePosixPath(file_path).suffix.lower() return suffix in SEMANTIC_EXTENSIONS # --------------------------------------------------------------------------- # Symbol extraction from a snapshot manifest # --------------------------------------------------------------------------- def _rekey_tree(tree: SymbolTree, file_path: str) -> SymbolTree: """Re-key *tree* so every address starts with *file_path*. The ``SymbolCache`` is keyed by content SHA-256. When two files have identical bytes they share one cache entry whose addresses carry the path of whichever file was parsed first. This function patches those addresses to the actual *file_path* being processed — an O(n) pass that is only triggered on a cache hit where the prefix differs. """ first = next(iter(tree), None) if first is None: return tree cached_prefix = first.split("::")[0] if "::" in first else None if cached_prefix == file_path: return tree # Addresses already correct — common case. return { f"{file_path}::{addr.split('::', 1)[1]}" if "::" in addr else addr: rec for addr, rec in tree.items() } def symbols_for_snapshot( root: pathlib.Path, manifest: Manifest, *, kind_filter: str | None = None, file_filter: str | None = None, language_filter: str | None = None, workdir: pathlib.Path | None = None, cache: SymbolCache | None = None, ) -> _SymbolTreeMap: """Extract symbol trees for all semantic files in *manifest*. Results are served from the persistent symbol cache when available, cutting a full-snapshot scan from ~1,300 ms to ~22 ms on a warm cache. The cache is keyed by the SHA-256 of the file bytes (``object_id`` for committed files; freshly computed SHA-256 for working-tree reads), so every hit is guaranteed correct. Args: root: Repository root (used to locate the object store and load/save the symbol cache). manifest: Snapshot manifest mapping file path → SHA-256. kind_filter: If set, only include symbols with this ``kind``. file_filter: If set, only include symbols from this exact file path. language_filter: If set, only include symbols from files of this language. workdir: When set, file content is read from *workdir* first (working-tree mode), falling back to the object store if the file does not exist on disk. Pass ``root`` to get live, uncommitted changes reflected immediately. cache: Optional pre-loaded ``SymbolCache`` instance. When ``None`` the cache is loaded from disk automatically and saved at the end of the call. Pass an explicit instance when batching multiple calls to share one cache load/save cycle. Returns: Dict mapping ``file_path → SymbolTree``; empty trees are omitted. """ if cache is None: active_cache: SymbolCache = load_symbol_cache(root) own_cache = True else: active_cache = cache own_cache = False result: _SymbolTreeMap = {} for file_path, object_id in sorted(manifest.items()): if not is_semantic(file_path): continue if file_filter and file_path != file_filter: continue if language_filter and language_of(file_path) != language_filter: continue # --- Determine cache key and read bytes if needed ---------------- # For working-tree files the object_id is the committed version; # the disk content may differ. We compute the SHA-256 of whatever # bytes we actually parse so the cache key is always content-addressed. cache_key = object_id raw: bytes | None = None if workdir is not None: disk_path = workdir / file_path if disk_path.is_file(): try: raw = disk_path.read_bytes() # Recompute the key for the actual disk content. cache_key = hashlib.sha256(raw).hexdigest() except OSError as exc: logger.debug("Could not read %s from workdir: %s", file_path, exc) # --- Cache lookup ------------------------------------------------ tree = active_cache.get(cache_key) if tree is None: # Cache miss — fetch bytes and parse. if raw is None: raw = read_object(root, object_id) if raw is None: logger.debug("Object %s missing for %s — skipping", object_id[:8], file_path) continue tree = parse_symbols(raw, file_path) active_cache.put(cache_key, tree) else: # Cache hit — re-key addresses when two files share the same # SHA-256 (identical content). The cache stores the tree built # for whichever file was parsed first; its addresses carry that # file's path prefix. Patch them to the current file_path so # cross-file clone detection and all downstream consumers get # correct symbol addresses. tree = _rekey_tree(tree, file_path) # --- Apply filters on top of the (potentially cached) full tree -- if kind_filter: tree = {addr: rec for addr, rec in tree.items() if rec["kind"] == kind_filter} if tree: result[file_path] = tree if own_cache: active_cache.save() return result # --------------------------------------------------------------------------- # Commit-graph walking # --------------------------------------------------------------------------- def walk_commits_bfs( root: pathlib.Path, start_commit_id: str, max_commits: int = 500, stop_at_commit_id: str | None = None, ) -> tuple[list[CommitRecord], bool]: """BFS walk of the commit DAG from *start_commit_id*, newest-first. Unlike the linear :func:`walk_commits`, this follows ``parent2_commit_id`` at merge commits so events on merged feature branches are never missed. Args: root: Repository root. start_commit_id: SHA-256 of the commit to start from. max_commits: Safety cap — returns ``truncated=True`` if reached. stop_at_commit_id: When set, stop BFS *before* entering this commit (exclusive lower bound, useful for range queries). Returns: ``(commits, truncated)`` — list sorted newest-first by ``committed_at``; ``truncated=True`` when ``max_commits`` was hit before exhausting the DAG. """ commits_by_id: _CommitIndex = {} queue: deque[str] = deque([start_commit_id]) seen: set[str] = set() truncated = False while queue: if len(commits_by_id) >= max_commits: truncated = True break commit_id = queue.popleft() if commit_id in seen: continue if stop_at_commit_id and commit_id == stop_at_commit_id: continue seen.add(commit_id) commit = read_commit(root, commit_id) if commit is None: continue commits_by_id[commit_id] = commit if commit.parent_commit_id: queue.append(commit.parent_commit_id) if commit.parent2_commit_id: queue.append(commit.parent2_commit_id) return ( sorted(commits_by_id.values(), key=lambda c: c.committed_at, reverse=True), truncated, ) # --------------------------------------------------------------------------- # Op traversal helpers # --------------------------------------------------------------------------- def flat_symbol_ops(ops: list[DomainOp]) -> Iterator[DomainOp]: """Yield all leaf ops, recursing into PatchOp.child_ops. Only yields ops that have a symbol-level address (i.e. contain ``::``). """ for op in ops: if op["op"] == "patch": for child in op["child_ops"]: if "::" in child["address"]: yield child elif "::" in op["address"]: yield op def touched_files(ops: list[DomainOp]) -> frozenset[str]: """Return the set of file paths that appear as PatchOp addresses in *ops*. Only counts files that had symbol-level child ops (semantic changes), not coarse file-level replace/insert/delete ops. """ files: set[str] = set() for op in ops: if op["op"] == "patch" and op["child_ops"]: files.add(op["address"]) return frozenset(files) def file_pairs(files: frozenset[str]) -> Iterator[tuple[str, str]]: """Yield all ordered pairs ``(a, b)`` with ``a < b`` from *files*.""" yield from itertools.combinations(sorted(files), 2)