[FEAT] Code-repair for candidates with unmatched test results

codeflash/discovery/functions_to_optimize.py

@@ -306,25 +306,43 @@ def levenshtein_distance(s1: str, s2: str) -> int:
     len1 = len(s1)
     len2 = len(s2)
     # Use a preallocated list instead of creating a new list every iteration
+
+    # Early exit for empty string cases
+    if len1 == 0:
+        return len2
+    if len2 == 0:
+        return len1
+
+    # Convert strings to lists for fast indexed access
+    s1_list = list(s1)
+    s2_list = list(s2)
+
+    # Preallocate and reuse arrays; avoid creating new ones every iteration
     previous = list(range(len1 + 1))
     current = [0] * (len1 + 1)
 
     for index2 in range(len2):
-        char2 = s2[index2]
+        char2 = s2_list[index2]
         current[0] = index2 + 1
+
+        # Remove redundant intermediate assignments for better cache locality
+        prev = previous
+        curr = current
+        s1_chars = s1_list
+        # Use local variables for frequently accessed values
         for index1 in range(len1):
-            char1 = s1[index1]
-            if char1 == char2:
-                current[index1 + 1] = previous[index1]
+            # Unrolling char1 assignment and equality check
+            if s1_chars[index1] == char2:
+                curr[index1 + 1] = prev[index1]
             else:
-                # Fast min calculation without tuple construct
-                a = previous[index1]
-                b = previous[index1 + 1]
-                c = current[index1]
-                min_val = min(b, a)
-                min_val = min(c, min_val)
-                current[index1 + 1] = 1 + min_val
-        # Swap references instead of copying
+                x = prev[index1]
+                y = prev[index1 + 1]
+                z = curr[index1]
+                min_xy = min(x, y)
+                min_xyz = min(z, min_xy)
+                curr[index1 + 1] = 1 + min_xyz
+
+        # Swap references rather than copying data
         previous, current = current, previous
     return previous[len1]
 

codeflash/models/models.py

@@ -3,6 +3,7 @@
 from collections import Counter, defaultdict
 from typing import TYPE_CHECKING
 
+import libcst as cst
 from rich.tree import Tree
 
 from codeflash.cli_cmds.console import DEBUG_MODE, lsp_log
@@ -505,6 +506,31 @@ def id(self) -> str:
             f"{self.function_getting_tested}:{self.iteration_id}"
         )
 
+    def find_func_in_class(self, class_node: cst.ClassDef, func_name: str) -> Optional[cst.FunctionDef]:
+        for stmt in class_node.body.body:
+            if isinstance(stmt, cst.FunctionDef) and stmt.name.value == func_name:
+                return stmt
+        return None
+
+    def get_src_code(self, test_path: Path) -> Optional[str]:
+        test_src = test_path.read_text(encoding="utf-8")
+        module_node = cst.parse_module(test_src)
+
+        if self.test_class_name:
+            for stmt in module_node.body:
+                if isinstance(stmt, cst.ClassDef) and stmt.name.value == self.test_class_name:
+                    func_node = self.find_func_in_class(stmt, self.test_function_name)
+                    if func_node:
+                        return module_node.code_for_node(func_node).strip()
+            # class not found
+            return None
+
+        # Otherwise, look for a top level function
+        for stmt in module_node.body:
+            if isinstance(stmt, cst.FunctionDef) and stmt.name.value == self.test_function_name:
+                return module_node.code_for_node(stmt).strip()
+        return None
+
     @staticmethod
     def from_str_id(string_id: str, iteration_id: str | None = None) -> InvocationId:
         components = string_id.split(":")
@@ -549,7 +575,10 @@ class TestResults(BaseModel):  # noqa: PLW1641
     # also we don't support deletion of test results elements - caution is advised
     test_results: list[FunctionTestInvocation] = []
     test_result_idx: dict[str, int] = {}
+
     perf_stdout: Optional[str] = None
+    # mapping between test function name and stdout failure message
+    test_failures: Optional[dict[str, str]] = None
 
     def add(self, function_test_invocation: FunctionTestInvocation) -> None:
         unique_id = function_test_invocation.unique_invocation_loop_id

codeflash/optimization/function_optimizer.py

@@ -1752,6 +1752,11 @@ def establish_original_code_baseline(
             )
         )
 
+    def get_results_not_matched_error(self) -> Failure:
+        logger.info("h4|Test results did not match the test results of the original code ❌")
+        console.rule()
+        return Failure("Test results did not match the test results of the original code.")
+
     def run_optimized_candidate(
         self,
         *,
@@ -1808,13 +1813,25 @@ def run_optimized_candidate(
                 )
             )
             console.rule()
