⚡️ Speed up method CommentMapper.visit_FunctionDef by 84% in PR #687 (granular-async-instrumentation)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #687

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

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

---

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

⏱️ Runtime : 3.03 milliseconds → 1.65 milliseconds (best of 148 runs)

📝 Explanation and details

The optimization replaces the expensive ast.walk() call with a targeted node traversal that only checks the immediate statement and its direct body children.

Key change: Instead of ast.walk(compound_line_node) which recursively traverses the entire AST subtree, the optimized code creates a focused list:

nodes_to_check = [compound_line_node]
nodes_to_check.extend(getattr(compound_line_node, 'body', []))

This dramatically reduces the number of nodes processed in the inner loop. The line profiler shows ast.walk() was the major bottleneck (46.2% of total time, 8.23ms), while the optimized version's equivalent loop takes only 1.9% of total time (180μs).

Why this works: The code only needs to check statements at the current level and one level deep (direct children in compound statement bodies like for, if, while, with). The original ast.walk() was doing unnecessary deep traversal of nested structures.

Performance impact: The optimization is most effective for test cases with compound statements (for/while/if/with blocks) containing multiple nested nodes, showing 73-156% speedups in those scenarios. Simple statement functions see smaller but consistent 1-3% improvements due to reduced overhead.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 54 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
# -- Patch the imported functions in the tested class --
import sys
import types
from pathlib import Path

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.edit_generated_tests import CommentMapper
from codeflash.code_utils.time_utils import format_perf, format_time
from codeflash.models.models import GeneratedTests
from codeflash.result.critic import performance_gain


# Minimal mock for GeneratedTests
class GeneratedTests:
    def __init__(self, behavior_file_path):
        self.behavior_file_path = Path(behavior_file_path)

# --- Helper functions for test setup ---

def make_funcdef(
    name="foo", 
    body=None, 
    lineno=1
):
    """Create a simple ast.FunctionDef node with the given body."""
    if body is None:
        body = [ast.Pass(lineno=lineno+1, col_offset=0)]
    return ast.FunctionDef(
        name=name,
        args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
        body=body,
        decorator_list=[],
        lineno=lineno,
        col_offset=0
    )

def make_with_block(body, lineno=2):
    """Create a simple ast.With node with the given body."""
    return ast.With(
        items=[],
        body=body,
        lineno=lineno,
        col_offset=0
    )

def make_for_block(body, lineno=2):
    """Create a simple ast.For node with the given body."""
    return ast.For(
        target=ast.Name(id="i", ctx=ast.Store(), lineno=lineno, col_offset=0),
        iter=ast.Name(id="range", ctx=ast.Load(), lineno=lineno, col_offset=0),
        body=body,
        orelse=[],
        lineno=lineno,
        col_offset=0
    )

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

def make_if_block(body, lineno=2):
    return ast.If(
        test=ast.Constant(value=True, lineno=lineno, col_offset=0),
        body=body,
        orelse=[],
        lineno=lineno,
        col_offset=0
    )

def make_while_block(body, lineno=2):
    return ast.While(
        test=ast.Constant(value=True, lineno=lineno, col_offset=0),
        body=body,
        orelse=[],
        lineno=lineno,
        col_offset=0
    )

# --- Test cases ---

# 1. BASIC TEST CASES

def test_single_statement_with_runtime_data():
    """Test that a single statement in a function gets a comment if runtime data is present."""
    func = make_funcdef(name="foo", body=[make_assign(lineno=2)], lineno=1)
    test = GeneratedTests("foo.py")
    # Compose the correct key
    key = "foo#foo#foo#2"
    original = {"foo#foo#foo#0": 100}
    optimized = {"foo#foo#foo#0": 50}
    # The keys are constructed as: qualified_name + "#" + abs_path + "#" + inv_id
    # In this case: context_stack = ["foo"], abs_path = Path("foo"), inv_id = "0"
    # So key should be: "foo#foo#0"
    original = {"foo#foo#0": 100}
    optimized = {"foo#foo#0": 50}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 10.5μs -> 10.4μs (0.488% faster)
    # Check comment format
    comment = mapper.results[2]

def test_multiple_statements_some_with_runtime_data():
    """Test that only statements with runtime data are commented."""
    func = make_funcdef(
        name="bar",
        body=[
            make_assign(lineno=2),
            make_assign(lineno=3),
            make_assign(lineno=4),
        ],
        lineno=1
    )
    test = GeneratedTests("bar.py")
    original = {"bar#bar#1": 80}
    optimized = {"bar#bar#1": 40}
    # Only the second statement (index 1) has runtime data
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 10.6μs -> 10.4μs (1.84% faster)
    comment = mapper.results[3]

