⚡️ Speed up method CommentMapper.visit_FunctionDef by 14% in PR #719 (async-support-for)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #719

If you approve this dependent PR, these changes will be merged into the original PR branch async-support-for.

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

---

📄 14% (0.14x) speedup for CommentMapper.visit_FunctionDef in codeflash/code_utils/edit_generated_tests.py

⏱️ Runtime : 4.56 milliseconds → 4.01 milliseconds (best of 241 runs)

📝 Explanation and details

The optimized code achieves a 13% speedup through several targeted micro-optimizations that reduce overhead in the hot loops:

Key optimizations applied:

1. Hoisted loop-invariant computations: Moved isinstance tuple constants (compound_types, valid_types) and frequently accessed attributes (get_comment, orig_rt, opt_rt) outside the loops to avoid repeated lookups.

2. Precomputed key prefix: Instead of repeatedly concatenating test_qualified_name + "#" + str(self.abs_path) inside loops, this is computed once as key_prefix and reused with f-string formatting.

3. Optimized getattr usage: Replaced the costly getattr(compound_line_node, "body", []) pattern with a single getattr(..., None) call, then conditionally building the nodes_to_check list using unpacking (*compound_line_node_body) when a body exists.

4. Reduced function call overhead: Cached the get_comment method reference and called it once per match_key, reusing the same comment for all nodes that share the same key, rather than calling it for each individual node.

5. String formatting optimization: Replaced string concatenation with f-string formatting for better performance.

Performance characteristics by test case:

• Large-scale tests show the best improvements (10-79% faster), particularly test_large_deeply_nested (78.8% faster) where the inner loop optimizations have maximum impact
• Basic cases show modest gains (1-4% faster) as there's less loop iteration overhead to optimize
• Edge cases with minimal computation show negligible or slightly negative impact due to the upfront setup cost of hoisted variables

The optimizations are most effective for functions with complex nested structures (for/while/if blocks) and many runtime entries, where the reduced per-iteration overhead compounds significantly.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import ast
from pathlib import Path

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.edit_generated_tests import CommentMapper


class GeneratedTests:
    # Minimal stub for the test object
    def __init__(self, behavior_file_path):
        self.behavior_file_path = behavior_file_path

# Helper to create an AST from code
def parse_func(code: str) -> ast.FunctionDef:
    tree = ast.parse(code)
    for node in tree.body:
        if isinstance(node, ast.FunctionDef):
            return node
    raise ValueError("No function definition found.")

# Helper to create a fake GeneratedTests object with a given file path
def make_test_obj(path: str = "foo.py"):
    return GeneratedTests(behavior_file_path=Path(path))

# ---------------- BASIC TEST CASES ----------------

def test_basic_single_statement():
    """
    Test a function with a single statement and a matching runtime entry.
    """
    code = "def foo():\n    x = 1"
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    key = "foo#foo#0"
    original = {key: 100_000_000}
    optimized = {key: 50_000_000}
    mapper = CommentMapper(test_obj, original, optimized)
    mapper.visit_FunctionDef(node) # 11.1μs -> 11.5μs (3.75% slower)
    # Check that the comment is correct
    comment = mapper.results[2]

def test_basic_two_statements():
    """
    Test a function with two statements, both with runtime entries.
    """
    code = "def foo():\n    x = 1\n    y = 2"
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    key = "foo#foo#"
    original = {key+"0": 100_000_000, key+"1": 80_000_000}
    optimized = {key+"0": 50_000_000, key+"1": 60_000_000}
    mapper = CommentMapper(test_obj, original, optimized)
    mapper.visit_FunctionDef(node) # 13.9μs -> 14.1μs (1.42% slower)

def test_basic_no_runtimes():
    """
    Test a function with statements but no runtime entries.
    """
    code = "def foo():\n    x = 1\n    y = 2"
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    mapper = CommentMapper(test_obj, {}, {})
    mapper.visit_FunctionDef(node) # 6.26μs -> 6.02μs (3.99% faster)

