Merge branch 'main' into feat/detached-worktrees

Author: Saga4

codeflash/api/cfapi.py

@@ -125,6 +125,8 @@ def suggest_changes(
     generated_tests: str,
     trace_id: str,
     coverage_message: str,
+    replay_tests: str = "",
+    concolic_tests: str = "",
 ) -> Response:
     """Suggest changes to a pull request.
 
@@ -148,6 +150,8 @@ def suggest_changes(
         "generatedTests": generated_tests,
         "traceId": trace_id,
         "coverage_message": coverage_message,
+        "replayTests": replay_tests,
+        "concolicTests": concolic_tests,
     }
     return make_cfapi_request(endpoint="/suggest-pr-changes", method="POST", payload=payload)
 
@@ -162,6 +166,8 @@ def create_pr(
     generated_tests: str,
     trace_id: str,
     coverage_message: str,
+    replay_tests: str = "",
+    concolic_tests: str = "",
 ) -> Response:
     """Create a pull request, targeting the specified branch. (usually 'main').
 
@@ -184,6 +190,8 @@ def create_pr(
         "generatedTests": generated_tests,
         "traceId": trace_id,
         "coverage_message": coverage_message,
+        "replayTests": replay_tests,
+        "concolicTests": concolic_tests,
     }
     return make_cfapi_request(endpoint="/create-pr", method="POST", payload=payload)
 

codeflash/code_utils/code_extractor.py

@@ -195,6 +195,64 @@ def visit_SimpleStatementLine(self, node: cst.SimpleStatementLine) -> None:
                 self.last_import_line = self.current_line
 
 
+class ConditionalImportCollector(cst.CSTVisitor):
+    """Collect imports inside top-level conditionals (e.g., if TYPE_CHECKING, try/except)."""
+
+    def __init__(self) -> None:
+        self.imports: set[str] = set()
+        self.depth = 0  # top-level
+
+    def get_full_dotted_name(self, expr: cst.BaseExpression) -> str:
+        if isinstance(expr, cst.Name):
+            return expr.value
+        if isinstance(expr, cst.Attribute):
+            return f"{self.get_full_dotted_name(expr.value)}.{expr.attr.value}"
+        return ""
+
+    def _collect_imports_from_block(self, block: cst.IndentedBlock) -> None:
+        for statement in block.body:
+            if isinstance(statement, cst.SimpleStatementLine):
+                for child in statement.body:
+                    if isinstance(child, cst.Import):
+                        for alias in child.names:
+                            module = self.get_full_dotted_name(alias.name)
+                            asname = alias.asname.name.value if alias.asname else alias.name.value
+                            self.imports.add(module if module == asname else f"{module}.{asname}")
+
+                    elif isinstance(child, cst.ImportFrom):
+                        if child.module is None:
+                            continue
+                        module = self.get_full_dotted_name(child.module)
+                        for alias in child.names:
+                            if isinstance(alias, cst.ImportAlias):
+                                name = alias.name.value
+                                asname = alias.asname.name.value if alias.asname else name
+                                self.imports.add(f"{module}.{asname}")
+
+    def visit_Module(self, node: cst.Module) -> None:
+        self.depth = 0
+
+    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
+        self.depth += 1
+
+    def leave_FunctionDef(self, node: cst.FunctionDef) -> None:
+        self.depth -= 1
+
+    def visit_ClassDef(self, node: cst.ClassDef) -> None:
+        self.depth += 1
+
+    def leave_ClassDef(self, node: cst.ClassDef) -> None:
+        self.depth -= 1
+
+    def visit_If(self, node: cst.If) -> None:
+        if self.depth == 0:
+            self._collect_imports_from_block(node.body)
+
+    def visit_Try(self, node: cst.Try) -> None:
+        if self.depth == 0:
+            self._collect_imports_from_block(node.body)
+
+
 class ImportInserter(cst.CSTTransformer):
     """Transformer that inserts global statements after the last import."""
 