-            if compare_test_results(baseline_results.behavior_test_results, candidate_behavior_results):
+            match, diffs = compare_test_results(baseline_results.behavior_test_results, candidate_behavior_results)
+            if match:
                 logger.info("h3|Test results matched ✅")
                 console.rule()
             else:
-                logger.info("h4|Test results did not match the test results of the original code ❌")
-                console.rule()
-                return Failure("Test results did not match the test results of the original code.")
+                result_unmatched_perc = len(diffs) / len(candidate_behavior_results)
+                if result_unmatched_perc > 0.5:
+                    # if the test unmatched percentage is greater than 50%, we can't fix it
+                    return self.get_results_not_matched_error()
+
+                # with the parsed test results diff ask the llm to fix the candidate to match the test results of the original code, and run again
+                # self.run_optimized_candidate(
+                #     optimization_candidate_index=optimization_candidate_index,
+                #     baseline_results=baseline_results,
+                #     original_helper_code=original_helper_code,
+                #     file_path_to_helper_classes=file_path_to_helper_classes,
+                # )
+                print(f"should try to fix it, diffs: {diffs}")
+                return self.get_results_not_matched_error()
 
             logger.info(f"loading|Running performance tests for candidate {optimization_candidate_index}...")
 

codeflash/verification/equivalence.py

@@ -1,4 +1,6 @@
 import sys
+from dataclasses import dataclass
+from enum import Enum
 
 from codeflash.cli_cmds.console import logger
 from codeflash.models.models import TestResults, TestType, VerificationType
@@ -7,21 +9,48 @@
 INCREASED_RECURSION_LIMIT = 5000
 
 
-def compare_test_results(original_results: TestResults, candidate_results: TestResults) -> bool:
+class TestDiffScope(Enum):
+    RETURN_VALUE = "return_value"
+    STDOUT = "stdout"
+    TIMED_OUT = "timed_out"
+    DID_PASS = "did_pass"  # noqa: S105
+
+
+@dataclass
+class TestDiff:
+    scope: TestDiffScope
+    test_src_code: str
+    pytest_error: str
+    original_value: any
+    candidate_value: any
+
+
+def compare_test_results(original_results: TestResults, candidate_results: TestResults) -> tuple[bool, list[TestDiff]]:
     # This is meant to be only called with test results for the first loop index
     if len(original_results) == 0 or len(candidate_results) == 0:
-        return False  # empty test results are not equal
+        return False, []  # empty test results are not equal
     original_recursion_limit = sys.getrecursionlimit()
     if original_recursion_limit < INCREASED_RECURSION_LIMIT:
         sys.setrecursionlimit(INCREASED_RECURSION_LIMIT)  # Increase recursion limit to avoid RecursionError
     test_ids_superset = original_results.get_all_unique_invocation_loop_ids().union(
         set(candidate_results.get_all_unique_invocation_loop_ids())
     )
-    are_equal: bool = True
+    test_diffs: list[TestDiff] = []
     did_all_timeout: bool = True
     for test_id in test_ids_superset:
         original_test_result = original_results.get_by_unique_invocation_loop_id(test_id)
         cdd_test_result = candidate_results.get_by_unique_invocation_loop_id(test_id)
+        candidate_test_failures = candidate_results.test_failures
+        # original_test_failures = original_results.test_failures
+        cdd_pytest_error = (
+            candidate_test_failures.get(original_test_result.id.test_function_name, "")
+            if candidate_test_failures
+            else ""
+        )
+        # original_pytest_error = (
+        #     original_test_failures.get(original_test_result.id.test_function_name, "") if original_test_failures else ""
+        # )
+
         if cdd_test_result is not None and original_test_result is None:
             continue
         # If helper function instance_state verification is not present, that's ok. continue
@@ -32,8 +61,7 @@ def compare_test_results(original_results: TestResults, candidate_results: TestR
         ):
             continue
         if original_test_result is None or cdd_test_result is None:
-            are_equal = False
-            break
+            return False, []
         did_all_timeout = did_all_timeout and original_test_result.timed_out
         if original_test_result.timed_out:
             continue
@@ -43,31 +71,42 @@ def compare_test_results(original_results: TestResults, candidate_results: TestR
             in {VerificationType.INIT_STATE_HELPER, VerificationType.INIT_STATE_FTO}
         ):
             superset_obj = True
+        test_src_code = original_test_result.id.get_src_code(original_test_result.file_name)
         if not comparator(original_test_result.return_value, cdd_test_result.return_value, superset_obj=superset_obj):
