⚡️ Speed up method InjectPerfOnly.find_and_update_line_node by 22% in PR #867 (inspect-signature-issue)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #867

If you approve this dependent PR, these changes will be merged into the original PR branch inspect-signature-issue.

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

---

📄 22% (0.22x) speedup for InjectPerfOnly.find_and_update_line_node in codeflash/code_utils/instrument_existing_tests.py

⏱️ Runtime : 2.01 milliseconds → 1.64 milliseconds (best of 113 runs)

📝 Explanation and details

The optimized code achieves a 22% speedup through two main optimizations that reduce overhead in AST traversal and attribute lookups:

1. Custom AST traversal replaces expensive ast.walk()
The original code uses ast.walk() which creates recursive stack frames for every AST node. The optimized version implements iter_ast_calls() - a manual iterative traversal that only visits ast.Call nodes using a single stack. This eliminates Python's recursion overhead and reduces the O(N) stack frame creation to a single stack operation.

2. Reduced attribute lookups in hot paths

• In node_in_call_position(): Uses getattr() with defaults to cache node attributes (node_lineno, node_end_lineno, etc.) instead of repeated hasattr() + attribute access
• In find_and_update_line_node(): Hoists frequently-accessed object attributes (fn_obj.qualified_name, self.mode, etc.) to local variables before the loop
• Pre-creates reusable AST nodes (codeflash_loop_index, codeflash_cur, codeflash_con) instead of recreating them in each iteration

Performance characteristics:

• Small AST trees (basic function calls): 5-28% faster due to reduced attribute lookups
• Large AST trees (deeply nested calls): 18-26% faster due to more efficient traversal avoiding ast.walk()
• Large call position lists: 26% faster due to optimized position checking with cached attributes

The optimizations are most effective for complex test instrumentation scenarios with large AST trees or many call positions to check, which is typical in code analysis and transformation workflows.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import ast
from typing import Any

# imports
import pytest
from codeflash.code_utils.instrument_existing_tests import InjectPerfOnly

# --- Minimal stubs for dependencies used in the function ---

class FunctionToOptimize:
    def __init__(self, function_name, qualified_name, is_async=False, parents=None, top_level_parent_name=None):
        self.function_name = function_name
        self.qualified_name = qualified_name
        self.is_async = is_async
        self.parents = parents or []
        self.top_level_parent_name = top_level_parent_name

class ParentStub:
    def __init__(self, type_, name):
        self.type = type_
        self.name = name

class CodePosition:
    def __init__(self, line_no=None, col_no=None, end_col_offset=None):
        self.line_no = line_no
        self.col_no = col_no
        self.end_col_offset = end_col_offset

class TestingMode:
    BEHAVIOR = "BEHAVIOR"
    PERFORMANCE = "PERFORMANCE"
from codeflash.code_utils.instrument_existing_tests import InjectPerfOnly

# --- Helper functions for test assertions ---

def get_first_call_node(stmt):
    """Recursively find the first ast.Call node in a statement."""
    for node in ast.walk(stmt):
        if isinstance(node, ast.Call):
            return node
    return None

def get_func_name(call_node):
    """Get function name from ast.Call node."""
    if isinstance(call_node.func, ast.Name):
        return call_node.func.id
    elif isinstance(call_node.func, ast.Attribute):
        return call_node.func.attr
    return None

def is_codeflash_wrap_call(call_node):
    return isinstance(call_node.func, ast.Name) and call_node.func.id == "codeflash_wrap"

def is_assign_to_bound_arguments(assign_node):
    return (
        isinstance(assign_node, ast.Assign)
        and any(isinstance(t, ast.Name) and t.id == "_call__bound__arguments" for t in assign_node.targets)
    )

def is_apply_defaults_expr(expr_node):
    return (
        isinstance(expr_node, ast.Expr)
        and isinstance(expr_node.value, ast.Call)
        and isinstance(expr_node.value.func, ast.Attribute)
        and expr_node.value.func.attr == "apply_defaults"
    )

# --- Test Cases ---

# 1. Basic Test Cases