def test_basic_with_for_loop():
    """
    Test a function with a for loop and runtime entries for the loop body.
    """
    code = (
        "def foo():\n"
        "    for i in range(2):\n"
        "        x = i\n"
        "    y = 3"
    )
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    key = "foo#foo#"
    # For loop is at index 0, body at index 0_0, then y=3 is at index 1
    original = {key+"0_0": 200_000_000, key+"1": 100_000_000}
    optimized = {key+"0_0": 100_000_000, key+"1": 50_000_000}
    mapper = CommentMapper(test_obj, original, optimized)
    mapper.visit_FunctionDef(node) # 15.0μs -> 14.7μs (1.90% faster)

# ---------------- EDGE TEST CASES ----------------

def test_edge_empty_function():
    """
    Test a function with an empty body.
    """
    code = "def foo():\n    pass"
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    mapper = CommentMapper(test_obj, {}, {})
    mapper.visit_FunctionDef(node) # 5.13μs -> 5.31μs (3.39% slower)

def test_edge_nested_if_in_for():
    """
    Test a function with a for loop containing an if, with runtime entries for the inner if body.
    """
    code = (
        "def foo():\n"
        "    for i in range(2):\n"
        "        if i == 1:\n"
        "            x = i\n"
        "    y = 3"
    )
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    key = "foo#foo#"
    # for is index 0, its body is index 0_0, if's body is index 0_0_0
    # But code only checks up to 2 levels: for (i=0), if (j=0), x=i (body[0])
    original = {key+"0_0": 300_000_000, key+"1": 100_000_000}
    optimized = {key+"0_0": 100_000_000, key+"1": 80_000_000}
    mapper = CommentMapper(test_obj, original, optimized)
    mapper.visit_FunctionDef(node) # 17.5μs -> 15.4μs (13.8% faster)

def test_edge_missing_optimized_runtime():
    """
    Test a function where only the original runtime is present for a statement.
    """
    code = "def foo():\n    x = 1"
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    key = "foo#foo#0"
    original = {key: 100_000_000}
    optimized = {}
    mapper = CommentMapper(test_obj, original, optimized)
    mapper.visit_FunctionDef(node) # 5.42μs -> 5.61μs (3.40% slower)

def test_edge_missing_original_runtime():
    """
    Test a function where only the optimized runtime is present for a statement.
    """
    code = "def foo():\n    x = 1"
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    key = "foo#foo#0"
    original = {}
    optimized = {key: 50_000_000}
    mapper = CommentMapper(test_obj, original, optimized)
    mapper.visit_FunctionDef(node) # 5.32μs -> 5.57μs (4.49% slower)

def test_edge_negative_perf_gain():
    """
    Test a function where the optimized runtime is slower (negative gain).
    """
    code = "def foo():\n    x = 1"
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    key = "foo#foo#0"
    original = {key: 50_000_000}
    optimized = {key: 100_000_000}
    mapper = CommentMapper(test_obj, original, optimized)
    mapper.visit_FunctionDef(node) # 10.5μs -> 10.7μs (1.50% slower)

def test_edge_multiple_functions():
    """
    Test that context_stack handles multiple functions independently.
    """
    code = (
        "def foo():\n"
        "    x = 1\n"
        "def bar():\n"
        "    y = 2"
    )
    tree = ast.parse(code)
    test_obj = make_test_obj("foo.py")
    foo_key = "foo#foo#0"
    bar_key = "bar#foo#0"
    original = {foo_key: 10_000_000, bar_key: 20_000_000}
    optimized = {foo_key: 5_000_000, bar_key: 10_000_000}
    mapper = CommentMapper(test_obj, original, optimized)
    for node in tree.body:
        if isinstance(node, ast.FunctionDef):
            mapper.visit_FunctionDef(node)

def test_edge_different_file_paths():
    """
    Test that the file path is included in the key and affects matching.
    """
    code = "def foo():\n    x = 1"
    node = parse_func(code)
    test_obj = make_test_obj("abc.py")
    key = "foo#abc#0"
    original = {key: 1_000_000}
    optimized = {key: 500_000}
    mapper = CommentMapper(test_obj, original, optimized)
    mapper.visit_FunctionDef(node) # 10.8μs -> 10.8μs (0.000% faster)