def test_with_block_with_runtime_data():
    """Test a function with a with-block containing a statement with runtime data."""
    inner = make_assign(lineno=3)
    with_block = make_with_block([inner], lineno=2)
    func = make_funcdef(name="baz", body=[with_block], lineno=1)
    test = GeneratedTests("baz.py")
    # The inv_id for with body is i_j: i=0 (with), j=0 (first in with body)
    original = {"baz#baz#0_0": 200}
    optimized = {"baz#baz#0_0": 100}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 18.3μs -> 10.1μs (81.6% faster)
    comment = mapper.results[3]

def test_for_block_with_runtime_data():
    """Test a function with a for-block containing a statement with runtime data."""
    inner = make_assign(lineno=4)
    for_block = make_for_block([inner], lineno=3)
    func = make_funcdef(name="qux", body=[for_block], lineno=2)
    test = GeneratedTests("qux.py")
    original = {"qux#qux#0_0": 300}
    optimized = {"qux#qux#0_0": 150}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 17.2μs -> 9.77μs (76.3% faster)
    comment = mapper.results[4]

# 2. EDGE TEST CASES

def test_no_statements():
    """Test a FunctionDef node with no statements."""
    func = make_funcdef(name="empty", body=[], lineno=1)
    test = GeneratedTests("empty.py")
    original = {}
    optimized = {}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 3.67μs -> 3.63μs (1.10% faster)

def test_statements_without_runtime_data():
    """Test that no comments are added if there is no runtime data for any statement."""
    func = make_funcdef(
        name="no_data",
        body=[make_assign(lineno=2), make_assign(lineno=3)],
        lineno=1
    )
    test = GeneratedTests("no_data.py")
    original = {}
    optimized = {}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 5.47μs -> 5.49μs (0.364% slower)

def test_deeply_nested_blocks():
    """Test function with nested compound statements (e.g., for inside with inside if)."""
    inner = make_assign(lineno=6)
    for_block = make_for_block([inner], lineno=5)
    with_block = make_with_block([for_block], lineno=4)
    if_block = make_if_block([with_block], lineno=3)
    func = make_funcdef(name="deep", body=[if_block], lineno=2)
    test = GeneratedTests("deep.py")
    # The inv_id for the assign is i=0 (if), j=0 (with), k=0 (for), l=0 (assign)
    # But only i_j is used: i=0 (if), j=0 (with body)
    # The code only supports one level of compound statement for inv_id
    # So the key for the assign in for inside with inside if is: "deep#deep#0_0"
    original = {"deep#deep#0_0": 1000}
    optimized = {"deep#deep#0_0": 500}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func)
    comment = mapper.results[6]

def test_block_with_missing_optimized_runtime():
    """Test that no comment is added if optimized runtime is missing."""
    func = make_funcdef(name="foo", body=[make_assign(lineno=2)], lineno=1)
    test = GeneratedTests("foo.py")
    original = {"foo#foo#0": 100}
    optimized = {}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 4.85μs -> 5.61μs (13.6% slower)

def test_block_with_missing_original_runtime():
    """Test that no comment is added if original runtime is missing."""
    func = make_funcdef(name="foo", body=[make_assign(lineno=2)], lineno=1)
    test = GeneratedTests("foo.py")
    original = {}
    optimized = {"foo#foo#0": 50}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 4.55μs -> 4.95μs (8.10% slower)

def test_block_with_slower_optimized_time():
    """Test that 'slower' is shown if optimized time is greater than original time."""
    func = make_funcdef(name="foo", body=[make_assign(lineno=2)], lineno=1)
    test = GeneratedTests("foo.py")
    original = {"foo#foo#0": 50}
    optimized = {"foo#foo#0": 100}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 9.11μs -> 9.31μs (2.14% slower)
    comment = mapper.results[2]

def test_function_with_multiple_compound_blocks():
    """Test a function with several compound blocks, each with runtime data."""
    with1 = make_with_block([make_assign(lineno=3)], lineno=2)
    for1 = make_for_block([make_assign(lineno=5)], lineno=4)
    func = make_funcdef(name="multi", body=[with1, for1], lineno=1)
    test = GeneratedTests("multi.py")
    original = {
        "multi#multi#0_0": 10,
        "multi#multi#1_0": 20
    }
    optimized = {
        "multi#multi#0_0": 5,
        "multi#multi#1_0": 10
    }
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 25.8μs -> 13.5μs (90.2% faster)

def test_function_with_while_block():
    """Test a function with a while-block containing a statement with runtime data."""
    inner = make_assign(lineno=3)
    while_block = make_while_block([inner], lineno=2)
    func = make_funcdef(name="whilefun", body=[while_block], lineno=1)
    test = GeneratedTests("whilefun.py")
    original = {"whilefun#whilefun#0_0": 400}
    optimized = {"whilefun#whilefun#0_0": 200}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 16.8μs -> 9.69μs (73.2% faster)

