⚡️ Speed up function `add_runtime_comments_to_generated_tests` by 239% in PR #488 (`fix-runtime-comments`) #489

codeflash-ai · 2025-07-02T23:50:33Z

⚡️ This pull request contains optimizations for PR #488

If you approve this dependent PR, these changes will be merged into the original PR branch fix-runtime-comments.

This PR will be automatically closed if the original PR is merged.

📄 239% (2.39x) speedup for `add_runtime_comments_to_generated_tests` in `codeflash/code_utils/edit_generated_tests.py`

⏱️ Runtime : 1.16 seconds → 341 milliseconds (best of 24 runs)

📝 Explanation and details

Here’s a heavily optimized rewrite of your function, focused on the main bottleneck: the tree.visit(transformer) call inside the main loop (~95% of your runtime!). Across the entire function, the following optimizations (all applied without changing any functional output) are used.

Precompute Data Structures: Several expensive operations (especially relative_to path gymnastics and conversions) are moved out of inner loops and stored as sensible lookups, since their results are almost invariant across tests.
Merge For Loops: The two near-identical for loops per invocation in leave_SimpleStatementLine are merged into one, halving search cost.
Optimize Invocation Matching: An indexed lookup is pre-built mapping the unique tuple keys (rel_path, qualified_name, cfo_loc) to their runtimes. This makes runtime-access O(1) instead of requiring a full scan per statement.
Avoid Deep AST/Normalized Source Walks: If possible, recommend optimizing find_codeflash_output_assignments to operate on the CST or directly on the parsed AST rather than reparsing source code. (The code preserves your current approach but this is a further large opportunity.)
Faster CST Name/Call detection: The leave_SimpleStatementLine’s _contains_myfunc_call is further micro-optimized by breaking as soon as a match is found (using exception for early escape), avoiding unnecessary traversal.
Minimize Object Creations: The GeneratedTests objects are only constructed once and appended.
Eliminating Minor Redundant Computation.
Reduce try/except Overhead: Only exceptions propagate; no functional change here.

Below is the optimized code, with comments kept as close as possible to your original code (apart from changed logic).

Summary of key gains:

The O(N*M) runtimes loop is now O(1) due to hash indexes.
All constant/cached values are precomputed outside the node visitor.
Deep tree walks and list traversals have early exits and critical-path logic is tightened.
No functional changes, all corner cases preserved.

Still slow?:
The biggest remaining hit will be the find_codeflash_output_assignments (which reparses source); move this to operate directly on CST if possible for further big wins.

Let me know your measured speedup! 🚀

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 23 Passed
🌀 Generated Regression Tests	✅ 23 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
`test_add_runtime_comments.py::TestAddRuntimeComments.test_add_runtime_comments_class_method`	1.14ms	1.10ms	✅3.57%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_add_runtime_comments_multiple_assignments`	2.36ms	2.17ms	✅8.52%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_add_runtime_comments_multiple_tests`	1.91ms	1.80ms	✅6.43%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_add_runtime_comments_no_codeflash_output`	985μs	945μs	✅4.16%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_add_runtime_comments_no_matching_runtimes`	921μs	916μs	✅0.553%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_add_runtime_comments_performance_regression`	1.61ms	1.46ms	✅10.4%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_add_runtime_comments_simple_function`	988μs	952μs	✅3.79%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_basic_runtime_comment_addition`	1.63ms	1.48ms	✅9.90%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_different_time_formats`	6.05ms	5.82ms	✅3.82%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_invalid_python_code_handling`	108μs	130μs	⚠️-17.1%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_missing_test_results`	1.53ms	1.43ms	✅7.25%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_multiple_generated_tests`	3.74ms	3.43ms	✅9.08%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_multiple_runtimes_uses_minimum`	1.63ms	1.50ms	✅8.77%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_multiple_test_functions`	3.42ms	3.09ms	✅10.7%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_multistatement_line_handling`	2.51ms	2.38ms	✅5.66%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_no_codeflash_output_assignment`	1.61ms	1.45ms	✅10.8%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_partial_test_results`	1.58ms	1.45ms	✅8.74%
`test_add_runtime_comments.py::TestAddRuntimeComments.test_preserved_test_attributes`	1.73ms	1.57ms	✅10.7%

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import ast
import os
from dataclasses import dataclass, field
from pathlib import Path
from textwrap import dedent
from typing import Dict, List, Optional