# ---------------- LARGE SCALE TEST CASES ----------------

def test_large_many_statements():
    """
    Test a function with many statements (up to 1000), all with runtime entries.
    """
    num_lines = 1000
    lines = [f"    x{i} = {i}" for i in range(num_lines)]
    code = "def foo():\n" + "\n".join(lines)
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    key = "foo#foo#"
    original = {key+str(i): 1_000_000 + i * 1000 for i in range(num_lines)}
    optimized = {key+str(i): 500_000 + i * 500 for i in range(num_lines)}
    mapper = CommentMapper(test_obj, original, optimized)
    mapper.visit_FunctionDef(node) # 1.89ms -> 1.71ms (10.5% faster)
    # All lines should be present in results
    for i in range(num_lines):
        lineno = i + 2  # first line is def, so statements start at line 2

def test_large_many_for_loops():
    """
    Test a function with many for loops, each with a single body statement and runtime entries.
    """
    num_loops = 500
    lines = []
    for i in range(num_loops):
        lines.append(f"    for i{i} in range(1):")
        lines.append(f"        x{i} = i{i}")
    code = "def foo():\n" + "\n".join(lines)
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    key = "foo#foo#"
    # for loops are at even indices, so body at 0_0, 2_0, ..., (2*num_loops-2)_0
    original = {f"{key}{2*i}_0": 2_000_000+i*1000 for i in range(num_loops)}
    optimized = {f"{key}{2*i}_0": 1_000_000+i*500 for i in range(num_loops)}
    mapper = CommentMapper(test_obj, original, optimized)
    mapper.visit_FunctionDef(node) # 830μs -> 700μs (18.6% faster)
    # Each x{i} = i{i} is at line 2+2*i+1
    for i in range(num_loops):
        lineno = 2 + 2*i + 1

def test_large_deeply_nested():
    """
    Test a function with a for loop containing an if, repeated 100 times.
    """
    num_blocks = 100
    lines = []
    for i in range(num_blocks):
        lines.append(f"    for i{i} in range(1):")
        lines.append(f"        if i{i} == 0:")
        lines.append(f"            x{i} = i{i}")
    code = "def foo():\n" + "\n".join(lines)
    node = parse_func(code)
    test_obj = make_test_obj("foo.py")
    key = "foo#foo#"
    # for at 0, if at 0_0, assign at 0_0_0, but code only checks up to 2 levels
    # So for each block, match key is {key}{i*3}_0
    original = {f"{key}{i*3}_0": 3_000_000+i*1000 for i in range(num_blocks)}
    optimized = {f"{key}{i*3}_0": 1_000_000+i*500 for i in range(num_blocks)}
    mapper = CommentMapper(test_obj, original, optimized)
    mapper.visit_FunctionDef(node) # 242μs -> 135μs (78.8% faster)
    # Each x{i} = i{i} is at line 2+3*i+2
    for i in range(num_blocks):
        lineno = 2 + 3*i + 2
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from __future__ import annotations

import ast
from pathlib import Path

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.edit_generated_tests import CommentMapper


class GeneratedTests:
    # Minimal stub for the test object
    def __init__(self, behavior_file_path):
        self.behavior_file_path = Path(behavior_file_path)

# ---------------------------
# Unit Tests for visit_FunctionDef
# ---------------------------

# Helper to build a minimal ast.FunctionDef
def make_funcdef(name, body):
    return ast.FunctionDef(
        name=name,
        args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
        body=body,
        decorator_list=[]
    )

# Helper to build a minimal ast.Assign node
def make_assign(lineno):
    return ast.Assign(
        targets=[ast.Name(id="x", ctx=ast.Store())],
        value=ast.Constant(value=1),
        lineno=lineno,
        col_offset=0
    )

