⚡️ Speed up method CommentMapper.visit_FunctionDef by 10% 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.

---

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

⏱️ Runtime : 1.31 milliseconds → 1.19 milliseconds (best of 153 runs)

📝 Explanation and details

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

Key optimizations:

1. Local variable caching for frequent lookups: Storing self.original_runtimes and self.optimized_runtimes in local variables (original_runtimes, optimized_runtimes) eliminates repeated attribute lookups in the inner loops, which is expensive due to Python's method resolution.

2. Eliminated unnecessary list allocations: The original code created nodes_to_check = [compound_line_node] and extended it with getattr(compound_line_node, "body", []) on every iteration. The optimized version uses hasattr() to check for a body attribute first, then iterates directly over compound_line_node.body when present, avoiding list creation and extension.

3. Reduced redundant attribute access: Added nbody = node.body and line_body = line_node.body to cache commonly accessed attributes, reducing repeated lookups.

4. Streamlined f-string usage: Consistent use of f-strings (f"{i}_{j}" instead of str(i) + "_" + str(j)) provides marginal performance improvements for string concatenation.

The optimizations are most effective for large-scale test cases - the annotated tests show the biggest improvements (13.4% faster) occur with functions containing many compound statements (50 for-loops with 10 statements each), where the eliminated allocations and attribute lookups compound significantly. Simple functions show minimal or slight performance regression due to the overhead of additional local variable assignments, but complex functions with nested loops see substantial gains.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 38 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
from types import SimpleNamespace

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


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

# --------- UNIT TESTS BELOW ---------

def make_test_mapper(
    src: str,
    original_runtimes: dict[str, int],
    optimized_runtimes: dict[str, int],
    filename="foo.py"
):
    # Parse the code, get the FunctionDef node, and build a CommentMapper
    tree = ast.parse(src)
    func_node = next((n for n in tree.body if isinstance(n, ast.FunctionDef)), None)
    test_obj = GeneratedTests(filename)
    mapper = CommentMapper(test_obj, original_runtimes, optimized_runtimes)
    return mapper, func_node

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

def test_single_statement_function():
    """Test a function with a single statement."""
    src = "def foo():\n    x = 1\n"
    # The key is: foo#foo (filename without .py) + #0
    key = "foo#foo#0"
    original = {key: 2_000_000_000}
    optimized = {key: 1_000_000_000}
    mapper, func_node = make_test_mapper(src, original, optimized, "foo.py")
    mapper.visit_FunctionDef(func_node) # 11.6μs -> 12.3μs (5.55% slower)
    # Check the comment format
    comment = mapper.results[2]

def test_multiple_statements():
    """Test a function with multiple top-level statements."""
    src = "def bar():\n    a = 1\n    b = 2\n    c = 3\n"
    key_base = "bar#bar"
    original = {
        f"{key_base}#0": 10,
        f"{key_base}#1": 20,
        f"{key_base}#2": 40,
    }
    optimized = {
        f"{key_base}#0": 5,
        f"{key_base}#1": 10,
        f"{key_base}#2": 40,
    }
    mapper, func_node = make_test_mapper(src, original, optimized, "bar.py")
    mapper.visit_FunctionDef(func_node) # 14.1μs -> 14.9μs (5.25% slower)

def test_with_compound_statement():
    """Test a function with a for-loop containing statements."""
    src = (
        "def baz():\n"
        "    for i in range(2):\n"
        "        a = i\n"
        "        b = i+1\n"
        "    c = 3\n"
    )
    key_base = "baz#baz"
    # for-loop is body[0], c=3 is body[1]
    # for-loop body: a=i (body[0]), b=i+1 (body[1])
    original = {
        f"{key_base}#0_0": 100,
        f"{key_base}#0_1": 200,
        f"{key_base}#1": 300,
    }
    optimized = {
        f"{key_base}#0_0": 50,
        f"{key_base}#0_1": 100,
        f"{key_base}#1": 150,
    }
    mapper, func_node = make_test_mapper(src, original, optimized, "baz.py")
    mapper.visit_FunctionDef(func_node) # 15.7μs -> 15.8μs (1.01% slower)

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

