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

---

📄 197% (1.97x) speedup for CommentMapper.visit_ClassDef in codeflash/code_utils/edit_generated_tests.py

⏱️ Runtime : 11.6 milliseconds → 3.90 milliseconds (best of 253 runs)

📝 Explanation and details

The optimization replaces ast.walk(node) with direct iteration over node.body in the visit_ClassDef method. This is a significant algorithmic improvement because:

What was changed:

• Changed for inner_node in ast.walk(node): to for inner_node in node.body:

Why this leads to a speedup:

• ast.walk(node) recursively traverses ALL descendant nodes in the AST subtree (classes, functions, statements, expressions, etc.), which creates unnecessary overhead
• node.body directly accesses only the immediate children of the class definition
• The line profiler shows the iteration went from 10,032 hits to just 409 hits - a 96% reduction in loop iterations
• The time spent on the iteration line dropped from 67.8% to 0.6% of total execution time

Performance characteristics:

• The optimization is most effective for classes with complex nested structures, as shown by the 196% speedup
• Large-scale test cases with 100+ methods and nested compound statements benefit significantly
• Basic test cases with simple class structures also see improvements due to reduced AST traversal overhead
• The optimization preserves exact functionality since we only need immediate class body elements (methods) anyway

This is a classic case of using the right data structure access pattern - direct indexing instead of tree traversal when you only need immediate children.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import ast
# Monkeypatch CommentMapper.get_comment for all tests
import types
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

# --- Minimal stubs for dependencies ---

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

# --- Helper functions for tests ---

def make_classdef_ast(class_name, methods=None):
    """
    Utility to create an ast.ClassDef node with given method(s).
    methods: list of ast.FunctionDef nodes
    """
    return ast.ClassDef(
        name=class_name,
        bases=[],
        keywords=[],
        body=methods if methods else [],
        decorator_list=[]
    )

def make_funcdef_ast(name, body=None):
    """
    Utility to create an ast.FunctionDef node.
    body: list of ast.stmt nodes
    """
    return ast.FunctionDef(
        name=name,
        args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
        body=body if body else [ast.Pass()],
        decorator_list=[],
        returns=None,
        type_comment=None
    )

def make_with_ast(body=None):
    """
    Utility to create an ast.With node with provided body.
    """
    return ast.With(
        items=[ast.withitem(context_expr=ast.Name(id='x', ctx=ast.Load()), optional_vars=None)],
        body=body if body else [ast.Pass()],
        type_comment=None
    )

def make_for_ast(body=None):
    """
    Utility to create an ast.For node with provided body.
    """
    return ast.For(
        target=ast.Name(id='i', ctx=ast.Store()),
        iter=ast.Name(id='xs', ctx=ast.Load()),
        body=body if body else [ast.Pass()],
        orelse=[]
    )

def make_if_ast(body=None):
    """
    Utility to create an ast.If node with provided body.
    """
    return ast.If(
        test=ast.NameConstant(value=True),
        body=body if body else [ast.Pass()],
        orelse=[]
    )

def make_assign_ast(lineno=1):
    """
    Utility to create an ast.Assign node with a specific lineno.
    """
    node = ast.Assign(
        targets=[ast.Name(id='a', ctx=ast.Store())],
        value=ast.Constant(value=1)
    )
    node.lineno = lineno
    return node

def make_pass_ast(lineno=1):
    node = ast.Pass()
    node.lineno = lineno
    return node

# --- Patch CommentMapper.get_comment for deterministic output ---

def fake_get_comment(self, key):
    # Return a deterministic comment string for testing
    return f"COMMENT:{key}"

CommentMapper.get_comment = fake_get_comment

# --- Unit Tests ---

# 1. BASIC TEST CASES

def test_single_class_no_methods():
    """
    Test a class with no methods: should not alter context_stack, results should remain empty.
    """
    test = GeneratedTests("foo.py")
    mapper = CommentMapper(test, {}, {})
    class_node = make_classdef_ast("MyClass")
    before_stack = list(mapper.context_stack)
    codeflash_output = mapper.visit_ClassDef(class_node); result = codeflash_output

def test_single_class_single_method_no_matches():
    """
    Class with one method, but no runtime keys match: results should be empty.
    """
    test = GeneratedTests("foo.py")
    orig = {}
    opt = {}
    mapper = CommentMapper(test, orig, opt)
    func = make_funcdef_ast("method", [make_pass_ast(lineno=5)])
    class_node = make_classdef_ast("C", [func])
    mapper.visit_ClassDef(class_node)