def make_simple_call_stmt(func_name="foo", args=None, keywords=None, lineno=1, col_offset=0):
    """Create a simple ast.Expr node with a function call."""
    args = args or []
    keywords = keywords or []
    call = ast.Call(func=ast.Name(id=func_name, ctx=ast.Load()), args=args, keywords=keywords)
    call.lineno = lineno
    call.col_offset = col_offset
    call.end_lineno = lineno
    call.end_col_offset = col_offset + 10
    expr = ast.Expr(value=call, lineno=lineno, col_offset=col_offset)
    return expr

def make_method_call_stmt(obj_name="obj", method_name="foo", args=None, keywords=None, lineno=1, col_offset=0):
    """Create a simple ast.Expr node with a method call."""
    args = args or []
    keywords = keywords or []
    attr = ast.Attribute(value=ast.Name(id=obj_name, ctx=ast.Load()), attr=method_name, ctx=ast.Load())
    call = ast.Call(func=attr, args=args, keywords=keywords)
    call.lineno = lineno
    call.col_offset = col_offset
    call.end_lineno = lineno
    call.end_col_offset = col_offset + 10
    expr = ast.Expr(value=call, lineno=lineno, col_offset=col_offset)
    return expr

@pytest.fixture
def default_function_obj():
    return FunctionToOptimize(function_name="foo", qualified_name="foo", is_async=False)

@pytest.fixture
def default_method_function_obj():
    return FunctionToOptimize(function_name="foo", qualified_name="Bar.foo", is_async=False, parents=[ParentStub("ClassDef", "Bar")], top_level_parent_name="Bar")

@pytest.fixture
def call_position():
    return [CodePosition(line_no=1, col_no=0, end_col_offset=10)]

def test_basic_function_call_transformation(default_function_obj, call_position):
    """Test a basic function call is wrapped and signature binding is inserted."""
    node = make_simple_call_stmt()
    inj = InjectPerfOnly(default_function_obj, "module.py", "pytest", call_position, mode=TestingMode.BEHAVIOR)
    codeflash_output = inj.find_and_update_line_node(node, "test_foo", "0"); result = codeflash_output
    call_node = get_first_call_node(result[2])
    # Should have codeflash_loop_index and codeflash_cur/con in args
    arg_names = [a.id for a in call_node.args if isinstance(a, ast.Name)]

def test_basic_method_call_transformation(default_method_function_obj, call_position):
    """Test a method call is wrapped and signature binding is inserted."""
    node = make_method_call_stmt(obj_name="self", method_name="foo")
    inj = InjectPerfOnly(default_method_function_obj, "module.py", "pytest", call_position, mode=TestingMode.BEHAVIOR)
    codeflash_output = inj.find_and_update_line_node(node, "test_bar", "1", test_class_name="Bar"); result = codeflash_output
    call_node = get_first_call_node(result[2])


def test_no_call_node_returns_none(default_function_obj, call_position):
    """Test that if there is no call node, None is returned."""
    node = ast.Pass(lineno=1, col_offset=0)
    inj = InjectPerfOnly(default_function_obj, "module.py", "pytest", call_position)
    codeflash_output = inj.find_and_update_line_node(node, "test_none", "0"); result = codeflash_output # 4.76μs -> 4.82μs (1.22% slower)

def test_call_node_not_in_position_returns_none(default_function_obj):
    """Test that if call node's position does not match, None is returned."""
    node = make_simple_call_stmt()
    # call_position is on a different line
    call_position = [CodePosition(line_no=2, col_no=0, end_col_offset=10)]
    inj = InjectPerfOnly(default_function_obj, "module.py", "pytest", call_position)
    codeflash_output = inj.find_and_update_line_node(node, "test_none", "0"); result = codeflash_output # 8.84μs -> 8.18μs (8.09% faster)

def test_async_function_returns_original(default_function_obj, call_position):
    """Test that if the function is async, the node is returned unchanged."""
    async_func = FunctionToOptimize(function_name="foo", qualified_name="foo", is_async=True)
    node = make_simple_call_stmt()
    inj = InjectPerfOnly(async_func, "module.py", "pytest", call_position)
    codeflash_output = inj.find_and_update_line_node(node, "test_async", "0"); result = codeflash_output # 9.82μs -> 7.63μs (28.6% faster)

