diff --git a/scripts/map.py b/scripts/map.py index b8f956c..56245ba 100644 --- a/scripts/map.py +++ b/scripts/map.py @@ -26,7 +26,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime from pathlib import Path -from dataclasses import dataclass, asdict +from dataclasses import dataclass, asdict, field from difflib import SequenceMatcher from collections import defaultdict from multiprocessing import cpu_count @@ -37,7 +37,7 @@ # Cache format version - bump when Symbol structure or file selection changes -CACHE_VERSION = 6 # v6: Added .mm (Objective-C++) and .metal file support +CACHE_VERSION = 7 # v7: Cache text elements (comments/strings) for FTS5 # Database schema version - bump when SQLite schema changes DB_VERSION = 1 # v1: Initial versioned schema @@ -87,6 +87,15 @@ class TextElement: content: str symbol_name: str | None = None # Symbol name if this is a docstring + def to_dict(self) -> dict: + """Convert to dictionary for JSON cache serialisation.""" + return asdict(self) + + @classmethod + def from_dict(cls, d: dict) -> "TextElement": + """Create from dictionary.""" + return cls(**d) + @dataclass class FileCache: @@ -94,12 +103,14 @@ class FileCache: mtime: float content_hash: str symbols: list[Symbol] + text_elements: list[TextElement] = field(default_factory=list) def to_dict(self) -> dict: return { "mtime": self.mtime, "content_hash": self.content_hash, "symbols": [s.to_dict() for s in self.symbols], + "text_elements": [e.to_dict() for e in self.text_elements], } @classmethod @@ -108,6 +119,9 @@ def from_dict(cls, d: dict) -> "FileCache": mtime=d["mtime"], content_hash=d["content_hash"], symbols=[Symbol.from_dict(s) for s in d["symbols"]], + # Older cache entries predate text extraction; CACHE_VERSION guards + # this, but stay tolerant in case of a hand-edited cache file. + text_elements=[TextElement.from_dict(e) for e in d.get("text_elements", [])], ) @@ -157,10 +171,12 @@ def save_if_needed(self) -> None: if self._dirty_count >= self.SAVE_INTERVAL: self.save() - def get_symbols(self, file_path: Path, rel_path: str) -> tuple[list[Symbol], bool]: + def get_symbols( + self, file_path: Path, rel_path: str + ) -> tuple[list[Symbol], list[TextElement], bool]: """ - Get symbols for a file, using cache if valid. - Returns (symbols, was_cached). + Get symbols and text elements for a file, using cache if valid. + Returns (symbols, text_elements, was_cached). """ cached = self.files.get(rel_path) @@ -168,32 +184,44 @@ def get_symbols(self, file_path: Path, rel_path: str) -> tuple[list[Symbol], boo if not file_path.exists(): if cached: del self.files[rel_path] - return [], False + return [], [], False current_mtime = file_path.stat().st_mtime # Fast path: mtime unchanged if cached and cached.mtime == current_mtime: - return cached.symbols, True + return cached.symbols, cached.text_elements, True # mtime changed - check content hash try: content = file_path.read_bytes() current_hash = hashlib.sha256(content).hexdigest() except IOError: - return [], False + return [], [], False # Content unchanged - just update mtime in cache if cached and cached.content_hash == current_hash: cached.mtime = current_mtime - return cached.symbols, True + return cached.symbols, cached.text_elements, True # Content changed - need to reparse - return [], False - - def update(self, rel_path: str, mtime: float, content_hash: str, symbols: list[Symbol]) -> None: - """Update cache with newly parsed symbols.""" - self.files[rel_path] = FileCache(mtime=mtime, content_hash=content_hash, symbols=symbols) + return [], [], False + + def update( + self, + rel_path: str, + mtime: float, + content_hash: str, + symbols: list[Symbol], + text_elements: list[TextElement], + ) -> None: + """Update cache with newly parsed symbols and text elements.""" + self.files[rel_path] = FileCache( + mtime=mtime, + content_hash=content_hash, + symbols=symbols, + text_elements=text_elements, + ) self._dirty_count += 1 def remove_stale(self, valid_paths: set[str]) -> None: @@ -215,13 +243,13 @@ def get_worker_count(percent: int = DEFAULT_WORKERS_PERCENT) -> int: return min(workers, MAX_WORKERS) -def parse_file_worker(args: tuple) -> tuple[str, float, str, list[dict], str]: +def parse_file_worker(args: tuple) -> tuple[str, float, str, list[dict], list[dict], str]: """ Worker function for parallel parsing. Takes (file_path_str, root_str, language) tuple. - Returns (rel_path, mtime, content_hash, symbols_as_dicts, language). + Returns (rel_path, mtime, content_hash, symbols_as_dicts, text_as_dicts, language). - Note: Returns dicts instead of Symbol objects for pickling. + Note: Returns dicts instead of Symbol/TextElement objects for pickling. """ file_path_str, root_str, language = args file_path = Path(file_path_str) @@ -233,21 +261,26 @@ def parse_file_worker(args: tuple) -> tuple[str, float, str, list[dict], str]: content = file_path.read_bytes() content_hash = hashlib.sha256(content).hexdigest() except IOError: - return (rel_path, 0, "", [], language) + return (rel_path, 0, "", [], [], language) # Parse based on language if language == "python": symbols = extract_symbols_from_python(file_path, root) + text_elements = extract_text_elements_from_python(file_path, root) elif language == "cpp": symbols = extract_symbols_from_cpp(file_path, root) + text_elements = extract_text_elements_from_cpp(file_path, root) elif language == "rust": symbols = extract_symbols_from_rust(file_path, root) + text_elements = extract_text_elements_from_rust(file_path, root) else: symbols = [] + text_elements = [] # Convert to dicts for pickling symbol_dicts = [s.to_dict() for s in symbols] - return (rel_path, mtime, content_hash, symbol_dicts, language) + text_dicts = [e.to_dict() for e in text_elements] + return (rel_path, mtime, content_hash, symbol_dicts, text_dicts, language) def get_function_signature(node: ast.FunctionDef | ast.AsyncFunctionDef) -> str: @@ -364,6 +397,171 @@ def get_rust_parser() -> Parser: return _rust_parser +# --------------------------------------------------------------------------- +# Text extraction for the FTS5 index +# +# These walk the same sources the symbol extractors do, but collect comments +# and string literals instead of declarations. Results land in code_text_fts so +# callers can answer "which file mentions this string?" without grepping the +# whole tree. +# --------------------------------------------------------------------------- + +# tree-sitter node type -> element_type. A matched node is not descended into, +# so the body of a comment or string never produces duplicate children. +CPP_TEXT_NODES: dict[str, str] = { + "comment": "comment", + "string_literal": "string_literal", + "raw_string_literal": "string_literal", + "char_literal": "string_literal", +} + +RUST_TEXT_NODES: dict[str, str] = { + "line_comment": "comment", + "block_comment": "comment", + "string_literal": "string_literal", + "raw_string_literal": "string_literal", + "char_literal": "string_literal", +} + +# Guard rails: generated headers can contain enormous string blobs, and a file +# with thousands of comments would bloat the index without adding signal. +MAX_TEXT_ELEMENT_CHARS = 2000 +MAX_TEXT_ELEMENTS_PER_FILE = 5000 + + +def clean_text_element(raw: str) -> str: + """Strip comment markers and surrounding quotes from raw source text.""" + text = raw.strip() + if not text: + return "" + + if text.startswith("/*"): + text = text[2:] + if text.endswith("*/"): + text = text[:-2] + else: + while text.startswith("//"): + text = text[2:] + if text[:1] in ("/", "!"): + text = text[1:] + + if len(text) >= 2 and text[0] in "\"'`" and text.endswith(text[0]): + text = text[1:-1] + + return text.strip()[:MAX_TEXT_ELEMENT_CHARS] + + +def walk_text_nodes( + node: Node, + source: bytes, + node_map: dict[str, str], + out: list[TextElement], + rel_path: str, +) -> None: + """Recursively collect comment/string nodes from a tree-sitter tree.""" + if len(out) >= MAX_TEXT_ELEMENTS_PER_FILE: + return + + element_type = node_map.get(node.type) + if element_type is not None: + raw = source[node.start_byte:node.end_byte].decode("utf-8", errors="replace") + content = clean_text_element(raw) + if content: + out.append(TextElement( + file_path=rel_path, + line_number=node.start_point[0] + 1, + element_type=element_type, + content=content, + )) + return + + for child in node.children: + walk_text_nodes(child, source, node_map, out, rel_path) + + +def extract_text_elements_from_tree( + file_path: Path, + relative_to: Path, + node_map: dict[str, str], + parser: Parser, +) -> list[TextElement]: + """Extract comments and string literals from a tree-sitter parsed file.""" + try: + source = file_path.read_bytes() + except IOError: + return [] + + try: + tree = parser.parse(source) + except Exception: + return [] + + out: list[TextElement] = [] + rel_path = str(file_path.relative_to(relative_to)) + walk_text_nodes(tree.root_node, source, node_map, out, rel_path) + return out + + +def extract_text_elements_from_cpp(file_path: Path, relative_to: Path) -> list[TextElement]: + """Extract comments and string literals from a C/C++/ObjC++/Metal file.""" + return extract_text_elements_from_tree( + file_path, relative_to, CPP_TEXT_NODES, get_cpp_parser() + ) + + +def extract_text_elements_from_rust(file_path: Path, relative_to: Path) -> list[TextElement]: + """Extract comments and string literals from a Rust file.""" + return extract_text_elements_from_tree( + file_path, relative_to, RUST_TEXT_NODES, get_rust_parser() + ) + + +def extract_text_elements_from_python(file_path: Path, relative_to: Path) -> list[TextElement]: + """ + Extract comments and string literals from a Python file. + + Uses the stdlib tokenizer rather than tree-sitter: Python is handled through + `ast` elsewhere in this module and no Python grammar is bundled. + """ + import tokenize + + try: + rel_path = str(file_path.relative_to(relative_to)) + except ValueError: + return [] + + out: list[TextElement] = [] + + try: + with open(file_path, "rb") as fh: + tokens = list(tokenize.tokenize(fh.readline)) + except (tokenize.TokenError, SyntaxError, OSError, UnicodeDecodeError): + return out + + for tok in tokens: + if len(out) >= MAX_TEXT_ELEMENTS_PER_FILE: + break + + if tok.type == tokenize.COMMENT: + content = tok.string.lstrip("#").strip()[:MAX_TEXT_ELEMENT_CHARS] + element_type = "comment" + elif tok.type == tokenize.STRING: + content = clean_text_element(tok.string) + element_type = "string_literal" + else: + continue + + if content: + out.append(TextElement( + file_path=rel_path, + line_number=tok.start[0], + element_type=element_type, + content=content, + )) + + return out + + def get_doc_comment(node: Node, source: bytes) -> str | None: """Extract doc comments (///, /**, //!) preceding a node.""" comments = [] @@ -768,7 +966,11 @@ def set_metadata(conn: sqlite3.Connection, key: str, value: str) -> None: conn.execute("INSERT OR REPLACE INTO metadata (key, value) VALUES (?, ?)", [key, value]) -def write_symbols_to_sqlite(symbols: list[Symbol], db_path: Path) -> None: +def write_symbols_to_sqlite( + symbols: list[Symbol], + text_elements: list[TextElement], + db_path: Path, +) -> None: """Write symbols to SQLite database for MCP server queries.""" # Connect directly - SQLite WAL mode + transactions handle atomicity and concurrency conn = sqlite3.connect(db_path, timeout=30.0) @@ -825,6 +1027,13 @@ def write_symbols_to_sqlite(symbols: list[Symbol], db_path: Path) -> None: [(s.name, s.kind, s.signature, s.docstring, s.file_path, s.line_number, s.end_line_number, s.parent) for s in symbols] ) + if text_elements: + conn.executemany( + """INSERT INTO code_text_fts (file_path, line_number, element_type, symbol_name, content) + VALUES (?, ?, ?, ?, ?)""", + [(e.file_path, e.line_number, e.element_type, e.symbol_name, e.content) for e in text_elements] + ) + # Set metadata to indicate successful indexing completion set_metadata(conn, 'status', 'completed') set_metadata(conn, 'db_version', str(DB_VERSION)) @@ -1062,7 +1271,8 @@ def main(): cache.save() # First pass: check cache and categorize files - all_symbols = [] + all_symbols: list[Symbol] = [] + all_text_elements: list[TextElement] = [] all_rel_paths = set() files_to_parse = [] # (file_path_str, root_str, language) @@ -1070,9 +1280,10 @@ def main(): for file_path in python_files: rel_path = str(file_path.relative_to(root)) all_rel_paths.add(rel_path) - symbols, was_cached = cache.get_symbols(file_path, rel_path) + symbols, text_elements, was_cached = cache.get_symbols(file_path, rel_path) if was_cached: all_symbols.extend(symbols) + all_text_elements.extend(text_elements) else: files_to_parse.append((str(file_path), str(root), "python")) @@ -1080,9 +1291,10 @@ def main(): for file_path in cpp_files: rel_path = str(file_path.relative_to(root)) all_rel_paths.add(rel_path) - symbols, was_cached = cache.get_symbols(file_path, rel_path) + symbols, text_elements, was_cached = cache.get_symbols(file_path, rel_path) if was_cached: all_symbols.extend(symbols) + all_text_elements.extend(text_elements) else: files_to_parse.append((str(file_path), str(root), "cpp")) @@ -1090,9 +1302,10 @@ def main(): for file_path in rust_files: rel_path = str(file_path.relative_to(root)) all_rel_paths.add(rel_path) - symbols, was_cached = cache.get_symbols(file_path, rel_path) + symbols, text_elements, was_cached = cache.get_symbols(file_path, rel_path) if was_cached: all_symbols.extend(symbols) + all_text_elements.extend(text_elements) else: files_to_parse.append((str(file_path), str(root), "rust")) @@ -1137,11 +1350,13 @@ def update_progress(status: str, completed: int = 0, total: int = 0, symbols: in completed = 0 for future in as_completed(futures): try: - rel_path, mtime, content_hash, symbol_dicts, lang = future.result() + rel_path, mtime, content_hash, symbol_dicts, text_dicts, lang = future.result() symbols = [Symbol.from_dict(d) for d in symbol_dicts] + text_elements = [TextElement.from_dict(d) for d in text_dicts] all_symbols.extend(symbols) + all_text_elements.extend(text_elements) if mtime > 0: # Valid result - cache.update(rel_path, mtime, content_hash, symbols) + cache.update(rel_path, mtime, content_hash, symbols, text_elements) completed += 1 if completed % update_interval == 0 or completed == len(files_to_parse): cache.save_if_needed() @@ -1153,11 +1368,13 @@ def update_progress(status: str, completed: int = 0, total: int = 0, symbols: in # Sequential parsing for small number of files completed = 0 for args in files_to_parse: - rel_path, mtime, content_hash, symbol_dicts, lang = parse_file_worker(args) + rel_path, mtime, content_hash, symbol_dicts, text_dicts, lang = parse_file_worker(args) symbols = [Symbol.from_dict(d) for d in symbol_dicts] + text_elements = [TextElement.from_dict(d) for d in text_dicts] all_symbols.extend(symbols) + all_text_elements.extend(text_elements) if mtime > 0: - cache.update(rel_path, mtime, content_hash, symbols) + cache.update(rel_path, mtime, content_hash, symbols, text_elements) cache.save_if_needed() completed += 1 if completed % update_interval == 0 or completed == len(files_to_parse): @@ -1170,7 +1387,7 @@ def update_progress(status: str, completed: int = 0, total: int = 0, symbols: in cache.save() # Write to SQLite database for MCP server queries - write_symbols_to_sqlite(all_symbols, db_path) + write_symbols_to_sqlite(all_symbols, all_text_elements, db_path) similar_classes = find_similar_classes(all_symbols) similar_functions = find_similar_functions(all_symbols)