perf: batch git cat-file reads, add pipeline timing

Author: nikolay-e

src/treemapper/diffctx/__init__.py

@@ -3,6 +3,7 @@
 import logging
 import re
 import subprocess
+import time
 from collections import defaultdict
 from pathlib import Path
 from typing import Any
@@ -18,6 +19,7 @@
 from .file_importance import compute_file_importance
 from .fragments import fragment_file  # type: ignore[attr-defined]
 from .git import (
+    CatFileBatch,
     GitError,
     get_changed_files,
     get_deleted_files,
@@ -82,6 +84,7 @@ def _read_file_content(
     file_path: Path,
     root_dir: Path,
     preferred_revs: list[str],
+    batch_reader: CatFileBatch | None = None,
 ) -> str | None:
     if file_path.suffix.lower() in KNOWN_BINARY_EXTENSIONS:
         return None
@@ -95,7 +98,10 @@ def _read_file_content(
 
     for rev in preferred_revs:
         try:
-            content = show_file_at_revision(root_dir, rev, rel)
+            if batch_reader is not None:
+                content = batch_reader.get(rev, rel)
+            else:
+                content = show_file_at_revision(root_dir, rev, rel)
             if _looks_binary(content):
                 return None
             return content
@@ -248,11 +254,12 @@ def _process_files_for_fragments(
     root_dir: Path,
     preferred_revs: list[str],
     seen_frag_ids: set[FragmentId],
+    batch_reader: CatFileBatch | None = None,
 ) -> list[Fragment]:
     max_frags = LIMITS.max_fragments
     fragments: list[Fragment] = []
     for file_path in files:
-        content = _read_file_content(file_path, root_dir, preferred_revs)
+        content = _read_file_content(file_path, root_dir, preferred_revs, batch_reader)
         if content is None:
             continue
         raw_frags = [f for f in fragment_file(file_path, content) if f.id not in seen_frag_ids]
@@ -477,8 +484,9 @@ def _create_whole_file_fragment(
     path: Path,
     root_dir: Path,
     preferred_revs: list[str],
+    batch_reader: CatFileBatch | None = None,
 ) -> Fragment | None:
-    content = _read_file_content(path, root_dir, preferred_revs)
+    content = _read_file_content(path, root_dir, preferred_revs, batch_reader)
     if not content or not content.strip():
         return None
     if _is_generated_file(path, content):
@@ -518,6 +526,7 @@ def _ensure_changed_files_represented(
     remaining_budget: int,
     root_dir: Path,
     preferred_revs: list[str],
+    batch_reader: CatFileBatch | None = None,
 ) -> list[Fragment]:
     selected_paths = {f.path for f in selected}
     missing_paths = set(changed_files) - selected_paths
@@ -537,7 +546,7 @@ def _ensure_changed_files_represented(
     for path in sorted(missing_paths):
         candidates = frags_by_path.get(path, [])
         if not candidates:
-            fallback = _create_whole_file_fragment(path, root_dir, preferred_revs)
+            fallback = _create_whole_file_fragment(path, root_dir, preferred_revs, batch_reader)
             candidates = [fallback] if fallback else []
 
         picked = _pick_smallest_fitting(candidates, selected_ids, budget_left)
@@ -652,36 +661,55 @@ def build_diff_context(
 
     preferred_revs = _build_preferred_revs(base_rev, head_rev)
 
-    seen_frag_ids: set[FragmentId] = set()
-    all_fragments = _process_files_for_fragments(changed_files, root_dir, preferred_revs, seen_frag_ids)
+    t0 = time.perf_counter()
 
-    all_candidate_files, is_large_repo = _collect_candidate_files(root_dir, set(changed_files), combined_spec)
-    all_candidate_files = _filter_whitelist(all_candidate_files, root_dir, wl_spec)
+    with CatFileBatch(root_dir) as batch_reader:
+        seen_frag_ids: set[FragmentId] = set()
+        all_fragments = _process_files_for_fragments(changed_files, root_dir, preferred_revs, seen_frag_ids, batch_reader)
 
-    edge_discovered = discover_all_related_files(changed_files, all_candidate_files, root_dir)
-    if len(edge_discovered) > _MAX_DISCOVERED_FILES:
-        logger.debug(
-            "diffctx: capping edge-discovered files from %d to %d",
-            len(edge_discovered),
-            _MAX_DISCOVERED_FILES,
-        )
-        edge_discovered = edge_discovered[:_MAX_DISCOVERED_FILES]
-    edge_discovered = [_normalize_path(p, root_dir) for p in edge_discovered]
-    all_fragments.extend(_process_files_for_fragments(edge_discovered, root_dir, preferred_revs, seen_frag_ids))
-
-    if not is_large_repo:
-        expanded_files = _expand_universe_by_rare_identifiers(
-            root_dir,
-            expansion_concepts,
-            changed_files + edge_discovered,
-            combined_spec,
-            candidate_files=all_candidate_files,
-            changed_files=changed_files,
-        )
-        expanded_files = [_normalize_path(p, root_dir) for p in expanded_files]
-        all_fragments.extend(_process_files_for_fragments(expanded_files, root_dir, preferred_revs, seen_frag_ids))
-    else:
-        logger.debug("diffctx: skipping rare-identifier expansion for large repo")
+        all_candidate_files, is_large_repo = _collect_candidate_files(root_dir, set(changed_files), combined_spec)
+        all_candidate_files = _filter_whitelist(all_candidate_files, root_dir, wl_spec)
+
+        t1 = time.perf_counter()
+
+        edge_discovered = discover_all_related_files(changed_files, all_candidate_files, root_dir)
+        if len(edge_discovered) > _MAX_DISCOVERED_FILES:
+            logger.debug(
+                "diffctx: capping edge-discovered files from %d to %d",
+                len(edge_discovered),
+                _MAX_DISCOVERED_FILES,
+            )
+            edge_discovered = edge_discovered[:_MAX_DISCOVERED_FILES]
+        edge_discovered = [_normalize_path(p, root_dir) for p in edge_discovered]
+        all_fragments.extend(_process_files_for_fragments(edge_discovered, root_dir, preferred_revs, seen_frag_ids, batch_reader))
+
+        t2 = time.perf_counter()
+
+        if not is_large_repo:
+            expanded_files = _expand_universe_by_rare_identifiers(
+                root_dir,
+                expansion_concepts,
+                changed_files + edge_discovered,
+                combined_spec,
+                candidate_files=all_candidate_files,
+                changed_files=changed_files,
+            )
+            expanded_files = [_normalize_path(p, root_dir) for p in expanded_files]
+            all_fragments.extend(
+                _process_files_for_fragments(expanded_files, root_dir, preferred_revs, seen_frag_ids, batch_reader)
+            )
+        else:
+            logger.debug("diffctx: skipping rare-identifier expansion for large repo")
+
+        t3 = time.perf_counter()
+
+    logger.debug(
+        "diffctx: timing — changed_files %.3fs, edge_discovery %.3fs, expansion %.3fs, total_io %.3fs",
+        t1 - t0,
+        t2 - t1,
+        t3 - t2,
+        t3 - t0,
+    )
 
     _assign_token_counts(all_fragments)
 
@@ -691,6 +719,8 @@ def build_diff_context(
     _assign_token_counts(signature_frags)
     all_fragments.extend(signature_frags)
 
+    t4 = time.perf_counter()
+
     if full:
         selected = _select_full_mode(all_fragments, changed_files)
         _log_full_mode(selected)
@@ -708,9 +738,15 @@ def build_diff_context(
         )
         effective_budget = budget_tokens if budget_tokens is not None else _UNLIMITED_BUDGET
         remaining = effective_budget - result.used_tokens
-        selected = _ensure_changed_files_represented(selected, all_fragments, changed_files, remaining, root_dir, preferred_revs)
+        with CatFileBatch(root_dir) as batch_reader:
+            selected = _ensure_changed_files_represented(
+                selected, all_fragments, changed_files, remaining, root_dir, preferred_revs, batch_reader
+            )
         _log_ppr_mode(selected, core_ids, budget_tokens, result, alpha, tau)
 
+    t5 = time.perf_counter()
+    logger.debug("diffctx: timing — graph+select %.3fs", t5 - t4)
+
     if no_content:
         for frag in selected:
             frag.content = ""

src/treemapper/diffctx/git.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
+import logging
 import re
 import subprocess
 from pathlib import Path
+from types import TracebackType
 
 from .types import DiffHunk
 
+logger = logging.getLogger(__name__)
+
 _HUNK_RE = re.compile(r"^@@ -(\d+)(?:,(\d+))? \+(\d+)(?:,(\d+))? @@")
 _RANGE_RE = re.compile(r"^\s*(\S+?)(\.\.\.?)(\S*?)\s*$")  # NOSONAR(S5852)
 
@@ -143,3 +147,66 @@ def get_renamed_old_paths(repo_root: Path, diff_range: str) -> set[Path]:
 def show_file_at_revision(repo_root: Path, rev: str, rel_path: Path) -> str:
     spec = f"{rev}:{rel_path.as_posix()}"
     return run_git(repo_root, ["show", spec])
+
+
+class CatFileBatch:
+    def __init__(self, repo_root: Path) -> None:
+        self._proc: subprocess.Popen[bytes] | None = None
+        self._repo_root = repo_root
+
+    def _ensure_started(self) -> subprocess.Popen[bytes]:
+        if self._proc is None or self._proc.poll() is not None:
+            self._proc = subprocess.Popen(
+                ["git", "-C", str(self._repo_root), "cat-file", "--batch"],
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+        return self._proc
+
+    def get(self, rev: str, rel_path: Path) -> str:
+        spec = f"{rev}:{rel_path.as_posix()}\n"
+        proc = self._ensure_started()
+        assert proc.stdin is not None and proc.stdout is not None
+        proc.stdin.write(spec.encode())
+        proc.stdin.flush()
+
+        header = proc.stdout.readline()
+        if not header:
+            raise GitError(f"cat-file: unexpected EOF for {spec.strip()}")
+
+        header_str = header.decode("utf-8", errors="replace").strip()
+        if header_str.endswith("missing"):
+            raise GitError(f"Path not found: {spec.strip()}")
+
+        parts = header_str.split()
+        if len(parts) < 3:
+            raise GitError(f"cat-file: malformed header: {header_str}")
+
+        size = int(parts[2])
+        content = proc.stdout.read(size)
+        proc.stdout.read(1)  # trailing LF
+
+        return content.decode("utf-8", errors="replace")
+
+    def close(self) -> None:
+        if self._proc is not None and self._proc.poll() is None:
+            assert self._proc.stdin is not None
+            self._proc.stdin.close()
+            try:
+                self._proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                self._proc.kill()
+                self._proc.wait()
+            self._proc = None
+
+    def __enter__(self) -> CatFileBatch:
+        return self
+
+    def __exit__(
+        self,
+        _exc_type: type[BaseException] | None,
+        _exc_val: BaseException | None,
+        _exc_tb: TracebackType | None,
+    ) -> None:
+        self.close()