import libcst as cst
# imports
import pytest
from codeflash.code_utils.edit_generated_tests import \
    add_runtime_comments_to_generated_tests

# --- Minimal stubs for dependencies ---

# InvocationId stub
@dataclass(frozen=True)
class InvocationId:
    test_module_path: str
    test_class_name: str
    test_function_name: str
    iteration_id: str

# GeneratedTests stub
@dataclass
class GeneratedTests:
    generated_original_test_source: str
    instrumented_behavior_test_source: str
    instrumented_perf_test_source: str
    behavior_file_path: Path
    perf_file_path: Path

# GeneratedTestsList stub
@dataclass
class GeneratedTestsList:
    generated_tests: List[GeneratedTests]

# TestConfig stub
@dataclass
class TestConfig:
    tests_root: Path
    project_root_path: Path
from codeflash.code_utils.edit_generated_tests import \
    add_runtime_comments_to_generated_tests

# ========== UNIT TESTS ==========

# Helper to build a minimal test source with a function call
def make_test_source(func_name: str, call_line: str = None, class_name: Optional[str] = None) -> str:
    if class_name:
        return f"""
class {class_name}:
    def test_func(self):
        codeflash_output = {func_name}()
"""
    else:
        return f"""
def test_func():
    codeflash_output = {func_name}()
"""

# Helper to create a GeneratedTestsList with one test
def make_generated_tests_list(source: str, behavior_path: Path, perf_path: Path) -> GeneratedTestsList:
    return GeneratedTestsList([
        GeneratedTests(
            generated_original_test_source=source,
            instrumented_behavior_test_source="",
            instrumented_perf_test_source="",
            behavior_file_path=behavior_path,
            perf_file_path=perf_path,
        )
    ])

# Helper to create InvocationId for a function
def make_invocation_id(test_module_path, test_class_name, test_function_name, iteration_id):
    return InvocationId(
        test_module_path=test_module_path,
        test_class_name=test_class_name,
        test_function_name=test_function_name,
        iteration_id=iteration_id
    )

# ========== BASIC TEST CASES ==========

def test_basic_runtime_comment_inserted():
    """Basic: Should add comment for a single function call in a function."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name)
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    # Matching InvocationId
    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    # Simulate original and optimized runtimes
    original_runtimes = {inv_id: [200_000_000]}  # 200ms
    optimized_runtimes = {inv_id: [100_000_000]} # 100ms

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    # The output should contain the runtime comment after the function call
    out_source = result.generated_tests[0].generated_original_test_source

def test_basic_runtime_comment_slower():
    """Basic: Should indicate 'slower' if optimized is slower."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name)
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id: [100_000_000]}  # 100ms
    optimized_runtimes = {inv_id: [200_000_000]} # 200ms

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_basic_no_matching_invocation():
    """Basic: If no matching invocation, no comment should be added."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name)
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    # InvocationId does not match function name
    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="other_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id: [100_000_000]}
    optimized_runtimes = {inv_id: [80_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_basic_multiple_calls():
    """Basic: Multiple calls in one function should each get a comment."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = """
def test_func():
    codeflash_output = myfunc()
    codeflash_output = myfunc()