def test_multiple_calls_only_first_transformed(default_function_obj, call_position):
    """Test that only the first matching call is transformed."""
    # Create a node with two calls: foo() and bar()
    call1 = ast.Call(func=ast.Name(id="foo", ctx=ast.Load()), args=[], keywords=[])
    call1.lineno = 1
    call1.col_offset = 0
    call1.end_lineno = 1
    call1.end_col_offset = 10
    call2 = ast.Call(func=ast.Name(id="bar", ctx=ast.Load()), args=[], keywords=[])
    call2.lineno = 1
    call2.col_offset = 11
    call2.end_lineno = 1
    call2.end_col_offset = 20
    expr = ast.Expr(
        value=ast.Tuple(elts=[call1, call2], ctx=ast.Load()), lineno=1, col_offset=0
    )
    inj = InjectPerfOnly(default_function_obj, "module.py", "pytest", call_position)
    codeflash_output = inj.find_and_update_line_node(expr, "test_multi", "0"); result = codeflash_output # 24.8μs -> 21.8μs (13.5% faster)
    # Only the first call node (foo) should be wrapped
    call_node = get_first_call_node(result[-1])
    # The second call (bar) should remain unchanged
    tuple_elts = result[-1].value.elts

def test_call_with_args_and_keywords(default_function_obj, call_position):
    """Test that calls with arguments and keywords are handled."""
    args = [ast.Constant(value=1), ast.Constant(value=2)]
    keywords = [ast.keyword(arg="x", value=ast.Constant(value=3))]
    node = make_simple_call_stmt(args=args, keywords=keywords)
    inj = InjectPerfOnly(default_function_obj, "module.py", "pytest", call_position)
    codeflash_output = inj.find_and_update_line_node(node, "test_args", "0"); result = codeflash_output # 22.9μs -> 19.1μs (19.9% faster)
    call_node = get_first_call_node(result[-1])

def test_method_call_with_different_method(default_method_function_obj, call_position):
    """Test that method calls with a different method name are not wrapped."""
    node = make_method_call_stmt(obj_name="self", method_name="notfoo")
    inj = InjectPerfOnly(default_method_function_obj, "module.py", "pytest", call_position)
    codeflash_output = inj.find_and_update_line_node(node, "test_bar", "1", test_class_name="Bar"); result = codeflash_output # 12.2μs -> 11.5μs (5.90% faster)

# 3. Large Scale Test Cases


def test_large_ast_tree_with_nested_calls(default_function_obj):
    """Test a large AST tree with deeply nested calls."""
    # Create a nested call: foo(bar(baz(1)))
    call_baz = ast.Call(func=ast.Name(id="baz", ctx=ast.Load()), args=[ast.Constant(value=1)], keywords=[])
    call_baz.lineno = 1
    call_baz.col_offset = 0
    call_baz.end_lineno = 1
    call_baz.end_col_offset = 5
    call_bar = ast.Call(func=ast.Name(id="bar", ctx=ast.Load()), args=[call_baz], keywords=[])
    call_bar.lineno = 1
    call_bar.col_offset = 6
    call_bar.end_lineno = 1
    call_bar.end_col_offset = 11
    call_foo = ast.Call(func=ast.Name(id="foo", ctx=ast.Load()), args=[call_bar], keywords=[])
    call_foo.lineno = 1
    call_foo.col_offset = 12
    call_foo.end_lineno = 1
    call_foo.end_col_offset = 17
    expr = ast.Expr(value=call_foo, lineno=1, col_offset=0)
    # Only the outermost call matches the position
    call_position = [CodePosition(line_no=1, col_no=12, end_col_offset=17)]
    inj = InjectPerfOnly(default_function_obj, "module.py", "pytest", call_position)
    codeflash_output = inj.find_and_update_line_node(expr, "test_nested", "0"); result = codeflash_output # 23.3μs -> 19.6μs (18.9% faster)
    call_node = get_first_call_node(result[-1])
    # The argument to foo should be bar(baz(1)), not wrapped
    arg = call_node.args[-1]
    if isinstance(arg, ast.Starred):
        # In BEHAVIOR mode, the argument is a Starred node
        pass
    else:
        pass