def test_single_class_single_method_with_match():
    """
    Class with one method, one statement, and runtimes match: should add a comment for that line.
    """
    test = GeneratedTests("foo.py")
    func_body = [make_assign_ast(lineno=10)]
    func = make_funcdef_ast("f", func_body)
    class_node = make_classdef_ast("C", [func])
    # Build the key as in _process_function_def
    key = "C.f#foo#C.f#foo#0"
    orig = {key: 1}
    opt = {key: 2}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_ClassDef(class_node)

def test_class_multiple_methods_some_matches():
    """
    Class with multiple methods, only some have runtime matches.
    """
    test = GeneratedTests("foo.py")
    func1 = make_funcdef_ast("a", [make_assign_ast(lineno=2)])
    func2 = make_funcdef_ast("b", [make_assign_ast(lineno=3)])
    class_node = make_classdef_ast("C", [func1, func2])
    key1 = "C.a#foo#C.a#foo#0"
    key2 = "C.b#foo#C.b#foo#0"
    orig = {key2: 5}
    opt = {key2: 6}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_ClassDef(class_node)

def test_nested_class_defs():
    """
    Class with a nested class: only outer class methods should be processed.
    """
    test = GeneratedTests("foo.py")
    inner_func = make_funcdef_ast("inner", [make_assign_ast(lineno=8)])
    inner_class = make_classdef_ast("Inner", [inner_func])
    outer_func = make_funcdef_ast("outer", [make_assign_ast(lineno=5)])
    class_node = make_classdef_ast("Outer", [outer_func, inner_class])
    key = "Outer.outer#foo#Outer.outer#foo#0"
    orig = {key: 1}
    opt = {key: 2}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_ClassDef(class_node)

# 2. EDGE TEST CASES

def test_empty_classdef():
    """
    ClassDef node with empty body.
    """
    test = GeneratedTests("foo.py")
    mapper = CommentMapper(test, {}, {})
    class_node = make_classdef_ast("Empty")
    mapper.visit_ClassDef(class_node)

def test_method_with_compound_statement_and_match():
    """
    Method with a With statement, and the inner statement matches runtime keys.
    """
    test = GeneratedTests("foo.py")
    assign = make_assign_ast(lineno=12)
    with_node = make_with_ast([assign])
    func = make_funcdef_ast("f", [with_node])
    class_node = make_classdef_ast("C", [func])
    key = "C.f#foo#C.f#foo#0_0"
    orig = {key: 11}
    opt = {key: 12}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_ClassDef(class_node)

def test_method_with_multiple_compound_statements():
    """
    Method with several compound statements (For, If, While), each with inner matches.
    """
    test = GeneratedTests("foo.py")
    assign1 = make_assign_ast(lineno=21)
    for_node = make_for_ast([assign1])
    assign2 = make_assign_ast(lineno=22)
    if_node = make_if_ast([assign2])
    func = make_funcdef_ast("f", [for_node, if_node])
    class_node = make_classdef_ast("C", [func])
    key1 = "C.f#foo#C.f#foo#0_0"
    key2 = "C.f#foo#C.f#foo#1_0"
    orig = {key1: 1, key2: 2}
    opt = {key1: 1, key2: 2}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_ClassDef(class_node)

def test_method_with_no_body():
    """
    Method with an empty body (should not error, no results).
    """
    test = GeneratedTests("foo.py")
    func = make_funcdef_ast("empty", [])
    class_node = make_classdef_ast("C", [func])
    mapper = CommentMapper(test, {}, {})
    mapper.visit_ClassDef(class_node)

def test_method_with_duplicate_linenos():
    """
    Method with two statements having the same lineno (should both be processed, last one wins in dict).
    """
    test = GeneratedTests("foo.py")
    assign1 = make_assign_ast(lineno=30)
    assign2 = make_assign_ast(lineno=30)
    func = make_funcdef_ast("f", [assign1, assign2])
    class_node = make_classdef_ast("C", [func])
    key1 = "C.f#foo#C.f#foo#0"
    key2 = "C.f#foo#C.f#foo#1"
    orig = {key1: 1, key2: 2}
    opt = {key1: 1, key2: 2}
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_ClassDef(class_node)

def test_context_stack_restoration_on_exception():
    """
    If an exception occurs in visit_FunctionDef, context_stack should be restored.
    """
    test = GeneratedTests("foo.py")
    func = make_funcdef_ast("f", [make_pass_ast(lineno=1)])
    class_node = make_classdef_ast("C", [func])
    mapper = CommentMapper(test, {}, {})
    # Patch visit_FunctionDef to raise
    orig_func = mapper.visit_FunctionDef
    def boom(node): raise RuntimeError("fail")
    mapper.visit_FunctionDef = boom
    before_stack = list(mapper.context_stack)
    with pytest.raises(RuntimeError):
        mapper.visit_ClassDef(class_node)

