⚡️ Speed up method CommentMapper.visit_AsyncFunctionDef by 399% in PR #769 (clean-async-branch)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #769

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

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

---

📄 399% (3.99x) speedup for CommentMapper.visit_AsyncFunctionDef in codeflash/code_utils/edit_generated_tests.py

⏱️ Runtime : 1.22 milliseconds → 244 microseconds (best of 178 runs)

📝 Explanation and details

The optimization achieves a 398% speedup by reducing repeated attribute lookups and string operations that were happening in tight loops.

Key optimizations applied:

1. Cached repeated attribute lookups: Instead of accessing self.context_stack, self.original_runtimes, self.optimized_runtimes, and self.get_comment repeatedly in loops, these are cached as local variables at the start of the method. This eliminates hundreds of attribute lookups during execution.

2. Pre-computed string concatenation: The original code was rebuilding test_qualified_name + "#" + str(self.abs_path) on every function call. The optimized version computes key_base once and reuses it, avoiding repeated string joins and str() calls.

3. Optimized loop structure: Replaced the while j >= 0 loop with for j in range(compound_body_len - 1, -1, -1) which has better performance characteristics and eliminates repeated len() calls.

4. Streamlined getattr usage: Instead of calling getattr(compound_line_node, "body", []) and then extending a list, the code now checks for the presence of a body attribute once and conditionally adds it, reducing function call overhead.

5. F-string optimization: Replaced string concatenation with f-strings (f"{i}_{j}" instead of str(i) + "_" + str(j)) which are more efficient in Python.

The line profiler shows the dramatic impact - the most expensive operations in the original (like self.get_comment() calls taking 23.9% and 18.8% of total time) now execute much faster due to cached attribute lookups. The optimizations are particularly effective for large-scale test cases with many statements or nested compound structures, showing 400%+ speedups in functions with 100+ statements.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import ast
from pathlib import Path

# imports
import pytest
from codeflash.code_utils.edit_generated_tests import CommentMapper


class GeneratedTests:
    def __init__(self, behavior_file_path: Path):
        self.behavior_file_path = behavior_file_path

# Helper to parse code and get the AsyncFunctionDef node
def get_async_func_node(code: str) -> ast.AsyncFunctionDef:
    tree = ast.parse(code)
    for node in tree.body:
        if isinstance(node, ast.AsyncFunctionDef):
            return node
    raise ValueError("No AsyncFunctionDef found")

# ========== Basic Test Cases ==========

def test_basic_single_statement():
    # Test with a simple async function with one statement
    code = "async def foo():\n    x = 1"
    node = get_async_func_node(code)
    # Prepare runtime dictionaries
    path = Path("testfile.py")
    test = GeneratedTests(path)
    key = "foo#testfile#0"
    original_runtimes = {key: 1000}
    optimized_runtimes = {key: 500}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_AsyncFunctionDef(node)
    comment = mapper.results[node.body[0].lineno]

def test_basic_multiple_statements():
    # Async function with multiple statements
    code = "async def foo():\n    x = 1\n    y = 2\n    z = 3"
    node = get_async_func_node(code)
    path = Path("testfile.py")
    test = GeneratedTests(path)
    # Provide runtime for each statement
    keys = [f"foo#testfile#{i}" for i in range(3)]
    original_runtimes = {keys[0]: 1000, keys[1]: 2000, keys[2]: 3000}
    optimized_runtimes = {keys[0]: 800, keys[1]: 1500, keys[2]: 2500}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_AsyncFunctionDef(node)
    # Each statement should get a comment
    for i, stmt in enumerate(node.body):
        comment = mapper.results[stmt.lineno]