-            are_equal = False
+            test_diffs.append(
+                TestDiff(
+                    scope=TestDiffScope.RETURN_VALUE,
+                    test_src_code=test_src_code,
+                    original_value=original_test_result.return_value,
+                    candidate_value=cdd_test_result.return_value,
+                    pytest_error=cdd_pytest_error,
+                )
+            )
+
             try:
-                logger.debug(
-                    "File Name: %s\n"
-                    "Test Type: %s\n"
-                    "Verification Type: %s\n"
-                    "Invocation ID: %s\n"
-                    "Original return value: %s\n"
-                    "Candidate return value: %s\n"
-                    "-------------------",
-                    original_test_result.file_name,
-                    original_test_result.test_type,
-                    original_test_result.verification_type,
-                    original_test_result.id,
-                    original_test_result.return_value,
-                    cdd_test_result.return_value,
+                print(
+                    f"File Name: {original_test_result.file_name}\n"
+                    f"Test Type: {original_test_result.test_type}\n"
+                    f"Verification Type: {original_test_result.verification_type}\n"
+                    f"Invocation ID: {original_test_result.id}\n"
+                    f"Original return value: {original_test_result.return_value}\n"
+                    f"Candidate return value: {cdd_test_result.return_value}\n"
                 )
             except Exception as e:
                 logger.error(e)
             break
         if (original_test_result.stdout and cdd_test_result.stdout) and not comparator(
             original_test_result.stdout, cdd_test_result.stdout
         ):
-            are_equal = False
+            test_diffs.append(
+                TestDiff(
+                    scope=TestDiffScope.STDOUT,
+                    test_src_code=test_src_code,
+                    original_value=original_test_result.stdout,
+                    candidate_value=cdd_test_result.stdout,
+                    pytest_error=cdd_pytest_error,
+                )
+            )
             break
 
         if original_test_result.test_type in {
@@ -76,9 +115,17 @@ def compare_test_results(original_results: TestResults, candidate_results: TestR
             TestType.GENERATED_REGRESSION,
             TestType.REPLAY_TEST,
         } and (cdd_test_result.did_pass != original_test_result.did_pass):
-            are_equal = False
+            test_diffs.append(
+                TestDiff(
+                    scope=TestDiffScope.DID_PASS,
+                    test_src_code=test_src_code,
+                    original_value=original_test_result.did_pass,
+                    candidate_value=cdd_test_result.did_pass,
+                    pytest_error=cdd_pytest_error,
+                )
+            )
             break
     sys.setrecursionlimit(original_recursion_limit)
     if did_all_timeout:
-        return False
-    return are_equal
+        return False, test_diffs
+    return len(test_diffs) == 0, test_diffs

codeflash/verification/parse_test_output.py

@@ -512,6 +512,43 @@ def merge_test_results(
     return merged_test_results
 
 
+def parse_test_failures_from_stdout(test_results: TestResults, stdout: str) -> TestResults:
+    stdout_lines = stdout.splitlines()
+    start_line = -1
+    end_line = -1
+    for i, line in enumerate(stdout_lines):
+        if start_line != -1 and end_line != -1:
+            break
+        if "FAILURES" in line:
+            start_line = i
+        elif "short test summary info" in line:
+            end_line = i
+    if start_line == -1 or end_line == -1:
+        return test_results
+
+    complete_failure_output_lines = stdout_lines[start_line:end_line]  # exclude last summary line
+
+    test_case_to_failure: dict[str, str] = {}
+
+    current_test_case: str | None = None
+    current_failure_lines: list[str] = []
+
+    for line in complete_failure_output_lines:
+        if line.startswith("_______"):
+            if current_test_case:
+                test_case_to_failure[current_test_case] = "".join(current_failure_lines)
+            current_test_case = line.strip("_ ").strip()
+            current_failure_lines = []
+        elif current_test_case:
+            current_failure_lines.append(line + "\n")
+
+    if current_test_case:
+        test_case_to_failure[current_test_case] = "".join(current_failure_lines)
+
+    test_results.test_failures = test_case_to_failure
+    return test_results
+
+
 def parse_test_results(
     test_xml_path: Path,
     test_files: TestFiles,
@@ -572,4 +609,9 @@ def parse_test_results(
             function_name=function_name,
         )
         coverage.log_coverage()
+    try:
+        parse_test_failures_from_stdout(results, run_result.stdout)
+    except Exception as e:
+        logger.exception(e)
+
     return results, coverage if all_args else None