# 3. LARGE SCALE TEST CASES

def test_large_number_of_statements():
    """Test that the function can handle a large number of statements efficiently."""
    N = 500
    body = [make_assign(lineno=i+2) for i in range(N)]
    func = make_funcdef(name="large", body=body, lineno=1)
    test = GeneratedTests("large.py")
    # Only even-indexed statements have runtime data
    original = {f"large#large#{i}": 100+i for i in range(0, N, 2)}
    optimized = {f"large#large#{i}": 50+i for i in range(0, N, 2)}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 496μs -> 490μs (1.12% faster)
    # Only even linenos should have comments
    for i in range(N):
        lineno = i+2
        if i % 2 == 0:
            comment = mapper.results[lineno]
        else:
            pass

def test_large_number_of_nested_blocks():
    """Test that the function can handle a large number of nested compound blocks."""
    N = 100
    # Each for-block contains one assign
    body = [make_for_block([make_assign(lineno=i*2+2)], lineno=i*2+1) for i in range(N)]
    func = make_funcdef(name="nested", body=body, lineno=1)
    test = GeneratedTests("nested.py")
    original = {f"nested#nested#{i}_0": 1000+i for i in range(N)}
    optimized = {f"nested#nested#{i}_0": 500+i for i in range(N)}
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 613μs -> 239μs (156% faster)
    for i in range(N):
        lineno = i*2+2
        comment = mapper.results[lineno]

def test_performance_with_mixed_blocks():
    """Test a function with a mix of compound and simple statements at scale."""
    N = 200
    body = []
    original = {}
    optimized = {}
    test = GeneratedTests("mix.py")
    # Alternate between assign and for-block
    for i in range(N):
        if i % 2 == 0:
            # Simple assign
            body.append(make_assign(lineno=i+2))
            original[f"mix#mix#{i}"] = 200+i
            optimized[f"mix#mix#{i}"] = 100+i
        else:
            # For block with one assign
            body.append(make_for_block([make_assign(lineno=i+2)], lineno=i+1))
            original[f"mix#mix#{i}_0"] = 300+i
            optimized[f"mix#mix#{i}_0"] = 150+i
    func = make_funcdef(name="mix", body=body, lineno=1)
    mapper = CommentMapper(test, original, optimized)
    mapper.visit_FunctionDef(func) # 757μs -> 366μs (106% faster)
    # Check all expected comments
    for i in range(N):
        if i % 2 == 0:
            lineno = i+2
        else:
            lineno = 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:
    def __init__(self, behavior_file_path):
        self.behavior_file_path = Path(behavior_file_path)

# unit tests

# Helper to parse code and extract the FunctionDef node
def get_func_node(code, func_name="foo"):
    module = ast.parse(code)
    for node in module.body:
        if isinstance(node, ast.FunctionDef) and node.name == func_name:
            return node
    raise ValueError("Function not found")

# Helper to build keys for runtime dicts
def build_key(func_name, abs_path, inv_id):
    return f"{func_name}#{abs_path}#{inv_id}"

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

def test_single_statement_function():
    # Test a function with a single statement
    code = "def foo():\n    x = 1"
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    key = build_key("foo", abs_path, "0")
    original_runtimes = {key: 100}
    optimized_runtimes = {key: 50}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_FunctionDef(node) # 9.86μs -> 9.83μs (0.295% faster)

def test_multiple_statements_function():
    # Test a function with several statements
    code = "def foo():\n    a = 1\n    b = 2\n    c = 3"
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    keys = [build_key("foo", abs_path, str(i)) for i in range(3)]
    original_runtimes = {k: 100*(i+1) for i, k in enumerate(keys)}
    optimized_runtimes = {k: 50*(i+1) for i, k in enumerate(keys)}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_FunctionDef(node) # 13.9μs -> 13.7μs (1.31% faster)
    # Should annotate lines 2, 3, 4
    for lineno in [2, 3, 4]:
        pass

def test_with_compound_statement():
    # Test a function with a for loop containing statements
    code = (
        "def foo():\n"
        "    for i in range(2):\n"
        "        x = i\n"
        "        y = i+1\n"
        "    z = 42"
    )
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    # for-loop is at index 0, z=42 is at index 1
    # x=i is at for.body[0], y=i+1 is at for.body[1]
    key_base = f"foo#{abs_path}"
    original_runtimes = {
        f"{key_base}#0_0": 100,
        f"{key_base}#0_1": 200,
        f"{key_base}#1": 300,
    }
    optimized_runtimes = {
        f"{key_base}#0_0": 80,
        f"{key_base}#0_1": 150,
        f"{key_base}#1": 250,
    }
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_FunctionDef(node) # 32.6μs -> 16.1μs (102% faster)

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

