getting worktree to work in both lsp and normal cli mode

Author: mohammedahmed18

codeflash/cli_cmds/cli.py

@@ -94,6 +94,7 @@ def parse_args() -> Namespace:
         help="Path to the directory of the project, where all the pytest-benchmark tests are located.",
     )
     parser.add_argument("--no-draft", default=False, action="store_true", help="Skip optimization for draft PRs")
+    parser.add_argument("--worktree", default=False, action="store_true", help="Use worktree for optimization")
 
     args, unknown_args = parser.parse_known_args()
     sys.argv[:] = [sys.argv[0], *unknown_args]

codeflash/code_utils/coverage_utils.py

@@ -39,16 +39,20 @@ def build_fully_qualified_name(function_name: str, code_context: CodeOptimizatio
     return full_name
 
 
-def generate_candidates(source_code_path: Path) -> list[str]:
+def generate_candidates(source_code_path: Path) -> set[str]:
     """Generate all the possible candidates for coverage data based on the source code path."""
-    candidates = [source_code_path.name]
+    candidates = set()
+    candidates.add(source_code_path.name)
     current_path = source_code_path.parent
 
+    last_added = source_code_path.name
     while current_path != current_path.parent:
-        candidate_path = str(Path(current_path.name) / candidates[-1])
-        candidates.append(candidate_path)
+        candidate_path = str(Path(current_path.name) / last_added)
+        candidates.add(candidate_path)
+        last_added = candidate_path
         current_path = current_path.parent
 
+    candidates.add(str(source_code_path))
     return candidates
 
 

codeflash/code_utils/git_utils.py

@@ -9,13 +9,14 @@
 from functools import cache
 from io import StringIO
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import git
 from rich.prompt import Confirm
 from unidiff import PatchSet
 
 from codeflash.cli_cmds.console import logger
+from codeflash.code_utils.compat import codeflash_cache_dir
 from codeflash.code_utils.config_consts import N_CANDIDATES
 
 if TYPE_CHECKING:
@@ -192,3 +193,75 @@ def get_last_commit_author_if_pr_exists(repo: Repo | None = None) -> str | None:
         return None
     else:
         return last_commit.author.name
+
+
+worktree_dirs = codeflash_cache_dir / "worktrees"
+patches_dir = codeflash_cache_dir / "patches"
+
+
+def create_worktree_snapshot_commit(worktree_dir: Path, commit_message: str) -> None:
+    repository = git.Repo(worktree_dir, search_parent_directories=True)
+    repository.git.commit("-am", commit_message, "--no-verify")
+
+
+def create_detached_worktree(module_root: Path) -> Optional[Path]:
+    if not check_running_in_git_repo(module_root):
+        logger.warning("Module is not in a git repository. Skipping worktree creation.")
+        return None
+    git_root = git_root_dir()
+    current_time_str = time.strftime("%Y%m%d-%H%M%S")
+    worktree_dir = worktree_dirs / f"{git_root.name}-{current_time_str}"
+
+    result = subprocess.run(["git", "worktree", "add", "-d", str(worktree_dir)], cwd=git_root, check=True)
+    if result.returncode != 0:
+        logger.error(f"Failed to create worktree: {result.stderr}")
+        return None
+
+    print(result.stdout)
+
+    # Get uncommitted diff from the original repo
+    repository = git.Repo(module_root, search_parent_directories=True)
+    uni_diff_text = repository.git.diff(None, "HEAD", ignore_blank_lines=True, ignore_space_at_eol=True)
+
+    if not uni_diff_text.strip():
+        logger.info("No uncommitted changes to copy to worktree.")
+        return worktree_dir
+
+    # Write the diff to a temporary file
+    with tempfile.NamedTemporaryFile(mode="w+", suffix=".codeflash.patch", delete=False) as tmp_patch_file:
+        tmp_patch_file.write(uni_diff_text + "\n")  # the new line here is a must otherwise the last hunk won't be valid
+        tmp_patch_file.flush()
+
+        patch_path = Path(tmp_patch_file.name).resolve()
+
+        # Apply the patch inside the worktree
+        try:
+            subprocess.run(["git", "apply", patch_path], cwd=worktree_dir, check=True)
+            create_worktree_snapshot_commit(worktree_dir, "Initial Snapshot")
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Failed to apply patch to worktree: {e}")
+
+        return worktree_dir
+
+
+def remove_worktree(worktree_dir: Path) -> None:
+    try:
+        repository = git.Repo(worktree_dir, search_parent_directories=True)
+        repository.git.worktree("remove", "--force", worktree_dir)
+    except Exception:
+        logger.exception(f"Failed to remove worktree: {worktree_dir}")
+
+
+def create_diff_from_worktree(worktree_dir: Path, files: list[str], fto_name: str) -> Path:
+    repository = git.Repo(worktree_dir, search_parent_directories=True)
+    uni_diff_text = repository.git.diff(None, "HEAD", *files, ignore_blank_lines=True, ignore_space_at_eol=True)
+    uni_diff_text = uni_diff_text.strip()
+    if not uni_diff_text:
+        logger.warning("No changes found in worktree.")
+        return None
+    # write to patches_dir
+    patches_dir.mkdir(parents=True, exist_ok=True)
+    patch_path = patches_dir / f"{worktree_dir.name}.{fto_name}.patch"
+    with patch_path.open("w", encoding="utf8") as f:
+        f.write(uni_diff_text + "\n")
+    return patch_path

codeflash/lsp/beta.py

@@ -100,6 +100,7 @@ def initialize_function_optimization(
         f"Args set - function: {server.optimizer.args.function}, file: {server.optimizer.args.file}", "Info"
     )
 
+    server.optimizer.worktree_mode()
     optimizable_funcs, _, _ = server.optimizer.get_optimizable_functions()
     if not optimizable_funcs:
         server.show_message_log(f"No optimizable functions found for {params.functionName}", "Warning")
@@ -334,9 +335,13 @@ def perform_function_optimization(  # noqa: PLR0911
     speedup = original_code_baseline.runtime / best_optimization.runtime
 
     server.show_message_log(f"Optimization completed for {params.functionName} with {speedup:.2f}x speedup", "Info")
+    diff_patch_files = server.optimizer.patch_files
 
     # CRITICAL: Clear the function filter after optimization to prevent state corruption
     server.optimizer.args.function = None
+    server.optimizer.patch_files = []
+    server.optimizer.current_worktree = None
+    server.optimizer.cleanup_temporary_paths()
     server.show_message_log("Cleared function filter to prevent state corruption", "Info")
 
     return {
@@ -345,4 +350,5 @@ def perform_function_optimization(  # noqa: PLR0911
         "message": "Optimization completed successfully",
         "extra": f"Speedup: {speedup:.2f}x faster",
         "optimization": optimized_source,
+        "diff_patch_files": diff_patch_files,
     }

codeflash/lsp/server.py

@@ -11,6 +11,8 @@
 if TYPE_CHECKING:
     from lsprotocol.types import InitializeParams, InitializeResult
 
+    from codeflash.optimization.optimizer import Optimizer
+
 
 class CodeflashLanguageServerProtocol(LanguageServerProtocol):
     _server: CodeflashLanguageServer
@@ -44,7 +46,7 @@ def _find_pyproject_toml(self, workspace_path: str) -> Path | None:
 class CodeflashLanguageServer(LanguageServer):
     def __init__(self, *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
         super().__init__(*args, **kwargs)
-        self.optimizer = None
+        self.optimizer: Optimizer | None = None
         self.args = None
 
     def prepare_optimizer_arguments(self, config_file: Path) -> None:

codeflash/optimization/function_optimizer.py

@@ -1239,7 +1239,9 @@ def process_review(
             "coverage_message": coverage_message,
         }
 
-        if not self.args.no_pr and not self.args.staging_review:
+        raise_pr = not self.args.no_pr
+
+        if raise_pr and not self.args.staging_review:
             data["git_remote"] = self.args.git_remote
             check_create_pr(**data)
         elif self.args.staging_review:
@@ -1250,12 +1252,24 @@ def process_review(
                 trace_id=self.function_trace_id, is_optimization_found=best_optimization is not None
             )
 
-        if ((not self.args.no_pr) or not self.args.staging_review) and (
-            self.args.all or env_utils.get_pr_number() or (self.args.file and not self.args.function)
+        if raise_pr and (
+            self.args.all
+            or env_utils.get_pr_number()
+            or self.args.replay_test
+            or (self.args.file and not self.args.function)
         ):
-            self.write_code_and_helpers(
-                self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
-            )
+            self.revert_code_and_helpers(original_helper_code)
+            return
+
+        if self.args.staging_review:
+            # always revert code and helpers when staging review
+            self.revert_code_and_helpers(original_helper_code)
+            return
+
+    def revert_code_and_helpers(self, original_helper_code: dict[Path, str]) -> None:
+        self.write_code_and_helpers(
+            self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
+        )
 
     def establish_original_code_baseline(
         self,

codeflash/optimization/optimizer.py

@@ -14,6 +14,13 @@
 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
+from codeflash.code_utils.git_utils import (
+    check_running_in_git_repo,
+    create_detached_worktree,
+    create_diff_from_worktree,
+    create_worktree_snapshot_commit,
+    remove_worktree,
+)
 from codeflash.either import is_successful
 from codeflash.models.models import ValidCode
 from codeflash.telemetry.posthog_cf import ph
@@ -48,6 +55,8 @@ def __init__(self, args: Namespace) -> None:
         self.functions_checkpoint: CodeflashRunCheckpoint | None = None
         self.current_function_being_optimized: FunctionToOptimize | None = None  # current only for the LSP
         self.current_function_optimizer: FunctionOptimizer | None = None
+        self.current_worktree: Path | None = None
+        self.patch_files: list[Path] = []
 
     def run_benchmarks(
         self, file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]], num_optimizable_functions: int
@@ -252,6 +261,10 @@ def run(self) -> None:
         if self.args.no_draft and is_pr_draft():
             logger.warning("PR is in draft mode, skipping optimization")
             return
+
+        if self.args.worktree:
+            self.worktree_mode()
+
         cleanup_paths(Optimizer.find_leftover_instrumented_test_files(self.test_cfg.tests_root))
 
         function_optimizer = None
@@ -260,7 +273,6 @@ def run(self) -> None:
             file_to_funcs_to_optimize, num_optimizable_functions
         )
         optimizations_found: int = 0
-        function_iterator_count: int = 0
         if self.args.test_framework == "pytest":
             self.test_cfg.concolic_test_root_dir = Path(
                 tempfile.mkdtemp(dir=self.args.tests_root, prefix="codeflash_concolic_")
@@ -296,8 +308,8 @@ def run(self) -> None:
                     except Exception as e:
                         logger.debug(f"Could not rank functions in {original_module_path}: {e}")
 
-                for function_to_optimize in functions_to_optimize:
-                    function_iterator_count += 1
+                for i, function_to_optimize in enumerate(functions_to_optimize):
+                    function_iterator_count = i + 1
                     logger.info(
                         f"Optimizing function {function_iterator_count} of {num_optimizable_functions}: "
                         f"{function_to_optimize.qualified_name}"
@@ -327,6 +339,22 @@ def run(self) -> None:
                             )
                         if is_successful(best_optimization):
                             optimizations_found += 1
+                            if self.current_worktree:
+                                read_writable_code = best_optimization.unwrap().code_context.read_writable_code
+                                relative_file_paths = [
+                                    code_string.file_path for code_string in read_writable_code.code_strings
+                                ]
+                                patch_path = create_diff_from_worktree(
+                                    self.current_worktree,
+                                    relative_file_paths,
+                                    self.current_function_optimizer.function_to_optimize.qualified_name,
+                                )
+                                self.patch_files.append(patch_path)
+                                if i < len(functions_to_optimize) - 1:
+                                    create_worktree_snapshot_commit(
+                                        self.current_worktree,
+                                        f"Optimizing {functions_to_optimize[i + 1].qualified_name}",
+                                    )
                         else:
                             logger.warning(best_optimization.failure())
                             console.rule()
@@ -337,6 +365,10 @@ def run(self) -> None:
                             function_optimizer.cleanup_generated_files()
 
             ph("cli-optimize-run-finished", {"optimizations_found": optimizations_found})
+            if len(self.patch_files) > 0:
+                logger.info(
+                    f"Created {len(self.patch_files)} patch(es) ({[str(patch_path) for patch_path in self.patch_files]})"
+                )
             if self.functions_checkpoint:
                 self.functions_checkpoint.cleanup()
             if hasattr(self.args, "command") and self.args.command == "optimize":
@@ -388,7 +420,46 @@ def cleanup_temporary_paths(self) -> None:
         if hasattr(get_run_tmp_file, "tmpdir"):
             get_run_tmp_file.tmpdir.cleanup()
             del get_run_tmp_file.tmpdir
-        cleanup_paths([self.test_cfg.concolic_test_root_dir, self.replay_tests_dir])
+        paths_to_clean = [self.test_cfg.concolic_test_root_dir, self.replay_tests_dir]
+        if self.current_worktree:
+            remove_worktree(self.current_worktree)
+        cleanup_paths(paths_to_clean)
+
+    def worktree_mode(self) -> None:
+        if self.current_worktree:
+            return
+        project_root = self.args.project_root
+        module_root = self.args.module_root
+
+        if check_running_in_git_repo(module_root):
+            relative_module_root = module_root.relative_to(project_root)
+            relative_optimized_file = self.args.file.relative_to(project_root) if self.args.file else None
+            relative_tests_root = self.test_cfg.tests_root.relative_to(project_root)
+            relative_benchmarks_root = (
+                self.args.benchmarks_root.relative_to(project_root) if self.args.benchmarks_root else None
+            )
+
+            worktree_dir = create_detached_worktree(module_root)
+            if worktree_dir is None:
+                logger.warning("Failed to create worktree. Skipping optimization.")
+                return
+            self.current_worktree = worktree_dir
+            # TODO: use a helper function to mutate self.args and self.test_cfg
+            self.args.module_root = worktree_dir / relative_module_root
+            self.args.project_root = worktree_dir
+            self.args.test_project_root = worktree_dir
+            self.args.tests_root = worktree_dir / relative_tests_root
+            if relative_benchmarks_root:
+                self.args.benchmarks_root = worktree_dir / relative_benchmarks_root
+
+            self.test_cfg.project_root_path = worktree_dir
+            self.test_cfg.tests_project_rootdir = worktree_dir
+            self.test_cfg.tests_root = worktree_dir / relative_tests_root
+            if relative_benchmarks_root:
+                self.test_cfg.benchmark_tests_root = worktree_dir / relative_benchmarks_root
+
+            if relative_optimized_file is not None:
+                self.args.file = worktree_dir / relative_optimized_file
 
 
 def run_with_args(args: Namespace) -> None:

codeflash/telemetry/sentry.py

@@ -14,7 +14,7 @@ def init_sentry(enabled: bool = False, exclude_errors: bool = False) -> None:  #
         )
 
         sentry_sdk.init(
-            dsn="https://4b9a1902f9361b48c04376df6483bc96@o4506833230561280.ingest.sentry.io/4506833262477312",
+            dsn="",
             integrations=[sentry_logging],
             # Set traces_sample_rate to 1.0 to capture 100%
             # of transactions for performance monitoring.

tests/test_code_utils.py

@@ -357,7 +357,7 @@ def test_cleanup_paths(multiple_existing_and_non_existing_files: list[Path]) ->
 
 def test_generate_candidates() -> None:
     source_code_path = Path("/Users/krrt7/Desktop/work/codeflash/cli/codeflash/code_utils/coverage_utils.py")
-    expected_candidates = [
+    expected_candidates = {
         "coverage_utils.py",
         "code_utils/coverage_utils.py",
         "codeflash/code_utils/coverage_utils.py",
@@ -367,7 +367,8 @@ def test_generate_candidates() -> None:
         "Desktop/work/codeflash/cli/codeflash/code_utils/coverage_utils.py",
         "krrt7/Desktop/work/codeflash/cli/codeflash/code_utils/coverage_utils.py",
         "Users/krrt7/Desktop/work/codeflash/cli/codeflash/code_utils/coverage_utils.py",
-    ]
+        "/Users/krrt7/Desktop/work/codeflash/cli/codeflash/code_utils/coverage_utils.py"
+    }
     assert generate_candidates(source_code_path) == expected_candidates