diff --git a/graphify/extract.py b/graphify/extract.py index d99d4b19..3c47ff1d 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -3650,7 +3650,167 @@ def _check_tree_sitter_version() -> None: ) -def extract(paths: list[Path], cache_root: Path | None = None) -> dict: +_DISPATCH: dict[str, Any] = { + ".py": extract_python, + ".js": extract_js, + ".jsx": extract_js, + ".mjs": extract_js, + ".ts": extract_js, + ".tsx": extract_js, + ".go": extract_go, + ".rs": extract_rust, + ".java": extract_java, + ".c": extract_c, + ".h": extract_c, + ".cpp": extract_cpp, + ".cc": extract_cpp, + ".cxx": extract_cpp, + ".hpp": extract_cpp, + ".rb": extract_ruby, + ".cs": extract_csharp, + ".kt": extract_kotlin, + ".kts": extract_kotlin, + ".scala": extract_scala, + ".php": extract_php, + ".swift": extract_swift, + ".lua": extract_lua, + ".toc": extract_lua, + ".zig": extract_zig, + ".ps1": extract_powershell, + ".ex": extract_elixir, + ".exs": extract_elixir, + ".m": extract_objc, + ".mm": extract_objc, + ".jl": extract_julia, + ".vue": extract_js, + ".svelte": extract_js, + ".dart": extract_dart, + ".v": extract_verilog, + ".sv": extract_verilog, + ".sql": extract_sql, +} + + +def _get_extractor(path: Path) -> Any | None: + """Return the correct extractor function for a file, or None if unsupported.""" + if path.name.endswith(".blade.php"): + return extract_blade + return _DISPATCH.get(path.suffix) + + +def _extract_single_file(args: tuple) -> tuple[int, dict]: + """Worker function for parallel extraction. Runs in a subprocess. + + Must be at module level (not a closure) so it can be pickled by + ProcessPoolExecutor. + + Args: + args: (index, path_str, cache_root_str) tuple + + Returns: + (index, result_dict) so results can be placed back in order. + """ + idx, path_str, cache_root_str = args + path = Path(path_str) + cache_root = Path(cache_root_str) + + # Check cache first (avoid re-extraction) + cached = load_cached(path, cache_root) + if cached is not None: + return idx, cached + + extractor = _get_extractor(path) + if extractor is None: + return idx, {"nodes": [], "edges": []} + + result = extractor(path) + if "error" not in result: + save_cached(path, result, cache_root) + return idx, result + + +def _extract_parallel( + uncached_work: list[tuple[int, Path]], + per_file: list[dict | None], + effective_root: Path, + max_workers: int | None, + total_files: int, +) -> None: + """Extract uncached files in parallel using ProcessPoolExecutor.""" + import concurrent.futures + + if max_workers is None: + max_workers = min(os.cpu_count() or 4, len(uncached_work), 8) + + root_str = str(effective_root) + work_items = [(idx, str(path), root_str) for idx, path in uncached_work] + + done_count = 0 + _PROGRESS_INTERVAL = 100 + with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as pool: + futures = { + pool.submit(_extract_single_file, item): item[0] for item in work_items + } + for future in concurrent.futures.as_completed(futures): + idx, result = future.result() + per_file[idx] = result + done_count += 1 + if ( + total_files >= _PROGRESS_INTERVAL + and done_count % _PROGRESS_INTERVAL == 0 + ): + print( + f" AST extraction: {done_count}/{len(uncached_work)} uncached files " + f"({done_count * 100 // len(uncached_work)}%) [{max_workers} workers]", + flush=True, + ) + if total_files >= _PROGRESS_INTERVAL: + print( + f" AST extraction: {total_files}/{total_files} files (100%) [{max_workers} workers]", + flush=True, + ) + + +def _extract_sequential( + uncached_work: list[tuple[int, Path]], + per_file: list[dict | None], + effective_root: Path, + total_files: int, +) -> None: + """Extract uncached files sequentially (fallback for small batches).""" + _PROGRESS_INTERVAL = 100 + for work_idx, (idx, path) in enumerate(uncached_work): + if ( + total_files >= _PROGRESS_INTERVAL + and work_idx % _PROGRESS_INTERVAL == 0 + and work_idx > 0 + ): + print( + f" AST extraction: {work_idx}/{len(uncached_work)} uncached files ({work_idx * 100 // len(uncached_work)}%)", + flush=True, + ) + extractor = _get_extractor(path) + if extractor is None: + per_file[idx] = {"nodes": [], "edges": []} + continue + result = extractor(path) + if "error" not in result: + save_cached(path, result, effective_root) + per_file[idx] = result + if total_files >= _PROGRESS_INTERVAL: + print(f" AST extraction: {total_files}/{total_files} files (100%)", flush=True) + + +_PARALLEL_THRESHOLD = 20 + + +def extract( + paths: list[Path], + cache_root: Path | None = None, + *, + parallel: bool = True, + max_workers: int | None = None, +) -> dict: """Extract AST nodes and edges from a list of code files. Two-pass process: @@ -3663,9 +3823,11 @@ def extract(paths: list[Path], cache_root: Path | None = None) -> dict: cache_root: explicit root for graphify-out/cache/ (overrides the inferred common path prefix). Pass Path('.') when running on a subdirectory so the cache stays at ./graphify-out/cache/. + parallel: if True and there are >= _PARALLEL_THRESHOLD uncached files, + use ProcessPoolExecutor for multi-core extraction. + max_workers: max subprocess count. Defaults to min(cpu_count, 8). """ _check_tree_sitter_version() - per_file: list[dict] = [] # Infer a common root for cache keys (use first diverging segment, not sum of all matches) try: @@ -3686,68 +3848,36 @@ def extract(paths: list[Path], cache_root: Path | None = None) -> dict: root = Path(".") root = root.resolve() - _DISPATCH: dict[str, Any] = { - ".py": extract_python, - ".js": extract_js, - ".jsx": extract_js, - ".mjs": extract_js, - ".ts": extract_js, - ".tsx": extract_js, - ".go": extract_go, - ".rs": extract_rust, - ".java": extract_java, - ".c": extract_c, - ".h": extract_c, - ".cpp": extract_cpp, - ".cc": extract_cpp, - ".cxx": extract_cpp, - ".hpp": extract_cpp, - ".rb": extract_ruby, - ".cs": extract_csharp, - ".kt": extract_kotlin, - ".kts": extract_kotlin, - ".scala": extract_scala, - ".php": extract_php, - ".swift": extract_swift, - ".lua": extract_lua, - ".toc": extract_lua, - ".zig": extract_zig, - ".ps1": extract_powershell, - ".ex": extract_elixir, - ".exs": extract_elixir, - ".m": extract_objc, - ".mm": extract_objc, - ".jl": extract_julia, - ".vue": extract_js, - ".svelte": extract_js, - ".dart": extract_dart, - ".v": extract_verilog, - ".sv": extract_verilog, - ".sql": extract_sql, - } - + effective_root = cache_root or root total = len(paths) - _PROGRESS_INTERVAL = 100 + + # Phase 1: separate cached hits from uncached work + per_file: list[dict | None] = [None] * total + uncached_work: list[tuple[int, Path]] = [] + for i, path in enumerate(paths): - if total >= _PROGRESS_INTERVAL and i % _PROGRESS_INTERVAL == 0 and i > 0: - print(f" AST extraction: {i}/{total} files ({i * 100 // total}%)", flush=True) - # .blade.php must be checked before suffix lookup since Path.suffix returns .php - if path.name.endswith(".blade.php"): - extractor = extract_blade - else: - extractor = _DISPATCH.get(path.suffix) - if extractor is None: + if _get_extractor(path) is None: + per_file[i] = {"nodes": [], "edges": []} continue - cached = load_cached(path, cache_root or root) + cached = load_cached(path, effective_root) if cached is not None: - per_file.append(cached) + per_file[i] = cached continue - result = extractor(path) - if "error" not in result: - save_cached(path, result, cache_root or root) - per_file.append(result) - if total >= _PROGRESS_INTERVAL: - print(f" AST extraction: {total}/{total} files (100%)", flush=True) + uncached_work.append((i, path)) + + # Phase 2: extract uncached files (parallel or sequential) + if uncached_work: + if parallel and len(uncached_work) >= _PARALLEL_THRESHOLD: + _extract_parallel( + uncached_work, per_file, effective_root, max_workers, total + ) + else: + _extract_sequential(uncached_work, per_file, effective_root, total) + + # Fill any remaining None slots (shouldn't happen, but defensive) + for i in range(total): + if per_file[i] is None: + per_file[i] = {"nodes": [], "edges": []} all_nodes: list[dict] = [] all_edges: list[dict] = [] diff --git a/tests/bench_extract.py b/tests/bench_extract.py new file mode 100644 index 00000000..b618d3d7 --- /dev/null +++ b/tests/bench_extract.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +"""Benchmark: sequential vs parallel AST extraction. + +Usage: + python tests/bench_extract.py [path-to-repo] + +Defaults to the current directory if no path is given. +Clears the AST cache between runs so every file is re-extracted. + +Example output: + === Graphify AST Extraction Benchmark === + Files: 1,247 + Languages: Python (412), TypeScript (389), Go (201), ... + + Sequential: 4.32s (8,934 nodes, 12,456 edges) + Parallel (8): 1.28s (8,934 nodes, 12,456 edges) + + Speedup: 3.38x + Results: ✓ identical +""" + +from __future__ import annotations + +import sys +import time +from collections import Counter +from pathlib import Path + +# Ensure the project root is importable +_project_root = Path(__file__).resolve().parent.parent +if str(_project_root) not in sys.path: + sys.path.insert(0, str(_project_root)) + +from graphify.extract import extract, collect_files +from graphify.cache import clear_cache + + +def _count_by_ext(paths: list[Path]) -> dict[str, int]: + """Count files by extension.""" + counter: Counter[str] = Counter() + for p in paths: + ext = p.suffix.lower() + counter[ext] += 1 + return dict(counter.most_common()) + + +_EXT_NAMES: dict[str, str] = { + ".py": "Python", + ".js": "JavaScript", + ".jsx": "JSX", + ".mjs": "MJS", + ".ts": "TypeScript", + ".tsx": "TSX", + ".go": "Go", + ".rs": "Rust", + ".java": "Java", + ".c": "C", + ".h": "C Header", + ".cpp": "C++", + ".cc": "C++", + ".cxx": "C++", + ".hpp": "C++ Header", + ".rb": "Ruby", + ".cs": "C#", + ".kt": "Kotlin", + ".kts": "Kotlin Script", + ".scala": "Scala", + ".php": "PHP", + ".swift": "Swift", + ".lua": "Lua", + ".toc": "Lua TOC", + ".zig": "Zig", + ".ps1": "PowerShell", + ".ex": "Elixir", + ".exs": "Elixir Script", + ".m": "Obj-C", + ".mm": "Obj-C++", + ".jl": "Julia", + ".vue": "Vue", + ".svelte": "Svelte", + ".dart": "Dart", + ".v": "Verilog", + ".sv": "SystemVerilog", + ".sql": "SQL", +} + + +def _format_languages(ext_counts: dict[str, int]) -> str: + parts = [] + for ext, count in ext_counts.items(): + name = _EXT_NAMES.get(ext, ext) + parts.append(f"{name} ({count})") + return ", ".join(parts) + + +def _run_extraction( + paths: list[Path], + cache_root: Path, + parallel: bool, + max_workers: int | None = None, +) -> tuple[float, int, int]: + """Run extraction, return (elapsed_seconds, node_count, edge_count).""" + clear_cache(cache_root) + t0 = time.perf_counter() + result = extract( + paths, cache_root=cache_root, parallel=parallel, max_workers=max_workers + ) + elapsed = time.perf_counter() - t0 + nodes = len(result.get("nodes", [])) + edges = len(result.get("edges", [])) + return elapsed, nodes, edges + + +def main() -> None: + target = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(".") + target = target.resolve() + + if not target.exists(): + print(f"Error: {target} does not exist", file=sys.stderr) + sys.exit(1) + + print("=== Graphify AST Extraction Benchmark ===\n") + print(f"Scanning {target} ...", flush=True) + + paths = collect_files(target) + if not paths: + print("No extractable files found.", file=sys.stderr) + sys.exit(1) + + ext_counts = _count_by_ext(paths) + print(f"Files: {len(paths):,}") + print(f"Languages: {_format_languages(ext_counts)}") + print() + + cache_root = target if target.is_dir() else target.parent + + # Workers count (same logic as _extract_parallel) + import os + + workers = min(os.cpu_count() or 4, len(paths), 8) + + # Run sequential + print("Running sequential extraction...", flush=True) + seq_time, seq_nodes, seq_edges = _run_extraction(paths, cache_root, parallel=False) + print(f"Sequential: {seq_time:.2f}s ({seq_nodes:,} nodes, {seq_edges:,} edges)") + + # Run parallel + print(f"\nRunning parallel extraction ({workers} workers)...", flush=True) + par_time, par_nodes, par_edges = _run_extraction( + paths, cache_root, parallel=True, max_workers=workers + ) + print( + f"Parallel ({workers}): {par_time:.2f}s ({par_nodes:,} nodes, {par_edges:,} edges)" + ) + + # Results + print() + if seq_time > 0: + speedup = seq_time / par_time if par_time > 0 else float("inf") + print(f"Speedup: {speedup:.2f}x") + print(f"Workers: {workers} (auto-detected)") + + # Validate correctness + if seq_nodes == par_nodes and seq_edges == par_edges: + print("Results: ✓ identical (node count, edge count match)") + else: + print("Results: ✗ MISMATCH!") + print(f" Sequential: {seq_nodes} nodes, {seq_edges} edges") + print(f" Parallel: {par_nodes} nodes, {par_edges} edges") + sys.exit(1) + + # Clean up cache after benchmark + clear_cache(cache_root) + print("\nCache cleared after benchmark.") + + +if __name__ == "__main__": + main()