def test_empty_function():
    """Test a function with no statements (just 'pass')."""
    src = "def empty():\n    pass\n"
    key = "empty#empty#0"
    # No runtimes for 'pass'
    original = {}
    optimized = {}
    mapper, func_node = make_test_mapper(src, original, optimized, "empty.py")
    mapper.visit_FunctionDef(func_node) # 5.02μs -> 5.36μs (6.34% slower)

def test_function_with_if_else():
    """Test a function with an if-else block."""
    src = (
        "def cond():\n"
        "    if True:\n"
        "        a = 1\n"
        "        b = 2\n"
        "    else:\n"
        "        c = 3\n"
        "    d = 4\n"
    )
    key_base = "cond#cond"
    # if is body[0], d=4 is body[1]
    # if body: a=1 (body[0]), b=2 (body[1])
    # else body: c=3 (body[0])
    original = {
        f"{key_base}#0_0": 10,
        f"{key_base}#0_1": 20,
        f"{key_base}#1": 30,
    }
    optimized = {
        f"{key_base}#0_0": 5,
        f"{key_base}#0_1": 10,
        f"{key_base}#1": 15,
    }
    mapper, func_node = make_test_mapper(src, original, optimized, "cond.py")
    mapper.visit_FunctionDef(func_node) # 15.5μs -> 15.6μs (0.705% slower)
    # c=3 (else) not included (no runtime for #0_2)

def test_missing_runtime_keys():
    """Test function where some statements have no runtime data."""
    src = "def foo():\n    x = 1\n    y = 2\n"
    key = "foo#foo#0"
    # Only x=1 has runtime data
    original = {key: 10}
    optimized = {key: 5}
    mapper, func_node = make_test_mapper(src, original, optimized, "foo.py")
    mapper.visit_FunctionDef(func_node) # 9.90μs -> 10.4μs (4.72% slower)

def test_function_with_nested_for():
    """Test a function with a nested for-loop."""
    src = (
        "def nest():\n"
        "    for i in range(2):\n"
        "        for j in range(2):\n"
        "            x = i + j\n"
        "    y = 99\n"
    )
    key_base = "nest#nest"
    # Outer for is body[0], y=99 is body[1]
    # Outer for body[0] is inner for
    # Inner for body[0] is x = i + j
    # inv_id for x = i+j: outer i=0, inner j=0: #0_0
    original = {
        f"{key_base}#0_0": 100,
        f"{key_base}#1": 200,
    }
    optimized = {
        f"{key_base}#0_0": 50,
        f"{key_base}#1": 100,
    }
    mapper, func_node = make_test_mapper(src, original, optimized, "nest.py")
    mapper.visit_FunctionDef(func_node) # 14.6μs -> 15.0μs (2.53% slower)

def test_function_with_no_body():
    """Test a function with an empty body (should not crash)."""
    func_node = ast.FunctionDef(
        name="nobody",
        args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
        body=[],
        decorator_list=[],
        lineno=1,
        col_offset=0
    )
    test_obj = GeneratedTests("nobody.py")
    mapper = CommentMapper(test_obj, {}, {})
    # Should not raise
    mapper.visit_FunctionDef(func_node) # 3.65μs -> 3.97μs (8.09% slower)

def test_function_with_non_assign_statements():
    """Test a function with statements that are not Assign or compound."""
    src = (
        "def foo():\n"
        "    print('hi')\n"
        "    return 42\n"
    )
    key_base = "foo#foo"
    original = {f"{key_base}#0": 10, f"{key_base}#1": 20}
    optimized = {f"{key_base}#0": 5, f"{key_base}#1": 10}
    mapper, func_node = make_test_mapper(src, original, optimized, "foo.py")
    mapper.visit_FunctionDef(func_node) # 11.9μs -> 12.6μs (5.26% slower)

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

def test_large_function_many_statements():
    """Test a function with 100 statements."""
    lines = ["def big():"] + [f"    x{i} = {i}" for i in range(100)]
    src = "\n".join(lines) + "\n"
    key_base = "big#big"
    original = {f"{key_base}#{i}": 1000 + i for i in range(100)}
    optimized = {f"{key_base}#{i}": 500 + i for i in range(100)}
    mapper, func_node = make_test_mapper(src, original, optimized, "big.py")
    mapper.visit_FunctionDef(func_node) # 189μs -> 185μs (2.29% faster)
    # Check a few random lines
    for i in [2, 51, 101]:
        pass