# 3. LARGE SCALE TEST CASES

def test_large_number_of_methods_and_statements():
    """
    Class with 100 methods, each with 5 statements, all with matches.
    """
    test = GeneratedTests("foo.py")
    orig = {}
    opt = {}
    methods = []
    for m in range(100):
        stmts = []
        for s in range(5):
            lineno = 1000 + m * 5 + s
            assign = make_assign_ast(lineno=lineno)
            stmts.append(assign)
            key = f"C.m{m}#foo#C.m{m}#foo#{s}"
            orig[key] = 1
            opt[key] = 2
        func = make_funcdef_ast(f"m{m}", stmts)
        methods.append(func)
    class_node = make_classdef_ast("C", methods)
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_ClassDef(class_node)
    # All lines should be present
    for m in range(100):
        for s in range(5):
            lineno = 1000 + m * 5 + s
            key = f"C.m{m}#foo#C.m{m}#foo#{s}"

def test_large_nested_compound_statements():
    """
    Method with 10 For statements, each with 10 assignments, all matched.
    """
    test = GeneratedTests("foo.py")
    orig = {}
    opt = {}
    for_nodes = []
    for i in range(10):
        assigns = []
        for j in range(10):
            lineno = 2000 + i * 10 + j
            assign = make_assign_ast(lineno=lineno)
            assigns.append(assign)
            key = f"C.f#foo#C.f#foo#{i}_{j}"
            orig[key] = 1
            opt[key] = 2
        for_node = make_for_ast(assigns)
        for_nodes.append(for_node)
    func = make_funcdef_ast("f", for_nodes)
    class_node = make_classdef_ast("C", [func])
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_ClassDef(class_node)
    for i in range(10):
        for j in range(10):
            lineno = 2000 + i * 10 + j
            key = f"C.f#foo#C.f#foo#{i}_{j}"

def test_large_class_many_methods_some_missing_matches():
    """
    Class with 50 methods, each with 10 statements, only even-numbered methods have matches.
    """
    test = GeneratedTests("foo.py")
    orig = {}
    opt = {}
    methods = []
    for m in range(50):
        stmts = []
        for s in range(10):
            lineno = 3000 + m * 10 + s
            assign = make_assign_ast(lineno=lineno)
            stmts.append(assign)
            if m % 2 == 0:
                key = f"C.m{m}#foo#C.m{m}#foo#{s}"
                orig[key] = 1
                opt[key] = 2
        func = make_funcdef_ast(f"m{m}", stmts)
        methods.append(func)
    class_node = make_classdef_ast("C", methods)
    mapper = CommentMapper(test, orig, opt)
    mapper.visit_ClassDef(class_node)
    for m in range(50):
        for s in range(10):
            lineno = 3000 + m * 10 + s
            if m % 2 == 0:
                key = f"C.m{m}#foo#C.m{m}#foo#{s}"
            else:
                pass
# 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
from types import SimpleNamespace

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


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

# Helper function to parse code and get the ClassDef node
def get_classdef_node(source: str) -> ast.ClassDef:
    tree = ast.parse(source)
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            return node
    raise ValueError("No ClassDef found in source!")

# Helper to generate a dummy path
def dummy_path(name="foo.py"):
    return Path(name)

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

def test_basic_single_method_no_body():
    # Test a class with a single empty method
    code = """
class MyClass:
    def foo(self):
        pass
    """
    node = get_classdef_node(code)
    test = GeneratedTests(dummy_path("foo.py"))
    cm = CommentMapper(test, {}, {})
    codeflash_output = cm.visit_ClassDef(node); result = codeflash_output

def test_basic_single_method_with_body_and_matching_keys():
    # Test a class with a method with statements, and a matching key in runtimes
    code = """
class MyClass:
    def foo(self):
        x = 1
        y = 2
        return x + y
    """
    node = get_classdef_node(code)
    test = GeneratedTests(dummy_path("foo.py"))
    key_prefix = "MyClass.foo#foo"
    # Only one statement will match (the last one, i=2)
    orig = {f"{key_prefix}#0": 1, f"{key_prefix}#1": 2, f"{key_prefix}#2": 3}
    opt = {f"{key_prefix}#0": 1, f"{key_prefix}#1": 2, f"{key_prefix}#2": 3}
    cm = CommentMapper(test, orig, opt)
    cm.abs_path = "foo"  # ensure str(self.abs_path) == "foo"
    cm.visit_ClassDef(node)
    # All three statements should be annotated
    # Find line numbers for x=1, y=2, return x+y
    lines = [stmt.lineno for stmt in node.body[0].body]
    for i, lineno in enumerate(lines):
        pass