def test_basic_with_compound_statement():
    # Async function with a for loop and assignments inside
    code = (
        "async def foo():\n"
        "    for i in range(2):\n"
        "        x = i\n"
        "        y = i + 1\n"
        "    z = 42"
    )
    node = get_async_func_node(code)
    path = Path("testfile.py")
    test = GeneratedTests(path)
    # Provide runtime for outer and inner statements
    key_base = "foo#testfile"
    original_runtimes = {
        f"{key_base}#0_0": 1000,  # x = i
        f"{key_base}#0_1": 2000,  # y = i + 1
        f"{key_base}#1": 3000     # z = 42
    }
    optimized_runtimes = {
        f"{key_base}#0_0": 900,
        f"{key_base}#0_1": 1800,
        f"{key_base}#1": 2500
    }
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_AsyncFunctionDef(node) # 18.6μs -> 10.2μs (81.8% faster)
    # Should annotate each inner statement and the one after the loop
    for i, stmt in enumerate(node.body):
        if isinstance(stmt, ast.For):
            for j, inner_stmt in enumerate(stmt.body):
                pass
        else:
            pass

# ========== Edge Test Cases ==========

def test_edge_empty_body():
    # Async function with no body
    code = "async def foo():\n    pass"
    node = get_async_func_node(code)
    path = Path("testfile.py")
    test = GeneratedTests(path)
    original_runtimes = {}
    optimized_runtimes = {}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_AsyncFunctionDef(node) # 5.59μs -> 5.86μs (4.62% slower)

def test_edge_missing_runtime_keys():
    # Async function with statements, but runtimes missing for some
    code = "async def foo():\n    x = 1\n    y = 2"
    node = get_async_func_node(code)
    path = Path("testfile.py")
    test = GeneratedTests(path)
    key = "foo#testfile#0"
    # Only one statement has runtime data
    original_runtimes = {key: 1000}
    optimized_runtimes = {key: 900}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_AsyncFunctionDef(node) # 11.2μs -> 6.34μs (76.6% faster)

def test_edge_nested_compound_statements():
    # Async function with nested for and if statements
    code = (
        "async def foo():\n"
        "    for i in range(2):\n"
        "        if i % 2 == 0:\n"
        "            x = i\n"
        "        else:\n"
        "            y = i\n"
        "    z = 42"
    )
    node = get_async_func_node(code)
    path = Path("testfile.py")
    test = GeneratedTests(path)
    key_base = "foo#testfile"
    original_runtimes = {
        f"{key_base}#0_0": 1000,  # x = i
        f"{key_base}#0_1": 1100,  # y = i
        f"{key_base}#1": 1200     # z = 42
    }
    optimized_runtimes = {
        f"{key_base}#0_0": 900,
        f"{key_base}#0_1": 1000,
        f"{key_base}#1": 1100
    }
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_AsyncFunctionDef(node) # 17.9μs -> 9.50μs (88.6% faster)
    # Should annotate both x = i and y = i, and z = 42
    found = 0
    for stmt in node.body:
        if isinstance(stmt, ast.For):
            for if_stmt in stmt.body:
                if isinstance(if_stmt, ast.If):
                    for j, branch in enumerate(if_stmt.body + if_stmt.orelse):
                        found += 1
        else:
            found += 1

def test_edge_non_assign_statements():
    # Async function with non-assign statements (e.g., return, expr)
    code = (
        "async def foo():\n"
        "    x = 1\n"
        "    print(x)\n"
        "    return x"
    )
    node = get_async_func_node(code)
    path = Path("testfile.py")
    test = GeneratedTests(path)
    key_base = "foo#testfile"
    original_runtimes = {
        f"{key_base}#0": 1000,  # x = 1
        f"{key_base}#1": 1100,  # print(x)
        f"{key_base}#2": 1200   # return x
    }
    optimized_runtimes = {
        f"{key_base}#0": 900,
        f"{key_base}#1": 1000,
        f"{key_base}#2": 1100
    }
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_AsyncFunctionDef(node) # 16.8μs -> 6.63μs (153% faster)
    # Should annotate all statements
    for stmt in node.body:
        pass

def test_edge_path_with_suffix():
    # Test that .with_suffix("") works as expected
    code = "async def foo():\n    x = 1"
    node = get_async_func_node(code)
    path = Path("testfile.py")
    test = GeneratedTests(path)
    key = "foo#testfile#0"
    original_runtimes = {key: 1000}
    optimized_runtimes = {key: 500}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_AsyncFunctionDef(node) # 10.1μs -> 5.39μs (87.4% faster)