"""
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    # Two invocation ids, one for each call (simulate iteration_id as 0 and 1)
    inv_id0 = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    inv_id1 = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="1_0"
    )
    original_runtimes = {inv_id0: [100_000_000], inv_id1: [200_000_000]}
    optimized_runtimes = {inv_id0: [80_000_000], inv_id1: [100_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_basic_class_method():
    """Basic: Should handle test method in a class."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name, class_name="TestClass")
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="TestClass",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id: [100_000_000]}
    optimized_runtimes = {inv_id: [90_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

# ========== EDGE TEST CASES ==========

def test_edge_zero_runtimes():
    """Edge: If either runtime is zero, no comment should be added."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name)
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    # One runtime is zero
    original_runtimes = {inv_id: [0]}
    optimized_runtimes = {inv_id: [100_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_edge_no_codeflash_output():
    """Edge: If there is no codeflash_output assignment, no comment should be added."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = """
def test_func():
    pass
"""
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    original_runtimes = {}
    optimized_runtimes = {}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_edge_multiple_classes():
    """Edge: Handles multiple classes and functions with same name in different classes."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = """
class TestClassA:
    def test_func(self):
        codeflash_output = myfunc()
class TestClassB:
    def test_func(self):
        codeflash_output = myfunc()
"""
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id_a = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="TestClassA",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    inv_id_b = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="TestClassB",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id_a: [100_000_000], inv_id_b: [200_000_000]}
    optimized_runtimes = {inv_id_a: [80_000_000], inv_id_b: [100_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_edge_non_matching_path():
    """Edge: If the test file path does not match InvocationId, no comment should be added."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name)
    behavior_path = Path("tests/other_module.py")
    perf_path = Path("tests/other_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id: [100_000_000]}
    optimized_runtimes = {inv_id: [80_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_edge_attribute_call():
    """Edge: Should detect calls like obj.myfunc() as well."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = """
def test_func():
    codeflash_output = obj.myfunc()
"""
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id: [100_000_000]}
    optimized_runtimes = {inv_id: [80_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source


def test_large_many_tests(monkeypatch):
    """Large Scale: Handles many test files and calls efficiently."""
    func_name = "myfunc"
    qualified_name = func_name
    N = 100  # 100 test functions
    test_sources = []
    behavior_paths = []
    perf_paths = []
    original_runtimes = {}
    optimized_runtimes = {}
    tests = []

    for i in range(N):
        source = f"""
def test_func_{i}():
    codeflash_output = myfunc()
"""
        behavior_path = Path(f"tests/test_module_{i}.py")
        perf_path = Path(f"tests/test_module_{i}.py")
        tests.append(GeneratedTests(
            generated_original_test_source=source,
            instrumented_behavior_test_source="",
            instrumented_perf_test_source="",
            behavior_file_path=behavior_path,
            perf_file_path=perf_path,
        ))
        inv_id = make_invocation_id(
            test_module_path=f"tests.test_module_{i}",
            test_class_name="",
            test_function_name=f"test_func_{i}",
            iteration_id="0_0"
        )
        # Each test has unique timings
        original_runtimes[inv_id] = [100_000_000 + i*1_000_000]
        optimized_runtimes[inv_id] = [90_000_000 + i*1_000_000]

    tests_list = GeneratedTestsList(tests)
    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    # Each output should have correct comment
    for i in range(N):
        out_source = result.generated_tests[i].generated_original_test_source
        expected_comment = f"# {100+i}ms -> {90+i}ms (10.00% faster)"

def test_large_many_calls_per_test():
    """Large Scale: Handles many calls in a single test function."""
    func_name = "myfunc"
    qualified_name = func_name
    N = 100  # 100 calls in one function
    call_lines = "\n".join([f"    codeflash_output = myfunc()" for _ in range(N)])
    test_source = f"""
def test_func():
{call_lines}
"""
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    original_runtimes = {}
    optimized_runtimes = {}
    for i in range(N):
        inv_id = make_invocation_id(
            test_module_path="tests.test_module",
            test_class_name="",
            test_function_name="test_func",
            iteration_id=f"{i}_0"
        )
        original_runtimes[inv_id] = [100_000_000 + i*1_000_000]
        optimized_runtimes[inv_id] = [90_000_000 + i*1_000_000]

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source
    # All 100 comments should appear
    for i in range(N):
        expected_comment = f"# {100+i}ms -> {90+i}ms (10.00% faster)"

def test_large_many_classes_and_methods():
    """Large Scale: Handles many classes and methods."""
    func_name = "myfunc"
    qualified_name = func_name
    N = 50  # 50 classes, each with 2 methods
    test_source = ""
    original_runtimes = {}
    optimized_runtimes = {}
    for i in range(N):
        test_source += f"""
class TestClass{i}:
    def test_func_a(self):
        codeflash_output = myfunc()
    def test_func_b(self):
        codeflash_output = myfunc()
"""
        inv_id_a = make_invocation_id(
            test_module_path=f"tests.test_module",
            test_class_name=f"TestClass{i}",
            test_function_name="test_func_a",
            iteration_id="0_0"
        )
        inv_id_b = make_invocation_id(
            test_module_path=f"tests.test_module",
            test_class_name=f"TestClass{i}",
            test_function_name="test_func_b",
            iteration_id="0_0"
        )
        original_runtimes[inv_id_a] = [100_000_000 + i*1_000_000]
        optimized_runtimes[inv_id_a] = [90_000_000 + i*1_000_000]
        original_runtimes[inv_id_b] = [200_000_000 + i*1_000_000]
        optimized_runtimes[inv_id_b] = [180_000_000 + i*1_000_000]

    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)
    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import ast
import os
from dataclasses import dataclass, field
from pathlib import Path
from textwrap import dedent
from typing import Dict, List

# Use libcst for parsing and transforming code
import libcst as cst
# imports
import pytest
from codeflash.code_utils.edit_generated_tests import \
    add_runtime_comments_to_generated_tests

# --- Minimal stub implementations for dependencies ---

# Simulate InvocationId
@dataclass(frozen=True)
class InvocationId:
    test_module_path: str
    test_class_name: str
    test_function_name: str
    iteration_id: str

# Simulate GeneratedTests
@dataclass
class GeneratedTests:
    generated_original_test_source: str
    instrumented_behavior_test_source: str = ""
    instrumented_perf_test_source: str = ""
    behavior_file_path: Path = Path("tests/test_file.py")
    perf_file_path: Path = Path("tests/test_file.py")

# Simulate GeneratedTestsList
@dataclass
class GeneratedTestsList:
    generated_tests: List[GeneratedTests]

# Simulate TestConfig
@dataclass
class TestConfig:
    tests_root: Path
    project_root_path: Path
from codeflash.code_utils.edit_generated_tests import \
    add_runtime_comments_to_generated_tests


def make_file_paths():
    project_root = Path("/project")
    tests_root = project_root / "tests"
    test_file = tests_root / "test_file.py"
    return project_root, tests_root, test_file

# Helper to build InvocationId
def make_invocation_id(
    test_module_path="tests.test_file",
    test_class_name="TestClass",
    test_function_name="test_func",
    iteration_id="0_0"
):
    return InvocationId(
        test_module_path=test_module_path,
        test_class_name=test_class_name,
        test_function_name=test_function_name,
        iteration_id=iteration_id
    )

# --- BASIC TEST CASES ---

def test_single_function_single_call_adds_comment():
    """
    Basic: Single function, single call, matching runtime data.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result = myfunc(1)
        codeflash_output = result
"""
    # Setup test config and test object
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])

    # Setup runtime data
    invocation_id = make_invocation_id()
    original_runtimes = {invocation_id: [1000]}
    optimized_runtimes = {invocation_id: [500]}

    # Call function
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output

    # Assert: comment is present and correct
    out_code = result.generated_tests[0].generated_original_test_source

def test_no_matching_invocation_no_comment():
    """
    Basic: No matching invocation id, no comment is added.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result = myfunc(2)
        codeflash_output = result
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])

    # Empty runtime data
    original_runtimes = {}
    optimized_runtimes = {}

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_multiple_calls_multiple_comments():
    """
    Basic: Multiple calls in one function, multiple comments.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result1 = myfunc(1)
        codeflash_output = result1
        result2 = myfunc(2)
        codeflash_output = result2
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])

    # Two calls, two invocation ids
    invocation_id1 = make_invocation_id(iteration_id="0_0")
    invocation_id2 = make_invocation_id(iteration_id="1_0")
    original_runtimes = {invocation_id1: [1000], invocation_id2: [2000]}
    optimized_runtimes = {invocation_id1: [800], invocation_id2: [1500]}

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_comment_shows_slower_when_optimized_slower():
    """
    Basic: Optimized is slower, comment says 'slower'.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result = myfunc(5)
        codeflash_output = result
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])

    invocation_id = make_invocation_id()
    original_runtimes = {invocation_id: [1000]}
    optimized_runtimes = {invocation_id: [1500]}

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

# --- EDGE TEST CASES ---

def test_function_with_no_calls_no_comment():
    """
    Edge: Function does not call the target function at all.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        x = 1 + 2
        y = x * 2
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes={},
        optimized_runtimes={}
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_zero_runtimes_no_comment():
    """
    Edge: Runtimes are zero, comment is not added.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result = myfunc(1)
        codeflash_output = result
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])
    invocation_id = make_invocation_id()
    original_runtimes = {invocation_id: [0]}
    optimized_runtimes = {invocation_id: [0]}
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_function_with_attribute_call():
    """
    Edge: Function call is as an attribute (obj.myfunc), comment is added.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        obj = SomeClass()
        result = obj.myfunc(7)
        codeflash_output = result
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])
    invocation_id = make_invocation_id()
    original_runtimes = {invocation_id: [100]}
    optimized_runtimes = {invocation_id: [50]}
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_multiple_classes_and_functions():
    """
    Edge: Multiple classes and functions, only correct function gets comment.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result = myfunc(1)
        codeflash_output = result