def test_basic_multiple_methods_and_classes():
    # Test a file with two classes, each with a method, and only one method has matching keys
    code = """
class A:
    def foo(self):
        x = 1

class B:
    def bar(self):
        y = 2
    """
    tree = ast.parse(code)
    test = GeneratedTests(dummy_path("bar.py"))
    cm = CommentMapper(test, {"A.foo#bar#0": 1}, {"A.foo#bar#0": 2})
    cm.abs_path = "bar"
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            cm.visit_ClassDef(node)
    # Only the 'x = 1' line in class A should be annotated
    foo_func = [n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == "foo"][0]
    x_lineno = foo_func.body[0].lineno

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

def test_edge_class_with_no_methods():
    # Test a class with no methods
    code = """
class Empty:
    pass
    """
    node = get_classdef_node(code)
    test = GeneratedTests(dummy_path("empty.py"))
    cm = CommentMapper(test, {}, {})
    cm.visit_ClassDef(node)

def test_edge_nested_classes():
    # Test a class with a nested class and both have methods
    code = """
class Outer:
    def foo(self):
        pass
    class Inner:
        def bar(self):
            z = 3
    """
    node = get_classdef_node(code)
    test = GeneratedTests(dummy_path("nested.py"))
    cm = CommentMapper(test, {"Outer.Inner.bar#nested#0": 1}, {"Outer.Inner.bar#nested#0": 2})
    cm.abs_path = "nested"
    cm.visit_ClassDef(node)
    # Only the 'z = 3' line in Inner.bar should be annotated
    # Find line number for z = 3
    inner = [n for n in ast.walk(node) if isinstance(n, ast.ClassDef) and n.name == "Inner"][0]
    bar_func = [n for n in ast.walk(inner) if isinstance(n, ast.FunctionDef) and n.name == "bar"][0]
    z_lineno = bar_func.body[0].lineno

def test_edge_function_with_control_structures():
    # Test a method with for/if/while/with blocks and matching keys for inner statements
    code = """
class Control:
    def baz(self):
        for i in range(2):
            a = i
        if True:
            b = 2
        while False:
            c = 3
        with open("f") as f:
            d = 4
    """
    node = get_classdef_node(code)
    test = GeneratedTests(dummy_path("ctrl.py"))
    key_prefix = "Control.baz#ctrl"
    orig = {
        f"{key_prefix}#0_0": 1,  # for: a = i
        f"{key_prefix}#1_0": 2,  # if: b = 2
        f"{key_prefix}#2_0": 3,  # while: c = 3
        f"{key_prefix}#3_0": 4,  # with: d = 4
    }
    opt = orig.copy()
    cm = CommentMapper(test, orig, opt)
    cm.abs_path = "ctrl"
    cm.visit_ClassDef(node)
    # Find line numbers for a = i, b = 2, c = 3, d = 4
    func = [n for n in ast.walk(node) if isinstance(n, ast.FunctionDef) and n.name == "baz"][0]
    block_lines = []
    for i, block in enumerate(func.body):
        block_lines.append(block.body[0].lineno)
    for i, lineno in enumerate(block_lines):
        pass

def test_edge_async_function():
    # Test a class with an async method
    code = """
class AsyncClass:
    async def afunc(self):
        x = 42
    """
    node = get_classdef_node(code)
    test = GeneratedTests(dummy_path("async.py"))
    key_prefix = "AsyncClass.afunc#async"
    orig = {f"{key_prefix}#0": 1}
    opt = {f"{key_prefix}#0": 1}
    cm = CommentMapper(test, orig, opt)
    cm.abs_path = "async"
    cm.visit_ClassDef(node)
    # Find line number for x = 42
    afunc = [n for n in ast.walk(node) if isinstance(n, ast.AsyncFunctionDef)][0]
    x_lineno = afunc.body[0].lineno

def test_edge_no_matching_keys():
    # Test a class with a method but no matching keys in either runtime dict
    code = """
class NoMatch:
    def foo(self):
        x = 1
    """
    node = get_classdef_node(code)
    test = GeneratedTests(dummy_path("nomatch.py"))
    cm = CommentMapper(test, {}, {})
    cm.abs_path = "nomatch"
    cm.visit_ClassDef(node)