# Helper to build a minimal ast.For node
def make_for(lineno, body):
    return ast.For(
        target=ast.Name(id="i", ctx=ast.Store()),
        iter=ast.Call(func=ast.Name(id="range", ctx=ast.Load()), args=[ast.Constant(value=3)], keywords=[]),
        body=body,
        orelse=[],
        lineno=lineno,
        col_offset=0
    )

# Helper to build a minimal ast.If node
def make_if(lineno, body):
    return ast.If(
        test=ast.Constant(value=True),
        body=body,
        orelse=[],
        lineno=lineno,
        col_offset=0
    )

# ---------------------------
# 1. Basic Test Cases
# ---------------------------

def test_single_simple_statement():
    # Function with one statement, one runtime entry
    assign = make_assign(lineno=2)
    func = make_funcdef("foo", [assign])
    test = GeneratedTests("file.py")
    key = "foo#file#0"
    orig = {key: 1000000}
    opt = {key: 500000}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_FunctionDef(func) # 11.4μs -> 11.5μs (0.777% slower)

def test_multiple_statements():
    # Function with several statements, all with runtime entries
    assigns = [make_assign(lineno=i+2) for i in range(3)]
    func = make_funcdef("bar", assigns)
    test = GeneratedTests("file.py")
    key_base = "bar#file"
    orig = {f"{key_base}#{i}": 1000000*(i+1) for i in range(3)}
    opt = {f"{key_base}#{i}": 500000*(i+1) for i in range(3)}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_FunctionDef(func) # 15.6μs -> 15.5μs (0.579% faster)
    for i, assign in enumerate(assigns):
        expected = f"# {1.00*(i+1):.2f} ms -> {0.50*(i+1):.2f} ms (50% faster)"

def test_no_runtime_data():
    # Function with statements, but no runtime data
    assigns = [make_assign(lineno=10), make_assign(lineno=11)]
    func = make_funcdef("baz", assigns)
    test = GeneratedTests("file.py")
    mapper = CommentMapper(test, {}, {})
    mapper.visit_FunctionDef(func) # 5.95μs -> 5.99μs (0.668% slower)

def test_partial_runtime_data():
    # Only some statements have runtime data
    assigns = [make_assign(lineno=5), make_assign(lineno=6)]
    func = make_funcdef("qux", assigns)
    test = GeneratedTests("file.py")
    key = "qux#file#0"
    orig = {key: 2000000}
    opt = {key: 1000000}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_FunctionDef(func) # 10.5μs -> 10.3μs (1.66% faster)

# ---------------------------
# 2. Edge Test Cases
# ---------------------------

def test_empty_function_body():
    # Function with empty body
    func = make_funcdef("empty", [])
    test = GeneratedTests("file.py")
    mapper = CommentMapper(test, {}, {})
    mapper.visit_FunctionDef(func) # 3.67μs -> 4.29μs (14.5% slower)

def test_nested_for_with_assigns():
    # Function with a for loop containing assigns
    inner_assign = make_assign(lineno=8)
    for_node = make_for(lineno=7, body=[inner_assign])
    func = make_funcdef("loopfunc", [for_node])
    test = GeneratedTests("file.py")
    key = "loopfunc#file#0_0"
    orig = {key: 3000000}
    opt = {key: 1500000}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_FunctionDef(func) # 10.7μs -> 10.7μs (0.187% faster)

def test_if_with_multiple_body_nodes():
    # If statement with multiple body nodes
    assign1 = make_assign(lineno=12)
    assign2 = make_assign(lineno=13)
    if_node = make_if(lineno=11, body=[assign1, assign2])
    func = make_funcdef("iffunc", [if_node])
    test = GeneratedTests("file.py")
    key1 = "iffunc#file#0_0"
    key2 = "iffunc#file#0_1"
    orig = {key1: 4000000, key2: 2000000}
    opt = {key1: 2000000, key2: 1000000}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_FunctionDef(func) # 13.6μs -> 13.3μs (2.49% faster)