# ========== Large Scale Test Cases ==========

def test_large_many_statements():
    # Async function with 100 statements
    code = "async def foo():\n" + "\n".join([f"    x{i} = {i}" for i in range(100)])
    node = get_async_func_node(code)
    path = Path("testfile.py")
    test = GeneratedTests(path)
    original_runtimes = {f"foo#testfile#{i}": 1000+i for i in range(100)}
    optimized_runtimes = {f"foo#testfile#{i}": 900+i for i in range(100)}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_AsyncFunctionDef(node) # 189μs -> 36.6μs (416% faster)
    # All 100 statements should be annotated
    for stmt in node.body:
        pass

def test_large_many_compound_statements():
    # Async function with 50 for loops, each with 2 assignments
    code = "async def foo():\n" + "\n".join([
        f"    for i{i} in range(2):\n        x{i} = i{i}\n        y{i} = i{i}+1"
        for i in range(50)
    ])
    node = get_async_func_node(code)
    path = Path("testfile.py")
    test = GeneratedTests(path)
    key_base = "foo#testfile"
    original_runtimes = {}
    optimized_runtimes = {}
    for i in range(50):
        original_runtimes[f"{key_base}#{i}_0"] = 1000 + i
        original_runtimes[f"{key_base}#{i}_1"] = 2000 + i
        optimized_runtimes[f"{key_base}#{i}_0"] = 900 + i
        optimized_runtimes[f"{key_base}#{i}_1"] = 1800 + i
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_AsyncFunctionDef(node) # 242μs -> 76.0μs (220% faster)
    # All inner assignments should be annotated
    count = 0
    for stmt in node.body:
        if isinstance(stmt, ast.For):
            for j, inner_stmt in enumerate(stmt.body):
                count += 1

def test_large_deeply_nested():
    # Async function with nested for/if, 10 levels deep
    code = "async def foo():\n"
    for i in range(10):
        code += "    " * (i+1) + f"for i{i} in range(1):\n"
    code += "    " * 11 + "x = 42\n"
    node = get_async_func_node(code)
    path = Path("testfile.py")
    test = GeneratedTests(path)
    # Only annotate the deepest assignment
    key = "foo#testfile#" + "_".join([str(i) for i in range(10)])  # Not used, but for illustration
    original_runtimes = {f"foo#testfile#0_0_0_0_0_0_0_0_0_0": 12345}
    optimized_runtimes = {f"foo#testfile#0_0_0_0_0_0_0_0_0_0": 12000}
    # Manually fix the key for the deepest assignment
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    # Patch the key to match what the function expects
    # In this case, the function will produce inv_id "0_0_0_0_0_0_0_0_0_0" for the deepest node
    # So we ensure our dicts match that
    mapper.original_runtimes = {"foo#testfile#0_0_0_0_0_0_0_0_0_0": 12345}
    mapper.optimized_runtimes = {"foo#testfile#0_0_0_0_0_0_0_0_0_0": 12000}
    mapper.visit_AsyncFunctionDef(node) # 7.67μs -> 8.10μs (5.19% slower)
    # Only the deepest assignment should be annotated
    found = False
    for stmt in node.body:
        # Traverse down to the deepest assignment
        current = stmt
        for _ in range(10):
            current = current.body[0]
        found = True
# 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 types import SimpleNamespace

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


class GeneratedTests:
    # Dummy class for compatibility
    def __init__(self, behavior_file_path):
        self.behavior_file_path = behavior_file_path

# Minimal Path mock for compatibility
class DummyPath:
    def __init__(self, name):
        self.name = name
    def with_suffix(self, suffix):
        return self.name

# Helper to create AsyncFunctionDef nodes
def make_async_func(name, body):
    return ast.AsyncFunctionDef(
        name=name,
        args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
        body=body,
        decorator_list=[],
        returns=None,
        type_comment=None,
    )

# Helper to create a dummy statement with lineno
def make_stmt(lineno):
    # Use ast.Pass as a simple statement
    node = ast.Pass()
    node.lineno = lineno
    return node

