fix: Java comparator key collision and add A/A baseline test

misrasaurabh1 · claude · misrasaurabh1 · commit 4ecb6260fc5d · 2026-03-01T02:05:07.000-08:00
The Java Comparator used only iteration_id as the map key when comparing
test results. Since every test method had iteration_id="1", all rows
collapsed to a single entry and only the last row survived. Because JUnit
test execution order is non-deterministic, the surviving row differed
between baseline and candidate runs, causing false correctness failures.

Fix the key to include test_module_path:test_class_name:test_function_name
for unique identification. Also add an A/A test that runs the original code
through the candidate pipeline to detect verification bugs like this.

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/codeflash-java-runtime/src/main/java/com/codeflash/Comparator.java b/codeflash-java-runtime/src/main/java/com/codeflash/Comparator.java
@@ -144,8 +144,11 @@ private static Map<String, byte[]> readTestResults(String dbPath) throws Excepti
         try (Connection conn = DriverManager.getConnection(url);
              Statement stmt = conn.createStatement();
              ResultSet rs = stmt.executeQuery(
-                 "SELECT iteration_id, return_value FROM test_results WHERE loop_index = 1")) {
+                 "SELECT test_module_path, test_class_name, test_function_name, iteration_id, return_value FROM test_results WHERE loop_index = 1")) {
             while (rs.next()) {
+                String testModulePath = rs.getString("test_module_path");
+                String testClassName = rs.getString("test_class_name");
+                String testFunctionName = rs.getString("test_function_name");
                 String iterationId = rs.getString("iteration_id");
                 byte[] returnValue = rs.getBytes("return_value");
                 // Strip the CODEFLASH_TEST_ITERATION suffix (e.g. "7_0" -> "7")
@@ -155,7 +158,10 @@ private static Map<String, byte[]> readTestResults(String dbPath) throws Excepti
                 if (lastUnderscore > 0) {
                     iterationId = iterationId.substring(0, lastUnderscore);
                 }
-                results.put(iterationId, returnValue);
+                // Use module:class:function:iteration as key to uniquely identify
+                // each invocation across different test files, classes, and methods
+                String key = testModulePath + ":" + testClassName + ":" + testFunctionName + "::" + iterationId;
+                results.put(key, returnValue);
             }
         }
         return results;
diff --git a/codeflash/languages/java/resources/codeflash-runtime-1.0.0.jar b/codeflash/languages/java/resources/codeflash-runtime-1.0.0.jar
diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py
@@ -1335,7 +1335,7 @@ def process_single_candidate(
                 optimized_code=candidate.source_code,
                 original_helper_code=original_helper_code,
             )
-            if not did_update:
+            if not did_update and candidate.optimization_id != "original-code-aa-test":
                 logger.info("No functions were replaced in the optimized code. Skipping optimization candidate.")
                 console.rule()
                 return None
@@ -1371,6 +1371,18 @@ def process_single_candidate(
         )
         console.rule()
 
+        # Handle A/A test: original code run through the candidate pipeline
+        if candidate.optimization_id == "original-code-aa-test":
+            if not is_successful(run_results):
+                logger.warning(
+                    "A/A test FAILED! Original code did not pass its own correctness check. "
+                    "This indicates a bug in the verification pipeline."
+                )
+                eval_ctx.record_failed_candidate(candidate.optimization_id)
+                return None
+            logger.info("A/A test passed: original code passed correctness verification as expected.")
+            return None
+
         if not is_successful(run_results):
             eval_ctx.record_failed_candidate(candidate.optimization_id)
             return None
@@ -1514,6 +1526,15 @@ def determine_best_candidate(
             language=self.function_to_optimize.language,
         )
 
+        # Prepend original code as an A/A test candidate (must pass correctness by definition)
+        aa_test_candidate = OptimizedCandidate(
+            source_code=code_context.read_writable_code,
+            explanation="Original code (A/A baseline test)",
+            optimization_id="original-code-aa-test",
+            source=OptimizedCandidateSource.OPTIMIZE,
+        )
+        candidates = [aa_test_candidate, *candidates]
+
         processor = CandidateProcessor(
             candidates,
             future_line_profile_results,
