diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 0a515076c..372d7892f 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Callable import libcst as cst +from git import Repo as GitRepo from rich.console import Group from rich.panel import Panel from rich.syntax import Syntax @@ -2222,11 +2223,11 @@ def process_review( console.print(Panel(panel_content, title="Optimization Review", border_style=display_info[1])) if raise_pr or staging_review: - data["root_dir"] = git_root_dir() + data["root_dir"] = git_root_dir(GitRepo(str(self.args.module_root), search_parent_directories=True)) if raise_pr and not staging_review and opt_review_result.review != "low": # Ensure root_dir is set for PR creation (needed for async functions that skip opt_review) if "root_dir" not in data: - data["root_dir"] = git_root_dir() + data["root_dir"] = git_root_dir(GitRepo(str(self.args.module_root), search_parent_directories=True)) data["git_remote"] = self.args.git_remote # Remove language from data dict as check_create_pr doesn't accept it pr_data = {k: v for k, v in data.items() if k != "language"} diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index b8f42010f..06540ca85 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -183,18 +183,54 @@ def get_optimizable_functions(self) -> tuple[dict[Path, list[FunctionToOptimize] """Discover functions to optimize.""" from codeflash.discovery.functions_to_optimize import get_functions_to_optimize - return get_functions_to_optimize( + # In worktree mode for git-diff discovery, file paths come from the original repo + # (via get_git_diff using cwd), but module_root/project_root have been mirrored to + # the worktree. Use the original roots for filtering so path comparisons match, + # then remap the discovered file paths to the worktree. + project_root = self.args.project_root + module_root = self.args.module_root + use_original_roots = ( + self.current_worktree and self.original_args_and_test_cfg and not self.args.all and not self.args.file + ) + if use_original_roots: + assert self.original_args_and_test_cfg is not None + original_args, _ = self.original_args_and_test_cfg + project_root = original_args.project_root + module_root = original_args.module_root + + result = get_functions_to_optimize( optimize_all=self.args.all, replay_test=self.args.replay_test, file=self.args.file, only_get_this_function=self.args.function, test_cfg=self.test_cfg, ignore_paths=self.args.ignore_paths, - project_root=self.args.project_root, - module_root=self.args.module_root, + project_root=project_root, + module_root=module_root, previous_checkpoint_functions=self.args.previous_checkpoint_functions, ) + # Remap discovered file paths from the original repo to the worktree so + # downstream optimization reads/writes happen in the worktree. + if use_original_roots: + import dataclasses + + assert self.current_worktree is not None + original_git_root = git_root_dir() + file_to_funcs, count, trace = result + remapped: dict[Path, list[FunctionToOptimize]] = {} + for file_path, funcs in file_to_funcs.items(): + new_path = mirror_path(Path(file_path), original_git_root, self.current_worktree) + remapped[new_path] = [ + dataclasses.replace( + func, file_path=mirror_path(func.file_path, original_git_root, self.current_worktree) + ) + for func in funcs + ] + return remapped, count, trace + + return result + def create_function_optimizer( self, function_to_optimize: FunctionToOptimize,