From a96918766fb9bf21eef4b744684f9f371b751cf5 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 10 Feb 2026 04:57:10 -0500 Subject: [PATCH 01/42] refactor: replace jedi_definition with definition_type on FunctionSource Store only the type string instead of the full Jedi Name object, removing the need for arbitrary_types_allowed and the runtime dependency on jedi in the model layer. --- codeflash/code_utils/code_replacer.py | 3 +- codeflash/context/code_context_extractor.py | 8 ++--- .../context/unused_definition_remover.py | 8 ++--- codeflash/models/models.py | 5 ++- codeflash/optimization/function_optimizer.py | 4 +-- tests/test_code_replacement.py | 34 +++++-------------- tests/test_function_dependencies.py | 5 ++- tests/test_unused_helper_revert.py | 4 +-- 8 files changed, 26 insertions(+), 45 deletions(-) diff --git a/codeflash/code_utils/code_replacer.py b/codeflash/code_utils/code_replacer.py index e543d184d..c19de3b7e 100644 --- a/codeflash/code_utils/code_replacer.py +++ b/codeflash/code_utils/code_replacer.py @@ -871,8 +871,7 @@ def replace_optimized_code( [ callee.qualified_name for callee in code_context.helper_functions - if callee.file_path == module_path - and (callee.jedi_definition is None or callee.jedi_definition.type != "class") + if callee.file_path == module_path and callee.definition_type != "class" ] ), candidate.source_code, diff --git a/codeflash/context/code_context_extractor.py b/codeflash/context/code_context_extractor.py index 61de73c32..5142ce518 100644 --- a/codeflash/context/code_context_extractor.py +++ b/codeflash/context/code_context_extractor.py @@ -265,7 +265,7 @@ def get_code_optimization_context_for_language( fully_qualified_name=helper.qualified_name, only_function_name=helper.name, source_code=helper.source_code, - jedi_definition=None, + definition_type=None, ) ) @@ -488,7 +488,7 @@ def get_function_to_optimize_as_function_source( fully_qualified_name=name.full_name, only_function_name=name.name, source_code=name.get_line_code(), - jedi_definition=name, + definition_type=name.type, ) except Exception as e: logger.exception(f"Error while getting function source: {e}") @@ -544,7 +544,7 @@ def get_function_sources_from_jedi( fully_qualified_name=definition.full_name, only_function_name=definition.name, source_code=definition.get_line_code(), - jedi_definition=definition, + definition_type=definition.type, ) file_path_to_function_source[definition_path].add(function_source) function_source_list.append(function_source) @@ -562,7 +562,7 @@ def get_function_sources_from_jedi( fully_qualified_name=f"{definition.full_name}.__init__", only_function_name="__init__", source_code=definition.get_line_code(), - jedi_definition=definition, + definition_type=definition.type, ) file_path_to_function_source[definition_path].add(function_source) function_source_list.append(function_source) diff --git a/codeflash/context/unused_definition_remover.py b/codeflash/context/unused_definition_remover.py index f4eec94e8..00b077f63 100644 --- a/codeflash/context/unused_definition_remover.py +++ b/codeflash/context/unused_definition_remover.py @@ -630,8 +630,8 @@ def _analyze_imports_in_optimized_code( helpers_by_file_and_func = defaultdict(dict) helpers_by_file = defaultdict(list) # preserved for "import module" for helper in code_context.helper_functions: - jedi_type = helper.jedi_definition.type if helper.jedi_definition else None - if jedi_type != "class": # Include when jedi_definition is None (non-Python) + jedi_type = helper.definition_type + if jedi_type != "class": # Include when definition_type is None (non-Python) func_name = helper.only_function_name module_name = helper.file_path.stem # Cache function lookup for this (module, func) @@ -789,8 +789,8 @@ def detect_unused_helper_functions( unused_helpers = [] entrypoint_file_path = function_to_optimize.file_path for helper_function in code_context.helper_functions: - jedi_type = helper_function.jedi_definition.type if helper_function.jedi_definition else None - if jedi_type != "class": # Include when jedi_definition is None (non-Python) + jedi_type = helper_function.definition_type + if jedi_type != "class": # Include when definition_type is None (non-Python) # Check if the helper function is called using multiple name variants helper_qualified_name = helper_function.qualified_name helper_simple_name = helper_function.only_function_name diff --git a/codeflash/models/models.py b/codeflash/models/models.py index d56672ba8..924d3a466 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -25,7 +25,6 @@ from re import Pattern from typing import Any, NamedTuple, Optional, cast -from jedi.api.classes import Name from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, ValidationError, model_validator from pydantic.dataclasses import dataclass @@ -136,14 +135,14 @@ class CoverReturnCode(IntEnum): ERROR = 2 -@dataclass(frozen=True, config={"arbitrary_types_allowed": True}) +@dataclass(frozen=True) class FunctionSource: file_path: Path qualified_name: str fully_qualified_name: str only_function_name: str source_code: str - jedi_definition: Name | None = None # None for non-Python languages + definition_type: str | None = None # e.g. "function", "class"; None for non-Python languages def __eq__(self, other: object) -> bool: if not isinstance(other, FunctionSource): diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index cac81fc92..e04e7626c 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -1501,8 +1501,8 @@ def replace_function_and_helpers_with_optimized_code( self.function_to_optimize.qualified_name ) for helper_function in code_context.helper_functions: - # Skip class definitions (jedi_definition may be None for non-Python languages) - if helper_function.jedi_definition is None or helper_function.jedi_definition.type != "class": + # Skip class definitions (definition_type may be None for non-Python languages) + if helper_function.definition_type != "class": read_writable_functions_by_file_path[helper_function.file_path].add(helper_function.qualified_name) for module_abspath, qualified_names in read_writable_functions_by_file_path.items(): did_update |= replace_function_definitions_in_module( diff --git a/tests/test_code_replacement.py b/tests/test_code_replacement.py index eccdc4e03..fbca6d71e 100644 --- a/tests/test_code_replacement.py +++ b/tests/test_code_replacement.py @@ -1,6 +1,5 @@ from __future__ import annotations -import dataclasses import os import re from collections import defaultdict @@ -19,28 +18,13 @@ replace_functions_in_file, ) from codeflash.discovery.functions_to_optimize import FunctionToOptimize -from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, FunctionParent +from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, FunctionParent, FunctionSource from codeflash.optimization.function_optimizer import FunctionOptimizer from codeflash.verification.verification_utils import TestConfig os.environ["CODEFLASH_API_KEY"] = "cf-test-key" -@dataclasses.dataclass -class JediDefinition: - type: str - - -@dataclasses.dataclass -class FakeFunctionSource: - file_path: Path - qualified_name: str - fully_qualified_name: str - only_function_name: str - source_code: str - jedi_definition: JediDefinition - - class Args: disable_imports_sorting = True formatter_cmds = ["disabled"] @@ -1137,7 +1121,7 @@ def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]: preexisting_objects = find_preexisting_objects(original_code) helper_functions = [ - FakeFunctionSource( + FunctionSource( file_path=Path( "/Users/saurabh/Library/CloudStorage/Dropbox/codeflash/cli/codeflash/verification/test_results.py" ), @@ -1145,7 +1129,7 @@ def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]: fully_qualified_name="codeflash.verification.test_results.TestType", only_function_name="TestType", source_code="", - jedi_definition=JediDefinition(type="class"), + definition_type="class", ) ] @@ -1160,7 +1144,7 @@ def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]: helper_functions_by_module_abspath = defaultdict(set) for helper_function in helper_functions: - if helper_function.jedi_definition.type != "class": + if helper_function.definition_type != "class": helper_functions_by_module_abspath[helper_function.file_path].add(helper_function.qualified_name) for module_abspath, qualified_names in helper_functions_by_module_abspath.items(): new_code: str = replace_functions_and_add_imports( @@ -1352,21 +1336,21 @@ def cosine_similarity_top_k( preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]] = find_preexisting_objects(original_code) helper_functions = [ - FakeFunctionSource( + FunctionSource( file_path=(Path(__file__).parent / "code_to_optimize" / "math_utils.py").resolve(), qualified_name="Matrix", fully_qualified_name="code_to_optimize.math_utils.Matrix", only_function_name="Matrix", source_code="", - jedi_definition=JediDefinition(type="class"), + definition_type="class", ), - FakeFunctionSource( + FunctionSource( file_path=(Path(__file__).parent / "code_to_optimize" / "math_utils.py").resolve(), qualified_name="cosine_similarity", fully_qualified_name="code_to_optimize.math_utils.cosine_similarity", only_function_name="cosine_similarity", source_code="", - jedi_definition=JediDefinition(type="function"), + definition_type="function", ), ] @@ -1425,7 +1409,7 @@ def cosine_similarity_top_k( ) helper_functions_by_module_abspath = defaultdict(set) for helper_function in helper_functions: - if helper_function.jedi_definition.type != "class": + if helper_function.definition_type != "class": helper_functions_by_module_abspath[helper_function.file_path].add(helper_function.qualified_name) for module_abspath, qualified_names in helper_functions_by_module_abspath.items(): new_helper_code: str = replace_functions_and_add_imports( diff --git a/tests/test_function_dependencies.py b/tests/test_function_dependencies.py index f51780f92..988f60b7b 100644 --- a/tests/test_function_dependencies.py +++ b/tests/test_function_dependencies.py @@ -151,10 +151,9 @@ def test_class_method_dependencies() -> None: # The code_context above should have the topologicalSortUtil function in it assert len(code_context.helper_functions) == 1 assert ( - code_context.helper_functions[0].jedi_definition.full_name - == "test_function_dependencies.Graph.topologicalSortUtil" + code_context.helper_functions[0].fully_qualified_name == "test_function_dependencies.Graph.topologicalSortUtil" ) - assert code_context.helper_functions[0].jedi_definition.name == "topologicalSortUtil" + assert code_context.helper_functions[0].only_function_name == "topologicalSortUtil" assert ( code_context.helper_functions[0].fully_qualified_name == "test_function_dependencies.Graph.topologicalSortUtil" ) diff --git a/tests/test_unused_helper_revert.py b/tests/test_unused_helper_revert.py index 18d21de32..c33138200 100644 --- a/tests/test_unused_helper_revert.py +++ b/tests/test_unused_helper_revert.py @@ -915,7 +915,7 @@ def local_helper(self, x): "only_function_name": "global_helper_1", "fully_qualified_name": "main.global_helper_1", "file_path": main_file, - "jedi_definition": type("MockJedi", (), {"type": "function"})(), + "definition_type": "function", }, )(), type( @@ -926,7 +926,7 @@ def local_helper(self, x): "only_function_name": "global_helper_2", "fully_qualified_name": "main.global_helper_2", "file_path": main_file, - "jedi_definition": type("MockJedi", (), {"type": "function"})(), + "definition_type": "function", }, )(), ] From b69a713fe7e054156dc60463f018b5c8dd9894bd Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 10 Feb 2026 04:57:21 -0500 Subject: [PATCH 02/42] feat: add persistent CallGraph class with SQLite caching Introduces CallGraph that uses Jedi infer()+goto() to build call edges, stores them in codeflash_cache.db with content-hash invalidation, and serves as a drop-in replacement for get_function_sources_from_jedi(). --- codeflash/context/call_graph.py | 295 ++++++++++++++++++++++++++++++++ 1 file changed, 295 insertions(+) create mode 100644 codeflash/context/call_graph.py diff --git a/codeflash/context/call_graph.py b/codeflash/context/call_graph.py new file mode 100644 index 000000000..15ff20119 --- /dev/null +++ b/codeflash/context/call_graph.py @@ -0,0 +1,295 @@ +from __future__ import annotations + +import hashlib +import os +import sqlite3 +from collections import defaultdict +from pathlib import Path +from typing import TYPE_CHECKING + +from codeflash.cli_cmds.console import logger +from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages +from codeflash.models.models import FunctionSource + +if TYPE_CHECKING: + from jedi.api.classes import Name + + +class CallGraph: + SCHEMA_VERSION = 1 + + def __init__(self, project_root: Path, db_path: Path | None = None) -> None: + import jedi + + self.project_root = project_root.resolve() + self.project_root_str = str(self.project_root) + self.jedi_project = jedi.Project(path=self.project_root) + + if db_path is None: + from codeflash.code_utils.compat import codeflash_cache_db + + db_path = codeflash_cache_db + + self.conn = sqlite3.connect(str(db_path)) + self.conn.execute("PRAGMA journal_mode=WAL") + self.indexed_file_hashes: dict[str, str] = {} + self._init_schema() + + def _init_schema(self) -> None: + cur = self.conn.cursor() + cur.execute("CREATE TABLE IF NOT EXISTS cg_schema_version (version INTEGER PRIMARY KEY)") + row = cur.execute("SELECT version FROM cg_schema_version LIMIT 1").fetchone() + if row is None: + cur.execute("INSERT INTO cg_schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,)) + elif row[0] != self.SCHEMA_VERSION: + # Schema mismatch — drop all cg_ tables and recreate + cur.execute("DROP TABLE IF EXISTS cg_call_edges") + cur.execute("DROP TABLE IF EXISTS cg_indexed_files") + cur.execute("DELETE FROM cg_schema_version") + cur.execute("INSERT INTO cg_schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,)) + + cur.execute( + """CREATE TABLE IF NOT EXISTS cg_indexed_files ( + project_root TEXT NOT NULL, + file_path TEXT NOT NULL, + file_hash TEXT NOT NULL, + PRIMARY KEY (project_root, file_path) + )""" + ) + cur.execute( + """CREATE TABLE IF NOT EXISTS cg_call_edges ( + project_root TEXT NOT NULL, + caller_file TEXT NOT NULL, + caller_qualified_name TEXT NOT NULL, + callee_file TEXT NOT NULL, + callee_qualified_name TEXT NOT NULL, + callee_fully_qualified_name TEXT NOT NULL, + callee_only_function_name TEXT NOT NULL, + callee_definition_type TEXT NOT NULL, + callee_source_line TEXT NOT NULL, + PRIMARY KEY (project_root, caller_file, caller_qualified_name, + callee_file, callee_qualified_name) + )""" + ) + cur.execute( + """CREATE INDEX IF NOT EXISTS idx_cg_edges_caller + ON cg_call_edges (project_root, caller_file, caller_qualified_name)""" + ) + self.conn.commit() + + def get_callees( + self, file_path_to_qualified_names: dict[Path, set[str]] + ) -> tuple[dict[Path, set[FunctionSource]], list[FunctionSource]]: + file_path_to_function_source: dict[Path, set[FunctionSource]] = defaultdict(set) + function_source_list: list[FunctionSource] = [] + + all_caller_keys: list[tuple[str, str]] = [] + for file_path, qualified_names in file_path_to_qualified_names.items(): + self.ensure_file_indexed(file_path) + fp_str = str(file_path.resolve()) + for qn in qualified_names: + all_caller_keys.append((fp_str, qn)) + + if not all_caller_keys: + return file_path_to_function_source, function_source_list + + cur = self.conn.cursor() + for caller_file, caller_qn in all_caller_keys: + rows = cur.execute( + """SELECT callee_file, callee_qualified_name, callee_fully_qualified_name, + callee_only_function_name, callee_definition_type, callee_source_line + FROM cg_call_edges + WHERE project_root = ? AND caller_file = ? AND caller_qualified_name = ?""", + (self.project_root_str, caller_file, caller_qn), + ).fetchall() + for callee_file, callee_qn, callee_fqn, callee_name, callee_type, callee_src in rows: + callee_path = Path(callee_file) + fs = FunctionSource( + file_path=callee_path, + qualified_name=callee_qn, + fully_qualified_name=callee_fqn, + only_function_name=callee_name, + source_code=callee_src, + definition_type=callee_type, + ) + file_path_to_function_source[callee_path].add(fs) + function_source_list.append(fs) + + return file_path_to_function_source, function_source_list + + def ensure_file_indexed(self, file_path: Path) -> None: + resolved = str(file_path.resolve()) + try: + content = file_path.read_text(encoding="utf-8") + except Exception: + return + file_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() + + cached_hash = self.indexed_file_hashes.get(resolved) + if cached_hash == file_hash: + return + + # Check DB for stored hash + row = self.conn.execute( + "SELECT file_hash FROM cg_indexed_files WHERE project_root = ? AND file_path = ?", + (self.project_root_str, resolved), + ).fetchone() + if row and row[0] == file_hash: + self.indexed_file_hashes[resolved] = file_hash + return + + self.index_file(file_path, file_hash) + + def index_file(self, file_path: Path, file_hash: str) -> None: + import jedi + + resolved = str(file_path.resolve()) + + # Delete stale data for this file + cur = self.conn.cursor() + cur.execute( + "DELETE FROM cg_call_edges WHERE project_root = ? AND caller_file = ?", (self.project_root_str, resolved) + ) + cur.execute( + "DELETE FROM cg_indexed_files WHERE project_root = ? AND file_path = ?", (self.project_root_str, resolved) + ) + + try: + script = jedi.Script(path=file_path, project=self.jedi_project) + refs = script.get_names(all_scopes=True, definitions=False, references=True) + except Exception: + logger.debug(f"CallGraph: failed to parse {file_path}") + cur.execute( + "INSERT OR REPLACE INTO cg_indexed_files (project_root, file_path, file_hash) VALUES (?, ?, ?)", + (self.project_root_str, resolved, file_hash), + ) + self.conn.commit() + self.indexed_file_hashes[resolved] = file_hash + return + + edges: set[tuple[str, str, str, str, str, str, str, str]] = set() + + for ref in refs: + try: + caller_qn = self._get_enclosing_function_qualified_name(ref) + if caller_qn is None: + continue + + definitions = self._resolve_definitions(ref) + if not definitions: + continue + + definition = definitions[0] + definition_path = definition.module_path + if definition_path is None: + continue + + if not self._is_valid_definition(definition, caller_qn): + continue + + if definition.type == "function": + callee_qn = get_qualified_name(definition.module_name, definition.full_name) + if len(callee_qn.split(".")) > 2: + continue + edges.add( + ( + resolved, + caller_qn, + str(definition_path), + callee_qn, + definition.full_name, + definition.name, + definition.type, + definition.get_line_code(), + ) + ) + elif definition.type == "class": + init_qn = get_qualified_name(definition.module_name, f"{definition.full_name}.__init__") + if len(init_qn.split(".")) > 2: + continue + edges.add( + ( + resolved, + caller_qn, + str(definition_path), + init_qn, + f"{definition.full_name}.__init__", + "__init__", + definition.type, + definition.get_line_code(), + ) + ) + except Exception: + continue + + cur.executemany( + """INSERT OR REPLACE INTO cg_call_edges + (project_root, caller_file, caller_qualified_name, + callee_file, callee_qualified_name, callee_fully_qualified_name, + callee_only_function_name, callee_definition_type, callee_source_line) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""", + [(self.project_root_str, *edge) for edge in edges], + ) + cur.execute( + "INSERT OR REPLACE INTO cg_indexed_files (project_root, file_path, file_hash) VALUES (?, ?, ?)", + (self.project_root_str, resolved, file_hash), + ) + self.conn.commit() + self.indexed_file_hashes[resolved] = file_hash + + def _resolve_definitions(self, ref: Name) -> list[Name]: + try: + inferred = ref.infer() + valid = [d for d in inferred if d.type in ("function", "class")] + if valid: + return valid + except Exception: + pass + + try: + return ref.goto(follow_imports=True, follow_builtin_imports=False) + except Exception: + return [] + + def _is_valid_definition(self, definition: Name, caller_qualified_name: str) -> bool: + definition_path = definition.module_path + if definition_path is None: + return False + if not str(definition_path).startswith(self.project_root_str + os.sep): + return False + if path_belongs_to_site_packages(definition_path): + return False + if not definition.full_name or not definition.full_name.startswith(definition.module_name): + return False + if definition.type not in ("function", "class"): + return False + # No self-edges + try: + def_qn = get_qualified_name(definition.module_name, definition.full_name) + if def_qn == caller_qualified_name: + return False + except ValueError: + return False + # Not an inner function of the caller + try: + from codeflash.optimization.function_context import belongs_to_function_qualified + + if belongs_to_function_qualified(definition, caller_qualified_name): + return False + except Exception: + pass + return True + + def _get_enclosing_function_qualified_name(self, ref: Name) -> str | None: + try: + parent = ref.parent() + if parent is None or parent.type != "function": + return None + if not parent.full_name or not parent.full_name.startswith(parent.module_name): + return None + return get_qualified_name(parent.module_name, parent.full_name) + except (ValueError, AttributeError): + return None + + def close(self) -> None: + self.conn.close() From 0078539fe20c1227889a6e5bee44cd0aafcf2cbc Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 10 Feb 2026 04:57:38 -0500 Subject: [PATCH 03/42] feat: wire CallGraph into the optimization pipeline Create CallGraph in Optimizer.run() for Python runs, pass it through FunctionOptimizer to code_context_extractor where it replaces get_function_sources_from_jedi() calls when available. --- codeflash/context/code_context_extractor.py | 19 +++++++++++++------ codeflash/optimization/function_optimizer.py | 5 ++++- codeflash/optimization/optimizer.py | 20 ++++++++++++++++++++ 3 files changed, 37 insertions(+), 7 deletions(-) diff --git a/codeflash/context/code_context_extractor.py b/codeflash/context/code_context_extractor.py index 5142ce518..4650de795 100644 --- a/codeflash/context/code_context_extractor.py +++ b/codeflash/context/code_context_extractor.py @@ -37,6 +37,7 @@ from jedi.api.classes import Name from libcst import CSTNode + from codeflash.context.call_graph import CallGraph from codeflash.context.unused_definition_remover import UsageInfo from codeflash.languages.base import HelperFunction @@ -78,6 +79,7 @@ def get_code_optimization_context( project_root_path: Path, optim_token_limit: int = OPTIMIZATION_CONTEXT_TOKEN_LIMIT, testgen_token_limit: int = TESTGEN_CONTEXT_TOKEN_LIMIT, + call_graph: CallGraph | None = None, ) -> CodeOptimizationContext: # Route to language-specific implementation for non-Python languages if not is_python(): @@ -86,9 +88,11 @@ def get_code_optimization_context( ) # Get FunctionSource representation of helpers of FTO - helpers_of_fto_dict, helpers_of_fto_list = get_function_sources_from_jedi( - {function_to_optimize.file_path: {function_to_optimize.qualified_name}}, project_root_path - ) + fto_input = {function_to_optimize.file_path: {function_to_optimize.qualified_name}} + if call_graph is not None: + helpers_of_fto_dict, helpers_of_fto_list = call_graph.get_callees(fto_input) + else: + helpers_of_fto_dict, helpers_of_fto_list = get_function_sources_from_jedi(fto_input, project_root_path) # Add function to optimize into helpers of FTO dict, as they'll be processed together fto_as_function_source = get_function_to_optimize_as_function_source(function_to_optimize, project_root_path) @@ -105,9 +109,12 @@ def get_code_optimization_context( qualified_names.update({f"{qn.rsplit('.', 1)[0]}.__init__" for qn in qualified_names if "." in qn}) # Get FunctionSource representation of helpers of helpers of FTO - helpers_of_helpers_dict, _helpers_of_helpers_list = get_function_sources_from_jedi( - helpers_of_fto_qualified_names_dict, project_root_path - ) + if call_graph is not None: + helpers_of_helpers_dict, _helpers_of_helpers_list = call_graph.get_callees(helpers_of_fto_qualified_names_dict) + else: + helpers_of_helpers_dict, _helpers_of_helpers_list = get_function_sources_from_jedi( + helpers_of_fto_qualified_names_dict, project_root_path + ) # Extract code context for optimization final_read_writable_code = extract_code_markdown_context_from_files( diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index e04e7626c..024e77c89 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -132,6 +132,7 @@ if TYPE_CHECKING: from argparse import Namespace + from codeflash.context.call_graph import CallGraph from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.either import Result from codeflash.models.models import ( @@ -437,6 +438,7 @@ def __init__( total_benchmark_timings: dict[BenchmarkKey, int] | None = None, args: Namespace | None = None, replay_tests_dir: Path | None = None, + call_graph: CallGraph | None = None, ) -> None: self.project_root = test_cfg.project_root_path self.test_cfg = test_cfg @@ -479,6 +481,7 @@ def __init__( self.function_benchmark_timings = function_benchmark_timings if function_benchmark_timings else {} self.total_benchmark_timings = total_benchmark_timings if total_benchmark_timings else {} self.replay_tests_dir = replay_tests_dir if replay_tests_dir else None + self.call_graph = call_graph n_tests = get_effort_value(EffortKeys.N_GENERATED_TESTS, self.effort) self.executor = concurrent.futures.ThreadPoolExecutor( max_workers=n_tests + 3 if self.experiment_id is None else n_tests + 4 @@ -1523,7 +1526,7 @@ def replace_function_and_helpers_with_optimized_code( def get_code_optimization_context(self) -> Result[CodeOptimizationContext, str]: try: new_code_ctx = code_context_extractor.get_code_optimization_context( - self.function_to_optimize, self.project_root + self.function_to_optimize, self.project_root, call_graph=self.call_graph ) except ValueError as e: return Failure(str(e)) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index b8f42010f..1ff956450 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -34,6 +34,7 @@ from codeflash.benchmarking.function_ranker import FunctionRanker from codeflash.code_utils.checkpoint import CodeflashRunCheckpoint + from codeflash.context.call_graph import CallGraph from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.models.models import BenchmarkKey, FunctionCalledInTest from codeflash.optimization.function_optimizer import FunctionOptimizer @@ -205,6 +206,7 @@ def create_function_optimizer( total_benchmark_timings: dict[BenchmarkKey, float] | None = None, original_module_ast: ast.Module | None = None, original_module_path: Path | None = None, + call_graph: CallGraph | None = None, ) -> FunctionOptimizer | None: from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast from codeflash.optimization.function_optimizer import FunctionOptimizer @@ -243,6 +245,7 @@ def create_function_optimizer( function_benchmark_timings=function_specific_timings, total_benchmark_timings=total_benchmark_timings if function_specific_timings else None, replay_tests_dir=self.replay_tests_dir, + call_graph=call_graph, ) def prepare_module_for_optimization( @@ -510,6 +513,19 @@ def run(self) -> None: function_benchmark_timings, total_benchmark_timings = self.run_benchmarks( file_to_funcs_to_optimize, num_optimizable_functions ) + + # Create persistent call graph for Python runs to cache Jedi analysis across functions + call_graph: CallGraph | None = None + from codeflash.languages import is_python + + if is_python(): + from codeflash.context.call_graph import CallGraph + + try: + call_graph = CallGraph(self.args.project_root) + except Exception: + logger.debug("Failed to initialize CallGraph, falling back to per-function Jedi analysis") + optimizations_found: int = 0 self.test_cfg.concolic_test_root_dir = Path( tempfile.mkdtemp(dir=self.args.tests_root, prefix="codeflash_concolic_") @@ -557,6 +573,7 @@ def run(self) -> None: total_benchmark_timings=total_benchmark_timings, original_module_ast=original_module_ast, original_module_path=original_module_path, + call_graph=call_graph, ) if function_optimizer is None: continue @@ -615,6 +632,9 @@ def run(self) -> None: else: logger.warning("⚠️ Failed to send completion email. Status") finally: + if call_graph is not None: + call_graph.close() + if function_optimizer: function_optimizer.cleanup_generated_files() From b604bf0a0ec6a8d23f6a4609b60f2e351f211bd5 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 10 Feb 2026 04:57:45 -0500 Subject: [PATCH 04/42] test: add unit and caching tests for CallGraph Covers same-file calls, cross-file calls, class instantiation, nested function exclusion, module-level exclusion, site-packages exclusion, empty/syntax-error files, and cache persistence. --- tests/test_call_graph.py | 286 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 286 insertions(+) create mode 100644 tests/test_call_graph.py diff --git a/tests/test_call_graph.py b/tests/test_call_graph.py new file mode 100644 index 000000000..945bf7bf2 --- /dev/null +++ b/tests/test_call_graph.py @@ -0,0 +1,286 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from pathlib import Path + +from codeflash.context.call_graph import CallGraph + + +@pytest.fixture +def project(tmp_path: Path) -> Path: + project_root = tmp_path / "project" + project_root.mkdir() + return project_root + + +@pytest.fixture +def db_path(tmp_path: Path) -> Path: + return tmp_path / "cache.db" + + +def write_file(project: Path, name: str, content: str) -> Path: + fp = project / name + fp.write_text(content, encoding="utf-8") + return fp + + +# --------------------------------------------------------------------------- +# Unit tests +# --------------------------------------------------------------------------- + + +def test_simple_function_call(project: Path, db_path: Path) -> None: + write_file( + project, + "mod.py", + """\ +def helper(): + return 1 + +def caller(): + return helper() +""", + ) + cg = CallGraph(project, db_path=db_path) + try: + _, result_list = cg.get_callees({project / "mod.py": {"caller"}}) + callee_qns = {fs.qualified_name for fs in result_list} + assert "helper" in callee_qns + finally: + cg.close() + + +def test_cross_file_call(project: Path, db_path: Path) -> None: + write_file( + project, + "utils.py", + """\ +def utility(): + return 42 +""", + ) + write_file( + project, + "main.py", + """\ +from utils import utility + +def caller(): + return utility() +""", + ) + cg = CallGraph(project, db_path=db_path) + try: + _, result_list = cg.get_callees({project / "main.py": {"caller"}}) + callee_qns = {fs.qualified_name for fs in result_list} + assert "utility" in callee_qns + # Should be in the utils.py file + callee_files = {fs.file_path.resolve() for fs in result_list if fs.qualified_name == "utility"} + assert (project / "utils.py").resolve() in callee_files + finally: + cg.close() + + +def test_class_instantiation(project: Path, db_path: Path) -> None: + write_file( + project, + "mod.py", + """\ +class MyClass: + def __init__(self): + pass + +def caller(): + obj = MyClass() + return obj +""", + ) + cg = CallGraph(project, db_path=db_path) + try: + _, result_list = cg.get_callees({project / "mod.py": {"caller"}}) + callee_types = {fs.definition_type for fs in result_list} + assert "class" in callee_types + finally: + cg.close() + + +def test_nested_function_excluded(project: Path, db_path: Path) -> None: + write_file( + project, + "mod.py", + """\ +def caller(): + def inner(): + return 1 + return inner() +""", + ) + cg = CallGraph(project, db_path=db_path) + try: + _, result_list = cg.get_callees({project / "mod.py": {"caller"}}) + assert len(result_list) == 0 + finally: + cg.close() + + +def test_module_level_not_tracked(project: Path, db_path: Path) -> None: + write_file( + project, + "mod.py", + """\ +def helper(): + return 1 + +x = helper() +""", + ) + cg = CallGraph(project, db_path=db_path) + try: + # Module level calls have no enclosing function, so no edges + _, result_list = cg.get_callees({project / "mod.py": {"helper"}}) + # helper itself doesn't call anything + assert len(result_list) == 0 + finally: + cg.close() + + +def test_site_packages_excluded(project: Path, db_path: Path) -> None: + write_file( + project, + "mod.py", + """\ +import os + +def caller(): + return os.path.join("a", "b") +""", + ) + cg = CallGraph(project, db_path=db_path) + try: + _, result_list = cg.get_callees({project / "mod.py": {"caller"}}) + # os.path.join is stdlib, should not appear + assert len(result_list) == 0 + finally: + cg.close() + + +def test_empty_file(project: Path, db_path: Path) -> None: + write_file(project, "mod.py", "") + cg = CallGraph(project, db_path=db_path) + try: + _, result_list = cg.get_callees({project / "mod.py": set()}) + assert len(result_list) == 0 + finally: + cg.close() + + +def test_syntax_error_file(project: Path, db_path: Path) -> None: + write_file(project, "mod.py", "def broken(\n") + cg = CallGraph(project, db_path=db_path) + try: + _, result_list = cg.get_callees({project / "mod.py": {"broken"}}) + assert len(result_list) == 0 + finally: + cg.close() + + +# --------------------------------------------------------------------------- +# Caching tests +# --------------------------------------------------------------------------- + + +def test_caching_no_reindex(project: Path, db_path: Path) -> None: + write_file( + project, + "mod.py", + """\ +def helper(): + return 1 + +def caller(): + return helper() +""", + ) + cg = CallGraph(project, db_path=db_path) + try: + cg.get_callees({project / "mod.py": {"caller"}}) + # Second call should use in-memory cache (hash unchanged) + resolved = str((project / "mod.py").resolve()) + assert resolved in cg.indexed_file_hashes + old_hash = cg.indexed_file_hashes[resolved] + cg.get_callees({project / "mod.py": {"caller"}}) + assert cg.indexed_file_hashes[resolved] == old_hash + finally: + cg.close() + + +def test_incremental_update_on_change(project: Path, db_path: Path) -> None: + fp = write_file( + project, + "mod.py", + """\ +def helper(): + return 1 + +def caller(): + return helper() +""", + ) + cg = CallGraph(project, db_path=db_path) + try: + _, result_list = cg.get_callees({project / "mod.py": {"caller"}}) + assert any(fs.qualified_name == "helper" for fs in result_list) + + # Modify the file — caller no longer calls helper + fp.write_text( + """\ +def helper(): + return 1 + +def new_helper(): + return 2 + +def caller(): + return new_helper() +""", + encoding="utf-8", + ) + _, result_list = cg.get_callees({project / "mod.py": {"caller"}}) + callee_qns = {fs.qualified_name for fs in result_list} + assert "new_helper" in callee_qns + finally: + cg.close() + + +def test_persistence_across_sessions(project: Path, db_path: Path) -> None: + write_file( + project, + "mod.py", + """\ +def helper(): + return 1 + +def caller(): + return helper() +""", + ) + # First session: index the file + cg1 = CallGraph(project, db_path=db_path) + try: + _, result_list = cg1.get_callees({project / "mod.py": {"caller"}}) + assert any(fs.qualified_name == "helper" for fs in result_list) + finally: + cg1.close() + + # Second session: should read from DB without re-indexing + cg2 = CallGraph(project, db_path=db_path) + try: + assert len(cg2.indexed_file_hashes) == 0 # in-memory cache is empty + _, result_list = cg2.get_callees({project / "mod.py": {"caller"}}) + assert any(fs.qualified_name == "helper" for fs in result_list) + finally: + cg2.close() From 5c4a65c183eccdee42b8a56071f61e09308c23ce Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 10 Feb 2026 05:27:59 -0500 Subject: [PATCH 05/42] feat: add Rich Live visualization for call graph indexing Replace the simple progress bar with a Live + Tree + Panel display that shows files being analyzed, call edges discovered, cache hits, and summary stats during call graph indexing. --- codeflash/cli_cmds/console.py | 84 +++++++++++++++++++++++++- codeflash/context/call_graph.py | 36 ++++++++--- codeflash/optimization/optimizer.py | 7 ++- tests/test_call_graph.py | 93 ++++++++++++++++++++++++++++- 4 files changed, 210 insertions(+), 10 deletions(-) diff --git a/codeflash/cli_cmds/console.py b/codeflash/cli_cmds/console.py index dab746c47..63c975a84 100644 --- a/codeflash/cli_cmds/console.py +++ b/codeflash/cli_cmds/console.py @@ -24,10 +24,11 @@ from codeflash.lsp.lsp_message import LspCodeMessage, LspTextMessage if TYPE_CHECKING: - from collections.abc import Generator + from collections.abc import Callable, Generator from rich.progress import TaskID + from codeflash.context.call_graph import IndexResult from codeflash.lsp.lsp_message import LspMessage DEBUG_MODE = logging.getLogger().getEffectiveLevel() == logging.DEBUG @@ -196,3 +197,84 @@ def test_files_progress_bar(total: int, description: str) -> Generator[tuple[Pro ) as progress: task_id = progress.add_task(description, total=total) yield progress, task_id + + +MAX_TREE_ENTRIES = 8 +MAX_EDGE_CHILDREN = 3 + + +@contextmanager +def call_graph_live_display(total: int) -> Generator[Callable[[IndexResult], None], None, None]: + from rich.console import Group + from rich.live import Live + from rich.panel import Panel + from rich.text import Text + from rich.tree import Tree + + if is_LSP_enabled(): + lsp_log(LspTextMessage(text="Building call graph", takes_time=True)) + yield lambda _result: None + return + + progress = Progress( + SpinnerColumn(next(spinners)), + TextColumn("[progress.description]{task.description}"), + BarColumn(complete_style="cyan", finished_style="green", pulse_style="yellow"), + MofNCompleteColumn(), + TimeElapsedColumn(), + TimeRemainingColumn(), + auto_refresh=False, + ) + task_id = progress.add_task("Analyzing files", total=total) + + results: list[IndexResult] = [] + stats_indexed = 0 + stats_cached = 0 + stats_edges = 0 + stats_errors = 0 + + def make_display() -> Panel: + tree = Tree("[bold]Call Graph[/bold]") + for result in results[-MAX_TREE_ENTRIES:]: + name = result.file_path.name + if result.error: + tree.add(f"[red]{name} (error)[/red]") + elif result.cached: + tree.add(f"[dim]{name} (cached)[/dim]") + else: + branch = tree.add(f"[cyan]{name}[/cyan] [dim]({result.num_edges} edges)[/dim]") + for caller, callee in result.edges[:MAX_EDGE_CHILDREN]: + branch.add(f"[dim]{caller} -> {callee}[/dim]") + remaining = len(result.edges) - MAX_EDGE_CHILDREN + if remaining > 0: + branch.add(f"[dim italic]...and {remaining} more[/dim italic]") + + parts: list[str] = [] + if stats_indexed: + parts.append(f"{stats_indexed} indexed") + if stats_cached: + parts.append(f"{stats_cached} cached") + if stats_errors: + parts.append(f"{stats_errors} errors") + parts.append(f"{stats_edges} edges") + stats_text = Text(" . ".join(parts), style="dim") + + return Panel( + Group(progress, Text(""), tree, Text(""), stats_text), title="Building Call Graph", border_style="cyan" + ) + + def update(result: IndexResult) -> None: + nonlocal stats_indexed, stats_cached, stats_edges, stats_errors + results.append(result) + if result.error: + stats_errors += 1 + elif result.cached: + stats_cached += 1 + else: + stats_indexed += 1 + stats_edges += result.num_edges + progress.advance(task_id) + live.update(make_display()) + + with Live(make_display(), console=console, transient=True, refresh_per_second=8) as live: + yield update diff --git a/codeflash/context/call_graph.py b/codeflash/context/call_graph.py index 15ff20119..7f2151203 100644 --- a/codeflash/context/call_graph.py +++ b/codeflash/context/call_graph.py @@ -4,6 +4,7 @@ import os import sqlite3 from collections import defaultdict +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING @@ -12,9 +13,20 @@ from codeflash.models.models import FunctionSource if TYPE_CHECKING: + from collections.abc import Callable, Iterable + from jedi.api.classes import Name +@dataclass(frozen=True, slots=True) +class IndexResult: + file_path: Path + cached: bool + num_edges: int + edges: tuple[tuple[str, str], ...] # (caller_qn, callee_name) pairs + error: bool + + class CallGraph: SCHEMA_VERSION = 1 @@ -117,17 +129,17 @@ def get_callees( return file_path_to_function_source, function_source_list - def ensure_file_indexed(self, file_path: Path) -> None: + def ensure_file_indexed(self, file_path: Path) -> IndexResult: resolved = str(file_path.resolve()) try: content = file_path.read_text(encoding="utf-8") except Exception: - return + return IndexResult(file_path=file_path, cached=False, num_edges=0, edges=(), error=True) file_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() cached_hash = self.indexed_file_hashes.get(resolved) if cached_hash == file_hash: - return + return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), error=False) # Check DB for stored hash row = self.conn.execute( @@ -136,11 +148,11 @@ def ensure_file_indexed(self, file_path: Path) -> None: ).fetchone() if row and row[0] == file_hash: self.indexed_file_hashes[resolved] = file_hash - return + return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), error=False) - self.index_file(file_path, file_hash) + return self.index_file(file_path, file_hash) - def index_file(self, file_path: Path, file_hash: str) -> None: + def index_file(self, file_path: Path, file_hash: str) -> IndexResult: import jedi resolved = str(file_path.resolve()) @@ -165,7 +177,7 @@ def index_file(self, file_path: Path, file_hash: str) -> None: ) self.conn.commit() self.indexed_file_hashes[resolved] = file_hash - return + return IndexResult(file_path=file_path, cached=False, num_edges=0, edges=(), error=True) edges: set[tuple[str, str, str, str, str, str, str, str]] = set() @@ -237,6 +249,9 @@ def index_file(self, file_path: Path, file_hash: str) -> None: self.conn.commit() self.indexed_file_hashes[resolved] = file_hash + edges_summary = tuple((caller_qn, callee_name) for (_, caller_qn, _, _, _, callee_name, _, _) in edges) + return IndexResult(file_path=file_path, cached=False, num_edges=len(edges), edges=edges_summary, error=False) + def _resolve_definitions(self, ref: Name) -> list[Name]: try: inferred = ref.infer() @@ -291,5 +306,12 @@ def _get_enclosing_function_qualified_name(self, ref: Name) -> str | None: except (ValueError, AttributeError): return None + def build_index(self, file_paths: Iterable[Path], on_progress: Callable[[IndexResult], None] | None = None) -> None: + """Pre-index a batch of files. Calls on_progress(result) after each file.""" + for file_path in file_paths: + result = self.ensure_file_indexed(file_path) + if on_progress is not None: + on_progress(result) + def close(self) -> None: self.conn.close() diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 1ff956450..2373dce31 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -11,7 +11,7 @@ from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient from codeflash.api.cfapi import send_completion_email -from codeflash.cli_cmds.console import console, logger, progress_bar +from codeflash.cli_cmds.console import call_graph_live_display, console, logger, progress_bar from codeflash.code_utils import env_utils from codeflash.code_utils.code_utils import cleanup_paths, get_run_tmp_file from codeflash.code_utils.env_utils import get_pr_number, is_pr_draft @@ -526,6 +526,11 @@ def run(self) -> None: except Exception: logger.debug("Failed to initialize CallGraph, falling back to per-function Jedi analysis") + if call_graph is not None and file_to_funcs_to_optimize: + source_files = list(file_to_funcs_to_optimize.keys()) + with call_graph_live_display(len(source_files)) as on_progress: + call_graph.build_index(source_files, on_progress=on_progress) + optimizations_found: int = 0 self.test_cfg.concolic_test_root_dir = Path( tempfile.mkdtemp(dir=self.args.tests_root, prefix="codeflash_concolic_") diff --git a/tests/test_call_graph.py b/tests/test_call_graph.py index 945bf7bf2..46b5f1135 100644 --- a/tests/test_call_graph.py +++ b/tests/test_call_graph.py @@ -7,7 +7,7 @@ if TYPE_CHECKING: from pathlib import Path -from codeflash.context.call_graph import CallGraph +from codeflash.context.call_graph import CallGraph, IndexResult @pytest.fixture @@ -284,3 +284,94 @@ def caller(): assert any(fs.qualified_name == "helper" for fs in result_list) finally: cg2.close() + + +def test_build_index_with_progress(project: Path, db_path: Path) -> None: + write_file( + project, + "a.py", + """\ +def helper_a(): + return 1 + +def caller_a(): + return helper_a() +""", + ) + write_file( + project, + "b.py", + """\ +from a import helper_a + +def caller_b(): + return helper_a() +""", + ) + + cg = CallGraph(project, db_path=db_path) + try: + progress_calls: list[IndexResult] = [] + files = [project / "a.py", project / "b.py"] + cg.build_index(files, on_progress=progress_calls.append) + + # Callback fired once per file + assert len(progress_calls) == 2 + + # Verify IndexResult fields for freshly indexed files + for result in progress_calls: + assert isinstance(result, IndexResult) + assert not result.error + assert not result.cached + assert result.num_edges > 0 + assert len(result.edges) == result.num_edges + + # Files are now indexed — get_callees should return correct results + _, result_list = cg.get_callees({project / "a.py": {"caller_a"}}) + callee_qns = {fs.qualified_name for fs in result_list} + assert "helper_a" in callee_qns + finally: + cg.close() + + +def test_build_index_cached_results(project: Path, db_path: Path) -> None: + write_file( + project, + "a.py", + """\ +def helper_a(): + return 1 + +def caller_a(): + return helper_a() +""", + ) + write_file( + project, + "b.py", + """\ +from a import helper_a + +def caller_b(): + return helper_a() +""", + ) + + cg = CallGraph(project, db_path=db_path) + try: + files = [project / "a.py", project / "b.py"] + # First pass — fresh indexing + cg.build_index(files) + + # Second pass — should all be cached + cached_results: list[IndexResult] = [] + cg.build_index(files, on_progress=cached_results.append) + + assert len(cached_results) == 2 + for result in cached_results: + assert result.cached + assert not result.error + assert result.num_edges == 0 + assert result.edges == () + finally: + cg.close() From 0c1397e8140185fed294e080a0be6039e282ff69 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 10 Feb 2026 06:12:32 -0500 Subject: [PATCH 06/42] feat: enrich call graph display with cross-file tracking and dependency summary Add cross-file edge detection to IndexResult, replace tree sub-entries with flat per-file dependency labels using plain language, and add a post-indexing summary panel showing per-function dependency stats. --- codeflash/cli_cmds/console.py | 74 +++++++++++++++++++++++------ codeflash/context/call_graph.py | 26 +++++++--- codeflash/optimization/optimizer.py | 4 +- tests/test_call_graph.py | 63 ++++++++++++++++++++++++ 4 files changed, 145 insertions(+), 22 deletions(-) diff --git a/codeflash/cli_cmds/console.py b/codeflash/cli_cmds/console.py index 63c975a84..6a5ae7997 100644 --- a/codeflash/cli_cmds/console.py +++ b/codeflash/cli_cmds/console.py @@ -25,10 +25,12 @@ if TYPE_CHECKING: from collections.abc import Callable, Generator + from pathlib import Path from rich.progress import TaskID - from codeflash.context.call_graph import IndexResult + from codeflash.context.call_graph import CallGraph, IndexResult + from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.lsp.lsp_message import LspMessage DEBUG_MODE = logging.getLogger().getEffectiveLevel() == logging.DEBUG @@ -200,7 +202,6 @@ def test_files_progress_bar(total: int, description: str) -> Generator[tuple[Pro MAX_TREE_ENTRIES = 8 -MAX_EDGE_CHILDREN = 3 @contextmanager @@ -231,10 +232,11 @@ def call_graph_live_display(total: int) -> Generator[Callable[[IndexResult], Non stats_indexed = 0 stats_cached = 0 stats_edges = 0 + stats_external = 0 stats_errors = 0 def make_display() -> Panel: - tree = Tree("[bold]Call Graph[/bold]") + tree = Tree("[bold]Dependencies[/bold]") for result in results[-MAX_TREE_ENTRIES:]: name = result.file_path.name if result.error: @@ -242,29 +244,35 @@ def make_display() -> Panel: elif result.cached: tree.add(f"[dim]{name} (cached)[/dim]") else: - branch = tree.add(f"[cyan]{name}[/cyan] [dim]({result.num_edges} edges)[/dim]") - for caller, callee in result.edges[:MAX_EDGE_CHILDREN]: - branch.add(f"[dim]{caller} -> {callee}[/dim]") - remaining = len(result.edges) - MAX_EDGE_CHILDREN - if remaining > 0: - branch.add(f"[dim italic]...and {remaining} more[/dim italic]") + local = result.num_edges - result.cross_file_edges + parts = [] + if local: + parts.append(f"{local} in same file") + if result.cross_file_edges: + parts.append(f"{result.cross_file_edges} from other modules") + label = ", ".join(parts) if parts else "no dependencies" + tree.add(f"[cyan]{name}[/cyan] [dim]{label}[/dim]") parts: list[str] = [] if stats_indexed: - parts.append(f"{stats_indexed} indexed") + parts.append(f"{stats_indexed} files analyzed") if stats_cached: parts.append(f"{stats_cached} cached") if stats_errors: parts.append(f"{stats_errors} errors") - parts.append(f"{stats_edges} edges") - stats_text = Text(" . ".join(parts), style="dim") + parts.append(f"{stats_edges} dependencies found") + if stats_external: + parts.append(f"{stats_external} from other modules") + stats_text = Text(" · ".join(parts), style="dim") return Panel( - Group(progress, Text(""), tree, Text(""), stats_text), title="Building Call Graph", border_style="cyan" + Group(progress, Text(""), tree, Text(""), stats_text), + title="Building Call Graph", + border_style="cyan", ) def update(result: IndexResult) -> None: - nonlocal stats_indexed, stats_cached, stats_edges, stats_errors + nonlocal stats_indexed, stats_cached, stats_edges, stats_external, stats_errors results.append(result) if result.error: stats_errors += 1 @@ -273,8 +281,46 @@ def update(result: IndexResult) -> None: else: stats_indexed += 1 stats_edges += result.num_edges + stats_external += result.cross_file_edges progress.advance(task_id) live.update(make_display()) with Live(make_display(), console=console, transient=True, refresh_per_second=8) as live: yield update + + +def call_graph_summary(call_graph: CallGraph, file_to_funcs: dict[Path, list[FunctionToOptimize]]) -> None: + from rich.panel import Panel + + total_functions = sum(len(funcs) for funcs in file_to_funcs.values()) + if total_functions == 0: + return + + total_callees = 0 + with_context = 0 + leaf_functions = 0 + + for file_path, funcs in file_to_funcs.items(): + for func in funcs: + _, func_callees = call_graph.get_callees({file_path: {func.qualified_name}}) + count = len(func_callees) + total_callees += count + if count > 0: + with_context += 1 + else: + leaf_functions += 1 + + avg_callees = total_callees / total_functions if total_functions > 0 else 0 + + summary = ( + f"{total_functions} functions ready for optimization · " + f"avg {avg_callees:.1f} dependencies/function\n" + f"{with_context} call other functions · " + f"{leaf_functions} are self-contained" + ) + + if is_LSP_enabled(): + lsp_log(LspTextMessage(text=summary)) + return + + console.print(Panel(summary, title="Dependency Summary", border_style="cyan")) diff --git a/codeflash/context/call_graph.py b/codeflash/context/call_graph.py index 7f2151203..3bd078992 100644 --- a/codeflash/context/call_graph.py +++ b/codeflash/context/call_graph.py @@ -23,7 +23,8 @@ class IndexResult: file_path: Path cached: bool num_edges: int - edges: tuple[tuple[str, str], ...] # (caller_qn, callee_name) pairs + edges: tuple[tuple[str, str, bool], ...] # (caller_qn, callee_name, is_cross_file) + cross_file_edges: int error: bool @@ -134,12 +135,12 @@ def ensure_file_indexed(self, file_path: Path) -> IndexResult: try: content = file_path.read_text(encoding="utf-8") except Exception: - return IndexResult(file_path=file_path, cached=False, num_edges=0, edges=(), error=True) + return IndexResult(file_path=file_path, cached=False, num_edges=0, edges=(), cross_file_edges=0, error=True) file_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() cached_hash = self.indexed_file_hashes.get(resolved) if cached_hash == file_hash: - return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), error=False) + return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), cross_file_edges=0, error=False) # Check DB for stored hash row = self.conn.execute( @@ -148,7 +149,7 @@ def ensure_file_indexed(self, file_path: Path) -> IndexResult: ).fetchone() if row and row[0] == file_hash: self.indexed_file_hashes[resolved] = file_hash - return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), error=False) + return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), cross_file_edges=0, error=False) return self.index_file(file_path, file_hash) @@ -177,7 +178,7 @@ def index_file(self, file_path: Path, file_hash: str) -> IndexResult: ) self.conn.commit() self.indexed_file_hashes[resolved] = file_hash - return IndexResult(file_path=file_path, cached=False, num_edges=0, edges=(), error=True) + return IndexResult(file_path=file_path, cached=False, num_edges=0, edges=(), cross_file_edges=0, error=True) edges: set[tuple[str, str, str, str, str, str, str, str]] = set() @@ -249,8 +250,19 @@ def index_file(self, file_path: Path, file_hash: str) -> IndexResult: self.conn.commit() self.indexed_file_hashes[resolved] = file_hash - edges_summary = tuple((caller_qn, callee_name) for (_, caller_qn, _, _, _, callee_name, _, _) in edges) - return IndexResult(file_path=file_path, cached=False, num_edges=len(edges), edges=edges_summary, error=False) + edges_summary = tuple( + (caller_qn, callee_name, caller_file != callee_file) + for (caller_file, caller_qn, callee_file, _, _, callee_name, _, _) in edges + ) + cross_file_count = sum(1 for e in edges_summary if e[2]) + return IndexResult( + file_path=file_path, + cached=False, + num_edges=len(edges), + edges=edges_summary, + cross_file_edges=cross_file_count, + error=False, + ) def _resolve_definitions(self, ref: Name) -> list[Name]: try: diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 2373dce31..80f4e7ce0 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -11,7 +11,7 @@ from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient from codeflash.api.cfapi import send_completion_email -from codeflash.cli_cmds.console import call_graph_live_display, console, logger, progress_bar +from codeflash.cli_cmds.console import call_graph_live_display, call_graph_summary, console, logger, progress_bar from codeflash.code_utils import env_utils from codeflash.code_utils.code_utils import cleanup_paths, get_run_tmp_file from codeflash.code_utils.env_utils import get_pr_number, is_pr_draft @@ -530,6 +530,8 @@ def run(self) -> None: source_files = list(file_to_funcs_to_optimize.keys()) with call_graph_live_display(len(source_files)) as on_progress: call_graph.build_index(source_files, on_progress=on_progress) + call_graph_summary(call_graph, file_to_funcs_to_optimize) + raise SystemExit optimizations_found: int = 0 self.test_cfg.concolic_test_root_dir = Path( diff --git a/tests/test_call_graph.py b/tests/test_call_graph.py index 46b5f1135..a2918434a 100644 --- a/tests/test_call_graph.py +++ b/tests/test_call_graph.py @@ -325,6 +325,7 @@ def caller_b(): assert not result.cached assert result.num_edges > 0 assert len(result.edges) == result.num_edges + assert result.cross_file_edges >= 0 # Files are now indexed — get_callees should return correct results _, result_list = cg.get_callees({project / "a.py": {"caller_a"}}) @@ -373,5 +374,67 @@ def caller_b(): assert not result.error assert result.num_edges == 0 assert result.edges == () + assert result.cross_file_edges == 0 + finally: + cg.close() + + +def test_cross_file_edges_tracked(project: Path, db_path: Path) -> None: + write_file( + project, + "utils.py", + """\ +def utility(): + return 42 +""", + ) + write_file( + project, + "main.py", + """\ +from utils import utility + +def caller(): + return utility() +""", + ) + + cg = CallGraph(project, db_path=db_path) + try: + progress_calls: list[IndexResult] = [] + cg.build_index([project / "utils.py", project / "main.py"], on_progress=progress_calls.append) + + # main.py should have cross-file edges (calls into utils.py) + main_result = next(r for r in progress_calls if r.file_path.name == "main.py") + assert main_result.cross_file_edges > 0 + # At least one edge tuple should have is_cross_file=True + assert any(is_cross_file for _, _, is_cross_file in main_result.edges) + finally: + cg.close() + + +def test_same_file_edges_not_cross_file(project: Path, db_path: Path) -> None: + write_file( + project, + "mod.py", + """\ +def helper(): + return 1 + +def caller(): + return helper() +""", + ) + + cg = CallGraph(project, db_path=db_path) + try: + progress_calls: list[IndexResult] = [] + cg.build_index([project / "mod.py"], on_progress=progress_calls.append) + + assert len(progress_calls) == 1 + result = progress_calls[0] + assert result.cross_file_edges == 0 + # All edges should have is_cross_file=False + assert all(not is_cross_file for _, _, is_cross_file in result.edges) finally: cg.close() From 92be009f29b806643ccfb2da951a2b2472d80327 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 10 Feb 2026 06:16:48 -0500 Subject: [PATCH 07/42] feat: rank functions by dependency count when no trace file is available Use the call graph to sort functions by callee count (most dependencies first) in --all mode without benchmarks, replacing arbitrary ordering. --- codeflash/cli_cmds/console.py | 4 +--- codeflash/optimization/optimizer.py | 25 +++++++++++++++++++++---- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/codeflash/cli_cmds/console.py b/codeflash/cli_cmds/console.py index 6a5ae7997..3ab6a9619 100644 --- a/codeflash/cli_cmds/console.py +++ b/codeflash/cli_cmds/console.py @@ -266,9 +266,7 @@ def make_display() -> Panel: stats_text = Text(" · ".join(parts), style="dim") return Panel( - Group(progress, Text(""), tree, Text(""), stats_text), - title="Building Call Graph", - border_style="cyan", + Group(progress, Text(""), tree, Text(""), stats_text), title="Building Call Graph", border_style="cyan" ) def update(result: IndexResult) -> None: diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 80f4e7ce0..a8e9de02c 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -389,7 +389,10 @@ def display_global_ranking( console.print(f"[dim]... and {len(globally_ranked) - display_count} more functions[/dim]") def rank_all_functions_globally( - self, file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]], trace_file_path: Path | None + self, + file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]], + trace_file_path: Path | None, + call_graph: CallGraph | None = None, ) -> list[tuple[Path, FunctionToOptimize]]: """Rank all functions globally across all files based on trace data. @@ -409,8 +412,10 @@ def rank_all_functions_globally( for file_path, functions in file_to_funcs_to_optimize.items(): all_functions.extend((file_path, func) for func in functions) - # If no trace file, return in original order + # If no trace file, rank by dependency count if call graph is available if not trace_file_path or not trace_file_path.exists(): + if call_graph is not None: + return self.rank_by_dependency_count(all_functions, call_graph) logger.debug("No trace file available, using original function order") return all_functions @@ -461,6 +466,17 @@ def rank_all_functions_globally( else: return globally_ranked + def rank_by_dependency_count( + self, all_functions: list[tuple[Path, FunctionToOptimize]], call_graph: CallGraph + ) -> list[tuple[Path, FunctionToOptimize]]: + counts: list[tuple[int, int, tuple[Path, FunctionToOptimize]]] = [] + for idx, (file_path, func) in enumerate(all_functions): + _, callee_list = call_graph.get_callees({file_path: {func.qualified_name}}) + counts.append((len(callee_list), idx, (file_path, func))) + counts.sort(key=lambda x: (-x[0], x[1])) + logger.debug(f"Ranked {len(counts)} functions by dependency count (most complex first)") + return [item for _, _, item in counts] + def run(self) -> None: from codeflash.code_utils.checkpoint import CodeflashRunCheckpoint @@ -531,7 +547,6 @@ def run(self) -> None: with call_graph_live_display(len(source_files)) as on_progress: call_graph.build_index(source_files, on_progress=on_progress) call_graph_summary(call_graph, file_to_funcs_to_optimize) - raise SystemExit optimizations_found: int = 0 self.test_cfg.concolic_test_root_dir = Path( @@ -548,7 +563,9 @@ def run(self) -> None: self.functions_checkpoint = CodeflashRunCheckpoint(self.args.module_root) # GLOBAL RANKING: Rank all functions together before optimizing - globally_ranked_functions = self.rank_all_functions_globally(file_to_funcs_to_optimize, trace_file_path) + globally_ranked_functions = self.rank_all_functions_globally( + file_to_funcs_to_optimize, trace_file_path, call_graph=call_graph + ) # Cache for module preparation (avoid re-parsing same files) prepared_modules: dict[Path, tuple[dict[Path, ValidCode], ast.Module | None]] = {} From 7d5638df0cc4bf3b8c7246e1aeec9c5487b1e899 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 10 Feb 2026 06:51:58 -0500 Subject: [PATCH 08/42] feat: parallelize call graph indexing with ProcessPoolExecutor Separate Jedi analysis (CPU-bound) from DB persistence so uncached files can be analyzed across multiple worker processes. Files are dispatched to a pool of up to 8 workers when >= 8 need indexing, with sequential fallback for small batches or on pool failure. --- codeflash/context/call_graph.py | 446 ++++++++++++++++++++++---------- 1 file changed, 308 insertions(+), 138 deletions(-) diff --git a/codeflash/context/call_graph.py b/codeflash/context/call_graph.py index 3bd078992..033f8399a 100644 --- a/codeflash/context/call_graph.py +++ b/codeflash/context/call_graph.py @@ -28,6 +28,166 @@ class IndexResult: error: bool +# --------------------------------------------------------------------------- +# Module-level helpers (must be top-level for ProcessPoolExecutor pickling) +# --------------------------------------------------------------------------- + +_PARALLEL_THRESHOLD = 8 + +# Per-worker state, initialised by _init_index_worker in child processes +_worker_jedi_project: object | None = None +_worker_project_root_str: str | None = None + + +def _init_index_worker(project_root: str) -> None: + import jedi + + global _worker_jedi_project, _worker_project_root_str + _worker_jedi_project = jedi.Project(path=project_root) + _worker_project_root_str = project_root + + +def _resolve_definitions(ref: Name) -> list[Name]: + try: + inferred = ref.infer() + valid = [d for d in inferred if d.type in ("function", "class")] + if valid: + return valid + except Exception: + pass + + try: + return ref.goto(follow_imports=True, follow_builtin_imports=False) + except Exception: + return [] + + +def _is_valid_definition(definition: Name, caller_qualified_name: str, project_root_str: str) -> bool: + definition_path = definition.module_path + if definition_path is None: + return False + + if not str(definition_path).startswith(project_root_str + os.sep): + return False + + if path_belongs_to_site_packages(definition_path): + return False + + if not definition.full_name or not definition.full_name.startswith(definition.module_name): + return False + + if definition.type not in ("function", "class"): + return False + + try: + def_qn = get_qualified_name(definition.module_name, definition.full_name) + if def_qn == caller_qualified_name: + return False + except ValueError: + return False + + try: + from codeflash.optimization.function_context import belongs_to_function_qualified + + if belongs_to_function_qualified(definition, caller_qualified_name): + return False + except Exception: + pass + + return True + + +def _get_enclosing_function_qn(ref: Name) -> str | None: + try: + parent = ref.parent() + if parent is None or parent.type != "function": + return None + if not parent.full_name or not parent.full_name.startswith(parent.module_name): + return None + return get_qualified_name(parent.module_name, parent.full_name) + except (ValueError, AttributeError): + return None + + +def _analyze_file(file_path: Path, jedi_project: object, project_root_str: str) -> tuple[set[tuple[str, ...]], bool]: + """Pure Jedi analysis — no DB access. Returns (edges, had_error).""" + import jedi + + resolved = str(file_path.resolve()) + + try: + script = jedi.Script(path=file_path, project=jedi_project) + refs = script.get_names(all_scopes=True, definitions=False, references=True) + except Exception: + return set(), True + + edges: set[tuple[str, str, str, str, str, str, str, str]] = set() + + for ref in refs: + try: + caller_qn = _get_enclosing_function_qn(ref) + if caller_qn is None: + continue + + definitions = _resolve_definitions(ref) + if not definitions: + continue + + definition = definitions[0] + definition_path = definition.module_path + if definition_path is None: + continue + + if not _is_valid_definition(definition, caller_qn, project_root_str): + continue + + # Extract common edge components + edge_base = (resolved, caller_qn, str(definition_path)) + + if definition.type == "function": + callee_qn = get_qualified_name(definition.module_name, definition.full_name) + if len(callee_qn.split(".")) > 2: + continue + edges.add( + ( + *edge_base, + callee_qn, + definition.full_name, + definition.name, + definition.type, + definition.get_line_code(), + ) + ) + elif definition.type == "class": + init_qn = get_qualified_name(definition.module_name, f"{definition.full_name}.__init__") + if len(init_qn.split(".")) > 2: + continue + edges.add( + ( + *edge_base, + init_qn, + f"{definition.full_name}.__init__", + "__init__", + definition.type, + definition.get_line_code(), + ) + ) + except Exception: + continue + + return edges, False + + +def _index_file_worker(args: tuple[str, str]) -> tuple[str, str, set[tuple[str, ...]], bool]: + """Worker entry point for ProcessPoolExecutor.""" + file_path_str, file_hash = args + edges, had_error = _analyze_file(Path(file_path_str), _worker_jedi_project, _worker_project_root_str) + return file_path_str, file_hash, edges, had_error + + +# --------------------------------------------------------------------------- + + class CallGraph: SCHEMA_VERSION = 1 @@ -96,16 +256,17 @@ def get_callees( file_path_to_function_source: dict[Path, set[FunctionSource]] = defaultdict(set) function_source_list: list[FunctionSource] = [] + # Build list of all caller keys all_caller_keys: list[tuple[str, str]] = [] for file_path, qualified_names in file_path_to_qualified_names.items(): self.ensure_file_indexed(file_path) - fp_str = str(file_path.resolve()) - for qn in qualified_names: - all_caller_keys.append((fp_str, qn)) + resolved = str(file_path.resolve()) + all_caller_keys.extend((resolved, qn) for qn in qualified_names) if not all_caller_keys: return file_path_to_function_source, function_source_list + # Query all callees cur = self.conn.cursor() for caller_file, caller_qn in all_caller_keys: rows = cur.execute( @@ -115,6 +276,7 @@ def get_callees( WHERE project_root = ? AND caller_file = ? AND caller_qualified_name = ?""", (self.project_root_str, caller_file, caller_qn), ).fetchall() + for callee_file, callee_qn, callee_fqn, callee_name, callee_type, callee_src in rows: callee_path = Path(callee_file) fs = FunctionSource( @@ -132,14 +294,16 @@ def get_callees( def ensure_file_indexed(self, file_path: Path) -> IndexResult: resolved = str(file_path.resolve()) + try: content = file_path.read_text(encoding="utf-8") except Exception: return IndexResult(file_path=file_path, cached=False, num_edges=0, edges=(), cross_file_edges=0, error=True) + file_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() - cached_hash = self.indexed_file_hashes.get(resolved) - if cached_hash == file_hash: + # Check in-memory cache first + if self.indexed_file_hashes.get(resolved) == file_hash: return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), cross_file_edges=0, error=False) # Check DB for stored hash @@ -147,6 +311,7 @@ def ensure_file_indexed(self, file_path: Path) -> IndexResult: "SELECT file_hash FROM cg_indexed_files WHERE project_root = ? AND file_path = ?", (self.project_root_str, resolved), ).fetchone() + if row and row[0] == file_hash: self.indexed_file_hashes[resolved] = file_hash return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), cross_file_edges=0, error=False) @@ -154,12 +319,18 @@ def ensure_file_indexed(self, file_path: Path) -> IndexResult: return self.index_file(file_path, file_hash) def index_file(self, file_path: Path, file_hash: str) -> IndexResult: - import jedi - resolved = str(file_path.resolve()) + edges, had_error = _analyze_file(file_path, self.jedi_project, self.project_root_str) + if had_error: + logger.debug(f"CallGraph: failed to parse {file_path}") + return self._persist_edges(file_path, resolved, file_hash, edges, had_error) - # Delete stale data for this file + def _persist_edges( + self, file_path: Path, resolved: str, file_hash: str, edges: set[tuple[str, ...]], had_error: bool + ) -> IndexResult: cur = self.conn.cursor() + + # Clear existing data for this file cur.execute( "DELETE FROM cg_call_edges WHERE project_root = ? AND caller_file = ?", (self.project_root_str, resolved) ) @@ -167,163 +338,162 @@ def index_file(self, file_path: Path, file_hash: str) -> IndexResult: "DELETE FROM cg_indexed_files WHERE project_root = ? AND file_path = ?", (self.project_root_str, resolved) ) - try: - script = jedi.Script(path=file_path, project=self.jedi_project) - refs = script.get_names(all_scopes=True, definitions=False, references=True) - except Exception: - logger.debug(f"CallGraph: failed to parse {file_path}") - cur.execute( - "INSERT OR REPLACE INTO cg_indexed_files (project_root, file_path, file_hash) VALUES (?, ?, ?)", - (self.project_root_str, resolved, file_hash), + # Insert new edges if parsing succeeded + if not had_error and edges: + cur.executemany( + """INSERT OR REPLACE INTO cg_call_edges + (project_root, caller_file, caller_qualified_name, + callee_file, callee_qualified_name, callee_fully_qualified_name, + callee_only_function_name, callee_definition_type, callee_source_line) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""", + [(self.project_root_str, *edge) for edge in edges], ) - self.conn.commit() - self.indexed_file_hashes[resolved] = file_hash - return IndexResult(file_path=file_path, cached=False, num_edges=0, edges=(), cross_file_edges=0, error=True) - edges: set[tuple[str, str, str, str, str, str, str, str]] = set() - - for ref in refs: - try: - caller_qn = self._get_enclosing_function_qualified_name(ref) - if caller_qn is None: - continue - - definitions = self._resolve_definitions(ref) - if not definitions: - continue - - definition = definitions[0] - definition_path = definition.module_path - if definition_path is None: - continue - - if not self._is_valid_definition(definition, caller_qn): - continue - - if definition.type == "function": - callee_qn = get_qualified_name(definition.module_name, definition.full_name) - if len(callee_qn.split(".")) > 2: - continue - edges.add( - ( - resolved, - caller_qn, - str(definition_path), - callee_qn, - definition.full_name, - definition.name, - definition.type, - definition.get_line_code(), - ) - ) - elif definition.type == "class": - init_qn = get_qualified_name(definition.module_name, f"{definition.full_name}.__init__") - if len(init_qn.split(".")) > 2: - continue - edges.add( - ( - resolved, - caller_qn, - str(definition_path), - init_qn, - f"{definition.full_name}.__init__", - "__init__", - definition.type, - definition.get_line_code(), - ) - ) - except Exception: - continue - - cur.executemany( - """INSERT OR REPLACE INTO cg_call_edges - (project_root, caller_file, caller_qualified_name, - callee_file, callee_qualified_name, callee_fully_qualified_name, - callee_only_function_name, callee_definition_type, callee_source_line) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""", - [(self.project_root_str, *edge) for edge in edges], - ) + # Record that this file has been indexed cur.execute( "INSERT OR REPLACE INTO cg_indexed_files (project_root, file_path, file_hash) VALUES (?, ?, ?)", (self.project_root_str, resolved, file_hash), ) + self.conn.commit() self.indexed_file_hashes[resolved] = file_hash + # Build summary for return value edges_summary = tuple( (caller_qn, callee_name, caller_file != callee_file) for (caller_file, caller_qn, callee_file, _, _, callee_name, _, _) in edges ) - cross_file_count = sum(1 for e in edges_summary if e[2]) + cross_file_count = sum(is_cross_file for _, _, is_cross_file in edges_summary) + return IndexResult( file_path=file_path, cached=False, num_edges=len(edges), edges=edges_summary, cross_file_edges=cross_file_count, - error=False, + error=had_error, ) - def _resolve_definitions(self, ref: Name) -> list[Name]: - try: - inferred = ref.infer() - valid = [d for d in inferred if d.type in ("function", "class")] - if valid: - return valid - except Exception: - pass + def build_index(self, file_paths: Iterable[Path], on_progress: Callable[[IndexResult], None] | None = None) -> None: + """Pre-index a batch of files, using multiprocessing for large uncached batches.""" + to_index: list[tuple[Path, str, str]] = [] - try: - return ref.goto(follow_imports=True, follow_builtin_imports=False) - except Exception: - return [] + for file_path in file_paths: + resolved = str(file_path.resolve()) - def _is_valid_definition(self, definition: Name, caller_qualified_name: str) -> bool: - definition_path = definition.module_path - if definition_path is None: - return False - if not str(definition_path).startswith(self.project_root_str + os.sep): - return False - if path_belongs_to_site_packages(definition_path): - return False - if not definition.full_name or not definition.full_name.startswith(definition.module_name): - return False - if definition.type not in ("function", "class"): - return False - # No self-edges - try: - def_qn = get_qualified_name(definition.module_name, definition.full_name) - if def_qn == caller_qualified_name: - return False - except ValueError: - return False - # Not an inner function of the caller - try: - from codeflash.optimization.function_context import belongs_to_function_qualified + try: + content = file_path.read_text(encoding="utf-8") + except Exception: + self._report_progress( + on_progress, + IndexResult( + file_path=file_path, cached=False, num_edges=0, edges=(), cross_file_edges=0, error=True + ), + ) + continue - if belongs_to_function_qualified(definition, caller_qualified_name): - return False - except Exception: - pass - return True + file_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() + + # Check if already cached (in-memory or DB) + if self._is_file_cached(resolved, file_hash): + self._report_progress( + on_progress, + IndexResult( + file_path=file_path, cached=True, num_edges=0, edges=(), cross_file_edges=0, error=False + ), + ) + continue + + to_index.append((file_path, resolved, file_hash)) + + if not to_index: + return + + # Index uncached files + if len(to_index) >= _PARALLEL_THRESHOLD: + self._build_index_parallel(to_index, on_progress) + else: + for file_path, _resolved, file_hash in to_index: + result = self.index_file(file_path, file_hash) + self._report_progress(on_progress, result) + + def _is_file_cached(self, resolved: str, file_hash: str) -> bool: + """Check if file is cached in memory or DB.""" + # Check in-memory cache + if self.indexed_file_hashes.get(resolved) == file_hash: + return True + + # Check DB cache + row = self.conn.execute( + "SELECT file_hash FROM cg_indexed_files WHERE project_root = ? AND file_path = ?", + (self.project_root_str, resolved), + ).fetchone() + + if row and row[0] == file_hash: + self.indexed_file_hashes[resolved] = file_hash + return True + + return False + + def _report_progress(self, on_progress: Callable[[IndexResult], None] | None, result: IndexResult) -> None: + """Report progress if callback provided.""" + if on_progress is not None: + on_progress(result) + + def _build_index_parallel( + self, to_index: list[tuple[Path, str, str]], on_progress: Callable[[IndexResult], None] | None + ) -> None: + from concurrent.futures import ProcessPoolExecutor, as_completed + + max_workers = min(os.cpu_count() or 1, len(to_index), 8) + path_info: dict[str, tuple[Path, str]] = {resolved: (fp, fh) for fp, resolved, fh in to_index} + worker_args = [(resolved, fh) for _fp, resolved, fh in to_index] + + logger.debug(f"CallGraph: indexing {len(to_index)} files across {max_workers} workers") - def _get_enclosing_function_qualified_name(self, ref: Name) -> str | None: try: - parent = ref.parent() - if parent is None or parent.type != "function": - return None - if not parent.full_name or not parent.full_name.startswith(parent.module_name): - return None - return get_qualified_name(parent.module_name, parent.full_name) - except (ValueError, AttributeError): - return None + with ProcessPoolExecutor( + max_workers=max_workers, initializer=_init_index_worker, initargs=(self.project_root_str,) + ) as executor: + futures = {executor.submit(_index_file_worker, args): args[0] for args in worker_args} + + for future in as_completed(futures): + resolved = futures[future] + file_path, file_hash = path_info[resolved] + + try: + _, _, edges, had_error = future.result() + except Exception: + logger.debug(f"CallGraph: worker failed for {file_path}") + self._persist_edges(file_path, resolved, file_hash, set(), had_error=True) + self._report_progress( + on_progress, + IndexResult( + file_path=file_path, cached=False, num_edges=0, edges=(), cross_file_edges=0, error=True + ), + ) + continue - def build_index(self, file_paths: Iterable[Path], on_progress: Callable[[IndexResult], None] | None = None) -> None: - """Pre-index a batch of files. Calls on_progress(result) after each file.""" - for file_path in file_paths: - result = self.ensure_file_indexed(file_path) - if on_progress is not None: - on_progress(result) + if had_error: + logger.debug(f"CallGraph: failed to parse {file_path}") + + result = self._persist_edges(file_path, resolved, file_hash, edges, had_error) + self._report_progress(on_progress, result) + + except Exception: + logger.debug("CallGraph: parallel indexing failed, falling back to sequential") + self._fallback_sequential_index(to_index, on_progress) + + def _fallback_sequential_index( + self, to_index: list[tuple[Path, str, str]], on_progress: Callable[[IndexResult], None] | None + ) -> None: + """Fallback to sequential indexing when parallel processing fails.""" + for file_path, resolved, file_hash in to_index: + # Skip files already persisted before the failure + if resolved in self.indexed_file_hashes: + continue + result = self.index_file(file_path, file_hash) + self._report_progress(on_progress, result) def close(self) -> None: self.conn.close() From 9cc10042ee1d203beb32237a97ffac6abd2f202c Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Wed, 11 Feb 2026 20:21:36 -0500 Subject: [PATCH 09/42] refactor: improve call graph live display and filter non-Python files Use bounded deque for results, batch updates every 8 results with manual refresh to reduce flicker, and filter source_files to Python-only before passing to the call graph indexer. --- codeflash/cli_cmds/console.py | 148 +++++++++++++++++----------- codeflash/optimization/optimizer.py | 2 +- 2 files changed, 89 insertions(+), 61 deletions(-) diff --git a/codeflash/cli_cmds/console.py b/codeflash/cli_cmds/console.py index 3ab6a9619..e9e921e91 100644 --- a/codeflash/cli_cmds/console.py +++ b/codeflash/cli_cmds/console.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +from collections import deque from contextlib import contextmanager from itertools import cycle from typing import TYPE_CHECKING, Optional @@ -214,7 +215,7 @@ def call_graph_live_display(total: int) -> Generator[Callable[[IndexResult], Non if is_LSP_enabled(): lsp_log(LspTextMessage(text="Building call graph", takes_time=True)) - yield lambda _result: None + yield lambda _: None return progress = Progress( @@ -228,87 +229,114 @@ def call_graph_live_display(total: int) -> Generator[Callable[[IndexResult], Non ) task_id = progress.add_task("Analyzing files", total=total) - results: list[IndexResult] = [] - stats_indexed = 0 - stats_cached = 0 - stats_edges = 0 - stats_external = 0 - stats_errors = 0 - - def make_display() -> Panel: - tree = Tree("[bold]Dependencies[/bold]") - for result in results[-MAX_TREE_ENTRIES:]: - name = result.file_path.name + results: deque[IndexResult] = deque(maxlen=MAX_TREE_ENTRIES) + stats = { + "indexed": 0, + "cached": 0, + "edges": 0, + "external": 0, + "errors": 0, + } + + tree = Tree("[bold]Dependencies[/bold]") + stats_text = Text("0 dependencies found", style="dim") + panel = Panel( + Group(progress, Text(""), tree, Text(""), stats_text), + title="Building Call Graph", + border_style="cyan", + ) + + def create_tree_node(result: IndexResult) -> Tree: + name = f"{result.file_path.parent.name}/{result.file_path.name}" + + if result.error: + return Tree(f"[red]{name} (error)[/red]") + + if result.cached: + return Tree(f"[dim]{name} (cached)[/dim]") + + local_edges = result.num_edges - result.cross_file_edges + edge_info = [] + + if local_edges: + edge_info.append(f"{local_edges} in same file") + if result.cross_file_edges: + edge_info.append(f"{result.cross_file_edges} from other modules") + + label = ", ".join(edge_info) if edge_info else "no dependencies" + return Tree(f"[cyan]{name}[/cyan] [dim]{label}[/dim]") + + def refresh_display() -> None: + tree.children = [create_tree_node(r) for r in results] + tree.children.extend([Tree(" ")] * (MAX_TREE_ENTRIES - len(results))) + + # Update stats + stat_parts = [] + if stats["indexed"]: + stat_parts.append(f"{stats['indexed']} files analyzed") + if stats["cached"]: + stat_parts.append(f"{stats['cached']} cached") + if stats["errors"]: + stat_parts.append(f"{stats['errors']} errors") + stat_parts.append(f"{stats['edges']} dependencies found") + if stats["external"]: + stat_parts.append(f"{stats['external']} from other modules") + + stats_text.truncate(0) + stats_text.append(" · ".join(stat_parts), style="dim") + + batch: list[IndexResult] = [] + + def process_batch() -> None: + for result in batch: + results.append(result) + if result.error: - tree.add(f"[red]{name} (error)[/red]") + stats["errors"] += 1 elif result.cached: - tree.add(f"[dim]{name} (cached)[/dim]") + stats["cached"] += 1 else: - local = result.num_edges - result.cross_file_edges - parts = [] - if local: - parts.append(f"{local} in same file") - if result.cross_file_edges: - parts.append(f"{result.cross_file_edges} from other modules") - label = ", ".join(parts) if parts else "no dependencies" - tree.add(f"[cyan]{name}[/cyan] [dim]{label}[/dim]") - - parts: list[str] = [] - if stats_indexed: - parts.append(f"{stats_indexed} files analyzed") - if stats_cached: - parts.append(f"{stats_cached} cached") - if stats_errors: - parts.append(f"{stats_errors} errors") - parts.append(f"{stats_edges} dependencies found") - if stats_external: - parts.append(f"{stats_external} from other modules") - stats_text = Text(" · ".join(parts), style="dim") - - return Panel( - Group(progress, Text(""), tree, Text(""), stats_text), title="Building Call Graph", border_style="cyan" - ) + stats["indexed"] += 1 + stats["edges"] += result.num_edges + stats["external"] += result.cross_file_edges + + progress.advance(task_id) + + batch.clear() + refresh_display() + live.refresh() def update(result: IndexResult) -> None: - nonlocal stats_indexed, stats_cached, stats_edges, stats_external, stats_errors - results.append(result) - if result.error: - stats_errors += 1 - elif result.cached: - stats_cached += 1 - else: - stats_indexed += 1 - stats_edges += result.num_edges - stats_external += result.cross_file_edges - progress.advance(task_id) - live.update(make_display()) - - with Live(make_display(), console=console, transient=True, refresh_per_second=8) as live: + batch.append(result) + if len(batch) >= 8: + process_batch() + + with Live(panel, console=console, transient=True, auto_refresh=False) as live: yield update + if batch: + process_batch() def call_graph_summary(call_graph: CallGraph, file_to_funcs: dict[Path, list[FunctionToOptimize]]) -> None: from rich.panel import Panel total_functions = sum(len(funcs) for funcs in file_to_funcs.values()) - if total_functions == 0: + if not total_functions: return total_callees = 0 with_context = 0 - leaf_functions = 0 for file_path, funcs in file_to_funcs.items(): for func in funcs: _, func_callees = call_graph.get_callees({file_path: {func.qualified_name}}) - count = len(func_callees) - total_callees += count - if count > 0: + callee_count = len(func_callees) + total_callees += callee_count + if callee_count > 0: with_context += 1 - else: - leaf_functions += 1 - avg_callees = total_callees / total_functions if total_functions > 0 else 0 + leaf_functions = total_functions - with_context + avg_callees = total_callees / total_functions summary = ( f"{total_functions} functions ready for optimization · " diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index a8e9de02c..c0bb5539e 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -543,7 +543,7 @@ def run(self) -> None: logger.debug("Failed to initialize CallGraph, falling back to per-function Jedi analysis") if call_graph is not None and file_to_funcs_to_optimize: - source_files = list(file_to_funcs_to_optimize.keys()) + source_files = [f for f in file_to_funcs_to_optimize if f.suffix in (".py", ".pyw")] with call_graph_live_display(len(source_files)) as on_progress: call_graph.build_index(source_files, on_progress=on_progress) call_graph_summary(call_graph, file_to_funcs_to_optimize) From f3f0b0e020784e63e97cdd604aaec9dc1f324c7f Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Wed, 11 Feb 2026 20:48:34 -0500 Subject: [PATCH 10/42] refactor: move CallGraph into Python language support layer Add DependencyResolver protocol and IndexResult to base.py, move call_graph.py to languages/python/, and use factory method in optimizer instead of is_python() gating. --- codeflash/cli_cmds/console.py | 16 ++----- codeflash/context/code_context_extractor.py | 5 +-- codeflash/languages/__init__.py | 4 ++ codeflash/languages/base.py | 45 ++++++++++++++++++- codeflash/languages/javascript/support.py | 3 ++ codeflash/languages/python/__init__.py | 3 +- .../python}/call_graph.py | 12 +---- codeflash/languages/python/support.py | 10 +++++ codeflash/optimization/function_optimizer.py | 4 +- codeflash/optimization/optimizer.py | 42 ++++++++--------- tests/test_call_graph.py | 3 +- 11 files changed, 92 insertions(+), 55 deletions(-) rename codeflash/{context => languages/python}/call_graph.py (98%) diff --git a/codeflash/cli_cmds/console.py b/codeflash/cli_cmds/console.py index e9e921e91..317d5feae 100644 --- a/codeflash/cli_cmds/console.py +++ b/codeflash/cli_cmds/console.py @@ -30,8 +30,8 @@ from rich.progress import TaskID - from codeflash.context.call_graph import CallGraph, IndexResult from codeflash.discovery.functions_to_optimize import FunctionToOptimize + from codeflash.languages.base import DependencyResolver, IndexResult from codeflash.lsp.lsp_message import LspMessage DEBUG_MODE = logging.getLogger().getEffectiveLevel() == logging.DEBUG @@ -230,20 +230,12 @@ def call_graph_live_display(total: int) -> Generator[Callable[[IndexResult], Non task_id = progress.add_task("Analyzing files", total=total) results: deque[IndexResult] = deque(maxlen=MAX_TREE_ENTRIES) - stats = { - "indexed": 0, - "cached": 0, - "edges": 0, - "external": 0, - "errors": 0, - } + stats = {"indexed": 0, "cached": 0, "edges": 0, "external": 0, "errors": 0} tree = Tree("[bold]Dependencies[/bold]") stats_text = Text("0 dependencies found", style="dim") panel = Panel( - Group(progress, Text(""), tree, Text(""), stats_text), - title="Building Call Graph", - border_style="cyan", + Group(progress, Text(""), tree, Text(""), stats_text), title="Building Call Graph", border_style="cyan" ) def create_tree_node(result: IndexResult) -> Tree: @@ -317,7 +309,7 @@ def update(result: IndexResult) -> None: process_batch() -def call_graph_summary(call_graph: CallGraph, file_to_funcs: dict[Path, list[FunctionToOptimize]]) -> None: +def call_graph_summary(call_graph: DependencyResolver, file_to_funcs: dict[Path, list[FunctionToOptimize]]) -> None: from rich.panel import Panel total_functions = sum(len(funcs) for funcs in file_to_funcs.values()) diff --git a/codeflash/context/code_context_extractor.py b/codeflash/context/code_context_extractor.py index 4650de795..3764f4544 100644 --- a/codeflash/context/code_context_extractor.py +++ b/codeflash/context/code_context_extractor.py @@ -37,9 +37,8 @@ from jedi.api.classes import Name from libcst import CSTNode - from codeflash.context.call_graph import CallGraph from codeflash.context.unused_definition_remover import UsageInfo - from codeflash.languages.base import HelperFunction + from codeflash.languages.base import DependencyResolver, HelperFunction def build_testgen_context( @@ -79,7 +78,7 @@ def get_code_optimization_context( project_root_path: Path, optim_token_limit: int = OPTIMIZATION_CONTEXT_TOKEN_LIMIT, testgen_token_limit: int = TESTGEN_CONTEXT_TOKEN_LIMIT, - call_graph: CallGraph | None = None, + call_graph: DependencyResolver | None = None, ) -> CodeOptimizationContext: # Route to language-specific implementation for non-Python languages if not is_python(): diff --git a/codeflash/languages/__init__.py b/codeflash/languages/__init__.py index 47136f4e7..daf33b43c 100644 --- a/codeflash/languages/__init__.py +++ b/codeflash/languages/__init__.py @@ -19,7 +19,9 @@ from codeflash.languages.base import ( CodeContext, + DependencyResolver, HelperFunction, + IndexResult, Language, LanguageSupport, ParentInfo, @@ -82,8 +84,10 @@ def __getattr__(name: str): __all__ = [ "CodeContext", + "DependencyResolver", "FunctionInfo", "HelperFunction", + "IndexResult", "Language", "LanguageSupport", "ParentInfo", diff --git a/codeflash/languages/base.py b/codeflash/languages/base.py index 99cefdf46..76357a4d1 100644 --- a/codeflash/languages/base.py +++ b/codeflash/languages/base.py @@ -11,10 +11,11 @@ from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Callable, Iterable, Sequence from pathlib import Path from codeflash.discovery.functions_to_optimize import FunctionToOptimize + from codeflash.models.models import FunctionSource from codeflash.languages.language_enum import Language from codeflash.models.function_types import FunctionParent @@ -34,6 +35,16 @@ def __getattr__(name: str) -> Any: raise AttributeError(msg) +@dataclass(frozen=True, slots=True) +class IndexResult: + file_path: Path + cached: bool + num_edges: int + edges: tuple[tuple[str, str, bool], ...] # (caller_qn, callee_name, is_cross_file) + cross_file_edges: int + error: bool + + @dataclass class HelperFunction: """A helper function that is a dependency of the target function. @@ -192,6 +203,29 @@ class ReferenceInfo: caller_function: str | None = None +@runtime_checkable +class DependencyResolver(Protocol): + """Protocol for language-specific dependency resolution. + + Implementations analyze source files to discover call-graph edges + between functions so the optimizer can extract richer context. + """ + + def build_index(self, file_paths: Iterable[Path], on_progress: Callable[[IndexResult], None] | None = None) -> None: + """Pre-index a batch of files.""" + ... + + def get_callees( + self, file_path_to_qualified_names: dict[Path, set[str]] + ) -> tuple[dict[Path, set[FunctionSource]], list[FunctionSource]]: + """Return callees for the given functions.""" + ... + + def close(self) -> None: + """Release resources (e.g. database connections).""" + ... + + @runtime_checkable class LanguageSupport(Protocol): """Protocol defining what a language implementation must provide. @@ -565,6 +599,15 @@ def ensure_runtime_environment(self, project_root: Path) -> bool: # Default implementation: just copy runtime files return False + def create_dependency_resolver(self, project_root: Path) -> DependencyResolver | None: + """Create a language-specific dependency resolver, if available. + + Returns: + A DependencyResolver instance, or None if not supported. + + """ + return None + def instrument_existing_test( self, test_path: Path, diff --git a/codeflash/languages/javascript/support.py b/codeflash/languages/javascript/support.py index d32cce001..9d41b9cd4 100644 --- a/codeflash/languages/javascript/support.py +++ b/codeflash/languages/javascript/support.py @@ -1942,6 +1942,9 @@ def ensure_runtime_environment(self, project_root: Path) -> bool: logger.error("Could not install codeflash. Please run: npm install --save-dev codeflash") return False + def create_dependency_resolver(self, project_root: Path) -> None: + return None + def instrument_existing_test( self, test_path: Path, diff --git a/codeflash/languages/python/__init__.py b/codeflash/languages/python/__init__.py index e599d1431..50eb1a51e 100644 --- a/codeflash/languages/python/__init__.py +++ b/codeflash/languages/python/__init__.py @@ -5,6 +5,7 @@ to the LanguageSupport protocol. """ +from codeflash.languages.python.call_graph import CallGraph from codeflash.languages.python.support import PythonSupport -__all__ = ["PythonSupport"] +__all__ = ["CallGraph", "PythonSupport"] diff --git a/codeflash/context/call_graph.py b/codeflash/languages/python/call_graph.py similarity index 98% rename from codeflash/context/call_graph.py rename to codeflash/languages/python/call_graph.py index 033f8399a..16002c5e2 100644 --- a/codeflash/context/call_graph.py +++ b/codeflash/languages/python/call_graph.py @@ -4,12 +4,12 @@ import os import sqlite3 from collections import defaultdict -from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING from codeflash.cli_cmds.console import logger from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages +from codeflash.languages.base import IndexResult from codeflash.models.models import FunctionSource if TYPE_CHECKING: @@ -18,16 +18,6 @@ from jedi.api.classes import Name -@dataclass(frozen=True, slots=True) -class IndexResult: - file_path: Path - cached: bool - num_edges: int - edges: tuple[tuple[str, str, bool], ...] # (caller_qn, callee_name, is_cross_file) - cross_file_edges: int - error: bool - - # --------------------------------------------------------------------------- # Module-level helpers (must be top-level for ProcessPoolExecutor pickling) # --------------------------------------------------------------------------- diff --git a/codeflash/languages/python/support.py b/codeflash/languages/python/support.py index 58f66d0b8..03c376478 100644 --- a/codeflash/languages/python/support.py +++ b/codeflash/languages/python/support.py @@ -9,6 +9,7 @@ from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.languages.base import ( CodeContext, + DependencyResolver, FunctionFilterCriteria, HelperFunction, Language, @@ -801,6 +802,15 @@ def ensure_runtime_environment(self, project_root: Path) -> bool: """ return True + def create_dependency_resolver(self, project_root: Path) -> DependencyResolver | None: + from codeflash.languages.python.call_graph import CallGraph + + try: + return CallGraph(project_root) + except Exception: + logger.debug("Failed to initialize CallGraph, falling back to per-function Jedi analysis") + return None + def instrument_existing_test( self, test_path: Path, diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 024e77c89..f2e1887db 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -132,9 +132,9 @@ if TYPE_CHECKING: from argparse import Namespace - from codeflash.context.call_graph import CallGraph from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.either import Result + from codeflash.languages.base import DependencyResolver from codeflash.models.models import ( BenchmarkKey, CodeStringsMarkdown, @@ -438,7 +438,7 @@ def __init__( total_benchmark_timings: dict[BenchmarkKey, int] | None = None, args: Namespace | None = None, replay_tests_dir: Path | None = None, - call_graph: CallGraph | None = None, + call_graph: DependencyResolver | None = None, ) -> None: self.project_root = test_cfg.project_root_path self.test_cfg = test_cfg diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index c0bb5539e..ff068f5c1 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -24,7 +24,7 @@ ) from codeflash.code_utils.time_utils import humanize_runtime from codeflash.either import is_successful -from codeflash.languages import is_javascript, set_current_language +from codeflash.languages import current_language_support, is_javascript, set_current_language from codeflash.models.models import ValidCode from codeflash.telemetry.posthog_cf import ph from codeflash.verification.verification_utils import TestConfig @@ -34,8 +34,8 @@ from codeflash.benchmarking.function_ranker import FunctionRanker from codeflash.code_utils.checkpoint import CodeflashRunCheckpoint - from codeflash.context.call_graph import CallGraph from codeflash.discovery.functions_to_optimize import FunctionToOptimize + from codeflash.languages.base import DependencyResolver from codeflash.models.models import BenchmarkKey, FunctionCalledInTest from codeflash.optimization.function_optimizer import FunctionOptimizer @@ -206,7 +206,7 @@ def create_function_optimizer( total_benchmark_timings: dict[BenchmarkKey, float] | None = None, original_module_ast: ast.Module | None = None, original_module_path: Path | None = None, - call_graph: CallGraph | None = None, + call_graph: DependencyResolver | None = None, ) -> FunctionOptimizer | None: from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast from codeflash.optimization.function_optimizer import FunctionOptimizer @@ -392,7 +392,7 @@ def rank_all_functions_globally( self, file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]], trace_file_path: Path | None, - call_graph: CallGraph | None = None, + call_graph: DependencyResolver | None = None, ) -> list[tuple[Path, FunctionToOptimize]]: """Rank all functions globally across all files based on trace data. @@ -467,7 +467,7 @@ def rank_all_functions_globally( return globally_ranked def rank_by_dependency_count( - self, all_functions: list[tuple[Path, FunctionToOptimize]], call_graph: CallGraph + self, all_functions: list[tuple[Path, FunctionToOptimize]], call_graph: DependencyResolver ) -> list[tuple[Path, FunctionToOptimize]]: counts: list[tuple[int, int, tuple[Path, FunctionToOptimize]]] = [] for idx, (file_path, func) in enumerate(all_functions): @@ -530,23 +530,17 @@ def run(self) -> None: file_to_funcs_to_optimize, num_optimizable_functions ) - # Create persistent call graph for Python runs to cache Jedi analysis across functions - call_graph: CallGraph | None = None - from codeflash.languages import is_python + # Create a language-specific dependency resolver (e.g. Jedi-based call graph for Python) + resolver: DependencyResolver | None = None + lang_support = current_language_support() + if hasattr(lang_support, "create_dependency_resolver"): + resolver = lang_support.create_dependency_resolver(self.args.project_root) - if is_python(): - from codeflash.context.call_graph import CallGraph - - try: - call_graph = CallGraph(self.args.project_root) - except Exception: - logger.debug("Failed to initialize CallGraph, falling back to per-function Jedi analysis") - - if call_graph is not None and file_to_funcs_to_optimize: - source_files = [f for f in file_to_funcs_to_optimize if f.suffix in (".py", ".pyw")] + if resolver is not None and file_to_funcs_to_optimize: + source_files = list(file_to_funcs_to_optimize.keys()) with call_graph_live_display(len(source_files)) as on_progress: - call_graph.build_index(source_files, on_progress=on_progress) - call_graph_summary(call_graph, file_to_funcs_to_optimize) + resolver.build_index(source_files, on_progress=on_progress) + call_graph_summary(resolver, file_to_funcs_to_optimize) optimizations_found: int = 0 self.test_cfg.concolic_test_root_dir = Path( @@ -564,7 +558,7 @@ def run(self) -> None: # GLOBAL RANKING: Rank all functions together before optimizing globally_ranked_functions = self.rank_all_functions_globally( - file_to_funcs_to_optimize, trace_file_path, call_graph=call_graph + file_to_funcs_to_optimize, trace_file_path, call_graph=resolver ) # Cache for module preparation (avoid re-parsing same files) prepared_modules: dict[Path, tuple[dict[Path, ValidCode], ast.Module | None]] = {} @@ -597,7 +591,7 @@ def run(self) -> None: total_benchmark_timings=total_benchmark_timings, original_module_ast=original_module_ast, original_module_path=original_module_path, - call_graph=call_graph, + call_graph=resolver, ) if function_optimizer is None: continue @@ -656,8 +650,8 @@ def run(self) -> None: else: logger.warning("⚠️ Failed to send completion email. Status") finally: - if call_graph is not None: - call_graph.close() + if resolver is not None: + resolver.close() if function_optimizer: function_optimizer.cleanup_generated_files() diff --git a/tests/test_call_graph.py b/tests/test_call_graph.py index a2918434a..22c306055 100644 --- a/tests/test_call_graph.py +++ b/tests/test_call_graph.py @@ -7,7 +7,8 @@ if TYPE_CHECKING: from pathlib import Path -from codeflash.context.call_graph import CallGraph, IndexResult +from codeflash.languages.base import IndexResult +from codeflash.languages.python.call_graph import CallGraph @pytest.fixture From ed31cd11eea78f485cda3a59e7a8e0a495036b7a Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Wed, 11 Feb 2026 21:03:31 -0500 Subject: [PATCH 11/42] refactor: filter files by language extensions and show project-relative paths in call graph Display file paths relative to project root in the call graph live display for easier navigation. Filter indexed files by the language support's file extensions to avoid processing irrelevant file types. --- codeflash/cli_cmds/console.py | 12 ++++++++++-- codeflash/languages/python/call_graph.py | 1 - codeflash/languages/python/support.py | 3 ++- codeflash/optimization/optimizer.py | 9 ++++----- 4 files changed, 16 insertions(+), 9 deletions(-) diff --git a/codeflash/cli_cmds/console.py b/codeflash/cli_cmds/console.py index 317d5feae..3e1aa64d7 100644 --- a/codeflash/cli_cmds/console.py +++ b/codeflash/cli_cmds/console.py @@ -206,7 +206,9 @@ def test_files_progress_bar(total: int, description: str) -> Generator[tuple[Pro @contextmanager -def call_graph_live_display(total: int) -> Generator[Callable[[IndexResult], None], None, None]: +def call_graph_live_display( + total: int, project_root: Path | None = None +) -> Generator[Callable[[IndexResult], None], None, None]: from rich.console import Group from rich.live import Live from rich.panel import Panel @@ -239,7 +241,13 @@ def call_graph_live_display(total: int) -> Generator[Callable[[IndexResult], Non ) def create_tree_node(result: IndexResult) -> Tree: - name = f"{result.file_path.parent.name}/{result.file_path.name}" + if project_root: + try: + name = str(result.file_path.resolve().relative_to(project_root.resolve())) + except ValueError: + name = f"{result.file_path.parent.name}/{result.file_path.name}" + else: + name = f"{result.file_path.parent.name}/{result.file_path.name}" if result.error: return Tree(f"[red]{name} (error)[/red]") diff --git a/codeflash/languages/python/call_graph.py b/codeflash/languages/python/call_graph.py index 16002c5e2..bc6dadb8a 100644 --- a/codeflash/languages/python/call_graph.py +++ b/codeflash/languages/python/call_graph.py @@ -131,7 +131,6 @@ def _analyze_file(file_path: Path, jedi_project: object, project_root_str: str) if not _is_valid_definition(definition, caller_qn, project_root_str): continue - # Extract common edge components edge_base = (resolved, caller_qn, str(definition_path)) if definition.type == "function": diff --git a/codeflash/languages/python/support.py b/codeflash/languages/python/support.py index 03c376478..429f9e1f9 100644 --- a/codeflash/languages/python/support.py +++ b/codeflash/languages/python/support.py @@ -9,7 +9,6 @@ from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.languages.base import ( CodeContext, - DependencyResolver, FunctionFilterCriteria, HelperFunction, Language, @@ -22,6 +21,8 @@ if TYPE_CHECKING: from collections.abc import Sequence + from codeflash.languages.base import DependencyResolver + logger = logging.getLogger(__name__) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index ff068f5c1..380a1d8be 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -531,14 +531,13 @@ def run(self) -> None: ) # Create a language-specific dependency resolver (e.g. Jedi-based call graph for Python) - resolver: DependencyResolver | None = None lang_support = current_language_support() - if hasattr(lang_support, "create_dependency_resolver"): - resolver = lang_support.create_dependency_resolver(self.args.project_root) + resolver = lang_support.create_dependency_resolver(self.args.project_root) if lang_support else None if resolver is not None and file_to_funcs_to_optimize: - source_files = list(file_to_funcs_to_optimize.keys()) - with call_graph_live_display(len(source_files)) as on_progress: + supported_exts = lang_support.file_extensions + source_files = [f for f in file_to_funcs_to_optimize if f.suffix in supported_exts] + with call_graph_live_display(len(source_files), project_root=self.args.project_root) as on_progress: resolver.build_index(source_files, on_progress=on_progress) call_graph_summary(resolver, file_to_funcs_to_optimize) From 5341ac8fdeab87b85c37f6535d659391a985f855 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Wed, 11 Feb 2026 21:13:32 -0500 Subject: [PATCH 12/42] fix: improve CLI output formatting for runtime estimate and call graph sections Split the runtime estimate and PR message into separate log lines to avoid awkward line wrapping. Add console rules between sections for clearer visual separation. --- codeflash/optimization/optimizer.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 380a1d8be..d58aad00a 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -519,12 +519,12 @@ def run(self) -> None: if self.args.all: three_min_in_ns = int(1.8e11) console.rule() - pr_message = ( - "\nCodeflash will keep opening pull requests as it finds optimizations." if not self.args.no_pr else "" - ) logger.info( - f"It might take about {humanize_runtime(num_optimizable_functions * three_min_in_ns)} to fully optimize this project.{pr_message}" + f"It might take about {humanize_runtime(num_optimizable_functions * three_min_in_ns)} to fully optimize this project." ) + if not self.args.no_pr: + logger.info("Codeflash will keep opening pull requests as it finds optimizations.") + console.rule() function_benchmark_timings, total_benchmark_timings = self.run_benchmarks( file_to_funcs_to_optimize, num_optimizable_functions @@ -539,6 +539,7 @@ def run(self) -> None: source_files = [f for f in file_to_funcs_to_optimize if f.suffix in supported_exts] with call_graph_live_display(len(source_files), project_root=self.args.project_root) as on_progress: resolver.build_index(source_files, on_progress=on_progress) + console.rule() call_graph_summary(resolver, file_to_funcs_to_optimize) optimizations_found: int = 0 From ee0da841eaa64a8cef0ffbfdf606ff5e995ad1ea Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Wed, 11 Feb 2026 21:50:36 -0500 Subject: [PATCH 13/42] refactor: simplify call graph DB schema to two flat human-readable tables MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the normalized relational hierarchy (cg_projects → cg_languages → cg_indexed_files/cg_call_edges) with two self-describing tables (indexed_files, call_edges) where every row includes project_root and language as text columns. --- codeflash/languages/python/call_graph.py | 94 +++++++++++++----------- codeflash/languages/python/support.py | 2 +- 2 files changed, 51 insertions(+), 45 deletions(-) diff --git a/codeflash/languages/python/call_graph.py b/codeflash/languages/python/call_graph.py index bc6dadb8a..596c8db4f 100644 --- a/codeflash/languages/python/call_graph.py +++ b/codeflash/languages/python/call_graph.py @@ -178,13 +178,14 @@ def _index_file_worker(args: tuple[str, str]) -> tuple[str, str, set[tuple[str, class CallGraph: - SCHEMA_VERSION = 1 + SCHEMA_VERSION = 2 - def __init__(self, project_root: Path, db_path: Path | None = None) -> None: + def __init__(self, project_root: Path, language: str = "python", db_path: Path | None = None) -> None: import jedi self.project_root = project_root.resolve() self.project_root_str = str(self.project_root) + self.language = language self.jedi_project = jedi.Project(path=self.project_root) if db_path is None: @@ -200,27 +201,35 @@ def __init__(self, project_root: Path, db_path: Path | None = None) -> None: def _init_schema(self) -> None: cur = self.conn.cursor() cur.execute("CREATE TABLE IF NOT EXISTS cg_schema_version (version INTEGER PRIMARY KEY)") + row = cur.execute("SELECT version FROM cg_schema_version LIMIT 1").fetchone() if row is None: cur.execute("INSERT INTO cg_schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,)) elif row[0] != self.SCHEMA_VERSION: - # Schema mismatch — drop all cg_ tables and recreate - cur.execute("DROP TABLE IF EXISTS cg_call_edges") - cur.execute("DROP TABLE IF EXISTS cg_indexed_files") + for table in [ + "cg_call_edges", "cg_indexed_files", "cg_languages", "cg_projects", "cg_project_meta", + "indexed_files", "call_edges", + ]: + cur.execute(f"DROP TABLE IF EXISTS {table}") cur.execute("DELETE FROM cg_schema_version") cur.execute("INSERT INTO cg_schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,)) cur.execute( - """CREATE TABLE IF NOT EXISTS cg_indexed_files ( + """ + CREATE TABLE IF NOT EXISTS indexed_files ( project_root TEXT NOT NULL, + language TEXT NOT NULL, file_path TEXT NOT NULL, file_hash TEXT NOT NULL, - PRIMARY KEY (project_root, file_path) - )""" + PRIMARY KEY (project_root, language, file_path) + ) + """ ) cur.execute( - """CREATE TABLE IF NOT EXISTS cg_call_edges ( + """ + CREATE TABLE IF NOT EXISTS call_edges ( project_root TEXT NOT NULL, + language TEXT NOT NULL, caller_file TEXT NOT NULL, caller_qualified_name TEXT NOT NULL, callee_file TEXT NOT NULL, @@ -229,13 +238,16 @@ def _init_schema(self) -> None: callee_only_function_name TEXT NOT NULL, callee_definition_type TEXT NOT NULL, callee_source_line TEXT NOT NULL, - PRIMARY KEY (project_root, caller_file, caller_qualified_name, + PRIMARY KEY (project_root, language, caller_file, caller_qualified_name, callee_file, callee_qualified_name) - )""" + ) + """ ) cur.execute( - """CREATE INDEX IF NOT EXISTS idx_cg_edges_caller - ON cg_call_edges (project_root, caller_file, caller_qualified_name)""" + """ + CREATE INDEX IF NOT EXISTS idx_call_edges_caller + ON call_edges (project_root, language, caller_file, caller_qualified_name) + """ ) self.conn.commit() @@ -259,11 +271,13 @@ def get_callees( cur = self.conn.cursor() for caller_file, caller_qn in all_caller_keys: rows = cur.execute( - """SELECT callee_file, callee_qualified_name, callee_fully_qualified_name, - callee_only_function_name, callee_definition_type, callee_source_line - FROM cg_call_edges - WHERE project_root = ? AND caller_file = ? AND caller_qualified_name = ?""", - (self.project_root_str, caller_file, caller_qn), + """ + SELECT callee_file, callee_qualified_name, callee_fully_qualified_name, + callee_only_function_name, callee_definition_type, callee_source_line + FROM call_edges + WHERE project_root = ? AND language = ? AND caller_file = ? AND caller_qualified_name = ? + """, + (self.project_root_str, self.language, caller_file, caller_qn), ).fetchall() for callee_file, callee_qn, callee_fqn, callee_name, callee_type, callee_src in rows: @@ -291,18 +305,7 @@ def ensure_file_indexed(self, file_path: Path) -> IndexResult: file_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() - # Check in-memory cache first - if self.indexed_file_hashes.get(resolved) == file_hash: - return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), cross_file_edges=0, error=False) - - # Check DB for stored hash - row = self.conn.execute( - "SELECT file_hash FROM cg_indexed_files WHERE project_root = ? AND file_path = ?", - (self.project_root_str, resolved), - ).fetchone() - - if row and row[0] == file_hash: - self.indexed_file_hashes[resolved] = file_hash + if self._is_file_cached(resolved, file_hash): return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), cross_file_edges=0, error=False) return self.index_file(file_path, file_hash) @@ -318,30 +321,35 @@ def _persist_edges( self, file_path: Path, resolved: str, file_hash: str, edges: set[tuple[str, ...]], had_error: bool ) -> IndexResult: cur = self.conn.cursor() + scope = (self.project_root_str, self.language) # Clear existing data for this file cur.execute( - "DELETE FROM cg_call_edges WHERE project_root = ? AND caller_file = ?", (self.project_root_str, resolved) + "DELETE FROM call_edges WHERE project_root = ? AND language = ? AND caller_file = ?", + (*scope, resolved), ) cur.execute( - "DELETE FROM cg_indexed_files WHERE project_root = ? AND file_path = ?", (self.project_root_str, resolved) + "DELETE FROM indexed_files WHERE project_root = ? AND language = ? AND file_path = ?", + (*scope, resolved), ) # Insert new edges if parsing succeeded if not had_error and edges: cur.executemany( - """INSERT OR REPLACE INTO cg_call_edges - (project_root, caller_file, caller_qualified_name, - callee_file, callee_qualified_name, callee_fully_qualified_name, - callee_only_function_name, callee_definition_type, callee_source_line) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""", - [(self.project_root_str, *edge) for edge in edges], + """ + INSERT OR REPLACE INTO call_edges + (project_root, language, caller_file, caller_qualified_name, + callee_file, callee_qualified_name, callee_fully_qualified_name, + callee_only_function_name, callee_definition_type, callee_source_line) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + [(*scope, *edge) for edge in edges], ) # Record that this file has been indexed cur.execute( - "INSERT OR REPLACE INTO cg_indexed_files (project_root, file_path, file_hash) VALUES (?, ?, ?)", - (self.project_root_str, resolved, file_hash), + "INSERT OR REPLACE INTO indexed_files (project_root, language, file_path, file_hash) VALUES (?, ?, ?, ?)", + (*scope, resolved, file_hash), ) self.conn.commit() @@ -408,14 +416,12 @@ def build_index(self, file_paths: Iterable[Path], on_progress: Callable[[IndexRe def _is_file_cached(self, resolved: str, file_hash: str) -> bool: """Check if file is cached in memory or DB.""" - # Check in-memory cache if self.indexed_file_hashes.get(resolved) == file_hash: return True - # Check DB cache row = self.conn.execute( - "SELECT file_hash FROM cg_indexed_files WHERE project_root = ? AND file_path = ?", - (self.project_root_str, resolved), + "SELECT file_hash FROM indexed_files WHERE project_root = ? AND language = ? AND file_path = ?", + (self.project_root_str, self.language, resolved), ).fetchone() if row and row[0] == file_hash: diff --git a/codeflash/languages/python/support.py b/codeflash/languages/python/support.py index 429f9e1f9..77a37759d 100644 --- a/codeflash/languages/python/support.py +++ b/codeflash/languages/python/support.py @@ -807,7 +807,7 @@ def create_dependency_resolver(self, project_root: Path) -> DependencyResolver | from codeflash.languages.python.call_graph import CallGraph try: - return CallGraph(project_root) + return CallGraph(project_root, language=self.language.value) except Exception: logger.debug("Failed to initialize CallGraph, falling back to per-function Jedi analysis") return None From 0a5e8141b46f4f4bb1d252c14b44d48566674f8c Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Wed, 11 Feb 2026 22:55:07 -0500 Subject: [PATCH 14/42] Update config_consts.py --- codeflash/code_utils/config_consts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codeflash/code_utils/config_consts.py b/codeflash/code_utils/config_consts.py index e344fad8a..32e5c7c54 100644 --- a/codeflash/code_utils/config_consts.py +++ b/codeflash/code_utils/config_consts.py @@ -4,8 +4,8 @@ from typing import Any, Union MAX_TEST_RUN_ITERATIONS = 5 -OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 16000 -TESTGEN_CONTEXT_TOKEN_LIMIT = 16000 +OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 100000 +TESTGEN_CONTEXT_TOKEN_LIMIT = 100000 INDIVIDUAL_TESTCASE_TIMEOUT = 15 MAX_FUNCTION_TEST_SECONDS = 60 MIN_IMPROVEMENT_THRESHOLD = 0.05 From 54aa7e1eabd37eb902961454beddf90d9fc30c29 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 12 Feb 2026 00:40:14 -0500 Subject: [PATCH 15/42] fix: skip call graph building in CI and fix ruff formatting Skip dependency resolver creation in CI environments where the cache DB doesn't persist between runs. Also apply ruff formatting to call_graph.py. --- codeflash/languages/python/call_graph.py | 15 +++++++++------ codeflash/optimization/optimizer.py | 5 ++++- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/codeflash/languages/python/call_graph.py b/codeflash/languages/python/call_graph.py index 596c8db4f..7d331d252 100644 --- a/codeflash/languages/python/call_graph.py +++ b/codeflash/languages/python/call_graph.py @@ -207,8 +207,13 @@ def _init_schema(self) -> None: cur.execute("INSERT INTO cg_schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,)) elif row[0] != self.SCHEMA_VERSION: for table in [ - "cg_call_edges", "cg_indexed_files", "cg_languages", "cg_projects", "cg_project_meta", - "indexed_files", "call_edges", + "cg_call_edges", + "cg_indexed_files", + "cg_languages", + "cg_projects", + "cg_project_meta", + "indexed_files", + "call_edges", ]: cur.execute(f"DROP TABLE IF EXISTS {table}") cur.execute("DELETE FROM cg_schema_version") @@ -325,12 +330,10 @@ def _persist_edges( # Clear existing data for this file cur.execute( - "DELETE FROM call_edges WHERE project_root = ? AND language = ? AND caller_file = ?", - (*scope, resolved), + "DELETE FROM call_edges WHERE project_root = ? AND language = ? AND caller_file = ?", (*scope, resolved) ) cur.execute( - "DELETE FROM indexed_files WHERE project_root = ? AND language = ? AND file_path = ?", - (*scope, resolved), + "DELETE FROM indexed_files WHERE project_root = ? AND language = ? AND file_path = ?", (*scope, resolved) ) # Insert new edges if parsing succeeded diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index d58aad00a..810c1758b 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -531,8 +531,11 @@ def run(self) -> None: ) # Create a language-specific dependency resolver (e.g. Jedi-based call graph for Python) + # Skip in CI — the cache DB doesn't persist between runs on ephemeral runners lang_support = current_language_support() - resolver = lang_support.create_dependency_resolver(self.args.project_root) if lang_support else None + resolver = None + if lang_support and not env_utils.is_ci(): + resolver = lang_support.create_dependency_resolver(self.args.project_root) if resolver is not None and file_to_funcs_to_optimize: supported_exts = lang_support.file_extensions From c096c82b50204c849db2cab5178b6fb1523270bc Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 12 Feb 2026 00:43:11 -0500 Subject: [PATCH 16/42] refactor: simplify compat.py by removing unnecessary class wrapper --- codeflash/code_utils/compat.py | 44 ++++++---------------------------- 1 file changed, 7 insertions(+), 37 deletions(-) diff --git a/codeflash/code_utils/compat.py b/codeflash/code_utils/compat.py index eb4e5b561..b73a6a5a7 100644 --- a/codeflash/code_utils/compat.py +++ b/codeflash/code_utils/compat.py @@ -2,46 +2,16 @@ import sys import tempfile from pathlib import Path -from typing import TYPE_CHECKING from platformdirs import user_config_dir -if TYPE_CHECKING: - codeflash_temp_dir: Path - codeflash_cache_dir: Path - codeflash_cache_db: Path +LF: str = os.linesep +IS_POSIX: bool = os.name != "nt" +SAFE_SYS_EXECUTABLE: str = Path(sys.executable).as_posix() +codeflash_cache_dir: Path = Path(user_config_dir(appname="codeflash", appauthor="codeflash-ai", ensure_exists=True)) -class Compat: - # os-independent newline - LF: str = os.linesep +codeflash_temp_dir: Path = Path(tempfile.gettempdir()) / "codeflash" +codeflash_temp_dir.mkdir(parents=True, exist_ok=True) - SAFE_SYS_EXECUTABLE: str = Path(sys.executable).as_posix() - - IS_POSIX: bool = os.name != "nt" - - @property - def codeflash_cache_dir(self) -> Path: - return Path(user_config_dir(appname="codeflash", appauthor="codeflash-ai", ensure_exists=True)) - - @property - def codeflash_temp_dir(self) -> Path: - temp_dir = Path(tempfile.gettempdir()) / "codeflash" - if not temp_dir.exists(): - temp_dir.mkdir(parents=True, exist_ok=True) - return temp_dir - - @property - def codeflash_cache_db(self) -> Path: - return self.codeflash_cache_dir / "codeflash_cache.db" - - -_compat = Compat() - - -codeflash_temp_dir = _compat.codeflash_temp_dir -codeflash_cache_dir = _compat.codeflash_cache_dir -codeflash_cache_db = _compat.codeflash_cache_db -LF = _compat.LF -SAFE_SYS_EXECUTABLE = _compat.SAFE_SYS_EXECUTABLE -IS_POSIX = _compat.IS_POSIX +codeflash_cache_db: Path = codeflash_cache_dir / "codeflash_cache.db" From 8555da065a7906a3c9ea2b40adc901dc9dc913cc Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Thu, 12 Feb 2026 05:48:27 +0000 Subject: [PATCH 17/42] fix: resolve mypy type errors in call_graph.py --- codeflash/languages/python/call_graph.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/codeflash/languages/python/call_graph.py b/codeflash/languages/python/call_graph.py index 7d331d252..bf4483d90 100644 --- a/codeflash/languages/python/call_graph.py +++ b/codeflash/languages/python/call_graph.py @@ -47,7 +47,8 @@ def _resolve_definitions(ref: Name) -> list[Name]: pass try: - return ref.goto(follow_imports=True, follow_builtin_imports=False) + result: list[Name] = ref.goto(follow_imports=True, follow_builtin_imports=False) + return result except Exception: return [] @@ -111,7 +112,7 @@ def _analyze_file(file_path: Path, jedi_project: object, project_root_str: str) except Exception: return set(), True - edges: set[tuple[str, str, str, str, str, str, str, str]] = set() + edges: set[tuple[str, ...]] = set() for ref in refs: try: @@ -170,6 +171,7 @@ def _analyze_file(file_path: Path, jedi_project: object, project_root_str: str) def _index_file_worker(args: tuple[str, str]) -> tuple[str, str, set[tuple[str, ...]], bool]: """Worker entry point for ProcessPoolExecutor.""" file_path_str, file_hash = args + assert _worker_project_root_str is not None edges, had_error = _analyze_file(Path(file_path_str), _worker_jedi_project, _worker_project_root_str) return file_path_str, file_hash, edges, had_error From 513e590b8b13562d40f41b037b4914d7d3eea0ba Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 12 Feb 2026 01:00:53 -0500 Subject: [PATCH 18/42] fix: revert token limits back to 16K from unintended 100K increase --- codeflash/code_utils/config_consts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codeflash/code_utils/config_consts.py b/codeflash/code_utils/config_consts.py index 32e5c7c54..e344fad8a 100644 --- a/codeflash/code_utils/config_consts.py +++ b/codeflash/code_utils/config_consts.py @@ -4,8 +4,8 @@ from typing import Any, Union MAX_TEST_RUN_ITERATIONS = 5 -OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 100000 -TESTGEN_CONTEXT_TOKEN_LIMIT = 100000 +OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 16000 +TESTGEN_CONTEXT_TOKEN_LIMIT = 16000 INDIVIDUAL_TESTCASE_TIMEOUT = 15 MAX_FUNCTION_TEST_SECONDS = 60 MIN_IMPROVEMENT_THRESHOLD = 0.05 From be4a2ca09ef56b56029bb7bb475c4b08c829c2d9 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 12 Feb 2026 01:01:45 -0500 Subject: [PATCH 19/42] feat: increase optimization and testgen token limits to 64K --- codeflash/code_utils/config_consts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codeflash/code_utils/config_consts.py b/codeflash/code_utils/config_consts.py index e344fad8a..7fd8814d6 100644 --- a/codeflash/code_utils/config_consts.py +++ b/codeflash/code_utils/config_consts.py @@ -4,8 +4,8 @@ from typing import Any, Union MAX_TEST_RUN_ITERATIONS = 5 -OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 16000 -TESTGEN_CONTEXT_TOKEN_LIMIT = 16000 +OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 64000 +TESTGEN_CONTEXT_TOKEN_LIMIT = 64000 INDIVIDUAL_TESTCASE_TIMEOUT = 15 MAX_FUNCTION_TEST_SECONDS = 60 MIN_IMPROVEMENT_THRESHOLD = 0.05 From 9e904483d85b19b5a544a53bfb43f28da372461a Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 12 Feb 2026 01:02:21 -0500 Subject: [PATCH 20/42] fix: use explicit token limits in tests to decouple from global constant --- tests/test_code_context_extractor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_code_context_extractor.py b/tests/test_code_context_extractor.py index c5009b898..f40ab1c8f 100644 --- a/tests/test_code_context_extractor.py +++ b/tests/test_code_context_extractor.py @@ -813,7 +813,7 @@ def helper_method(self): ending_line=None, ) - code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root) + code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root, optim_token_limit=16000) read_write_context, read_only_context = code_ctx.read_writable_code, code_ctx.read_only_context_code hashing_context = code_ctx.hashing_code_context # In this scenario, the read-only code context is too long, so the read-only docstrings are removed. @@ -1006,7 +1006,7 @@ def helper_method(self): ) # In this scenario, the read-writable code is too long, so we abort. with pytest.raises(ValueError, match="Read-writable code has exceeded token limit, cannot proceed"): - code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root) + code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root, optim_token_limit=16000) def test_example_class_token_limit_4(tmp_path: Path) -> None: @@ -1059,7 +1059,7 @@ def helper_method(self): # In this scenario, the read-writable code context becomes too large because the __init__ function is referencing the global x variable instead of the class attribute self.x, so we abort. with pytest.raises(ValueError, match="Read-writable code has exceeded token limit, cannot proceed"): - code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root) + code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root, optim_token_limit=16000) def test_example_class_token_limit_5(tmp_path: Path) -> None: From fc42548f9f49870864e54effd7395826206bce25 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 12 Feb 2026 01:03:34 -0500 Subject: [PATCH 21/42] test: update token limit tests for 64K default --- tests/test_code_context_extractor.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_code_context_extractor.py b/tests/test_code_context_extractor.py index f40ab1c8f..6f8cd77cd 100644 --- a/tests/test_code_context_extractor.py +++ b/tests/test_code_context_extractor.py @@ -768,7 +768,7 @@ def helper_method(self): def test_example_class_token_limit_1(tmp_path: Path) -> None: docstring_filler = " ".join( - ["This is a long docstring that will be used to fill up the token limit." for _ in range(1000)] + ["This is a long docstring that will be used to fill up the token limit." for _ in range(4000)] ) code = f""" class MyClass: @@ -813,7 +813,7 @@ def helper_method(self): ending_line=None, ) - code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root, optim_token_limit=16000) + code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root) read_write_context, read_only_context = code_ctx.read_writable_code, code_ctx.read_only_context_code hashing_context = code_ctx.hashing_code_context # In this scenario, the read-only code context is too long, so the read-only docstrings are removed. @@ -961,7 +961,7 @@ def helper_method(self): def test_example_class_token_limit_3(tmp_path: Path) -> None: string_filler = " ".join( - ["This is a long string that will be used to fill up the token limit." for _ in range(1000)] + ["This is a long string that will be used to fill up the token limit." for _ in range(4000)] ) code = f""" class MyClass: @@ -1006,12 +1006,12 @@ def helper_method(self): ) # In this scenario, the read-writable code is too long, so we abort. with pytest.raises(ValueError, match="Read-writable code has exceeded token limit, cannot proceed"): - code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root, optim_token_limit=16000) + code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root) def test_example_class_token_limit_4(tmp_path: Path) -> None: string_filler = " ".join( - ["This is a long string that will be used to fill up the token limit." for _ in range(1000)] + ["This is a long string that will be used to fill up the token limit." for _ in range(4000)] ) code = f""" class MyClass: @@ -1059,7 +1059,7 @@ def helper_method(self): # In this scenario, the read-writable code context becomes too large because the __init__ function is referencing the global x variable instead of the class attribute self.x, so we abort. with pytest.raises(ValueError, match="Read-writable code has exceeded token limit, cannot proceed"): - code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root, optim_token_limit=16000) + code_ctx = get_code_optimization_context(function_to_optimize, opt.args.project_root) def test_example_class_token_limit_5(tmp_path: Path) -> None: From 80759c905e085fddb98e06c21165baede218ec55 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 12 Feb 2026 01:05:49 -0500 Subject: [PATCH 22/42] fix: add None guard for lang_support before accessing file_extensions --- codeflash/optimization/optimizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 810c1758b..b22f0df9c 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -537,7 +537,7 @@ def run(self) -> None: if lang_support and not env_utils.is_ci(): resolver = lang_support.create_dependency_resolver(self.args.project_root) - if resolver is not None and file_to_funcs_to_optimize: + if resolver is not None and lang_support is not None and file_to_funcs_to_optimize: supported_exts = lang_support.file_extensions source_files = [f for f in file_to_funcs_to_optimize if f.suffix in supported_exts] with call_graph_live_display(len(source_files), project_root=self.args.project_root) as on_progress: From d7edef82b7d91c9eef8d88bbe8bc6aff6245a552 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 12 Feb 2026 01:07:41 -0500 Subject: [PATCH 23/42] refactor: batch callee counting in call_graph_summary with new count_callees_per_function method --- codeflash/cli_cmds/console.py | 14 ++++++------- codeflash/languages/base.py | 6 ++++++ codeflash/languages/python/call_graph.py | 26 ++++++++++++++++++++++++ 3 files changed, 39 insertions(+), 7 deletions(-) diff --git a/codeflash/cli_cmds/console.py b/codeflash/cli_cmds/console.py index 3e1aa64d7..ddfa04771 100644 --- a/codeflash/cli_cmds/console.py +++ b/codeflash/cli_cmds/console.py @@ -327,13 +327,13 @@ def call_graph_summary(call_graph: DependencyResolver, file_to_funcs: dict[Path, total_callees = 0 with_context = 0 - for file_path, funcs in file_to_funcs.items(): - for func in funcs: - _, func_callees = call_graph.get_callees({file_path: {func.qualified_name}}) - callee_count = len(func_callees) - total_callees += callee_count - if callee_count > 0: - with_context += 1 + callee_counts = call_graph.count_callees_per_function( + {file_path: {func.qualified_name for func in funcs} for file_path, funcs in file_to_funcs.items()} + ) + for count in callee_counts.values(): + total_callees += count + if count > 0: + with_context += 1 leaf_functions = total_functions - with_context avg_callees = total_callees / total_functions diff --git a/codeflash/languages/base.py b/codeflash/languages/base.py index 76357a4d1..1e2254b1d 100644 --- a/codeflash/languages/base.py +++ b/codeflash/languages/base.py @@ -221,6 +221,12 @@ def get_callees( """Return callees for the given functions.""" ... + def count_callees_per_function( + self, file_path_to_qualified_names: dict[Path, set[str]] + ) -> dict[str, int]: + """Return the number of callees for each caller qualified name.""" + ... + def close(self) -> None: """Release resources (e.g. database connections).""" ... diff --git a/codeflash/languages/python/call_graph.py b/codeflash/languages/python/call_graph.py index bf4483d90..aea6d696a 100644 --- a/codeflash/languages/python/call_graph.py +++ b/codeflash/languages/python/call_graph.py @@ -302,6 +302,32 @@ def get_callees( return file_path_to_function_source, function_source_list + def count_callees_per_function( + self, file_path_to_qualified_names: dict[Path, set[str]] + ) -> dict[str, int]: + all_caller_keys: list[tuple[str, str]] = [] + for file_path, qualified_names in file_path_to_qualified_names.items(): + self.ensure_file_indexed(file_path) + resolved = str(file_path.resolve()) + all_caller_keys.extend((resolved, qn) for qn in qualified_names) + + if not all_caller_keys: + return {} + + counts: dict[str, int] = {} + cur = self.conn.cursor() + for caller_file, caller_qn in all_caller_keys: + row = cur.execute( + """ + SELECT COUNT(*) FROM call_edges + WHERE project_root = ? AND language = ? AND caller_file = ? AND caller_qualified_name = ? + """, + (self.project_root_str, self.language, caller_file, caller_qn), + ).fetchone() + counts[caller_qn] = row[0] if row else 0 + + return counts + def ensure_file_indexed(self, file_path: Path) -> IndexResult: resolved = str(file_path.resolve()) From c3fdf31a96affa86dd40f86b28c03732022e2a24 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 12 Feb 2026 01:11:39 -0500 Subject: [PATCH 24/42] refactor: use batch count_callees_per_function for dependency ranking and summary --- codeflash/languages/base.py | 4 +-- codeflash/languages/python/call_graph.py | 4 +-- codeflash/optimization/optimizer.py | 14 +++++----- tests/test_call_graph.py | 33 ++++++++++++++++++++++++ 4 files changed, 42 insertions(+), 13 deletions(-) diff --git a/codeflash/languages/base.py b/codeflash/languages/base.py index 1e2254b1d..7e7f5f127 100644 --- a/codeflash/languages/base.py +++ b/codeflash/languages/base.py @@ -221,9 +221,7 @@ def get_callees( """Return callees for the given functions.""" ... - def count_callees_per_function( - self, file_path_to_qualified_names: dict[Path, set[str]] - ) -> dict[str, int]: + def count_callees_per_function(self, file_path_to_qualified_names: dict[Path, set[str]]) -> dict[str, int]: """Return the number of callees for each caller qualified name.""" ... diff --git a/codeflash/languages/python/call_graph.py b/codeflash/languages/python/call_graph.py index aea6d696a..5dfa2d330 100644 --- a/codeflash/languages/python/call_graph.py +++ b/codeflash/languages/python/call_graph.py @@ -302,9 +302,7 @@ def get_callees( return file_path_to_function_source, function_source_list - def count_callees_per_function( - self, file_path_to_qualified_names: dict[Path, set[str]] - ) -> dict[str, int]: + def count_callees_per_function(self, file_path_to_qualified_names: dict[Path, set[str]]) -> dict[str, int]: all_caller_keys: list[tuple[str, str]] = [] for file_path, qualified_names in file_path_to_qualified_names.items(): self.ensure_file_indexed(file_path) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index b22f0df9c..89530ebf3 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -469,13 +469,13 @@ def rank_all_functions_globally( def rank_by_dependency_count( self, all_functions: list[tuple[Path, FunctionToOptimize]], call_graph: DependencyResolver ) -> list[tuple[Path, FunctionToOptimize]]: - counts: list[tuple[int, int, tuple[Path, FunctionToOptimize]]] = [] - for idx, (file_path, func) in enumerate(all_functions): - _, callee_list = call_graph.get_callees({file_path: {func.qualified_name}}) - counts.append((len(callee_list), idx, (file_path, func))) - counts.sort(key=lambda x: (-x[0], x[1])) - logger.debug(f"Ranked {len(counts)} functions by dependency count (most complex first)") - return [item for _, _, item in counts] + file_to_qns: dict[Path, set[str]] = defaultdict(set) + for file_path, func in all_functions: + file_to_qns[file_path].add(func.qualified_name) + callee_counts = call_graph.count_callees_per_function(dict(file_to_qns)) + ranked = sorted(enumerate(all_functions), key=lambda x: (-callee_counts.get(x[1][1].qualified_name, 0), x[0])) + logger.debug(f"Ranked {len(ranked)} functions by dependency count (most complex first)") + return [item for _, item in ranked] def run(self) -> None: from codeflash.code_utils.checkpoint import CodeflashRunCheckpoint diff --git a/tests/test_call_graph.py b/tests/test_call_graph.py index 22c306055..d2fcc104f 100644 --- a/tests/test_call_graph.py +++ b/tests/test_call_graph.py @@ -414,6 +414,39 @@ def caller(): cg.close() +def test_count_callees_per_function(project: Path, db_path: Path) -> None: + write_file( + project, + "mod.py", + """\ +def helper_a(): + return 1 + +def helper_b(): + return 2 + +def caller_one(): + return helper_a() + helper_b() + +def caller_two(): + return helper_a() + +def leaf(): + return 42 +""", + ) + + cg = CallGraph(project, db_path=db_path) + try: + cg.build_index([project / "mod.py"]) + counts = cg.count_callees_per_function({project / "mod.py": {"caller_one", "caller_two", "leaf"}}) + assert counts["caller_one"] == 2 + assert counts["caller_two"] == 1 + assert counts["leaf"] == 0 + finally: + cg.close() + + def test_same_file_edges_not_cross_file(project: Path, db_path: Path) -> None: write_file( project, From 457278331d455ae7a7e8566a71025d04d16969d1 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 12 Feb 2026 01:14:36 -0500 Subject: [PATCH 25/42] fix: remove slots=True from dataclass for Python 3.9 compatibility --- codeflash/languages/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/languages/base.py b/codeflash/languages/base.py index 7e7f5f127..b300b1950 100644 --- a/codeflash/languages/base.py +++ b/codeflash/languages/base.py @@ -35,7 +35,7 @@ def __getattr__(name: str) -> Any: raise AttributeError(msg) -@dataclass(frozen=True, slots=True) +@dataclass(frozen=True) class IndexResult: file_path: Path cached: bool From 267dff9702c746dd6b9130bc73bfe93816702322 Mon Sep 17 00:00:00 2001 From: "codeflash-ai[bot]" <148906541+codeflash-ai[bot]@users.noreply.github.com> Date: Thu, 12 Feb 2026 06:40:52 +0000 Subject: [PATCH 26/42] Optimize call_graph_summary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The optimized code achieves a **151% speedup** (from 7.86 to 3.12 microseconds) primarily through three key optimizations: ## 1. Module-Level Import Hoisting Moving `from rich.panel import Panel` from inside `call_graph_summary()` to the top-level module imports eliminates repeated import overhead on every function call. The line profiler shows this import took ~30,000 ns in the original (0.5% of total time). While seemingly small, this overhead is completely eliminated in the optimized version. ## 2. C-Level Aggregation with Built-in `sum()` The optimization replaces Python-level accumulation loops with native `sum()` calls that execute at C speed: **Original approach** (manual accumulation): ```python total_callees = 0 with_context = 0 for count in callee_counts.values(): total_callees += count if count > 0: with_context += 1 ``` This loop incurred ~828,000 ns across 2,005 iterations (234,973 + 301,448 + 292,402 ns). **Optimized approach** (C-level sum): ```python total_callees = sum(callee_counts.values()) with_context = sum(1 for count in callee_counts.values() if count > 0) ``` The new approach completes in ~405,000 ns total (16,399 + 389,145 ns) - nearly **2x faster** for the aggregation logic alone. ## 3. Leveraging `map()` for Initial Summation Using `sum(map(len, file_to_funcs.values()))` instead of a generator expression provides a minor efficiency gain by pushing the iteration into C-level code, though the improvement here is marginal (34,533 ns → 24,396 ns). ## Performance Characteristics Based on the annotated tests, these optimizations excel when: - **Large-scale scenarios**: The `test_large_scale_many_functions_single_file` (1000 functions) and `test_large_scale_multiple_files_distribution` (1000 functions across 10 files) benefit most from reduced per-iteration overhead - **Frequent invocations**: If `call_graph_summary()` is called multiple times in a session, the eliminated import overhead compounds savings - **Non-empty function sets**: The optimization's impact is proportional to the number of callees being aggregated The changes preserve all behavior - same summary text, same Panel display, same LSP handling - while delivering substantial runtime improvements through strategic use of Python's built-in functions that leverage optimized C implementations. --- codeflash/cli_cmds/console.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/codeflash/cli_cmds/console.py b/codeflash/cli_cmds/console.py index ddfa04771..bed8bc6c3 100644 --- a/codeflash/cli_cmds/console.py +++ b/codeflash/cli_cmds/console.py @@ -23,6 +23,7 @@ from codeflash.lsp.helpers import is_LSP_enabled from codeflash.lsp.lsp_logger import enhanced_log from codeflash.lsp.lsp_message import LspCodeMessage, LspTextMessage +from rich.panel import Panel if TYPE_CHECKING: from collections.abc import Callable, Generator @@ -318,22 +319,20 @@ def update(result: IndexResult) -> None: def call_graph_summary(call_graph: DependencyResolver, file_to_funcs: dict[Path, list[FunctionToOptimize]]) -> None: - from rich.panel import Panel - - total_functions = sum(len(funcs) for funcs in file_to_funcs.values()) + total_functions = sum(map(len, file_to_funcs.values())) if not total_functions: return - total_callees = 0 - with_context = 0 + # Build the mapping expected by the dependency resolver + file_items = file_to_funcs.items() + mapping = {file_path: {func.qualified_name for func in funcs} for file_path, funcs in file_items} + + callee_counts = call_graph.count_callees_per_function(mapping) + + # Use built-in sum for C-level loops to reduce Python overhead + total_callees = sum(callee_counts.values()) + with_context = sum(1 for count in callee_counts.values() if count > 0) - callee_counts = call_graph.count_callees_per_function( - {file_path: {func.qualified_name for func in funcs} for file_path, funcs in file_to_funcs.items()} - ) - for count in callee_counts.values(): - total_callees += count - if count > 0: - with_context += 1 leaf_functions = total_functions - with_context avg_callees = total_callees / total_functions From 182c1b002dd79fc9be55daebb53386d1aedabc6d Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Thu, 12 Feb 2026 06:43:39 +0000 Subject: [PATCH 27/42] style: auto-fix linting issues --- codeflash/cli_cmds/console.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/codeflash/cli_cmds/console.py b/codeflash/cli_cmds/console.py index bed8bc6c3..305ad458b 100644 --- a/codeflash/cli_cmds/console.py +++ b/codeflash/cli_cmds/console.py @@ -8,6 +8,7 @@ from rich.console import Console from rich.logging import RichHandler +from rich.panel import Panel from rich.progress import ( BarColumn, MofNCompleteColumn, @@ -23,7 +24,6 @@ from codeflash.lsp.helpers import is_LSP_enabled from codeflash.lsp.lsp_logger import enhanced_log from codeflash.lsp.lsp_message import LspCodeMessage, LspTextMessage -from rich.panel import Panel if TYPE_CHECKING: from collections.abc import Callable, Generator @@ -333,7 +333,6 @@ def call_graph_summary(call_graph: DependencyResolver, file_to_funcs: dict[Path, total_callees = sum(callee_counts.values()) with_context = sum(1 for count in callee_counts.values() if count > 0) - leaf_functions = total_functions - with_context avg_callees = total_callees / total_functions From 4523ac2d0c12842e32331a678cea5b4b05e3065b Mon Sep 17 00:00:00 2001 From: "codeflash-ai[bot]" <148906541+codeflash-ai[bot]@users.noreply.github.com> Date: Thu, 12 Feb 2026 07:03:35 +0000 Subject: [PATCH 28/42] Optimize _analyze_imports_in_optimized_code MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The optimized code achieves a **38% runtime improvement** (10.3ms → 7.42ms) by replacing the inefficient `ast.walk()` traversal with a targeted `ast.NodeVisitor` pattern. **Key Optimization:** The original code used `ast.walk(optimized_ast)` which visits **every node in the AST** (4,466 nodes in the profiled example), performing `isinstance()` checks on each one to find Import/ImportFrom nodes. This resulted in 18.77ms spent just traversing the tree (46.9% of total runtime). The optimized version introduces an `_ImportCollector` class that uses Python's `ast.NodeVisitor` pattern to selectively visit only Import and ImportFrom nodes. By defining `visit_Import()` and `visit_ImportFrom()` methods, the collector automatically skips irrelevant nodes during traversal. This reduces the collection phase to just 2.88ms (12% of runtime), saving approximately 15.89ms. **Performance Profile:** - The line profiler shows the `collector.visit()` call takes 2.88ms vs. the original loop's 18.77ms - The subsequent processing loop over collected nodes runs faster (1.69k iterations vs. 4.42k), eliminating 62% of unnecessary `isinstance()` checks - All other operations (helper preprocessing, dictionary lookups, set operations) remain essentially unchanged **Test Case Behavior:** The optimization is most effective for: - **Large ASTs with many nodes**: The `test_large_scale_many_import_statements_with_helpers` shows 62.6% speedup (662μs → 407μs) when processing 200 import statements, demonstrating the benefit of selective traversal - **Complex code with deep nesting**: ASTs with more non-import nodes see greater relative gains Smaller test cases show 30-40% slower runtimes due to the overhead of instantiating the collector class, but these are measuring microsecond differences (8-25μs) that are negligible in real-world usage where the function processes larger ASTs. **Practical Impact:** This function analyzes import statements in optimized code to map names to helper functions. Given its role in code optimization workflows, it likely processes many ASTs repeatedly. The 38% runtime reduction directly improves the optimization pipeline's throughput, especially when analyzing codebases with numerous import statements. --- .../context/unused_definition_remover.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/codeflash/context/unused_definition_remover.py b/codeflash/context/unused_definition_remover.py index 00b077f63..9883b8e84 100644 --- a/codeflash/context/unused_definition_remover.py +++ b/codeflash/context/unused_definition_remover.py @@ -638,7 +638,23 @@ def _analyze_imports_in_optimized_code( helpers_by_file_and_func[module_name].setdefault(func_name, []).append(helper) helpers_by_file[module_name].append(helper) - for node in ast.walk(optimized_ast): + # Collect only import nodes to avoid per-node isinstance checks across the whole AST + class _ImportCollector(ast.NodeVisitor): + def __init__(self) -> None: + self.nodes: list[ast.AST] = [] + + def visit_Import(self, node: ast.Import) -> None: + self.nodes.append(node) + # No need to recurse further for import nodes + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + self.nodes.append(node) + # No need to recurse further for import-from nodes + + collector = _ImportCollector() + collector.visit(optimized_ast) + + for node in collector.nodes: if isinstance(node, ast.ImportFrom): # Handle "from module import function" statements module_name = node.module @@ -655,6 +671,7 @@ def _analyze_imports_in_optimized_code( imported_set.add(helper.qualified_name) imported_set.add(helper.fully_qualified_name) + elif isinstance(node, ast.Import): # Handle "import module" statements for alias in node.names: From 0e284ad1803c042cc4c8538c0e8365eb61cd1405 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Thu, 12 Feb 2026 07:06:28 +0000 Subject: [PATCH 29/42] style: auto-fix linting issues --- codeflash/context/unused_definition_remover.py | 1 - 1 file changed, 1 deletion(-) diff --git a/codeflash/context/unused_definition_remover.py b/codeflash/context/unused_definition_remover.py index 9883b8e84..3547623ae 100644 --- a/codeflash/context/unused_definition_remover.py +++ b/codeflash/context/unused_definition_remover.py @@ -671,7 +671,6 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None: imported_set.add(helper.qualified_name) imported_set.add(helper.fully_qualified_name) - elif isinstance(node, ast.Import): # Handle "import module" statements for alias in node.names: From 11543f0d0b1ee93ec5c56a24032d991a23250515 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Mon, 16 Feb 2026 23:03:05 -0500 Subject: [PATCH 30/42] fix: use (file_path, qualified_name) key in count_callees_per_function Bare qualified_name keys could collide across files (e.g. `helper` in both a.py and b.py), causing counts to be silently overwritten. --- codeflash/languages/base.py | 6 ++++-- codeflash/languages/python/call_graph.py | 14 ++++++++------ codeflash/optimization/optimizer.py | 4 +++- tests/test_call_graph.py | 9 +++++---- 4 files changed, 20 insertions(+), 13 deletions(-) diff --git a/codeflash/languages/base.py b/codeflash/languages/base.py index 480ea01c1..edbfe1eae 100644 --- a/codeflash/languages/base.py +++ b/codeflash/languages/base.py @@ -221,8 +221,10 @@ def get_callees( """Return callees for the given functions.""" ... - def count_callees_per_function(self, file_path_to_qualified_names: dict[Path, set[str]]) -> dict[str, int]: - """Return the number of callees for each caller qualified name.""" + def count_callees_per_function( + self, file_path_to_qualified_names: dict[Path, set[str]] + ) -> dict[tuple[Path, str], int]: + """Return the number of callees for each (file_path, qualified_name) pair.""" ... def close(self) -> None: diff --git a/codeflash/languages/python/call_graph.py b/codeflash/languages/python/call_graph.py index 5dfa2d330..8d99a26bd 100644 --- a/codeflash/languages/python/call_graph.py +++ b/codeflash/languages/python/call_graph.py @@ -302,19 +302,21 @@ def get_callees( return file_path_to_function_source, function_source_list - def count_callees_per_function(self, file_path_to_qualified_names: dict[Path, set[str]]) -> dict[str, int]: - all_caller_keys: list[tuple[str, str]] = [] + def count_callees_per_function( + self, file_path_to_qualified_names: dict[Path, set[str]] + ) -> dict[tuple[Path, str], int]: + all_caller_keys: list[tuple[Path, str, str]] = [] for file_path, qualified_names in file_path_to_qualified_names.items(): self.ensure_file_indexed(file_path) resolved = str(file_path.resolve()) - all_caller_keys.extend((resolved, qn) for qn in qualified_names) + all_caller_keys.extend((file_path, resolved, qn) for qn in qualified_names) if not all_caller_keys: return {} - counts: dict[str, int] = {} + counts: dict[tuple[Path, str], int] = {} cur = self.conn.cursor() - for caller_file, caller_qn in all_caller_keys: + for orig_path, caller_file, caller_qn in all_caller_keys: row = cur.execute( """ SELECT COUNT(*) FROM call_edges @@ -322,7 +324,7 @@ def count_callees_per_function(self, file_path_to_qualified_names: dict[Path, se """, (self.project_root_str, self.language, caller_file, caller_qn), ).fetchone() - counts[caller_qn] = row[0] if row else 0 + counts[(orig_path, caller_qn)] = row[0] if row else 0 return counts diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 89530ebf3..682e8bb36 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -473,7 +473,9 @@ def rank_by_dependency_count( for file_path, func in all_functions: file_to_qns[file_path].add(func.qualified_name) callee_counts = call_graph.count_callees_per_function(dict(file_to_qns)) - ranked = sorted(enumerate(all_functions), key=lambda x: (-callee_counts.get(x[1][1].qualified_name, 0), x[0])) + ranked = sorted( + enumerate(all_functions), key=lambda x: (-callee_counts.get((x[1][0], x[1][1].qualified_name), 0), x[0]) + ) logger.debug(f"Ranked {len(ranked)} functions by dependency count (most complex first)") return [item for _, item in ranked] diff --git a/tests/test_call_graph.py b/tests/test_call_graph.py index d2fcc104f..9dda10e97 100644 --- a/tests/test_call_graph.py +++ b/tests/test_call_graph.py @@ -439,10 +439,11 @@ def leaf(): cg = CallGraph(project, db_path=db_path) try: cg.build_index([project / "mod.py"]) - counts = cg.count_callees_per_function({project / "mod.py": {"caller_one", "caller_two", "leaf"}}) - assert counts["caller_one"] == 2 - assert counts["caller_two"] == 1 - assert counts["leaf"] == 0 + mod_path = project / "mod.py" + counts = cg.count_callees_per_function({mod_path: {"caller_one", "caller_two", "leaf"}}) + assert counts[(mod_path, "caller_one")] == 2 + assert counts[(mod_path, "caller_two")] == 1 + assert counts[(mod_path, "leaf")] == 0 finally: cg.close() From 43d74a8f9e436d5791696586ea236c75692be280 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Wed, 18 Feb 2026 17:10:32 -0500 Subject: [PATCH 31/42] fix: use iterative DFS for function discovery to avoid RecursionError Deeply nested expression trees (e.g. large dict/list literals) at module or class level caused the recursive ast.NodeVisitor to exceed Python's default recursion limit. Replace the FunctionWithReturnStatement visitor class with an iterative stack-based traversal. --- codeflash/discovery/functions_to_optimize.py | 56 ++++++++++---------- 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index 86d574af1..5951b137b 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -114,32 +114,34 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> None: ) -class FunctionWithReturnStatement(ast.NodeVisitor): - def __init__(self, file_path: Path) -> None: - self.functions: list[FunctionToOptimize] = [] - self.ast_path: list[FunctionParent] = [] - self.file_path: Path = file_path - - def visit_FunctionDef(self, node: FunctionDef) -> None: - if function_has_return_statement(node) and not function_is_a_property(node): - self.functions.append( - FunctionToOptimize(function_name=node.name, file_path=self.file_path, parents=self.ast_path[:]) - ) - - def visit_AsyncFunctionDef(self, node: AsyncFunctionDef) -> None: - if function_has_return_statement(node) and not function_is_a_property(node): - self.functions.append( - FunctionToOptimize( - function_name=node.name, file_path=self.file_path, parents=self.ast_path[:], is_async=True +def find_functions_with_return_statement( + ast_module: ast.Module, file_path: Path +) -> list[FunctionToOptimize]: + results: list[FunctionToOptimize] = [] + # (node, parent_path) — iterative DFS avoids RecursionError on deeply nested ASTs + stack: list[tuple[ast.AST, list[FunctionParent]]] = [(ast_module, [])] + while stack: + node, ast_path = stack.pop() + if isinstance(node, (FunctionDef, AsyncFunctionDef)): + if function_has_return_statement(node) and not function_is_a_property(node): + results.append( + FunctionToOptimize( + function_name=node.name, + file_path=file_path, + parents=ast_path[:], + is_async=isinstance(node, AsyncFunctionDef), + ) ) - ) - - def generic_visit(self, node: ast.AST) -> None: - if isinstance(node, (FunctionDef, AsyncFunctionDef, ClassDef)): - self.ast_path.append(FunctionParent(node.name, node.__class__.__name__)) - super().generic_visit(node) - if isinstance(node, (FunctionDef, AsyncFunctionDef, ClassDef)): - self.ast_path.pop() + # Don't recurse into function bodies (matches original visitor behaviour) + continue + child_path = ( + ast_path + [FunctionParent(node.name, node.__class__.__name__)] + if isinstance(node, ClassDef) + else ast_path + ) + for child in reversed(list(ast.iter_child_nodes(node))): + stack.append((child, child_path)) + return results # ============================================================================= @@ -237,9 +239,7 @@ def _find_all_functions_in_python_file(file_path: Path) -> dict[Path, list[Funct if DEBUG_MODE: logger.exception(e) return functions - function_name_visitor = FunctionWithReturnStatement(file_path) - function_name_visitor.visit(ast_module) - functions[file_path] = function_name_visitor.functions + functions[file_path] = find_functions_with_return_statement(ast_module, file_path) return functions From 5663985712a264833fce8abae2205f68c98c9502 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Wed, 18 Feb 2026 22:13:57 +0000 Subject: [PATCH 32/42] style: auto-fix linting issues --- codeflash/discovery/functions_to_optimize.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index 5951b137b..8821e0e9a 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -114,9 +114,7 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> None: ) -def find_functions_with_return_statement( - ast_module: ast.Module, file_path: Path -) -> list[FunctionToOptimize]: +def find_functions_with_return_statement(ast_module: ast.Module, file_path: Path) -> list[FunctionToOptimize]: results: list[FunctionToOptimize] = [] # (node, parent_path) — iterative DFS avoids RecursionError on deeply nested ASTs stack: list[tuple[ast.AST, list[FunctionParent]]] = [(ast_module, [])] @@ -135,9 +133,7 @@ def find_functions_with_return_statement( # Don't recurse into function bodies (matches original visitor behaviour) continue child_path = ( - ast_path + [FunctionParent(node.name, node.__class__.__name__)] - if isinstance(node, ClassDef) - else ast_path + [*ast_path, FunctionParent(node.name, node.__class__.__name__)] if isinstance(node, ClassDef) else ast_path ) for child in reversed(list(ast.iter_child_nodes(node))): stack.append((child, child_path)) From 890c466b1afd402568d36eef502a4b2d9592612b Mon Sep 17 00:00:00 2001 From: "codeflash-ai[bot]" <148906541+codeflash-ai[bot]@users.noreply.github.com> Date: Wed, 18 Feb 2026 22:22:40 +0000 Subject: [PATCH 33/42] Optimize find_functions_with_return_statement The optimized code achieves a **26% runtime improvement** by making the AST traversal in `function_has_return_statement` more targeted and efficient. **Key Optimization:** The critical change is in how `function_has_return_statement` traverses the AST when searching for `Return` nodes: **Original approach:** ```python stack.extend(ast.iter_child_nodes(node)) ``` This visits *all* child nodes including expressions, names, constants, and other non-statement nodes. **Optimized approach:** ```python for child in ast.iter_child_nodes(node): if isinstance(child, ast.stmt): stack.append(child) ``` This only pushes statement nodes onto the stack, since `Return` is a statement type (`ast.stmt`). **Why This Is Faster:** 1. **Reduced Node Traversal**: In typical Python functions, there are many more expression nodes (variable references, literals, operators, etc.) than statement nodes. For example, a simple `return x + y` has 1 Return statement but multiple Name and BinOp expression nodes underneath. The optimization skips all the expression-level nodes. 2. **Lower Python Overhead**: Fewer nodes in the stack means fewer loop iterations, fewer `isinstance` checks on non-Return nodes, and less list manipulation overhead. 3. **Preserved Correctness**: Since `Return` nodes are always statements in Python's AST (they inherit from `ast.stmt`), filtering to only statement nodes cannot miss any Return nodes. **Performance Impact by Test Case:** The optimization shows particularly strong gains for: - **Functions without returns** (up to 91% faster): Early termination without traversing deep expression trees - **Large codebases** (34-41% faster on tests with 1000+ functions): The cumulative effect across many function bodies - **Functions with complex expressions but no returns** (82% faster): Avoiding expensive traversal of unused expression subtrees - **Generator functions without explicit returns** (64% faster): Skipping yield expression internals The optimization maintains correctness across all test cases including nested classes, async functions, properties, and various control structures, while delivering consistent runtime improvements. --- codeflash/discovery/functions_to_optimize.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index 8821e0e9a..bb8b2f902 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -961,7 +961,11 @@ def function_has_return_statement(function_node: FunctionDef | AsyncFunctionDef) node = stack.pop() if isinstance(node, ast.Return): return True - stack.extend(ast.iter_child_nodes(node)) + # Only push child nodes that are statements; Return nodes are statements, + # so this preserves correctness while avoiding unnecessary traversal into expr/Name/etc. + for child in ast.iter_child_nodes(node): + if isinstance(child, ast.stmt): + stack.append(child) return False From 88b0ee50cfef8273f8170ed0880e4cd8daccfb14 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Wed, 18 Feb 2026 17:24:39 -0500 Subject: [PATCH 34/42] fix: batch SQL queries and deduplicate Path.resolve() in call graph Replace per-function SQL loops in get_callees() and count_callees_per_function() with temp table JOINs, and thread resolved path strings through to avoid redundant resolve() calls. --- .codex/skills/.gitignore | 1 + .gemini/skills/.gitignore | 1 + codeflash/languages/python/call_graph.py | 108 +++++++++++++---------- 3 files changed, 65 insertions(+), 45 deletions(-) diff --git a/.codex/skills/.gitignore b/.codex/skills/.gitignore index b1cda282a..4e7af4bc6 100644 --- a/.codex/skills/.gitignore +++ b/.codex/skills/.gitignore @@ -1,2 +1,3 @@ # Managed by Tessl tessl:* +tessl__* diff --git a/.gemini/skills/.gitignore b/.gemini/skills/.gitignore index b1cda282a..4e7af4bc6 100644 --- a/.gemini/skills/.gitignore +++ b/.gemini/skills/.gitignore @@ -1,2 +1,3 @@ # Managed by Tessl tessl:* +tessl__* diff --git a/codeflash/languages/python/call_graph.py b/codeflash/languages/python/call_graph.py index 8d99a26bd..f9641d2e0 100644 --- a/codeflash/languages/python/call_graph.py +++ b/codeflash/languages/python/call_graph.py @@ -264,41 +264,44 @@ def get_callees( file_path_to_function_source: dict[Path, set[FunctionSource]] = defaultdict(set) function_source_list: list[FunctionSource] = [] - # Build list of all caller keys all_caller_keys: list[tuple[str, str]] = [] for file_path, qualified_names in file_path_to_qualified_names.items(): - self.ensure_file_indexed(file_path) resolved = str(file_path.resolve()) + self.ensure_file_indexed(file_path, resolved) all_caller_keys.extend((resolved, qn) for qn in qualified_names) if not all_caller_keys: return file_path_to_function_source, function_source_list - # Query all callees cur = self.conn.cursor() - for caller_file, caller_qn in all_caller_keys: - rows = cur.execute( - """ - SELECT callee_file, callee_qualified_name, callee_fully_qualified_name, - callee_only_function_name, callee_definition_type, callee_source_line - FROM call_edges - WHERE project_root = ? AND language = ? AND caller_file = ? AND caller_qualified_name = ? - """, - (self.project_root_str, self.language, caller_file, caller_qn), - ).fetchall() - - for callee_file, callee_qn, callee_fqn, callee_name, callee_type, callee_src in rows: - callee_path = Path(callee_file) - fs = FunctionSource( - file_path=callee_path, - qualified_name=callee_qn, - fully_qualified_name=callee_fqn, - only_function_name=callee_name, - source_code=callee_src, - definition_type=callee_type, - ) - file_path_to_function_source[callee_path].add(fs) - function_source_list.append(fs) + cur.execute("CREATE TEMP TABLE IF NOT EXISTS _caller_keys (caller_file TEXT, caller_qualified_name TEXT)") + cur.execute("DELETE FROM _caller_keys") + cur.executemany("INSERT INTO _caller_keys VALUES (?, ?)", all_caller_keys) + + rows = cur.execute( + """ + SELECT ce.callee_file, ce.callee_qualified_name, ce.callee_fully_qualified_name, + ce.callee_only_function_name, ce.callee_definition_type, ce.callee_source_line + FROM call_edges ce + INNER JOIN _caller_keys ck + ON ce.caller_file = ck.caller_file AND ce.caller_qualified_name = ck.caller_qualified_name + WHERE ce.project_root = ? AND ce.language = ? + """, + (self.project_root_str, self.language), + ).fetchall() + + for callee_file, callee_qn, callee_fqn, callee_name, callee_type, callee_src in rows: + callee_path = Path(callee_file) + fs = FunctionSource( + file_path=callee_path, + qualified_name=callee_qn, + fully_qualified_name=callee_fqn, + only_function_name=callee_name, + source_code=callee_src, + definition_type=callee_type, + ) + file_path_to_function_source[callee_path].add(fs) + function_source_list.append(fs) return file_path_to_function_source, function_source_list @@ -307,30 +310,44 @@ def count_callees_per_function( ) -> dict[tuple[Path, str], int]: all_caller_keys: list[tuple[Path, str, str]] = [] for file_path, qualified_names in file_path_to_qualified_names.items(): - self.ensure_file_indexed(file_path) resolved = str(file_path.resolve()) + self.ensure_file_indexed(file_path, resolved) all_caller_keys.extend((file_path, resolved, qn) for qn in qualified_names) if not all_caller_keys: return {} - counts: dict[tuple[Path, str], int] = {} cur = self.conn.cursor() - for orig_path, caller_file, caller_qn in all_caller_keys: - row = cur.execute( - """ - SELECT COUNT(*) FROM call_edges - WHERE project_root = ? AND language = ? AND caller_file = ? AND caller_qualified_name = ? - """, - (self.project_root_str, self.language, caller_file, caller_qn), - ).fetchone() - counts[(orig_path, caller_qn)] = row[0] if row else 0 + cur.execute("CREATE TEMP TABLE IF NOT EXISTS _count_keys (caller_file TEXT, caller_qualified_name TEXT)") + cur.execute("DELETE FROM _count_keys") + cur.executemany( + "INSERT INTO _count_keys VALUES (?, ?)", [(resolved, qn) for _, resolved, qn in all_caller_keys] + ) + + rows = cur.execute( + """ + SELECT ck.caller_file, ck.caller_qualified_name, COUNT(ce.rowid) + FROM _count_keys ck + LEFT JOIN call_edges ce + ON ce.caller_file = ck.caller_file AND ce.caller_qualified_name = ck.caller_qualified_name + AND ce.project_root = ? AND ce.language = ? + GROUP BY ck.caller_file, ck.caller_qualified_name + """, + (self.project_root_str, self.language), + ).fetchall() + + resolved_to_path: dict[str, Path] = {resolved: fp for fp, resolved, _ in all_caller_keys} + counts: dict[tuple[Path, str], int] = {} + for caller_file, caller_qn, cnt in rows: + counts[(resolved_to_path[caller_file], caller_qn)] = cnt return counts - def ensure_file_indexed(self, file_path: Path) -> IndexResult: - resolved = str(file_path.resolve()) + def ensure_file_indexed(self, file_path: Path, resolved: str | None = None) -> IndexResult: + if resolved is None: + resolved = str(file_path.resolve()) + # Always read and hash the file before checking the cache so we detect on-disk changes try: content = file_path.read_text(encoding="utf-8") except Exception: @@ -341,10 +358,11 @@ def ensure_file_indexed(self, file_path: Path) -> IndexResult: if self._is_file_cached(resolved, file_hash): return IndexResult(file_path=file_path, cached=True, num_edges=0, edges=(), cross_file_edges=0, error=False) - return self.index_file(file_path, file_hash) + return self.index_file(file_path, file_hash, resolved) - def index_file(self, file_path: Path, file_hash: str) -> IndexResult: - resolved = str(file_path.resolve()) + def index_file(self, file_path: Path, file_hash: str, resolved: str | None = None) -> IndexResult: + if resolved is None: + resolved = str(file_path.resolve()) edges, had_error = _analyze_file(file_path, self.jedi_project, self.project_root_str) if had_error: logger.debug(f"CallGraph: failed to parse {file_path}") @@ -441,8 +459,8 @@ def build_index(self, file_paths: Iterable[Path], on_progress: Callable[[IndexRe if len(to_index) >= _PARALLEL_THRESHOLD: self._build_index_parallel(to_index, on_progress) else: - for file_path, _resolved, file_hash in to_index: - result = self.index_file(file_path, file_hash) + for file_path, resolved, file_hash in to_index: + result = self.index_file(file_path, file_hash, resolved) self._report_progress(on_progress, result) def _is_file_cached(self, resolved: str, file_hash: str) -> bool: @@ -518,7 +536,7 @@ def _fallback_sequential_index( # Skip files already persisted before the failure if resolved in self.indexed_file_hashes: continue - result = self.index_file(file_path, file_hash) + result = self.index_file(file_path, file_hash, resolved) self._report_progress(on_progress, result) def close(self) -> None: From f7b04013ca683c59e9e836aff819bbe0bc7e22af Mon Sep 17 00:00:00 2001 From: "codeflash-ai[bot]" <148906541+codeflash-ai[bot]@users.noreply.github.com> Date: Wed, 18 Feb 2026 22:34:59 +0000 Subject: [PATCH 35/42] Optimize function_has_return_statement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The optimized code achieves a **146% speedup** (from 1.47ms to 595μs) by eliminating the overhead of `ast.iter_child_nodes()` and replacing it with direct field access on AST nodes. **Key optimizations:** 1. **Direct stack initialization**: Instead of starting with `[function_node]` and then traversing into its body, the stack is initialized directly with `list(function_node.body)`. This skips one iteration and avoids processing the function definition wrapper itself. 2. **Manual field traversal**: Rather than calling `ast.iter_child_nodes(node)` which is a generator that yields all child nodes, the code directly accesses `node._fields` and uses `getattr()` to inspect each field. This eliminates the generator overhead and function call costs associated with `ast.iter_child_nodes()`. 3. **Targeted statement filtering**: By checking `isinstance(child, ast.stmt)` or `isinstance(item, ast.stmt)` only on relevant fields (handling both single statements and lists of statements), the traversal focuses on statement nodes where `ast.Return` can appear, avoiding unnecessary checks on expression nodes. **Why this is faster:** - **Reduced function call overhead**: `ast.iter_child_nodes()` is a generator function that incurs call/yield overhead on every iteration. Direct attribute access via `getattr()` is faster for small numbers of fields. - **Fewer iterations**: The line profiler shows the original code's `ast.iter_child_nodes()` line hit 5,453 times (69% of runtime), while the optimized version's field iteration hits only 3,290 times (17.4% of runtime). - **Better cache locality**: Direct field access patterns may benefit from better CPU cache utilization compared to generator state management. **Test case performance:** The optimization shows dramatic improvements particularly for: - **Functions with many sequential statements** (2365% faster for 1000 statements, 1430% faster for 1000 nested functions) - **Simple functions** (234-354% faster for basic return detection) - **Moderately complex control flow** (80-125% faster for nested conditionals/loops) The speedup is consistent across all test cases, with early-return scenarios benefiting the most as the optimization allows faster discovery of the return statement before processing unnecessary nodes. --- codeflash/discovery/functions_to_optimize.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index bb8b2f902..abad0d85e 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -956,15 +956,20 @@ def filter_files_optimized(file_path: Path, tests_root: Path, ignore_paths: list def function_has_return_statement(function_node: FunctionDef | AsyncFunctionDef) -> bool: # Custom DFS, return True as soon as a Return node is found - stack: list[ast.AST] = [function_node] + stack: list[ast.AST] = list(function_node.body) while stack: node = stack.pop() if isinstance(node, ast.Return): return True # Only push child nodes that are statements; Return nodes are statements, # so this preserves correctness while avoiding unnecessary traversal into expr/Name/etc. - for child in ast.iter_child_nodes(node): - if isinstance(child, ast.stmt): + for field in getattr(node, "_fields", ()): + child = getattr(node, field, None) + if isinstance(child, list): + for item in child: + if isinstance(item, ast.stmt): + stack.append(item) + elif isinstance(child, ast.stmt): stack.append(child) return False From 522969baadbf3a283ac96e314c0f2a37f7e8fc04 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 19 Feb 2026 01:17:11 -0500 Subject: [PATCH 36/42] fix: restore call_graph parameter to get_code_optimization_context Add DependencyResolver parameter back to get_code_optimization_context() that was lost during file move from codeflash/context/ to codeflash/languages/python/context/. When call_graph is available, use it for helper discovery instead of Jedi-based fallback. --- .../python/context/code_context_extractor.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index 9a4daf726..5ef02d788 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -36,7 +36,7 @@ if TYPE_CHECKING: from jedi.api.classes import Name - from codeflash.languages.base import HelperFunction + from codeflash.languages.base import DependencyResolver, HelperFunction from codeflash.languages.python.context.unused_definition_remover import UsageInfo # Error message constants @@ -80,6 +80,7 @@ def get_code_optimization_context( project_root_path: Path, optim_token_limit: int = OPTIMIZATION_CONTEXT_TOKEN_LIMIT, testgen_token_limit: int = TESTGEN_CONTEXT_TOKEN_LIMIT, + call_graph: DependencyResolver | None = None, ) -> CodeOptimizationContext: # Route to language-specific implementation for non-Python languages if not is_python(): @@ -88,9 +89,11 @@ def get_code_optimization_context( ) # Get FunctionSource representation of helpers of FTO - helpers_of_fto_dict, helpers_of_fto_list = get_function_sources_from_jedi( - {function_to_optimize.file_path: {function_to_optimize.qualified_name}}, project_root_path - ) + fto_input = {function_to_optimize.file_path: {function_to_optimize.qualified_name}} + if call_graph is not None: + helpers_of_fto_dict, helpers_of_fto_list = call_graph.get_callees(fto_input) + else: + helpers_of_fto_dict, helpers_of_fto_list = get_function_sources_from_jedi(fto_input, project_root_path) # Add function to optimize into helpers of FTO dict, as they'll be processed together fto_as_function_source = get_function_to_optimize_as_function_source(function_to_optimize, project_root_path) From f1c707a7c4ea82f8e6f1e854e8a986e54bfc5717 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 19 Feb 2026 01:51:49 -0500 Subject: [PATCH 37/42] Simplify dependency summary output --- codeflash/cli_cmds/console.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/codeflash/cli_cmds/console.py b/codeflash/cli_cmds/console.py index 305ad458b..f44aff44e 100644 --- a/codeflash/cli_cmds/console.py +++ b/codeflash/cli_cmds/console.py @@ -336,11 +336,12 @@ def call_graph_summary(call_graph: DependencyResolver, file_to_funcs: dict[Path, leaf_functions = total_functions - with_context avg_callees = total_callees / total_functions + function_label = "function" if total_functions == 1 else "functions" + summary = ( - f"{total_functions} functions ready for optimization · " - f"avg {avg_callees:.1f} dependencies/function\n" - f"{with_context} call other functions · " - f"{leaf_functions} are self-contained" + f"{total_functions} {function_label} ready for optimization\n" + f"Uses other functions: {with_context} · " + f"Standalone: {leaf_functions}" ) if is_LSP_enabled(): From 12f36fb064b6015181079aeed571f70991e78bfa Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 19 Feb 2026 01:57:41 -0500 Subject: [PATCH 38/42] Clarify call graph UI text --- codeflash/cli_cmds/console.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/codeflash/cli_cmds/console.py b/codeflash/cli_cmds/console.py index f44aff44e..fdc5a420a 100644 --- a/codeflash/cli_cmds/console.py +++ b/codeflash/cli_cmds/console.py @@ -235,8 +235,8 @@ def call_graph_live_display( results: deque[IndexResult] = deque(maxlen=MAX_TREE_ENTRIES) stats = {"indexed": 0, "cached": 0, "edges": 0, "external": 0, "errors": 0} - tree = Tree("[bold]Dependencies[/bold]") - stats_text = Text("0 dependencies found", style="dim") + tree = Tree("[bold]Recent Files[/bold]") + stats_text = Text("0 calls found", style="dim") panel = Panel( Group(progress, Text(""), tree, Text(""), stats_text), title="Building Call Graph", border_style="cyan" ) @@ -260,11 +260,11 @@ def create_tree_node(result: IndexResult) -> Tree: edge_info = [] if local_edges: - edge_info.append(f"{local_edges} in same file") + edge_info.append(f"{local_edges} calls in same file") if result.cross_file_edges: - edge_info.append(f"{result.cross_file_edges} from other modules") + edge_info.append(f"{result.cross_file_edges} calls from other modules") - label = ", ".join(edge_info) if edge_info else "no dependencies" + label = ", ".join(edge_info) if edge_info else "no calls" return Tree(f"[cyan]{name}[/cyan] [dim]{label}[/dim]") def refresh_display() -> None: @@ -279,9 +279,9 @@ def refresh_display() -> None: stat_parts.append(f"{stats['cached']} cached") if stats["errors"]: stat_parts.append(f"{stats['errors']} errors") - stat_parts.append(f"{stats['edges']} dependencies found") + stat_parts.append(f"{stats['edges']} calls found") if stats["external"]: - stat_parts.append(f"{stats['external']} from other modules") + stat_parts.append(f"{stats['external']} cross-file calls") stats_text.truncate(0) stats_text.append(" · ".join(stat_parts), style="dim") @@ -312,7 +312,7 @@ def update(result: IndexResult) -> None: if len(batch) >= 8: process_batch() - with Live(panel, console=console, transient=True, auto_refresh=False) as live: + with Live(panel, console=console, transient=False, auto_refresh=False) as live: yield update if batch: process_batch() @@ -348,4 +348,4 @@ def call_graph_summary(call_graph: DependencyResolver, file_to_funcs: dict[Path, lsp_log(LspTextMessage(text=summary)) return - console.print(Panel(summary, title="Dependency Summary", border_style="cyan")) + console.print(Panel(summary, title="Call Graph Summary", border_style="cyan")) From 09ad2f59d1555e4383eefd7c00a30c12ce6a5326 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 19 Feb 2026 02:26:20 -0500 Subject: [PATCH 39/42] Rename Python call graph to reference graph --- codeflash/languages/python/__init__.py | 4 ++-- .../python/{call_graph.py => reference_graph.py} | 13 +++++++------ codeflash/languages/python/support.py | 6 +++--- 3 files changed, 12 insertions(+), 11 deletions(-) rename codeflash/languages/python/{call_graph.py => reference_graph.py} (97%) diff --git a/codeflash/languages/python/__init__.py b/codeflash/languages/python/__init__.py index 50eb1a51e..939d5941f 100644 --- a/codeflash/languages/python/__init__.py +++ b/codeflash/languages/python/__init__.py @@ -5,7 +5,7 @@ to the LanguageSupport protocol. """ -from codeflash.languages.python.call_graph import CallGraph +from codeflash.languages.python.reference_graph import ReferenceGraph from codeflash.languages.python.support import PythonSupport -__all__ = ["CallGraph", "PythonSupport"] +__all__ = ["PythonSupport", "ReferenceGraph"] diff --git a/codeflash/languages/python/call_graph.py b/codeflash/languages/python/reference_graph.py similarity index 97% rename from codeflash/languages/python/call_graph.py rename to codeflash/languages/python/reference_graph.py index f9641d2e0..4f389fd66 100644 --- a/codeflash/languages/python/call_graph.py +++ b/codeflash/languages/python/reference_graph.py @@ -21,6 +21,7 @@ # --------------------------------------------------------------------------- # Module-level helpers (must be top-level for ProcessPoolExecutor pickling) # --------------------------------------------------------------------------- +# TODO: create call graph. _PARALLEL_THRESHOLD = 8 @@ -179,7 +180,7 @@ def _index_file_worker(args: tuple[str, str]) -> tuple[str, str, set[tuple[str, # --------------------------------------------------------------------------- -class CallGraph: +class ReferenceGraph: SCHEMA_VERSION = 2 def __init__(self, project_root: Path, language: str = "python", db_path: Path | None = None) -> None: @@ -365,7 +366,7 @@ def index_file(self, file_path: Path, file_hash: str, resolved: str | None = Non resolved = str(file_path.resolve()) edges, had_error = _analyze_file(file_path, self.jedi_project, self.project_root_str) if had_error: - logger.debug(f"CallGraph: failed to parse {file_path}") + logger.debug(f"ReferenceGraph: failed to parse {file_path}") return self._persist_edges(file_path, resolved, file_hash, edges, had_error) def _persist_edges( @@ -493,7 +494,7 @@ def _build_index_parallel( path_info: dict[str, tuple[Path, str]] = {resolved: (fp, fh) for fp, resolved, fh in to_index} worker_args = [(resolved, fh) for _fp, resolved, fh in to_index] - logger.debug(f"CallGraph: indexing {len(to_index)} files across {max_workers} workers") + logger.debug(f"ReferenceGraph: indexing {len(to_index)} files across {max_workers} workers") try: with ProcessPoolExecutor( @@ -508,7 +509,7 @@ def _build_index_parallel( try: _, _, edges, had_error = future.result() except Exception: - logger.debug(f"CallGraph: worker failed for {file_path}") + logger.debug(f"ReferenceGraph: worker failed for {file_path}") self._persist_edges(file_path, resolved, file_hash, set(), had_error=True) self._report_progress( on_progress, @@ -519,13 +520,13 @@ def _build_index_parallel( continue if had_error: - logger.debug(f"CallGraph: failed to parse {file_path}") + logger.debug(f"ReferenceGraph: failed to parse {file_path}") result = self._persist_edges(file_path, resolved, file_hash, edges, had_error) self._report_progress(on_progress, result) except Exception: - logger.debug("CallGraph: parallel indexing failed, falling back to sequential") + logger.debug("ReferenceGraph: parallel indexing failed, falling back to sequential") self._fallback_sequential_index(to_index, on_progress) def _fallback_sequential_index( diff --git a/codeflash/languages/python/support.py b/codeflash/languages/python/support.py index 1480e690c..b026e99e5 100644 --- a/codeflash/languages/python/support.py +++ b/codeflash/languages/python/support.py @@ -753,12 +753,12 @@ def ensure_runtime_environment(self, project_root: Path) -> bool: return True def create_dependency_resolver(self, project_root: Path) -> DependencyResolver | None: - from codeflash.languages.python.call_graph import CallGraph + from codeflash.languages.python.reference_graph import ReferenceGraph try: - return CallGraph(project_root, language=self.language.value) + return ReferenceGraph(project_root, language=self.language.value) except Exception: - logger.debug("Failed to initialize CallGraph, falling back to per-function Jedi analysis") + logger.debug("Failed to initialize ReferenceGraph, falling back to per-function Jedi analysis") return None def instrument_existing_test( From cb9115831251a83967233558276ea12a0c82afb4 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 19 Feb 2026 02:30:27 -0500 Subject: [PATCH 40/42] refactor: rename test file and imports to match reference graph rename --- ..._call_graph.py => test_reference_graph.py} | 36 +++++++++---------- 1 file changed, 18 insertions(+), 18 deletions(-) rename tests/{test_call_graph.py => test_reference_graph.py} (92%) diff --git a/tests/test_call_graph.py b/tests/test_reference_graph.py similarity index 92% rename from tests/test_call_graph.py rename to tests/test_reference_graph.py index 9dda10e97..6e3ab0c65 100644 --- a/tests/test_call_graph.py +++ b/tests/test_reference_graph.py @@ -8,7 +8,7 @@ from pathlib import Path from codeflash.languages.base import IndexResult -from codeflash.languages.python.call_graph import CallGraph +from codeflash.languages.python.reference_graph import ReferenceGraph @pytest.fixture @@ -46,7 +46,7 @@ def caller(): return helper() """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: _, result_list = cg.get_callees({project / "mod.py": {"caller"}}) callee_qns = {fs.qualified_name for fs in result_list} @@ -74,7 +74,7 @@ def caller(): return utility() """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: _, result_list = cg.get_callees({project / "main.py": {"caller"}}) callee_qns = {fs.qualified_name for fs in result_list} @@ -100,7 +100,7 @@ def caller(): return obj """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: _, result_list = cg.get_callees({project / "mod.py": {"caller"}}) callee_types = {fs.definition_type for fs in result_list} @@ -120,7 +120,7 @@ def inner(): return inner() """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: _, result_list = cg.get_callees({project / "mod.py": {"caller"}}) assert len(result_list) == 0 @@ -139,7 +139,7 @@ def helper(): x = helper() """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: # Module level calls have no enclosing function, so no edges _, result_list = cg.get_callees({project / "mod.py": {"helper"}}) @@ -160,7 +160,7 @@ def caller(): return os.path.join("a", "b") """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: _, result_list = cg.get_callees({project / "mod.py": {"caller"}}) # os.path.join is stdlib, should not appear @@ -171,7 +171,7 @@ def caller(): def test_empty_file(project: Path, db_path: Path) -> None: write_file(project, "mod.py", "") - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: _, result_list = cg.get_callees({project / "mod.py": set()}) assert len(result_list) == 0 @@ -181,7 +181,7 @@ def test_empty_file(project: Path, db_path: Path) -> None: def test_syntax_error_file(project: Path, db_path: Path) -> None: write_file(project, "mod.py", "def broken(\n") - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: _, result_list = cg.get_callees({project / "mod.py": {"broken"}}) assert len(result_list) == 0 @@ -206,7 +206,7 @@ def caller(): return helper() """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: cg.get_callees({project / "mod.py": {"caller"}}) # Second call should use in-memory cache (hash unchanged) @@ -231,7 +231,7 @@ def caller(): return helper() """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: _, result_list = cg.get_callees({project / "mod.py": {"caller"}}) assert any(fs.qualified_name == "helper" for fs in result_list) @@ -270,7 +270,7 @@ def caller(): """, ) # First session: index the file - cg1 = CallGraph(project, db_path=db_path) + cg1 = ReferenceGraph(project, db_path=db_path) try: _, result_list = cg1.get_callees({project / "mod.py": {"caller"}}) assert any(fs.qualified_name == "helper" for fs in result_list) @@ -278,7 +278,7 @@ def caller(): cg1.close() # Second session: should read from DB without re-indexing - cg2 = CallGraph(project, db_path=db_path) + cg2 = ReferenceGraph(project, db_path=db_path) try: assert len(cg2.indexed_file_hashes) == 0 # in-memory cache is empty _, result_list = cg2.get_callees({project / "mod.py": {"caller"}}) @@ -310,7 +310,7 @@ def caller_b(): """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: progress_calls: list[IndexResult] = [] files = [project / "a.py", project / "b.py"] @@ -359,7 +359,7 @@ def caller_b(): """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: files = [project / "a.py", project / "b.py"] # First pass — fresh indexing @@ -400,7 +400,7 @@ def caller(): """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: progress_calls: list[IndexResult] = [] cg.build_index([project / "utils.py", project / "main.py"], on_progress=progress_calls.append) @@ -436,7 +436,7 @@ def leaf(): """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: cg.build_index([project / "mod.py"]) mod_path = project / "mod.py" @@ -461,7 +461,7 @@ def caller(): """, ) - cg = CallGraph(project, db_path=db_path) + cg = ReferenceGraph(project, db_path=db_path) try: progress_calls: list[IndexResult] = [] cg.build_index([project / "mod.py"], on_progress=progress_calls.append) From 82783f8c92a911dbe28787a2b801f263559ed974 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 19 Feb 2026 02:40:03 -0500 Subject: [PATCH 41/42] disable it for now --- codeflash/optimization/optimizer.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 680efdf1e..941502a4d 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -572,16 +572,17 @@ def run(self) -> None: # Skip in CI — the cache DB doesn't persist between runs on ephemeral runners lang_support = current_language_support() resolver = None - if lang_support and not env_utils.is_ci(): - resolver = lang_support.create_dependency_resolver(self.args.project_root) - - if resolver is not None and lang_support is not None and file_to_funcs_to_optimize: - supported_exts = lang_support.file_extensions - source_files = [f for f in file_to_funcs_to_optimize if f.suffix in supported_exts] - with call_graph_live_display(len(source_files), project_root=self.args.project_root) as on_progress: - resolver.build_index(source_files, on_progress=on_progress) - console.rule() - call_graph_summary(resolver, file_to_funcs_to_optimize) + # CURRENTLY DISABLED: The resolver is currently not used for anything until i clean up the repo structure for python + # if lang_support and not env_utils.is_ci(): + # resolver = lang_support.create_dependency_resolver(self.args.project_root) + + # if resolver is not None and lang_support is not None and file_to_funcs_to_optimize: + # supported_exts = lang_support.file_extensions + # source_files = [f for f in file_to_funcs_to_optimize if f.suffix in supported_exts] + # with call_graph_live_display(len(source_files), project_root=self.args.project_root) as on_progress: + # resolver.build_index(source_files, on_progress=on_progress) + # console.rule() + # call_graph_summary(resolver, file_to_funcs_to_optimize) optimizations_found: int = 0 self.test_cfg.concolic_test_root_dir = Path( From d2dea5c1c13089333664daecced168b3b7970f6a Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 19 Feb 2026 02:44:38 -0500 Subject: [PATCH 42/42] fix: remove stale jedi_definition argument from FunctionSource calls --- .../languages/python/context/code_context_extractor.py | 3 --- codeflash/optimization/optimizer.py | 8 +++++++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index 5ef02d788..0e42022f6 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -255,7 +255,6 @@ def get_code_optimization_context_for_language( fully_qualified_name=helper.qualified_name, only_function_name=helper.name, source_code=helper.source_code, - jedi_definition=None, ) ) @@ -477,7 +476,6 @@ def get_function_to_optimize_as_function_source( fully_qualified_name=name.full_name, only_function_name=name.name, source_code=name.get_line_code(), - jedi_definition=name, ) except Exception as e: logger.exception(f"Error while getting function source: {e}") @@ -545,7 +543,6 @@ def get_function_sources_from_jedi( fully_qualified_name=fqn, only_function_name=func_name, source_code=definition.get_line_code(), - jedi_definition=definition, ) file_path_to_function_source[definition_path].add(function_source) function_source_list.append(function_source) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 941502a4d..3211ab59b 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -11,7 +11,13 @@ from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient from codeflash.api.cfapi import send_completion_email -from codeflash.cli_cmds.console import call_graph_live_display, call_graph_summary, console, logger, progress_bar +from codeflash.cli_cmds.console import ( # noqa: F401 + call_graph_live_display, + call_graph_summary, + console, + logger, + progress_bar, +) from codeflash.code_utils import env_utils from codeflash.code_utils.code_utils import cleanup_paths, get_run_tmp_file from codeflash.code_utils.env_utils import get_pr_number, is_pr_draft