Merge remote-tracking branch 'origin/multi-language' into multi-language

# Conflicts:
#	codeflash/verification/verification_utils.py

Author: misrasaurabh1

codeflash/context/code_context_extractor.py

@@ -207,6 +207,154 @@ def get_code_optimization_context(
         preexisting_objects=preexisting_objects,
     )
 
+def get_code_optimization_context_for_language(
+    function_to_optimize: FunctionToOptimize,
+    project_root_path: Path,
+    optim_token_limit: int = OPTIMIZATION_CONTEXT_TOKEN_LIMIT,
+    testgen_token_limit: int = TESTGEN_CONTEXT_TOKEN_LIMIT,
+) -> CodeOptimizationContext:
+    """Extract code optimization context for non-Python languages.
+
+    Uses the language support abstraction to extract code context and converts
+    it to the CodeOptimizationContext format expected by the pipeline.
+
+    This function supports multi-file context extraction, grouping helpers by file
+    and creating proper CodeStringsMarkdown with file paths for multi-file replacement.
+
+    Args:
+        function_to_optimize: The function to extract context for.
+        project_root_path: Root of the project.
+        optim_token_limit: Token limit for optimization context.
+        testgen_token_limit: Token limit for testgen context.
+
+    Returns:
+        CodeOptimizationContext with target code and dependencies.
+
+    """
+    from codeflash.languages import get_language_support
+    from codeflash.languages.base import FunctionInfo, ParentInfo
+
+    # Get language support for this function
+    language = Language(function_to_optimize.language)
+    lang_support = get_language_support(language)
+
+    # Convert FunctionToOptimize to FunctionInfo for language support
+    parents = tuple(ParentInfo(name=p.name, type=p.type) for p in function_to_optimize.parents)
+    func_info = FunctionInfo(
+        name=function_to_optimize.function_name,
+        file_path=function_to_optimize.file_path,
+        start_line=function_to_optimize.starting_line or 1,
+        end_line=function_to_optimize.ending_line or 1,
+        parents=parents,
+        is_async=function_to_optimize.is_async,
+        is_method=len(function_to_optimize.parents) > 0,
+        language=language,
+    )
+
+    # Extract code context using language support
+    code_context = lang_support.extract_code_context(func_info, project_root_path, project_root_path)
+
+    # Build imports string if available
+    imports_code = "\n".join(code_context.imports) if code_context.imports else ""
+
+    # Get relative path for target file
+    try:
+        target_relative_path = function_to_optimize.file_path.resolve().relative_to(project_root_path.resolve())
+    except ValueError:
+        target_relative_path = function_to_optimize.file_path
+
+    # Group helpers by file path
+    helpers_by_file: dict[Path, list] = defaultdict(list)
+    helper_function_sources = []
+
+    for helper in code_context.helper_functions:
+        helpers_by_file[helper.file_path].append(helper)
+
+        # Convert to FunctionSource for pipeline compatibility
+        helper_function_sources.append(
+            FunctionSource(
+                file_path=helper.file_path,
+                qualified_name=helper.qualified_name,
+                fully_qualified_name=helper.qualified_name,
+                only_function_name=helper.name,
+                source_code=helper.source_code,
+                jedi_definition=None,
+            )
+        )
+
+    # Build read-writable code (target file + same-file helpers + global variables)
+    read_writable_code_strings = []
+
+    # Combine target code with same-file helpers
+    target_file_code = code_context.target_code
+    same_file_helpers = helpers_by_file.get(function_to_optimize.file_path, [])
+    if same_file_helpers:
+        helper_code = "\n\n".join(h.source_code for h in same_file_helpers)
+        target_file_code = target_file_code + "\n\n" + helper_code
+
+    # Add global variables (module-level declarations) referenced by the function and helpers
+    # These should be included in read-writable context so AI can modify them if needed
+    if code_context.read_only_context:
+        target_file_code = code_context.read_only_context + "\n\n" + target_file_code
+
+    # Add imports to target file code
+    if imports_code:
+        target_file_code = imports_code + "\n\n" + target_file_code
+
+    read_writable_code_strings.append(
+        CodeString(code=target_file_code, file_path=target_relative_path, language=function_to_optimize.language)
+    )
+
+    # Add helper files (cross-file helpers)
+    for file_path, file_helpers in helpers_by_file.items():
+        if file_path == function_to_optimize.file_path:
+            continue  # Already included in target file
+
+        try:
+            helper_relative_path = file_path.resolve().relative_to(project_root_path.resolve())
+        except ValueError:
+            helper_relative_path = file_path
+
+        # Combine all helpers from this file
+        combined_helper_code = "\n\n".join(h.source_code for h in file_helpers)
+
+        read_writable_code_strings.append(
+            CodeString(
+                code=combined_helper_code, file_path=helper_relative_path, language=function_to_optimize.language
+            )
+        )
+
+    read_writable_code = CodeStringsMarkdown(
+        code_strings=read_writable_code_strings, language=function_to_optimize.language
+    )
+
+    # Build testgen context (same as read_writable for non-Python)
+    testgen_context = CodeStringsMarkdown(
+        code_strings=read_writable_code_strings.copy(), language=function_to_optimize.language
+    )
+
+    # Check token limits
+    read_writable_tokens = encoded_tokens_len(read_writable_code.markdown)
+    if read_writable_tokens > optim_token_limit:
+        raise ValueError("Read-writable code has exceeded token limit, cannot proceed")
+
+    testgen_tokens = encoded_tokens_len(testgen_context.markdown)
+    if testgen_tokens > testgen_token_limit:
+        raise ValueError("Testgen code context has exceeded token limit, cannot proceed")
+
+    # Generate code hash from all read-writable code
+    code_hash = hashlib.sha256(read_writable_code.flat.encode("utf-8")).hexdigest()
+
+    return CodeOptimizationContext(
+        testgen_context=testgen_context,
+        read_writable_code=read_writable_code,
+        # Global variables are now included in read-writable code, so don't duplicate in read-only
+        read_only_context_code="",
+        hashing_code_context=read_writable_code.flat,
+        hashing_code_context_hash=code_hash,
+        helper_functions=helper_function_sources,
+        preexisting_objects=set(),  # Not implemented for non-Python yet
+    )
 
 def extract_code_markdown_context_from_files(
     helpers_of_fto: dict[Path, set[FunctionSource]],

codeflash/discovery/functions_to_optimize.py

@@ -155,7 +155,7 @@ class FunctionToOptimize:
     starting_line: Optional[int] = None
     ending_line: Optional[int] = None
     starting_col: Optional[int] = None  # Column offset for precise location
-    ending_col: Optional[int] = None    # Column offset for precise location
+    ending_col: Optional[int] = None  # Column offset for precise location
     is_async: bool = False
     language: str = "python"  # Language identifier for multi-language support
 
@@ -186,11 +186,14 @@ def qualified_name_with_modules_from_root(self, project_root_path: Path) -> str:
 # =============================================================================
 
 
-def get_files_for_language(module_root_path: Path, language: Language | None = None) -> list[Path]:
+def get_files_for_language(
+    module_root_path: Path, ignore_paths: list[Path], language: Language | None = None
+) -> list[Path]:
     """Get all source files for supported languages.
 
     Args:
         module_root_path: Root path to search for source files.
+        ignore_paths: List of paths to ignore (can be files or directories).
         language: Optional specific language to filter for. If None, includes all supported languages.
 
     Returns:
@@ -206,7 +209,10 @@ def get_files_for_language(module_root_path: Path, language: Language | None = N
     files = []
     for ext in extensions:
         pattern = f"*{ext}"
-        files.extend(module_root_path.rglob(pattern))
+        for file_path in module_root_path.rglob(pattern):
+            if any(file_path.is_relative_to(ignore_path) for ignore_path in ignore_paths):
+                continue
+            files.append(file_path)
     return files
 
 
@@ -289,7 +295,7 @@ def get_functions_to_optimize(
         if optimize_all:
             logger.info("!lsp|Finding all functions in the module '%s'…", optimize_all)
             console.rule()
-            functions = get_all_files_and_functions(Path(optimize_all))
+            functions = get_all_files_and_functions(Path(optimize_all), ignore_paths)
         elif replay_test:
             functions, trace_file_path = get_all_replay_test_functions(
                 replay_test=replay_test, test_cfg=test_cfg, project_root_path=project_root
@@ -452,20 +458,21 @@ def get_functions_within_lines(modified_lines: dict[str, list[int]]) -> dict[str
 
 
 def get_all_files_and_functions(
-    module_root_path: Path, language: Language | None = None
+    module_root_path: Path, ignore_paths: list[Path], language: Language | None = None
 ) -> dict[str, list[FunctionToOptimize]]:
     """Get all optimizable functions from files in the module root.
 
     Args:
         module_root_path: Root path to search for source files.
+        ignore_paths: List of paths to ignore.
         language: Optional specific language to filter for. If None, includes all supported languages.
 
     Returns:
         Dictionary mapping file paths to lists of FunctionToOptimize.
 
     """
     functions: dict[str, list[FunctionToOptimize]] = {}
-    for file_path in get_files_for_language(module_root_path, language):
+    for file_path in get_files_for_language(module_root_path, ignore_paths, language):
         # Find all the functions in the file
         functions.update(find_all_functions_in_file(file_path).items())
     # Randomize the order of the files to optimize to avoid optimizing the same file in the same order every time.

codeflash/languages/javascript/edit_tests.py

@@ -7,7 +7,6 @@
 from __future__ import annotations
 
 import re
-from pathlib import Path
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.time_utils import format_perf, format_time
@@ -23,6 +22,7 @@ def format_runtime_comment(original_time: int, optimized_time: int) -> str:
 
     Returns:
         Formatted comment string with // prefix.
+
     """
     perf_gain = format_perf(
         abs(performance_gain(original_runtime_ns=original_time, optimized_runtime_ns=optimized_time) * 100)
@@ -31,11 +31,7 @@ def format_runtime_comment(original_time: int, optimized_time: int) -> str:
     return f"// {format_time(original_time)} -> {format_time(optimized_time)} ({perf_gain}% {status})"
 
 
-def add_runtime_comments(
-    source: str,
-    original_runtimes: dict[str, int],
-    optimized_runtimes: dict[str, int],
-) -> str:
+def add_runtime_comments(source: str, original_runtimes: dict[str, int], optimized_runtimes: dict[str, int]) -> str:
     """Add runtime comments to JavaScript test source code.
 
     For JavaScript, we match timing data by test function name and add comments
@@ -48,6 +44,7 @@ def add_runtime_comments(
 
     Returns:
         Source code with runtime comments added.
+
     """
     logger.debug(f"[js-annotations] original_runtimes has {len(original_runtimes)} entries")
     logger.debug(f"[js-annotations] optimized_runtimes has {len(optimized_runtimes)} entries")
@@ -144,6 +141,7 @@ def remove_test_functions(source: str, functions_to_remove: list[str]) -> str:
 
     Returns:
         Source code with specified functions removed.
+
     """
     if not functions_to_remove:
         return source
@@ -152,8 +150,7 @@ def remove_test_functions(source: str, functions_to_remove: list[str]) -> str:
         # Pattern to match test('name', ...) or it('name', ...) blocks
         # This handles nested callbacks and multi-line test bodies
         test_pattern = re.compile(
-            rf"(?:test|it)\s*\(\s*['\"]" + re.escape(func_name) + rf"['\"].*?\)\s*;?\s*\n?",
-            re.DOTALL,
+            r"(?:test|it)\s*\(\s*['\"]" + re.escape(func_name) + r"['\"].*?\)\s*;?\s*\n?", re.DOTALL
         )
 
         # Try to find and remove matching test blocks
@@ -180,6 +177,7 @@ def _find_block_end(source: str, start: int) -> int:
 
     Returns:
         Position after the closing brace, or start if not found.
+
     """
     # Find the opening brace
     brace_pos = source.find("{", start)
@@ -230,21 +228,3 @@ def _find_block_end(source: str, start: int) -> int:
         i += 1
 
     return start
-
-
-def get_comment_prefix() -> str:
-    """Get the comment prefix for JavaScript.
-
-    Returns:
-        The JavaScript single-line comment prefix.
-    """
-    return "//"
-
-
-def get_test_file_suffix() -> str:
-    """Get the test file suffix for JavaScript.
-
-    Returns:
-        The Jest test file suffix.
-    """
-    return ".test.js"

codeflash/languages/javascript/instrument.py

@@ -119,7 +119,7 @@ def _instrument_js_test_code(code: str, func_name: str, test_file_path: str, mod
     """
     # Add codeflash helper import if not already present
     # Support both npm package (codeflash) and legacy local file (codeflash-jest-helper)
-    has_codeflash_import = "codeflash" in code or "codeflash-jest-helper" in code
+    has_codeflash_import = "codeflash" in code
     if not has_codeflash_import:
         # Detect module system: ESM uses "import ... from", CommonJS uses "require()"
         is_esm = bool(re.search(r"^\s*import\s+.+\s+from\s+['\"]", code, re.MULTILINE))
@@ -128,11 +128,7 @@ def _instrument_js_test_code(code: str, func_name: str, test_file_path: str, mod
             # ESM: Use import statement at the top of the file (after any other imports)
             helper_import = "import codeflash from 'codeflash';\n"
             # Find the last import statement to add after
-            import_matches = list(re.finditer(
-                r"^import\s+.+\s+from\s+['\"][^'\"]+['\"]\s*;?\s*\n",
-                code,
-                re.MULTILINE,
-            ))
+            import_matches = list(re.finditer(r"^import\s+.+\s+from\s+['\"][^'\"]+['\"]\s*;?\s*\n", code, re.MULTILINE))
             if import_matches:
                 # Add after the last import
                 last_import = import_matches[-1]
@@ -145,11 +141,7 @@ def _instrument_js_test_code(code: str, func_name: str, test_file_path: str, mod
             # CommonJS: Use require statement
             helper_require = "const codeflash = require('codeflash');\n"
             # Find the first require statement to add after
-            import_match = re.search(
-                r"^((?:const|let|var)\s+.+?require\([^)]+\).*;?\s*\n)",
-                code,
-                re.MULTILINE,
-            )
+            import_match = re.search(r"^((?:const|let|var)\s+.+?require\([^)]+\).*;?\s*\n)", code, re.MULTILINE)
             if import_match:
                 insert_pos = import_match.end()
                 code = code[:insert_pos] + helper_require + code[insert_pos:]