def test_large_function_many_compound_statements():
    """Test a function with 50 for-loops, each with 10 statements."""
    lines = ["def bigcomp():"] 
    for i in range(50):
        lines.append(f"    for i{i} in range(2):")
        for j in range(10):
            lines.append(f"        x{i}_{j} = {j}")
    src = "\n".join(lines) + "\n"
    key_base = "bigcomp#bigcomp"
    original = {}
    optimized = {}
    # for-loop is body[i], its body[j] is x{i}_{j}
    for i in range(50):
        for j in range(10):
            original[f"{key_base}#{i}_"+str(j)] = 1000 + i*10 + j
            optimized[f"{key_base}#{i}_"+str(j)] = 500 + i*10 + j
    mapper, func_node = make_test_mapper(src, original, optimized, "bigcomp.py")
    mapper.visit_FunctionDef(func_node) # 1.02ms -> 896μs (13.4% faster)
    # Each x{i}_{j} is at line 2 + i*11 + j
    expected_lines = set()
    line_num = 2
    for i in range(50):
        for j in range(10):
            expected_lines.add(line_num + 1 + j)
        line_num += 11
    # Spot check a few lines
    for i in [2, 100, 501]:
        pass


#------------------------------------------------
import ast
from pathlib import Path
from types import SimpleNamespace

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


# Mocks for codeflash.models.models
class GeneratedTests:
    def __init__(self, behavior_file_path):
        self.behavior_file_path = behavior_file_path

# --- Unit Tests ---

# Helper to create a fake path object
class FakePath:
    def __init__(self, path):
        self._path = path
    def with_suffix(self, s):
        # returns self with suffix removed, as string
        if "." in self._path:
            return self._path.rsplit(".", 1)[0]
        return self._path

# Helper to run CommentMapper.visit_FunctionDef and return results
def run_mapper_on_funcdef(
    func_src: str,
    original_runtimes: dict[str, int],
    optimized_runtimes: dict[str, int],
    file_path="foo.py",
):
    """
    Parses func_src, runs CommentMapper, and returns the results dict.
    """
    tree = ast.parse(func_src)
    # Find the function def node
    func_node = None
    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef):
            func_node = node
            break
    test = GeneratedTests(FakePath(file_path))
    mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
    mapper.visit_FunctionDef(func_node)
    return mapper.results

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

def test_basic_single_statement():
    """
    Test a function with a single statement.
    """
    src = "def foo():\n    x = 1"
    # The key is: foo#foo#0
    original = {"foo#foo#0": 100}
    optimized = {"foo#foo#0": 50}
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_basic_multiple_statements():
    """
    Test a function with several statements.
    """
    src = "def foo():\n    x = 1\n    y = 2\n    z = 3"
    original = {
        "foo#foo#0": 100,
        "foo#foo#1": 200,
        "foo#foo#2": 300,
    }
    optimized = {
        "foo#foo#0": 80,
        "foo#foo#1": 180,
        "foo#foo#2": 150,
    }
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_basic_with_compound_statement_for():
    """
    Test a function with a for loop containing assignments.
    """
    src = (
        "def foo():\n"
        "    for i in range(2):\n"
        "        x = i\n"
        "        y = i+1"
    )
    # The keys for for-loop body are foo#foo#0_0 and foo#foo#0_1
    original = {
        "foo#foo#0_0": 100,
        "foo#foo#0_1": 200,
    }
    optimized = {
        "foo#foo#0_0": 90,
        "foo#foo#0_1": 180,
    }
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_basic_with_if_statement():
    """
    Test a function with an if statement.
    """
    src = (
        "def foo():\n"
        "    if True:\n"
        "        x = 1\n"
        "    y = 2"
    )
    # The keys: foo#foo#0_0 for x=1, foo#foo#1 for y=2
    original = {
        "foo#foo#0_0": 100,
        "foo#foo#1": 200,
    }
    optimized = {
        "foo#foo#0_0": 80,
        "foo#foo#1": 150,
    }
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

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