@@ -329,8 +387,19 @@ def add_needed_imports_from_module(
     except Exception as e:
         logger.error(f"Error parsing source module code: {e}")
         return dst_module_code
+
+    cond_import_collector = ConditionalImportCollector()
+    try:
+        parsed_dst_module = cst.parse_module(dst_module_code)
+        parsed_dst_module.visit(cond_import_collector)
+    except cst.ParserSyntaxError as e:
+        logger.exception(f"Syntax error in destination module code: {e}")
+        return dst_module_code  # Return the original code if there's a syntax error
+
     try:
         for mod in gatherer.module_imports:
+            if mod in cond_import_collector.imports:
+                continue
             AddImportsVisitor.add_needed_import(dst_context, mod)
             RemoveImportsVisitor.remove_unused_import(dst_context, mod)
         for mod, obj_seq in gatherer.object_mapping.items():
@@ -339,28 +408,29 @@ def add_needed_imports_from_module(
                     f"{mod}.{obj}" in helper_functions_fqn or dst_context.full_module_name == mod  # avoid circular deps
                 ):
                     continue  # Skip adding imports for helper functions already in the context
+                if f"{mod}.{obj}" in cond_import_collector.imports:
+                    continue
                 AddImportsVisitor.add_needed_import(dst_context, mod, obj)
                 RemoveImportsVisitor.remove_unused_import(dst_context, mod, obj)
     except Exception as e:
         logger.exception(f"Error adding imports to destination module code: {e}")
         return dst_module_code
     for mod, asname in gatherer.module_aliases.items():
+        if f"{mod}.{asname}" in cond_import_collector.imports:
+            continue
         AddImportsVisitor.add_needed_import(dst_context, mod, asname=asname)
         RemoveImportsVisitor.remove_unused_import(dst_context, mod, asname=asname)
     for mod, alias_pairs in gatherer.alias_mapping.items():
         for alias_pair in alias_pairs:
             if f"{mod}.{alias_pair[0]}" in helper_functions_fqn:
                 continue
+            if f"{mod}.{alias_pair[1]}" in cond_import_collector.imports:
+                continue
             AddImportsVisitor.add_needed_import(dst_context, mod, alias_pair[0], asname=alias_pair[1])
             RemoveImportsVisitor.remove_unused_import(dst_context, mod, alias_pair[0], asname=alias_pair[1])
 
     try:
-        parsed_module = cst.parse_module(dst_module_code)
-    except cst.ParserSyntaxError as e:
-        logger.exception(f"Syntax error in destination module code: {e}")
-        return dst_module_code  # Return the original code if there's a syntax error
-    try:
-        transformed_module = AddImportsVisitor(dst_context).transform_module(parsed_module)
+        transformed_module = AddImportsVisitor(dst_context).transform_module(parsed_dst_module)
         transformed_module = RemoveImportsVisitor(dst_context).transform_module(transformed_module)
         return transformed_module.code.lstrip("\n")
     except Exception as e:

codeflash/optimization/function_optimizer.py

@@ -1196,7 +1196,7 @@ def process_review(
         if concolic_test_str:
             generated_tests_str += "\n#------------------------------------------------\n" + concolic_test_str
 
-        existing_tests = existing_tests_source_for(
+        existing_tests, replay_tests, concolic_tests = existing_tests_source_for(
             self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root),
             function_to_all_tests,
             test_cfg=self.test_cfg,
@@ -1237,6 +1237,8 @@ def process_review(
             if self.experiment_id
             else self.function_trace_id,
             "coverage_message": coverage_message,
+            "replay_tests": replay_tests,
+            "concolic_tests": concolic_tests,
         }
 
         raise_pr = not self.args.no_pr

codeflash/result/create_pr.py

@@ -34,12 +34,16 @@ def existing_tests_source_for(
     test_cfg: TestConfig,
     original_runtimes_all: dict[InvocationId, list[int]],
     optimized_runtimes_all: dict[InvocationId, list[int]],
-) -> str:
+) -> tuple[str, str, str]:
     test_files = function_to_tests.get(function_qualified_name_with_modules_from_root)
     if not test_files:
-        return ""
-    output: str = ""
-    rows = []
+        return "", "", ""
+    output_existing: str = ""
+    output_concolic: str = ""
+    output_replay: str = ""
+    rows_existing = []
+    rows_concolic = []
+    rows_replay = []
     headers = ["Test File::Test Function", "Original ⏱️", "Optimized ⏱️", "Speedup"]
     tests_root = test_cfg.tests_root
     original_tests_to_runtimes: dict[Path, dict[str, int]] = {}
@@ -99,28 +103,79 @@ def existing_tests_source_for(
                     * 100
                 )
                 if greater:
-                    rows.append(
+                    if "__replay_test_" in str(print_filename):
+                        rows_replay.append(
+                            [
+                                f"`{print_filename}::{qualified_name}`",
+                                f"{print_original_runtime}",
+                                f"{print_optimized_runtime}",
+                                f"{perf_gain}%⚠️",
+                            ]
+                        )
+                    elif "codeflash_concolic" in str(print_filename):
+                        rows_concolic.append(
+                            [
+                                f"`{print_filename}::{qualified_name}`",
+                                f"{print_original_runtime}",
+                                f"{print_optimized_runtime}",
+                                f"{perf_gain}%⚠️",
+                            ]
+                        )
+                    else:
+                        rows_existing.append(
+                            [
+                                f"`{print_filename}::{qualified_name}`",
+                                f"{print_original_runtime}",
+                                f"{print_optimized_runtime}",
+                                f"{perf_gain}%⚠️",
+                            ]
+                        )
+                elif "__replay_test_" in str(print_filename):
+                    rows_replay.append(
                         [
                             f"`{print_filename}::{qualified_name}`",
                             f"{print_original_runtime}",
                             f"{print_optimized_runtime}",
-                            f"⚠️{perf_gain}%",
+                            f"{perf_gain}%✅",
+                        ]
+                    )
+                elif "codeflash_concolic" in str(print_filename):
+                    rows_concolic.append(
+                        [
+                            f"`{print_filename}::{qualified_name}`",
+                            f"{print_original_runtime}",
+                            f"{print_optimized_runtime}",
+                            f"{perf_gain}%✅",
                         ]
                     )
                 else:
-                    rows.append(
+                    rows_existing.append(
                         [
                             f"`{print_filename}::{qualified_name}`",
                             f"{print_original_runtime}",
                             f"{print_optimized_runtime}",
-                            f"✅{perf_gain}%",
+                            f"{perf_gain}%✅",
                         ]
                     )
-    output += tabulate(  # type: ignore[no-untyped-call]
-        headers=headers, tabular_data=rows, tablefmt="pipe", colglobalalign=None, preserve_whitespace=True
+    output_existing += tabulate(  # type: ignore[no-untyped-call]
+        headers=headers, tabular_data=rows_existing, tablefmt="pipe", colglobalalign=None, preserve_whitespace=True
+    )
+    output_existing += "\n"
+    if len(rows_existing) == 0:
+        output_existing = ""
+    output_concolic += tabulate(  # type: ignore[no-untyped-call]
+        headers=headers, tabular_data=rows_concolic, tablefmt="pipe", colglobalalign=None, preserve_whitespace=True
+    )
+    output_concolic += "\n"
+    if len(rows_concolic) == 0:
+        output_concolic = ""
+    output_replay += tabulate(  # type: ignore[no-untyped-call]
+        headers=headers, tabular_data=rows_replay, tablefmt="pipe", colglobalalign=None, preserve_whitespace=True
     )
-    output += "\n"
-    return output
+    output_replay += "\n"
+    if len(rows_replay) == 0:
+        output_replay = ""
+    return output_existing, output_replay, output_concolic
 
 
 def check_create_pr(
@@ -131,6 +186,8 @@ def check_create_pr(
     generated_original_test_source: str,
     function_trace_id: str,
     coverage_message: str,
+    replay_tests: str,
+    concolic_tests: str,
     git_remote: Optional[str] = None,
 ) -> None:
     pr_number: Optional[int] = env_utils.get_pr_number()
@@ -171,6 +228,8 @@ def check_create_pr(
             generated_tests=generated_original_test_source,
             trace_id=function_trace_id,
             coverage_message=coverage_message,
+            replay_tests=replay_tests,
+            concolic_tests=concolic_tests,
         )
         if response.ok:
             logger.info(f"Suggestions were successfully made to PR #{pr_number}")
@@ -218,6 +277,8 @@ def check_create_pr(
             generated_tests=generated_original_test_source,
             trace_id=function_trace_id,
             coverage_message=coverage_message,
+            replay_tests=replay_tests,
+            concolic_tests=concolic_tests,
         )
         if response.ok:
             pr_id = response.text