From 606cca5c2f645a4bf62cb3d5a43d2a1034cfd77a Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Wed, 6 May 2026 21:50:18 -0500 Subject: [PATCH 1/2] perf: avoid double dict lookup for line profiler results --- codeflash/models/models.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 640e5230a..af90e7bf3 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -547,10 +547,9 @@ def handle_duplicate_candidate( self.optimized_runtimes[candidate.optimization_id] = self.optimized_runtimes.get(past_opt_id) # Line profiler results only available for successful runs - if past_opt_id in self.optimized_line_profiler_results: - self.optimized_line_profiler_results[candidate.optimization_id] = self.optimized_line_profiler_results[ - past_opt_id - ] + line_profiler_result = self.optimized_line_profiler_results.get(past_opt_id) + if line_profiler_result is not None: + self.optimized_line_profiler_results[candidate.optimization_id] = line_profiler_result self.optimizations_post[candidate.optimization_id] = self.ast_code_to_id[normalized_code][ "shorter_source_code" From 5e579ca7629fcf396fb1db7c39f1058e9672137b Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Wed, 6 May 2026 21:57:00 -0500 Subject: [PATCH 2/2] test: add unit tests for line profiler result copying --- tests/test_candidate_evaluation_context.py | 78 ++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 tests/test_candidate_evaluation_context.py diff --git a/tests/test_candidate_evaluation_context.py b/tests/test_candidate_evaluation_context.py new file mode 100644 index 000000000..e5b81a0b3 --- /dev/null +++ b/tests/test_candidate_evaluation_context.py @@ -0,0 +1,78 @@ +from unittest.mock import MagicMock, patch + +from codeflash.models.models import ( + CandidateEvaluationContext, + CodeString, + CodeStringsMarkdown, + OptimizedCandidate, + OptimizedCandidateSource, +) + + +def make_source_code(code: str = "pass") -> CodeStringsMarkdown: + return CodeStringsMarkdown(code_strings=[CodeString(code=code, file_path=None)]) + + +def make_candidate(optimization_id: str, code: str = "pass") -> OptimizedCandidate: + return OptimizedCandidate( + source_code=make_source_code(code), + explanation="test", + optimization_id=optimization_id, + source=OptimizedCandidateSource.OPTIMIZE, + ) + + +def make_context_with_past_candidate(past_opt_id: str, normalized_code: str, code: str = "pass") -> CandidateEvaluationContext: + ctx = CandidateEvaluationContext() + ctx.ast_code_to_id[normalized_code] = { + "optimization_id": past_opt_id, + "shorter_source_code": make_source_code(code), + "diff_len": 10, + } + ctx.speedup_ratios[past_opt_id] = 2.0 + ctx.is_correct[past_opt_id] = True + ctx.optimized_runtimes[past_opt_id] = 0.5 + return ctx + + +@patch("codeflash.code_utils.code_utils.diff_length", return_value=10) +def test_copy_line_profiler_results_existing_key(_mock_diff: MagicMock) -> None: + past_opt_id = "past-123" + normalized_code = "normalized_code_abc" + ctx = make_context_with_past_candidate(past_opt_id, normalized_code) + ctx.optimized_line_profiler_results[past_opt_id] = "line_profile_output_data" + + candidate = make_candidate("new-456") + ctx.handle_duplicate_candidate(candidate, normalized_code, "original_flat") + + assert ctx.optimized_line_profiler_results["new-456"] == "line_profile_output_data" + + +@patch("codeflash.code_utils.code_utils.diff_length", return_value=10) +def test_copy_line_profiler_results_missing_key(_mock_diff: MagicMock) -> None: + past_opt_id = "past-789" + normalized_code = "normalized_code_xyz" + ctx = make_context_with_past_candidate(past_opt_id, normalized_code) + # No line profiler result for past_opt_id + + candidate = make_candidate("new-012") + ctx.handle_duplicate_candidate(candidate, normalized_code, "original_flat") + + assert "new-012" not in ctx.optimized_line_profiler_results + + +@patch("codeflash.code_utils.code_utils.diff_length", return_value=10) +def test_copy_line_profiler_results_does_not_corrupt(_mock_diff: MagicMock) -> None: + past_opt_id = "past-aaa" + normalized_code = "normalized_code_zzz" + ctx = make_context_with_past_candidate(past_opt_id, normalized_code) + original_value = "original_profile_data" + ctx.optimized_line_profiler_results[past_opt_id] = original_value + + candidate = make_candidate("new-bbb") + ctx.handle_duplicate_candidate(candidate, normalized_code, "original_flat") + + # Original entry unchanged + assert ctx.optimized_line_profiler_results[past_opt_id] == original_value + # New entry is equal but does not affect original on mutation + assert ctx.optimized_line_profiler_results["new-bbb"] == original_value