feat: optimize current diff

Author: mohammedahmed18

codeflash/discovery/functions_to_optimize.py

@@ -208,7 +208,7 @@ def get_functions_to_optimize(
             logger.info("Finding all functions modified in the current git diff ...")
             console.rule()
             ph("cli-optimizing-git-diff")
-            functions = get_functions_within_git_diff()
+            functions = get_functions_within_git_diff(uncommitted_changes=False)
         filtered_modified_functions, functions_count = filter_functions(
             functions, test_cfg.tests_root, ignore_paths, project_root, module_root, previous_checkpoint_functions
         )
@@ -224,8 +224,8 @@ def get_functions_to_optimize(
         return filtered_modified_functions, functions_count, trace_file_path
 
 
-def get_functions_within_git_diff() -> dict[str, list[FunctionToOptimize]]:
-    modified_lines: dict[str, list[int]] = get_git_diff(uncommitted_changes=False)
+def get_functions_within_git_diff(uncommitted_changes: bool) -> dict[str, list[FunctionToOptimize]]:  # noqa: FBT001
+    modified_lines: dict[str, list[int]] = get_git_diff(uncommitted_changes=uncommitted_changes)
     modified_functions: dict[str, list[FunctionToOptimize]] = {}
     for path_str, lines_in_file in modified_lines.items():
         path = Path(path_str)

codeflash/lsp/beta.py

@@ -11,6 +11,7 @@
 from codeflash.api.cfapi import get_codeflash_api_key, get_user_id
 from codeflash.code_utils.git_utils import create_diff_patch_from_worktree
 from codeflash.code_utils.shell_utils import save_api_key_to_rc
+from codeflash.discovery.functions_to_optimize import filter_functions, get_functions_within_git_diff
 from codeflash.either import is_successful
 from codeflash.lsp.server import CodeflashLanguageServer, CodeflashLanguageServerProtocol
 from codeflash.result.explanation import Explanation
@@ -40,6 +41,23 @@ class ProvideApiKeyParams:
 server = CodeflashLanguageServer("codeflash-language-server", "v1.0", protocol_cls=CodeflashLanguageServerProtocol)
 
 
+@server.feature("getOptimizableFunctionsInCurrentDiff")
+def get_functions_in_current_git_diff(
+    server: CodeflashLanguageServer, _params: OptimizableFunctionsParams
+) -> dict[str, str | list[str]]:
+    functions = get_functions_within_git_diff(uncommitted_changes=True)
+    file_to_funcs_to_optimize, _ = filter_functions(
+        modified_functions=functions,
+        tests_root=server.optimizer.test_cfg.tests_root,
+        ignore_paths=[],
+        project_root=server.optimizer.args.project_root,
+        module_root=server.optimizer.args.module_root,
+        previous_checkpoint_functions={},
+    )
+    qualified_names: list[str] = [func.qualified_name for funcs in file_to_funcs_to_optimize.values() for func in funcs]
+    return {"functions": qualified_names, "status": "success"}
+
+
 @server.feature("getOptimizableFunctions")
 def get_optimizable_functions(
     server: CodeflashLanguageServer, params: OptimizableFunctionsParams