# Helper to create a compound statement (e.g., ast.For) with body and lineno
def make_for(lineno, body):
    node = ast.For(target=ast.Name(id="i", ctx=ast.Store()), iter=ast.Name(id="lst", ctx=ast.Load()), body=body, orelse=[])
    node.lineno = lineno
    return node

def make_with(lineno, body):
    node = ast.With(items=[], body=body)
    node.lineno = lineno
    return node

def make_if(lineno, body):
    node = ast.If(test=ast.Name(id="cond", ctx=ast.Load()), body=body, orelse=[])
    node.lineno = lineno
    return node

def make_assign(lineno):
    node = ast.Assign(targets=[ast.Name(id="x", ctx=ast.Store())], value=ast.Constant(value=1))
    node.lineno = lineno
    return node

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

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

def test_basic_single_stmt():
    # Single statement async function
    func = make_async_func("foo", [make_stmt(10)])
    # Prepare runtime keys
    key = "foo#dummy_path#0"
    orig = {key: 2_000_000_000}
    opt = {key: 1_000_000_000}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 8.40μs -> 3.01μs (179% faster)

def test_basic_multiple_stmts():
    # Multiple statements
    func = make_async_func("bar", [make_stmt(5), make_stmt(6)])
    key1 = "bar#dummy_path#1"
    key0 = "bar#dummy_path#0"
    orig = {key0: 3_000_000_000, key1: 2_000_000_000}
    opt = {key0: 1_000_000_000, key1: 1_500_000_000}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 11.1μs -> 3.35μs (231% faster)

def test_basic_compound_for():
    # Async function with a for loop containing statements
    for_body = [make_stmt(20), make_stmt(21)]
    func = make_async_func("baz", [make_for(15, for_body)])
    key_0_1 = "baz#dummy_path#0_1"
    key_0_0 = "baz#dummy_path#0_0"
    orig = {key_0_0: 5_000_000_000, key_0_1: 4_000_000_000}
    opt = {key_0_0: 2_000_000_000, key_0_1: 1_000_000_000}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 12.4μs -> 5.51μs (124% faster)

def test_basic_compound_with_assign():
    # Async function with a with block containing assignments
    with_body = [make_assign(30)]
    func = make_async_func("qux", [make_with(25, with_body)])
    key_0_0 = "qux#dummy_path#0_0"
    orig = {key_0_0: 1_000_000_000}
    opt = {key_0_0: 2_000_000_000}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 8.48μs -> 4.58μs (85.1% faster)

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

def test_edge_no_body():
    # Async function with empty body
    func = make_async_func("empty", [])
    orig = {}
    opt = {}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 1.70μs -> 1.84μs (7.54% slower)

def test_edge_missing_runtime_keys():
    # Async function with statements but no runtime keys
    func = make_async_func("foo", [make_stmt(42)])
    orig = {}
    opt = {}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 2.62μs -> 2.79μs (6.10% slower)

def test_edge_compound_missing_keys():
    # Compound statement with some missing runtime keys
    for_body = [make_stmt(50), make_stmt(51)]
    func = make_async_func("edge", [make_for(49, for_body)])
    key_0_1 = "edge#dummy_path#0_1"
    orig = {key_0_1: 4_000_000_000}
    opt = {key_0_1: 2_000_000_000}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 9.85μs -> 5.26μs (87.2% faster)

def test_edge_nested_compound():
    # Nested compound statements (e.g., for inside if)
    inner_body = [make_stmt(61)]
    if_body = [make_for(60, inner_body)]
    func = make_async_func("nest", [make_if(59, if_body)])
    key_0_0_0 = "nest#dummy_path#0_0"
    orig = {key_0_0_0: 10_000_000_000}
    opt = {key_0_0_0: 5_000_000_000}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 11.4μs -> 5.33μs (115% faster)

def test_edge_assign_in_compound():
    # Assignment inside compound statement
    with_body = [make_assign(70)]
    func = make_async_func("assign", [make_with(69, with_body)])
    key_0_0 = "assign#dummy_path#0_0"
    orig = {key_0_0: 2_000_000_000}
    opt = {key_0_0: 1_000_000_000}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 8.02μs -> 4.49μs (78.6% faster)