class OtherClass:
    def test_func(self):
        result = myfunc(2)
        codeflash_output = result
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])
    # Only first class/function gets runtime data
    invocation_id = make_invocation_id(test_class_name="TestClass")
    original_runtimes = {invocation_id: [1000]}
    optimized_runtimes = {invocation_id: [900]}
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source


def test_large_number_of_functions_and_calls():
    """
    Large Scale: Many functions and calls, all get comments.
    """
    project_root, tests_root, test_file = make_file_paths()
    # Generate 100 functions, each with 2 calls
    num_funcs = 100
    test_lines = []
    original_runtimes = {}
    optimized_runtimes = {}
    for i in range(num_funcs):
        test_lines.append(f"class TestClass{i}:")
        test_lines.append(f"    def test_func(self):")
        for j in range(2):
            test_lines.append(f"        result{j} = myfunc({i*2+j})")
            test_lines.append(f"        codeflash_output = result{j}")
            # Build invocation id for each call
            inv_id = InvocationId(
                test_module_path="tests.test_file",
                test_class_name=f"TestClass{i}",
                test_function_name="test_func",
                iteration_id=f"{j}_0"
            )
            original_runtimes[inv_id] = [1000 + i*10 + j]
            optimized_runtimes[inv_id] = [500 + i*5 + j]
    test_code = "\n".join(test_lines)
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_large_number_of_tests_in_list():
    """
    Large Scale: Many test files in GeneratedTestsList, all get processed.
    """
    project_root, tests_root, test_file = make_file_paths()
    num_tests = 50
    generated_tests = []
    original_runtimes = {}
    optimized_runtimes = {}
    for i in range(num_tests):
        code = f"""
class TestClass{i}:
    def test_func(self):
        result = myfunc({i})
        codeflash_output = result
"""
        test_obj = GeneratedTests(
            generated_original_test_source=code,
            behavior_file_path=test_file,
            perf_file_path=test_file
        )
        generated_tests.append(test_obj)
        inv_id = InvocationId(
            test_module_path="tests.test_file",
            test_class_name=f"TestClass{i}",
            test_function_name="test_func",
            iteration_id="0_0"
        )
        original_runtimes[inv_id] = [1000 + i]
        optimized_runtimes[inv_id] = [900 + i]
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    generated_tests_list = GeneratedTestsList(generated_tests)
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests_list,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    # All 50 tests should have a comment
    for idx in range(num_tests):
        out_code = result.generated_tests[idx].generated_original_test_source
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr488-2025-07-02T23.50.27 and push.

