perf: backport libcst visitor dispatch cache from codeflash-python

Cache the visitor dispatch tables that libcst rebuilds on every
MatcherDecoratableTransformer/Visitor instantiation. The tables
depend only on the class, not the instance, so caching by type is
safe. Saves ~27ms per visitor instantiation (24x faster).

Also fix pre-existing ruff F821 in cli.py (missing exit_with_message
import in process_pyproject_config).

Author: KRRT7

codeflash/benchmarking/instrument_codeflash_trace.py

@@ -4,6 +4,7 @@
 
 import libcst as cst
 
+import codeflash.code_utils._libcst_cache  # noqa: F401
 from codeflash.code_utils.formatter import sort_imports
 
 if TYPE_CHECKING:

codeflash/cli_cmds/cli.py

@@ -82,7 +82,7 @@ def process_and_validate_cmd_args(args: Namespace) -> Namespace:
 
 def process_pyproject_config(args: Namespace) -> Namespace:
     from codeflash.code_utils import env_utils
-    from codeflash.code_utils.code_utils import normalize_ignore_paths
+    from codeflash.code_utils.code_utils import exit_with_message, normalize_ignore_paths
     from codeflash.code_utils.config_parser import parse_config_file
     from codeflash.languages.test_framework import set_current_test_framework
     from codeflash.lsp.helpers import is_LSP_enabled

codeflash/code_utils/_libcst_cache.py

@@ -0,0 +1,64 @@
+"""Cache libcst visitor dispatch table construction.
+
+libcst's ``MatcherDecoratableTransformer`` and
+``MatcherDecoratableVisitor`` rebuild visitor dispatch tables on
+every instantiation by iterating ``dir(self)`` (~600 attributes)
+and calling ``getattr`` + ``inspect.ismethod`` on each.  The
+results depend only on the class, not the instance, so caching
+by ``type(obj)`` is safe.
+
+Import this module before any libcst visitors are instantiated
+to install the cache.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import libcst.matchers._visitors as _mv
+
+_visit_cache: dict[type, Any] = {}
+_leave_cache: dict[type, Any] = {}
+_matchers_cache: dict[type, Any] = {}
+
+_original_visit = _mv._gather_constructed_visit_funcs  # noqa: SLF001
+_original_leave = _mv._gather_constructed_leave_funcs  # noqa: SLF001
+_original_matchers = _mv._gather_matchers  # noqa: SLF001
+
+
+def _cached_visit(obj: object) -> Any:
+    """Return cached visit-function dispatch table for the object's class."""
+    cls = type(obj)
+    try:
+        return _visit_cache[cls]
+    except KeyError:
+        result = _original_visit(obj)
+        _visit_cache[cls] = result
+        return result
+
+
+def _cached_leave(obj: object) -> Any:
+    """Return cached leave-function dispatch table for the object's class."""
+    cls = type(obj)
+    try:
+        return _leave_cache[cls]
+    except KeyError:
+        result = _original_leave(obj)
+        _leave_cache[cls] = result
+        return result
+
+
+def _cached_matchers(obj: object) -> Any:
+    """Return cached matcher dispatch table for the object's class."""
+    cls = type(obj)
+    try:
+        return dict(_matchers_cache[cls])
+    except KeyError:
+        result = _original_matchers(obj)
+        _matchers_cache[cls] = result
+        return dict(result)
+
+
+_mv._gather_constructed_visit_funcs = _cached_visit  # noqa: SLF001
+_mv._gather_constructed_leave_funcs = _cached_leave  # noqa: SLF001
+_mv._gather_matchers = _cached_matchers  # noqa: SLF001

codeflash/code_utils/instrument_existing_tests.py

@@ -7,6 +7,7 @@
 
 import libcst as cst
 
+import codeflash.code_utils._libcst_cache  # noqa: F401
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import get_run_tmp_file, module_name_from_file_path
 from codeflash.code_utils.formatter import sort_imports

codeflash/languages/function_optimizer.py

@@ -20,6 +20,7 @@
 from rich.text import Text
 from rich.tree import Tree
 
+import codeflash.code_utils._libcst_cache  # noqa: F401
 from codeflash.api.aiservice import AiServiceClient, AIServiceRefinerRequest, LocalAiServiceClient
 from codeflash.api.cfapi import add_code_context_hash, create_staging, get_cfapi_base_urls, mark_optimization_success
 from codeflash.benchmarking.utils import process_benchmark_data

codeflash/languages/python/context/code_context_extractor.py

@@ -11,6 +11,7 @@
 
 import libcst as cst
 
+import codeflash.code_utils._libcst_cache  # noqa: F401
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import encoded_tokens_len, get_qualified_name, path_belongs_to_site_packages
 from codeflash.code_utils.config_consts import (

codeflash/languages/python/context/unused_definition_remover.py

@@ -9,6 +9,7 @@
 
 import libcst as cst
 
+import codeflash.code_utils._libcst_cache  # noqa: F401
 from codeflash.cli_cmds.console import logger
 from codeflash.languages import current_language
 from codeflash.languages.base import Language

codeflash/languages/python/static_analysis/code_extractor.py

@@ -11,6 +11,7 @@
 from libcst.codemod.visitors import AddImportsVisitor, GatherImportsVisitor, RemoveImportsVisitor
 from libcst.helpers import calculate_module_and_package
 
+import codeflash.code_utils._libcst_cache  # noqa: F401
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.config_consts import MAX_CONTEXT_LEN_REVIEW
 from codeflash.languages.base import Language

codeflash/languages/python/static_analysis/code_replacer.py

@@ -9,6 +9,7 @@
 import libcst as cst
 from libcst.metadata import PositionProvider
 
+import codeflash.code_utils._libcst_cache  # noqa: F401
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.config_parser import find_conftest_files
 from codeflash.code_utils.formatter import sort_imports

codeflash/languages/python/static_analysis/edit_generated_tests.py

@@ -10,6 +10,7 @@
 from libcst import MetadataWrapper
 from libcst.metadata import PositionProvider
 
+import codeflash.code_utils._libcst_cache  # noqa: F401
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.time_utils import format_perf, format_time
 from codeflash.models.models import GeneratedTests, GeneratedTestsList