def test_edge_multiple_compounds():
    # Multiple compound statements in one function
    for_body1 = [make_stmt(80)]
    for_body2 = [make_stmt(81)]
    func = make_async_func("multi", [make_for(78, for_body1), make_for(79, for_body2)])
    key_0_0 = "multi#dummy_path#0_0"
    key_1_0 = "multi#dummy_path#1_0"
    orig = {key_0_0: 8_000_000_000, key_1_0: 4_000_000_000}
    opt = {key_0_0: 4_000_000_000, key_1_0: 2_000_000_000}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 11.8μs -> 5.96μs (97.3% faster)

def test_edge_if_else_branch():
    # If statement with both body and orelse
    if_body = [make_stmt(90)]
    else_body = [make_stmt(91)]
    node = ast.If(test=ast.Name(id="cond", ctx=ast.Load()), body=if_body, orelse=else_body)
    node.lineno = 89
    func = make_async_func("branch", [node])
    key_0_0 = "branch#dummy_path#0_0"
    key_0_1 = "branch#dummy_path#0_1"
    orig = {key_0_0: 1_000_000_000, key_0_1: 2_000_000_000}
    opt = {key_0_0: 500_000_000, key_0_1: 1_000_000_000}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 8.19μs -> 4.36μs (87.8% faster)

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

def test_large_many_statements():
    # Large async function with many statements
    N = 500
    body = [make_stmt(lineno) for lineno in range(1000, 1000 + N)]
    func = make_async_func("bigfunc", body)
    orig = {}
    opt = {}
    # Annotate every 10th statement
    for i in range(N):
        if i % 10 == 0:
            key = f"bigfunc#dummy_path#{i}"
            orig[key] = 1_000_000_000 + i * 1_000_000
            opt[key] = 500_000_000 + i * 500_000
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func)
    # Only every 10th line should be annotated
    for i in range(N):
        lineno = 1000 + i
        if i % 10 == 0:
            # Check comment format
            comment = mapper.results[lineno]
        else:
            pass

def test_large_many_compounds():
    # Large async function with many compound statements
    N = 100
    func_body = []
    orig = {}
    opt = {}
    for i in range(N):
        # Each compound has a statement with unique lineno
        stmt = make_stmt(2000 + i)
        compound = make_for(1999 + i, [stmt])
        func_body.append(compound)
        key = f"huge#dummy_path#{i}_0"
        orig[key] = 2_000_000_000 + i * 1_000_000
        opt[key] = 1_000_000_000 + i * 500_000
    func = make_async_func("huge", func_body)
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func)
    # All N statements should be annotated
    for i in range(N):
        lineno = 2000 + i
        comment = mapper.results[lineno]

def test_large_nested_compounds():
    # Large async function with nested compounds
    N = 50
    func_body = []
    orig = {}
    opt = {}
    for i in range(N):
        inner_stmt = make_stmt(3000 + i)
        inner_for = make_for(2999 + i, [inner_stmt])
        outer_if = make_if(2998 + i, [inner_for])
        func_body.append(outer_if)
        key = f"deep#dummy_path#{i}_0_0"
        orig[key] = 5_000_000_000 + i * 1_000_000
        opt[key] = 2_500_000_000 + i * 500_000
    func = make_async_func("deep", func_body)
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func)
    # All N inner statements should be annotated
    for i in range(N):
        lineno = 3000 + i
        comment = mapper.results[lineno]

def test_large_no_runtime_keys():
    # Large async function with no runtime keys
    N = 100
    body = [make_stmt(lineno) for lineno in range(4000, 4000 + N)]
    func = make_async_func("norun", body)
    orig = {}
    opt = {}
    mapper = CommentMapper(GeneratedTests(DummyPath("dummy_path")), orig, opt)
    mapper.visit_AsyncFunctionDef(func) # 38.4μs -> 32.6μs (17.9% faster)
# 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-pr769-2025-09-27T01.01.48 and push.

[KRRT7]

still don't like these

[KRRT7]

will manually implement this later.