perf: cache file reads and AST parses during discovery pass

Introduces a discovery-scoped cache (`discovery_cache()` context
manager) that ensures each file is read from disk at most once and
parsed into an AST at most once within a single discovery run.

Key changes:
- `read_file_cached()`: returns cached file content when
  `discovery_cache()` is active, falls back to normal read otherwise
- `parse_ast_cached()`: returns cached `ast.Module` when active
- `get_functions_to_optimize()`: wraps its body with `discovery_cache()`
- `find_all_functions_in_file()`: uses `read_file_cached`
- `inspect_top_level_functions_or_methods()`: uses `parse_ast_cached`
- `get_all_replay_test_functions()`: uses `parse_ast_cached`
- JS/TS export helpers: use `read_file_cached`
- Removed dead code: `_find_all_functions_via_language_support()` (never
  called, had a type error passing Path as source str)

Signature changes: None. All public function signatures are unchanged.
The cache is transparent -- when not inside `discovery_cache()`, all
functions behave identically to before (direct reads/parses).

Author: KRRT7

codeflash/discovery/functions_to_optimize.py

@@ -37,6 +37,7 @@
 
 if TYPE_CHECKING:
     from argparse import Namespace
+    from collections.abc import Generator
 
     from codeflash.models.models import CodeOptimizationContext
     from codeflash.verification.verification_utils import TestConfig
@@ -51,6 +52,46 @@ class FunctionProperties:
     staticmethod_class_name: Optional[str]
 
 
+# =============================================================================
+# Discovery-scoped file/AST cache
+# =============================================================================
+
+_active_discovery_cache: dict[Path, str] | None = None
+_active_ast_cache: dict[Path, ast.Module] | None = None
+
+
+@contextlib.contextmanager
+def discovery_cache() -> Generator[None, None, None]:
+    global _active_discovery_cache, _active_ast_cache
+    _active_discovery_cache = {}
+    _active_ast_cache = {}
+    try:
+        yield
+    finally:
+        _active_discovery_cache = None
+        _active_ast_cache = None
+
+
+def read_file_cached(file_path: Path) -> str:
+    if _active_discovery_cache is not None:
+        if file_path not in _active_discovery_cache:
+            _active_discovery_cache[file_path] = file_path.read_text(encoding="utf-8")
+        return _active_discovery_cache[file_path]
+    return file_path.read_text(encoding="utf-8")
+
+
+def parse_ast_cached(file_path: Path, source: str | None = None) -> ast.Module:
+    if _active_ast_cache is not None:
+        if file_path not in _active_ast_cache:
+            if source is None:
+                source = read_file_cached(file_path)
+            _active_ast_cache[file_path] = ast.parse(source)
+        return _active_ast_cache[file_path]
+    if source is None:
+        source = read_file_cached(file_path)
+    return ast.parse(source)
+
+
 # =============================================================================
 # Multi-language support helpers
 # =============================================================================