…% in PR #488 (`fix-runtime-comments`) Here’s a heavily optimized rewrite of your function, focused on the main bottleneck: the `tree.visit(transformer)` call inside the main loop (~95% of your runtime!). Across the entire function, the following optimizations (all applied **without changing any functional output**) are used. 1. **Precompute Data Structures:** Several expensive operations (especially `relative_to` path gymnastics and conversions) are moved out of inner loops and stored as sensible lookups, since their results are almost invariant across tests. 2. **Merge For Loops:** The two near-identical `for` loops per invocation in `leave_SimpleStatementLine` are merged into one, halving search cost. 3. **Optimize Invocation Matching:** An indexed lookup is pre-built mapping the unique tuple keys `(rel_path, qualified_name, cfo_loc)` to their runtimes. This makes runtime-access O(1) instead of requiring a full scan per statement. 4. **Avoid Deep AST/Normalized Source Walks:** If possible, recommend optimizing `find_codeflash_output_assignments` to operate on the CST or directly on the parsed AST rather than reparsing source code. (**The code preserves your current approach but this is a further large opportunity.**) 5. **Faster CST Name/Call detection:** The `leave_SimpleStatementLine`’s `_contains_myfunc_call` is further micro-optimized by breaking as soon as a match is found (using exception for early escape), avoiding unnecessary traversal. 6. **Minimize Object Creations:** The `GeneratedTests` objects are only constructed once and appended. 7. **Eliminating Minor Redundant Computation.** 8. **Reduce try/except Overhead:** Only exceptions propagate; no functional change here. Below is the optimized code, with comments kept as close as possible to your original code (apart from changed logic). **Summary of key gains:** - The O(N*M) runtimes loop is now O(1) due to hash indexes. - All constant/cached values are precomputed outside the node visitor. - Deep tree walks and list traversals have early exits and critical-path logic is tightened. - No functional changes, all corner cases preserved. **Still slow?**: The biggest remaining hit will be the `find_codeflash_output_assignments` (which reparses source); move this to operate directly on CST if possible for further big wins. Let me know your measured speedup! 🚀