def test_empty_function_body():
    """
    Test a function with an empty body (pass).
    """
    src = "def foo():\n    pass"
    original = {}
    optimized = {}
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_function_with_no_matching_keys():
    """
    Test a function with statements but no matching keys in runtimes.
    """
    src = "def foo():\n    x = 1\n    y = 2"
    original = {}
    optimized = {}
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_function_with_partial_matching_keys():
    """
    Test a function with only one statement having matching keys.
    """
    src = "def foo():\n    x = 1\n    y = 2"
    original = {"foo#foo#1": 100}
    optimized = {"foo#foo#1": 50}
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_function_with_nested_compound_statements():
    """
    Test function with nested for and if statements.
    """
    src = (
        "def foo():\n"
        "    for i in range(2):\n"
        "        if i == 1:\n"
        "            x = i\n"
        "        y = i\n"
    )
    # x = i is foo#foo#0_0, y = i is foo#foo#0_1
    original = {
        "foo#foo#0_0": 100,
        "foo#foo#0_1": 200,
    }
    optimized = {
        "foo#foo#0_0": 50,
        "foo#foo#0_1": 180,
    }
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_function_with_slower_optimized_time():
    """
    Test a function where the optimized time is slower.
    """
    src = "def foo():\n    x = 1"
    original = {"foo#foo#0": 100}
    optimized = {"foo#foo#0": 150}
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_function_with_zero_original_time():
    """
    If original time is zero, performance_gain returns 0.0, so 0% faster/slower.
    """
    src = "def foo():\n    x = 1"
    original = {"foo#foo#0": 0}
    optimized = {"foo#foo#0": 0}
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_function_with_non_py_suffix():
    """
    Test that .with_suffix works for non-.py files.
    """
    src = "def foo():\n    x = 1"
    original = {"foo#foo#0": 100}
    optimized = {"foo#foo#0": 50}
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.txt")

def test_function_with_multiple_functions():
    """
    Test that only the first FunctionDef is processed.
    """
    src = (
        "def foo():\n"
        "    x = 1\n"
        "def bar():\n"
        "    y = 2"
    )
    # Only foo's assignment should be processed
    original = {"foo#foo#0": 100}
    optimized = {"foo#foo#0": 50}
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

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

def test_large_function_many_statements():
    """
    Test a function with 500 statements.
    """
    lines = ["def foo():"]
    original = {}
    optimized = {}
    for i in range(500):
        lines.append(f"    x{i} = {i}")
        key = f"foo#foo#{i}"
        original[key] = 100 + i
        optimized[key] = 50 + i
    src = "\n".join(lines)
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_large_function_many_for_loops():
    """
    Test a function with 100 for loops, each with 2 assignments.
    """
    lines = ["def foo():"]
    original = {}
    optimized = {}
    for i in range(100):
        lines.append(f"    for i{i} in range(2):")
        lines.append(f"        x{i} = {i}")
        lines.append(f"        y{i} = {i}+1")
        # For-loop body: index is i=loop, j=body index
        key_x = f"foo#foo#{i}_0"
        key_y = f"foo#foo#{i}_1"
        original[key_x] = 100 + i
        original[key_y] = 200 + i
        optimized[key_x] = 90 + i
        optimized[key_y] = 180 + i
    src = "\n".join(lines)
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_large_function_deeply_nested_compounds():
    """
    Test a function with 10 nested for-loops, each with an if and an assignment.
    """
    indent = ""
    src = "def foo():\n"
    original = {}
    optimized = {}
    for i in range(10):
        indent += "    "
        src += f"{indent}for i{i} in range(2):\n"
    indent += "    "
    src += f"{indent}if True:\n"
    indent += "    "
    src += f"{indent}x = 1\n"
    # The innermost assignment is at line 13
    key = "foo#foo#" + "0" * 10 + "_0"
    # But the actual key for the innermost for-loop is foo#foo#9_0
    key = "foo#foo#9_0"
    original[key] = 1000
    optimized[key] = 900
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")

def test_large_function_with_mixed_statements():
    """
    Test a function with a mix of 100 assignments and 100 for-loops.
    """
    lines = ["def foo():"]
    original = {}
    optimized = {}
    for i in range(100):
        lines.append(f"    x{i} = {i}")
        key = f"foo#foo#{i}"
        original[key] = 100 + i
        optimized[key] = 90 + i
    for i in range(100, 200):
        lines.append(f"    for j{i} in range(2):")
        lines.append(f"        y{i} = {i}")
        key = f"foo#foo#{i}_0"
        original[key] = 200 + i
        optimized[key] = 180 + i
    src = "\n".join(lines)
    results = run_mapper_on_funcdef(src, original, optimized, file_path="foo.py")
# 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-pr687-2025-09-13T00.38.14 and push.