def test_for_with_nested_body():
    # For loop with a nested for loop inside
    inner_assign = make_assign(lineno=22)
    inner_for = make_for(lineno=21, body=[inner_assign])
    outer_for = make_for(lineno=20, body=[inner_for])
    func = make_funcdef("nested_for", [outer_for])
    test = GeneratedTests("file.py")
    key = "nested_for#file#0_0"
    orig = {key: 10000000}
    opt = {key: 8000000}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_FunctionDef(func) # 13.5μs -> 10.9μs (23.7% faster)

def test_statement_with_no_lineno():
    # Statement missing lineno attribute (should not crash)
    expr = ast.Expr(value=ast.Constant(value="no lineno"))  # no lineno set
    func = make_funcdef("nolineno", [expr])
    test = GeneratedTests("file.py")
    mapper = CommentMapper(test, {}, {})
    # Should not raise
    mapper.visit_FunctionDef(func) # 4.66μs -> 4.88μs (4.51% slower)

def test_faster_and_slower_status():
    # Test both 'faster' and 'slower' status in comment
    assign = make_assign(lineno=30)
    func = make_funcdef("statusfunc", [assign])
    test = GeneratedTests("file.py")
    key = "statusfunc#file#0"
    orig = {key: 1000000}
    opt = {key: 2000000}  # slower
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_FunctionDef(func) # 9.50μs -> 9.78μs (2.86% slower)

def test_zero_original_runtime():
    # Test when original runtime is zero (avoid division by zero)
    assign = make_assign(lineno=40)
    func = make_funcdef("zerofunc", [assign])
    test = GeneratedTests("file.py")
    key = "zerofunc#file#0"
    orig = {key: 0}
    opt = {key: 0}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_FunctionDef(func) # 8.58μs -> 8.70μs (1.38% slower)

# ---------------------------
# 3. Large Scale Test Cases
# ---------------------------

def test_large_function_many_statements():
    # Function with 100 statements, all with runtime data
    assigns = [make_assign(lineno=100+i) for i in range(100)]
    func = make_funcdef("bigfunc", assigns)
    test = GeneratedTests("file.py")
    key_base = "bigfunc#file"
    orig = {f"{key_base}#{i}": 1000000 for i in range(100)}
    opt = {f"{key_base}#{i}": 900000 for i in range(100)}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_FunctionDef(func) # 198μs -> 180μs (9.75% faster)
    for assign in assigns:
        pass

def test_large_nested_for():
    # Function with a for loop containing 500 assigns
    assigns = [make_assign(lineno=200+i) for i in range(500)]
    for_node = make_for(lineno=199, body=assigns)
    func = make_funcdef("hugefor", [for_node])
    test = GeneratedTests("file.py")
    key_base = "hugefor#file#0"
    orig = {f"{key_base}_{i}": 2000000 for i in range(500)}
    opt = {f"{key_base}_{i}": 1000000 for i in range(500)}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_FunctionDef(func) # 1.06ms -> 954μs (10.6% faster)
    for assign in assigns:
        pass

def test_large_if_with_many_body_nodes():
    # If statement with 50 assigns
    assigns = [make_assign(lineno=300+i) for i in range(50)]
    if_node = make_if(lineno=299, body=assigns)
    func = make_funcdef("bigif", [if_node])
    test = GeneratedTests("file.py")
    key_base = "bigif#file#0"
    orig = {f"{key_base}_{i}": 5000000 for i in range(50)}
    opt = {f"{key_base}_{i}": 4000000 for i in range(50)}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_FunctionDef(func) # 115μs -> 106μs (8.14% faster)
    for assign in assigns:
        pass
# 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-pr719-2025-09-08T23.46.13 and push.

[KRRT7]

through several targeted micro-optimizations that reduce overhead in the hot loops:

@mohammedahmed18 is it possible the prompt is inadvertently looking for micro-optimizations? I'm still seeing a couple of those.

ultimately closing because my dependent PR is a bit messed up, will rerun opt later

[mohammedahmed18]

@KRRT7 the prompt doesn’t block micro opts on purpose, for not holding codeflash back from generating possible optimizations