def test_large_call_position_list(default_function_obj):
    """Test with a large call_positions list, only matching positions wrap calls."""
    node = make_simple_call_stmt(lineno=50)
    # 999 positions that don't match, and 1 that does
    positions = [CodePosition(line_no=i, col_no=0, end_col_offset=10) for i in range(1, 1000)]
    positions.append(CodePosition(line_no=50, col_no=0, end_col_offset=10))
    inj = InjectPerfOnly(default_function_obj, "module.py", "pytest", positions)
    codeflash_output = inj.find_and_update_line_node(node, "test_largepos", "0"); result = codeflash_output # 28.4μs -> 22.5μs (26.2% faster)
    call_node = get_first_call_node(result[-1])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import ast
from types import SimpleNamespace

# imports
import pytest
from codeflash.code_utils.instrument_existing_tests import InjectPerfOnly


# function to test
def find_and_update_line_node(
    test_node: ast.stmt, node_name: str, index: str, test_class_name: str | None = None
) -> list[ast.stmt] | None:
    """
    Given an AST statement node, finds the first ast.Call node within it that matches a call position,
    and rewrites it (and possibly adds argument binding statements) as described in the original code.
    Returns a list of AST statement nodes (possibly including the original test_node and new nodes),
    or None if no matching call is found.
    """
    # For testability, we'll simulate a minimal context for the function:
    # We'll use a global variable for call_positions, function_object, module_path, mode, etc.
    # (In production, these are attributes of the class, but for testing we inject them globally.)
    global _find_and_update_ctx
    ctx = _find_and_update_ctx

    def get_call_arguments(call_node: ast.Call):
        return SimpleNamespace(args=call_node.args, keywords=call_node.keywords)

    def node_in_call_position(node: ast.AST, call_positions: list):
        if isinstance(node, ast.Call) and hasattr(node, "lineno") and hasattr(node, "col_offset"):
            for pos in call_positions:
                if (
                    pos.line_no is not None
                    and getattr(node, "end_lineno", None) is not None
                    and node.lineno <= pos.line_no <= node.end_lineno
                ):
                    if pos.line_no == node.lineno and node.col_offset <= pos.col_no:
                        return True
                    if (
                        pos.line_no == node.end_lineno
                        and pos.end_col_offset is not None
                        and pos.end_col_offset >= pos.col_no
                    ):
                        return True
                    if node.lineno < pos.line_no < node.end_lineno:
                        return True
        return False

    return_statement = [test_node]
    call_node = None
    for node in ast.walk(test_node):
        if isinstance(node, ast.Call) and node_in_call_position(node, ctx.call_positions):
            call_node = node
            all_args = get_call_arguments(call_node)
            if isinstance(node.func, ast.Name):
                function_name = node.func.id

                if getattr(ctx.function_object, "is_async", False):
                    return [test_node]

                # Create the signature binding statements
                bind_call = ast.Assign(
                    targets=[ast.Name(id="_call__bound__arguments", ctx=ast.Store())],
                    value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Call(
                                func=ast.Attribute(
                                    value=ast.Name(id="inspect", ctx=ast.Load()), attr="signature", ctx=ast.Load()
                                ),
                                args=[ast.Name(id=function_name, ctx=ast.Load())],
                                keywords=[],
                            ),
                            attr="bind",
                            ctx=ast.Load(),
                        ),
                        args=all_args.args,
                        keywords=all_args.keywords,
                    ),
                    lineno=test_node.lineno,
                    col_offset=test_node.col_offset,
                )

                apply_defaults = ast.Expr(
                    value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                            attr="apply_defaults",
                            ctx=ast.Load(),
                        ),
                        args=[],
                        keywords=[],
                    ),
                    lineno=test_node.lineno + 1,
                    col_offset=test_node.col_offset,
                )

                node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
                node.args = [
                    ast.Name(id=function_name, ctx=ast.Load()),
                    ast.Constant(value=ctx.module_path),
                    ast.Constant(value=test_class_name or None),
                    ast.Constant(value=node_name),
                    ast.Constant(value=getattr(ctx.function_object, "qualified_name", "")),
                    ast.Constant(value=index),
                    ast.Name(id="codeflash_loop_index", ctx=ast.Load()),
                    *(
                        [ast.Name(id="codeflash_cur", ctx=ast.Load()), ast.Name(id="codeflash_con", ctx=ast.Load())]
                        if ctx.mode == "BEHAVIOR"
                        else []
                    ),
                    *(
                        call_node.args
                        if ctx.mode == "PERFORMANCE"
                        else [
                            ast.Starred(
                                value=ast.Attribute(
                                    value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                    attr="args",
                                    ctx=ast.Load(),
                                ),
                                ctx=ast.Load(),
                            )
                        ]
                    ),
                ]
                node.keywords = (
                    [
                        ast.keyword(
                            value=ast.Attribute(
                                value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                attr="kwargs",
                                ctx=ast.Load(),
                            )
                        )
                    ]
                    if ctx.mode == "BEHAVIOR"
                    else call_node.keywords
                )

                return_statement = (
                    [bind_call, apply_defaults, test_node] if ctx.mode == "BEHAVIOR" else [test_node]
                )
                break
            if isinstance(node.func, ast.Attribute):
                function_to_test = node.func.attr
                if function_to_test == getattr(ctx.function_object, "function_name", ""):
                    if getattr(ctx.function_object, "is_async", False):
                        return [test_node]

                    function_name = ast.unparse(node.func)

                    bind_call = ast.Assign(
                        targets=[ast.Name(id="_call__bound__arguments", ctx=ast.Store())],
                        value=ast.Call(
                            func=ast.Attribute(
                                value=ast.Call(
                                    func=ast.Attribute(
                                        value=ast.Name(id="inspect", ctx=ast.Load()),
                                        attr="signature",
                                        ctx=ast.Load(),
                                    ),
                                    args=[ast.parse(function_name, mode="eval").body],
                                    keywords=[],
                                ),
                                attr="bind",
                                ctx=ast.Load(),
                            ),
                            args=all_args.args,
                            keywords=all_args.keywords,
                        ),
                        lineno=test_node.lineno,
                        col_offset=test_node.col_offset,
                    )

                    apply_defaults = ast.Expr(
                        value=ast.Call(
                            func=ast.Attribute(
                                value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                attr="apply_defaults",
                                ctx=ast.Load(),
                            ),
                            args=[],
                            keywords=[],
                        ),
                        lineno=test_node.lineno + 1,
                        col_offset=test_node.col_offset,
                    )

                    node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
                    node.args = [
                        ast.parse(function_name, mode="eval").body,
                        ast.Constant(value=ctx.module_path),
                        ast.Constant(value=test_class_name or None),
                        ast.Constant(value=node_name),
                        ast.Constant(value=getattr(ctx.function_object, "qualified_name", "")),
                        ast.Constant(value=index),
                        ast.Name(id="codeflash_loop_index", ctx=ast.Load()),
                        *(
                            [
                                ast.Name(id="codeflash_cur", ctx=ast.Load()),
                                ast.Name(id="codeflash_con", ctx=ast.Load()),
                            ]
                            if ctx.mode == "BEHAVIOR"
                            else []
                        ),
                        *(
                            call_node.args
                            if ctx.mode == "PERFORMANCE"
                            else [
                                ast.Starred(
                                    value=ast.Attribute(
                                        value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                        attr="args",
                                        ctx=ast.Load(),
                                    ),
                                    ctx=ast.Load(),
                                )
                            ]
                        ),
                    ]
                    node.keywords = (
                        [
                            ast.keyword(
                                value=ast.Attribute(
                                    value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                    attr="kwargs",
                                    ctx=ast.Load(),
                                )
                            )
                        ]
                        if ctx.mode == "BEHAVIOR"
                        else call_node.keywords
                    )

                    return_statement = (
                        [bind_call, apply_defaults, test_node] if ctx.mode == "BEHAVIOR" else [test_node]
                    )
                    break

    if call_node is None:
        return None
    return return_statement

# Helper classes for simulating context and positions
class DummyFunctionObj:
    def __init__(self, function_name="foo", qualified_name="foo", is_async=False):
        self.function_name = function_name
        self.qualified_name = qualified_name
        self.is_async = is_async
        self.parents = []
        self.top_level_parent_name = None

class DummyCodePosition:
    def __init__(self, line_no, col_no, end_col_offset=None):
        self.line_no = line_no
        self.col_no = col_no
        self.end_col_offset = end_col_offset

# 1. Basic Test Cases

To edit these changes git checkout codeflash/optimize-pr867-2025-11-01T00.02.02 and push.