misrasaurabh1 · 2025-07-02T23:57:38Z

@aseembits93 i don't see the runtime info in the generated tests here

codeflash-ai · 2025-07-03T20:13:39Z

This PR has been automatically closed because the original PR #488 by aseembits93 was closed.

codeflash-ai bot added the ⚡️ codeflash Optimization PR opened by Codeflash AI label Jul 2, 2025

codeflash-ai bot mentioned this pull request Jul 2, 2025

runtime timing info even if codeflash_output assignment is missing #488

Closed

codeflash-ai bot closed this Jul 3, 2025

codeflash-ai bot deleted the codeflash/optimize-pr488-2025-07-02T23.50.27 branch July 3, 2025 20:13

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

⚡️ Speed up function `add_runtime_comments_to_generated_tests` by 239% in PR #488 (`fix-runtime-comments`) #489

⚡️ Speed up function `add_runtime_comments_to_generated_tests` by 239% in PR #488 (`fix-runtime-comments`) #489

Uh oh!

codeflash-ai bot commented Jul 2, 2025

Uh oh!

misrasaurabh1 commented Jul 2, 2025

Uh oh!

codeflash-ai bot commented Jul 3, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

⚡️ Speed up function add_runtime_comments_to_generated_tests by 239% in PR #488 (fix-runtime-comments) #489

⚡️ Speed up function add_runtime_comments_to_generated_tests by 239% in PR #488 (fix-runtime-comments) #489

Uh oh!

Conversation

codeflash-ai bot commented Jul 2, 2025

⚡️ This pull request contains optimizations for PR #488

📄 239% (2.39x) speedup for add_runtime_comments_to_generated_tests in codeflash/code_utils/edit_generated_tests.py

📝 Explanation and details

Uh oh!

misrasaurabh1 commented Jul 2, 2025

Uh oh!

codeflash-ai bot commented Jul 3, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

⚡️ Speed up function `add_runtime_comments_to_generated_tests` by 239% in PR #488 (`fix-runtime-comments`) #489

⚡️ Speed up function `add_runtime_comments_to_generated_tests` by 239% in PR #488 (`fix-runtime-comments`) #489

📄 239% (2.39x) speedup for `add_runtime_comments_to_generated_tests` in `codeflash/code_utils/edit_generated_tests.py`