def test_no_runtime_data():
    # Test a function where no runtime data is present
    code = "def foo():\n    x = 1\n    y = 2"
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    mapper = CommentMapper(test, {}, {})
    mapper.visit_FunctionDef(node) # 5.78μs -> 5.80μs (0.328% slower)

def test_missing_optimized_runtime():
    # Test a function where only original runtime is present
    code = "def foo():\n    x = 1"
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    key = build_key("foo", abs_path, "0")
    original_runtimes = {key: 100}
    optimized_runtimes = {}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_FunctionDef(node) # 5.21μs -> 5.08μs (2.54% faster)

def test_nested_compound_statements():
    # Test a function with nested compound statements (for inside if)
    code = (
        "def foo():\n"
        "    if True:\n"
        "        for i in range(2):\n"
        "            x = i\n"
        "    y = 5"
    )
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    key_base = f"foo#{abs_path}"
    # if is at index 0, for is at if.body[0], x=i is at for.body[0]
    original_runtimes = {
        f"{key_base}#0_0": 100,  # x = i
        f"{key_base}#1": 200,    # y = 5
    }
    optimized_runtimes = {
        f"{key_base}#0_0": 80,
        f"{key_base}#1": 150,
    }
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_FunctionDef(node) # 30.7μs -> 15.2μs (102% faster)

def test_zero_runtime():
    # Test with original time zero (edge case for division)
    code = "def foo():\n    x = 1"
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    key = build_key("foo", abs_path, "0")
    original_runtimes = {key: 0}
    optimized_runtimes = {key: 0}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_FunctionDef(node) # 9.11μs -> 8.98μs (1.46% faster)

def test_slower_optimized_runtime():
    # Test when optimized is slower than original
    code = "def foo():\n    x = 1"
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    key = build_key("foo", abs_path, "0")
    original_runtimes = {key: 100}
    optimized_runtimes = {key: 200}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_FunctionDef(node) # 9.10μs -> 9.06μs (0.442% faster)

def test_function_with_no_body():
    # Test a function with just 'pass'
    code = "def foo():\n    pass"
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    mapper = CommentMapper(test, {}, {})
    mapper.visit_FunctionDef(node) # 4.81μs -> 4.87μs (1.23% slower)

def test_function_with_docstring():
    # Test a function with a docstring and a statement
    code = 'def foo():\n    """Docstring"""\n    x = 1'
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    key = build_key("foo", abs_path, "1")  # x=1 is at index 1 (after docstring)
    original_runtimes = {key: 100}
    optimized_runtimes = {key: 50}
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_FunctionDef(node) # 9.93μs -> 9.91μs (0.202% faster)

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

def test_large_function_many_statements():
    # Test a function with 100 statements
    lines = ["def foo():"]
    for i in range(100):
        lines.append(f"    x{i} = {i}")
    code = "\n".join(lines)
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    original_runtimes = {}
    optimized_runtimes = {}
    for i in range(100):
        key = build_key("foo", abs_path, str(i))
        original_runtimes[key] = 100 + i
        optimized_runtimes[key] = 50 + i
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_FunctionDef(node) # 158μs -> 159μs (0.546% slower)
    # Should annotate lines 2-101
    for lineno in range(2, 102):
        pass

def test_large_function_many_compound_statements():
    # Test a function with 50 for-loops each with 2 statements
    lines = ["def foo():"]
    for i in range(50):
        lines.append(f"    for i{i} in range(2):")
        lines.append(f"        x{i} = i{i}")
        lines.append(f"        y{i} = i{i}+1")
    code = "\n".join(lines)
    node = get_func_node(code)
    abs_path = "file"
    test = GeneratedTests(abs_path)
    original_runtimes = {}
    optimized_runtimes = {}
    for i in range(50):
        # for-loop is at index i
        key_base = build_key("foo", abs_path, f"{i}_0")
        original_runtimes[key_base] = 100 + i
        optimized_runtimes[key_base] = 50 + i
        key_base2 = build_key("foo", abs_path, f"{i}_1")
        original_runtimes[key_base2] = 200 + i
        optimized_runtimes[key_base2] = 100 + i
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_FunctionDef(node) # 712μs -> 198μs (259% faster)
    # All inner lines should be annotated
    # Each for-loop adds 2 lines, starting at line 3, 6, 9, ..., up to 150
    for i in range(50):
        lineno1 = 2 + i*3 + 1
        lineno2 = 2 + i*3 + 2

To edit these changes git checkout codeflash/optimize-pr687-2025-09-03T05.27.10 and push.

[KRRT7]

this is good too, why aren't we finding these in original PRs?

[misrasaurabh1]

maybe codeflash wasn't as smart back then 😄 @KRRT7