def test_edge_multiple_statements_in_compound_block():
    # Test a function with multiple statements in a for block, only some have matching keys
    code = """
class MultiBlock:
    def foo(self):
        for i in range(2):
            a = 1
            b = 2
            c = 3
    """
    node = get_classdef_node(code)
    test = GeneratedTests(dummy_path("multi.py"))
    key_prefix = "MultiBlock.foo#multi"
    orig = {f"{key_prefix}#0_1": 1, f"{key_prefix}#0_2": 2}
    opt = orig.copy()
    cm = CommentMapper(test, orig, opt)
    cm.abs_path = "multi"
    cm.visit_ClassDef(node)
    # Only b = 2 and c = 3 should be annotated
    func = [n for n in ast.walk(node) if isinstance(n, ast.FunctionDef) and n.name == "foo"][0]
    for_block = func.body[0]
    # for block: a = 1 (j=0), b = 2 (j=1), c = 3 (j=2)
    lines = [stmt.lineno for stmt in for_block.body]

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

def test_large_class_many_methods():
    # Test a class with 100 methods, each with one statement, and all have matching keys
    method_count = 100
    code_lines = ["class Large:"]
    for i in range(method_count):
        code_lines.append(f"    def foo_{i}(self):")
        code_lines.append(f"        x_{i} = {i}")
    code = "\n".join(code_lines)
    node = get_classdef_node(code)
    test = GeneratedTests(dummy_path("large.py"))
    orig = {}
    opt = {}
    for i in range(method_count):
        key = f"Large.foo_{i}#large#0"
        orig[key] = i
        opt[key] = i
    cm = CommentMapper(test, orig, opt)
    cm.abs_path = "large"
    cm.visit_ClassDef(node)
    # Each method's single statement should be annotated
    for func in [n for n in ast.walk(node) if isinstance(n, ast.FunctionDef)]:
        stmt = func.body[0]
        key = f"Large.{func.name}#large#0"

def test_large_class_with_large_compound_blocks():
    # Test a class with a method containing a for loop with 500 statements, all have matching keys
    stmt_count = 500
    code_lines = ["class BigCompound:", "    def foo(self):", "        for i in range(2):"]
    for j in range(stmt_count):
        code_lines.append(f"            x_{j} = {j}")
    code = "\n".join(code_lines)
    node = get_classdef_node(code)
    test = GeneratedTests(dummy_path("big.py"))
    orig = {}
    opt = {}
    key_prefix = "BigCompound.foo#big"
    for j in range(stmt_count):
        orig[f"{key_prefix}#0_{j}"] = j
        opt[f"{key_prefix}#0_{j}"] = j
    cm = CommentMapper(test, orig, opt)
    cm.abs_path = "big"
    cm.visit_ClassDef(node)
    # All 500 statements in the for block should be annotated
    func = [n for n in ast.walk(node) if isinstance(n, ast.FunctionDef) and n.name == "foo"][0]
    for_block = func.body[0]
    for j, stmt in enumerate(for_block.body):
        pass

def test_large_multiple_classes_and_methods():
    # Test 10 classes, each with 10 methods, each with 5 statements, all have matching keys
    class_count = 10
    method_count = 10
    stmt_count = 5
    code_lines = []
    orig = {}
    opt = {}
    for c in range(class_count):
        cname = f"C{c}"
        code_lines.append(f"class {cname}:")
        for m in range(method_count):
            code_lines.append(f"    def foo_{m}(self):")
            for s in range(stmt_count):
                code_lines.append(f"        x_{s} = {s}")
            # Prepare keys for each statement
            key_prefix = f"{cname}.foo_{m}#all"
            for s in range(stmt_count):
                orig[f"{key_prefix}#{s}"] = s
                opt[f"{key_prefix}#{s}"] = s
    code = "\n".join(code_lines)
    tree = ast.parse(code)
    test = GeneratedTests(dummy_path("all.py"))
    cm = CommentMapper(test, orig, opt)
    cm.abs_path = "all"
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            cm.visit_ClassDef(node)
    # Check that all statements are annotated
    for func in [n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)]:
        cname = func.parent.name if hasattr(func, "parent") else func.name.split("_")[0]
        key_prefix = f"{func.parent.name}.{func.name}#all" if hasattr(func, "parent") else f"{cname}.{func.name}#all"
        for s, stmt in enumerate(func.body):
            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-pr687-2025-09-03T05.07.18 and push.

[misrasaurabh1]

this looks good? We don't need to traverse all - we only support direct methods of a class.

[KRRT7]

it does look good, making a slight edit and merging