@@ -152,7 +193,7 @@ def _is_js_ts_function_exported(file_path: Path, function_name: str) -> tuple[bo
     from codeflash.languages.javascript.treesitter import get_analyzer_for_file
 
     try:
-        source = file_path.read_text(encoding="utf-8")
+        source = read_file_cached(file_path)
         analyzer = get_analyzer_for_file(file_path)
         return analyzer.is_function_exported(source, function_name)
     except Exception as e:
@@ -170,7 +211,7 @@ def _is_js_ts_function_exists_but_not_exported(file_path: Path, function_name: s
     from codeflash.languages.javascript.treesitter import get_analyzer_for_file
 
     try:
-        source = file_path.read_text(encoding="utf-8")
+        source = read_file_cached(file_path)
         analyzer = get_analyzer_for_file(file_path)
         all_funcs = analyzer.find_functions(
             source, include_methods=True, include_arrow_functions=True, require_name=True
@@ -183,28 +224,6 @@ def _is_js_ts_function_exists_but_not_exported(file_path: Path, function_name: s
     return False
 
 
-def _find_all_functions_via_language_support(file_path: Path) -> dict[Path, list[FunctionToOptimize]]:
-    """Find all optimizable functions using the language support abstraction.
-
-    This function uses the registered language support for the file's language
-    to discover functions, then converts them to FunctionToOptimize instances.
-    """
-    from codeflash.languages.base import FunctionFilterCriteria
-
-    functions: dict[Path, list[FunctionToOptimize]] = {}
-
-    try:
-        lang_support = get_language_support(file_path)
-        require_return = lang_support.language != Language.JAVA
-        criteria = FunctionFilterCriteria(require_return=require_return)
-        source = file_path.read_text(encoding="utf-8")
-        functions[file_path] = lang_support.discover_functions(source, file_path, criteria)
-    except Exception as e:
-        logger.debug(f"Failed to discover functions in {file_path}: {e}")
-
-    return functions
-
-
 def get_functions_to_optimize(
     optimize_all: str | None,
     replay_test: list[Path] | None,
@@ -222,7 +241,7 @@ def get_functions_to_optimize(
     functions: dict[Path, list[FunctionToOptimize]]
     trace_file_path: Path | None = None
     is_lsp = is_LSP_enabled()
-    with warnings.catch_warnings():
+    with discovery_cache(), warnings.catch_warnings():
         warnings.simplefilter(action="ignore", category=SyntaxWarning)
         if optimize_all:
             logger.info("!lsp|Finding all functions in the module '%s'…", optimize_all)
@@ -489,7 +508,7 @@ def find_all_functions_in_file(file_path: Path) -> dict[Path, list[FunctionToOpt
         lang_support = get_language_support(file_path)
         require_return = lang_support.language != Language.JAVA
         criteria = FunctionFilterCriteria(require_return=require_return)
-        source = file_path.read_text(encoding="utf-8")
+        source = read_file_cached(file_path)
         return {file_path: lang_support.discover_functions(source, file_path, criteria)}
     except Exception as e:
         logger.debug(f"Failed to discover functions in {file_path}: {e}")
@@ -506,21 +525,20 @@ def get_all_replay_test_functions(
     trace_file_path: Path | None = None
     for replay_test_file in replay_test:
         try:
-            with replay_test_file.open("r", encoding="utf8") as f:
-                tree = ast.parse(f.read())
-                for node in ast.walk(tree):
-                    if isinstance(node, ast.Assign):
-                        for target in node.targets:
-                            if (
-                                isinstance(target, ast.Name)
-                                and target.id == "trace_file_path"
-                                and isinstance(node.value, ast.Constant)
-                                and isinstance(node.value.value, str)
-                            ):
-                                trace_file_path = Path(node.value.value)
-                                break
-                        if trace_file_path:
+            tree = parse_ast_cached(replay_test_file)
+            for node in ast.walk(tree):
+                if isinstance(node, ast.Assign):
+                    for target in node.targets:
+                        if (
+                            isinstance(target, ast.Name)
+                            and target.id == "trace_file_path"
+                            and isinstance(node.value, ast.Constant)
+                            and isinstance(node.value.value, str)
+                        ):
+                            trace_file_path = Path(node.value.value)
                             break
+                    if trace_file_path:
+                        break
             if trace_file_path:
                 break
         except Exception as e:
@@ -634,7 +652,7 @@ def _get_java_replay_test_functions(
         from codeflash.languages.registry import get_language_support
 
         lang_support = get_language_support(source_file)
-        source_code = source_file.read_text(encoding="utf-8")
+        source_code = read_file_cached(source_file)
         all_functions = lang_support.discover_functions(source_code, source_file)
 
         for func in all_functions:
@@ -762,11 +780,10 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None:
 def inspect_top_level_functions_or_methods(
     file_name: Path, function_or_method_name: str, class_name: str | None = None, line_no: int | None = None
 ) -> FunctionProperties | None:
-    with file_name.open(encoding="utf8") as file:
-        try:
-            ast_module = ast.parse(file.read())
-        except Exception:
-            return None
+    try:
+        ast_module = parse_ast_cached(file_name)
+    except Exception:
+        return None
     visitor = TopLevelFunctionOrMethodVisitor(
         file_name=file_name, function_or_method_name=function_or_method_name, class_name=class_name, line_no=line_no
     )

tests/test_discovery_cache.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import ast
+import tempfile
+from pathlib import Path
+from unittest.mock import patch
+
+from codeflash.discovery.functions_to_optimize import (
+    discovery_cache,
+    find_all_functions_in_file,
+    inspect_top_level_functions_or_methods,
+    parse_ast_cached,
+    read_file_cached,
+)
+
+
+def test_read_file_cached_without_context_manager(tmp_path: Path) -> None:
+    f = tmp_path / "sample.py"
+    f.write_text("x = 1\n", encoding="utf-8")
+    assert read_file_cached(f) == "x = 1\n"
+
+
+def test_read_file_cached_returns_same_object_within_context(tmp_path: Path) -> None:
+    f = tmp_path / "sample.py"
+    f.write_text("x = 1\n", encoding="utf-8")
+    with discovery_cache():
+        result1 = read_file_cached(f)
+        result2 = read_file_cached(f)
+    assert result1 is result2
+
+
+def test_read_file_cached_does_not_persist_across_contexts(tmp_path: Path) -> None:
+    f = tmp_path / "sample.py"
+    f.write_text("x = 1\n", encoding="utf-8")
+    with discovery_cache():
+        result1 = read_file_cached(f)
+    f.write_text("x = 2\n", encoding="utf-8")
+    with discovery_cache():
+        result2 = read_file_cached(f)
+    assert result1 != result2
+
+
+def test_parse_ast_cached_returns_same_object_within_context(tmp_path: Path) -> None:
+    f = tmp_path / "sample.py"
+    f.write_text("def foo():\n    return 1\n", encoding="utf-8")
+    with discovery_cache():
+        tree1 = parse_ast_cached(f)
+        tree2 = parse_ast_cached(f)
+    assert tree1 is tree2
+    assert isinstance(tree1, ast.Module)
+
+
+def test_parse_ast_cached_uses_provided_source(tmp_path: Path) -> None:
+    f = tmp_path / "sample.py"
+    f.write_text("x = 1\n", encoding="utf-8")
+    source = "y = 2\n"
+    with discovery_cache():
+        tree = parse_ast_cached(f, source=source)
+    assert any(
+        isinstance(n, ast.Assign)
+        and isinstance(n.targets[0], ast.Name)
+        and n.targets[0].id == "y"
+        for n in ast.walk(tree)
+    )
+
+
+def test_discovery_cache_avoids_redundant_reads(tmp_path: Path) -> None:
+    f = tmp_path / "module.py"
+    f.write_text("def bar():\n    return 42\n", encoding="utf-8")
+    with discovery_cache():
+        with patch.object(Path, "read_text", wraps=f.read_text) as mock_read:
+            read_file_cached(f)
+            read_file_cached(f)
+            read_file_cached(f)
+            assert mock_read.call_count == 1
+
+
+def test_find_all_functions_in_file_uses_cache(tmp_path: Path) -> None:
+    f = tmp_path / "module.py"
+    f.write_text("def compute(x):\n    return x * 2\n", encoding="utf-8")
+    with discovery_cache():
+        result = find_all_functions_in_file(f)
+        assert f in result
+        assert result[f][0].function_name == "compute"
+
+
+def test_inspect_top_level_functions_uses_cache(tmp_path: Path) -> None:
+    f = tmp_path / "module.py"
+    f.write_text("def top_func(a, b):\n    return a + b\n", encoding="utf-8")
+    with discovery_cache():
+        props = inspect_top_level_functions_or_methods(f, "top_func")
+    assert props is not None
+    assert props.is_top_level
+    assert props.has_args
+
+
+def test_find_and_inspect_share_cached_content(tmp_path: Path) -> None:
+    f = tmp_path / "module.py"
+    f.write_text(
+        "class MyClass:\n    def method(self):\n        return 1\n\ndef standalone():\n    return 2\n",
+        encoding="utf-8",
+    )
+    with discovery_cache():
+        with patch.object(Path, "read_text", wraps=f.read_text) as mock_read:
+            find_all_functions_in_file(f)
+            props = inspect_top_level_functions_or_methods(f, "method", class_name="MyClass")
+            assert mock_read.call_count == 1
+    assert props is not None
+    assert props.is_top_level
+
+
+def test_discovery_results_correct_with_multiple_files(tmp_path: Path) -> None:
+    f1 = tmp_path / "a.py"
+    f1.write_text("def alpha():\n    return 'a'\n", encoding="utf-8")
+    f2 = tmp_path / "b.py"
+    f2.write_text("def beta(x):\n    return x + 1\n", encoding="utf-8")
+
+    with discovery_cache():
+        r1 = find_all_functions_in_file(f1)
+        r2 = find_all_functions_in_file(f2)
+
+    assert r1[f1][0].function_name == "alpha"
+    assert r2[f2][0].function_name == "beta"
+
+
+def test_cache_handles_invalid_syntax_gracefully(tmp_path: Path) -> None:
+    f = tmp_path / "broken.py"
+    f.write_text("def incomplete(:\n", encoding="utf-8")
+    with discovery_cache():
+        result = find_all_functions_in_file(f)
+    assert result == {}
+
+
+def test_cache_handles_nonexistent_file_in_parse_ast(tmp_path: Path) -> None:
+    f = tmp_path / "nonexistent.py"
+    with discovery_cache():
+        try:
+            parse_ast_cached(f)
+            assert False, "Should have raised"
+        except (FileNotFoundError, OSError):
+            pass