⚡️ Speed up method AsyncCallInstrumenter._process_test_function by 30% 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.

---

📄 30% (0.30x) speedup for AsyncCallInstrumenter._process_test_function in codeflash/code_utils/instrument_existing_tests.py

⏱️ Runtime : 2.08 milliseconds → 1.61 milliseconds (best of 14 runs)

📝 Explanation and details

The optimized code achieves a 29% speedup through three key optimizations:

1. Replaced ast.walk() with manual stack traversal in _instrument_statement()
The original code used ast.walk() which creates a generator and recursively yields nodes. The optimized version uses an explicit stack with ast.iter_child_nodes(), eliminating generator overhead. This is the primary performance gain, as shown in the line profiler where ast.walk() took 81.9% of execution time in the original vs the new manual traversal being more efficient.

2. Optimized timeout decorator check with early exit
Instead of using any() with a generator expression that always evaluates all decorators, the optimized version uses a manual loop with break when the timeout decorator is found. This avoids unnecessary iterations when the decorator is found early, particularly beneficial for unittest frameworks.

3. Minor micro-optimizations

• Cached self.async_call_counter to a local variable to reduce attribute lookups
• Replaced hasattr(stmt, "lineno") with getattr(stmt, "lineno", 1) to avoid double attribute access
• Cached node.decorator_list reference to avoid repeated attribute access

Performance characteristics by test type:

• Large-scale tests (500+ async calls): The stack-based traversal shows significant gains due to reduced generator overhead
• unittest framework tests: Early exit optimization provides 33-99% speedup when timeout decorators are found quickly
• Mixed target/non-target calls: Manual traversal avoids unnecessary deep walks through non-matching nodes
• Small functions: Minor but consistent 10-25% improvements from micro-optimizations

The optimizations are most effective for codebases with many async calls or complex AST structures where the reduced generator overhead and early exits provide compound benefits.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import ast
import types

# imports
import pytest
from codeflash.code_utils.instrument_existing_tests import \
    AsyncCallInstrumenter


# Dummy imports/classes for dependencies (since we can't import actual codeflash modules)
class CodePosition:
    def __init__(self, lineno, col_offset):
        self.lineno = lineno
        self.col_offset = col_offset

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

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

# The AsyncCallInstrumenter class as provided above (already included in prompt)

# Helper to parse code and get function node
def get_func_node(code, func_name):
    tree = ast.parse(code)
    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name:
            return node
    raise ValueError(f"Function {func_name} not found")

# Helper to get all env assignments in a function node
def get_env_assignments(node):
    env_assigns = []
    for stmt in node.body:
        if isinstance(stmt, ast.Assign):
            if (
                isinstance(stmt.targets[0], ast.Subscript)
                and isinstance(stmt.targets[0].value, ast.Attribute)
                and stmt.targets[0].value.attr == "environ"
                and stmt.targets[0].slice.value == "CODEFLASH_CURRENT_LINE_ID"
            ):
                env_assigns.append(stmt)
    return env_assigns

# Helper to check for timeout_decorator in decorator_list
def has_timeout_decorator(node):
    for d in node.decorator_list:
        if (
            isinstance(d, ast.Call)
            and isinstance(d.func, ast.Name)
            and d.func.id == "timeout_decorator.timeout"
            and len(d.args) == 1
            and isinstance(d.args[0], ast.Constant)
            and d.args[0].value == 15
        ):
            return True
    return False

# -------------------------------
# Basic Test Cases
# -------------------------------

def test_no_async_calls_no_instrumentation():
    """Test: Function with no async calls should not be instrumented."""
    code = """
def test_simple():
    x = 1
    y = x + 2
    return y
"""
    func_node = get_func_node(code, "test_simple")
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="pytest",
        call_positions=[],
    )
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output # 21.4μs -> 17.7μs (20.8% faster)

def test_async_call_in_target_position_instruments_env_assignment():
    """Test: Awaiting target function at target position triggers env assignment."""
    code = """
async def test_async():
    await foo()
    return 42
"""
    func_node = get_func_node(code, "test_async")
    # Simulate call position at line 3, col 4 (await foo())
    call_pos = [CodePosition(3, 4)]
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="pytest",
        call_positions=call_pos,
    )
    # Patch ast.Call node with lineno/col_offset for test
    for node in ast.walk(func_node):
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            node.value.lineno = 3
            node.value.col_offset = 4
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output
    env_assigns = get_env_assignments(new_node)

def test_unittest_framework_adds_timeout_decorator():
    """Test: When using 'unittest', timeout_decorator is added if not present."""
    code = """
def test_unittest():
    pass
"""
    func_node = get_func_node(code, "test_unittest")
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="unittest",
        call_positions=[],
    )
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output # 11.9μs -> 8.13μs (46.7% faster)

def test_unittest_framework_does_not_duplicate_timeout_decorator():
    """Test: If timeout_decorator already present, do not add again."""
    code = """
@timeout_decorator.timeout(15)
def test_unittest():
    pass
"""
    func_node = get_func_node(code, "test_unittest")
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="unittest",
        call_positions=[],
    )
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output # 10.0μs -> 6.54μs (52.9% faster)

def test_multiple_awaits_instrument_multiple_env_assignments():
    """Test: Multiple awaits of target function at different positions are instrumented."""
    code = """
async def test_multi():
    await foo()
    await foo()
    await foo()
"""
    func_node = get_func_node(code, "test_multi")
    positions = [CodePosition(3, 4), CodePosition(4, 4), CodePosition(5, 4)]
    # Patch nodes
    i = 0
    for node in ast.walk(func_node):
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            node.value.lineno = positions[i].lineno
            node.value.col_offset = positions[i].col_offset
            i += 1
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="pytest",
        call_positions=positions,
    )
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output
    env_assigns = get_env_assignments(new_node)
    # Values should be "0", "1", "2"
    values = [assign.value.value for assign in env_assigns]

# -------------------------------
# Edge Test Cases
# -------------------------------

def test_async_call_not_in_target_position_no_instrumentation():
    """Test: Awaiting target function NOT at target position does not instrument."""
    code = """
async def test_async():
    await foo()
    return 42
"""
    func_node = get_func_node(code, "test_async")
    # Target position is not where the call is
    call_pos = [CodePosition(100, 100)]
    for node in ast.walk(func_node):
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            node.value.lineno = 3
            node.value.col_offset = 4
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="pytest",
        call_positions=call_pos,
    )
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output

def test_non_target_function_not_instrumented():
    """Test: Awaiting a non-target function is not instrumented."""
    code = """
async def test_async():
    await bar()
"""
    func_node = get_func_node(code, "test_async")
    call_pos = [CodePosition(3, 4)]
    for node in ast.walk(func_node):
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            node.value.lineno = 3
            node.value.col_offset = 4
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),  # target is 'foo', but call is 'bar'
        module_path="dummy.py",
        test_framework="pytest",
        call_positions=call_pos,
    )
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output # 8.60μs -> 7.88μs (9.19% faster)

def test_attribute_call_target_function_instrumented():
    """Test: Awaiting method call (obj.foo()) is instrumented if target is 'foo'."""
    code = """
async def test_async():
    await obj.foo()
"""
    func_node = get_func_node(code, "test_async")
    call_pos = [CodePosition(3, 10)]
    # Patch call node
    for node in ast.walk(func_node):
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            node.value.lineno = 3
            node.value.col_offset = 10
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="pytest",
        call_positions=call_pos,
    )
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output
    env_assigns = get_env_assignments(new_node)

def test_env_assignment_lineno_matches_stmt():
    """Test: Env assignment's lineno matches original statement's lineno if present."""
    code = """
async def test_async():
    await foo()
"""
    func_node = get_func_node(code, "test_async")
    call_pos = [CodePosition(3, 4)]
    # Patch call node and statement with lineno
    for node in ast.walk(func_node):
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            node.value.lineno = 3
            node.value.col_offset = 4
            node.lineno = 3
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="pytest",
        call_positions=call_pos,
    )
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output
    env_assigns = get_env_assignments(new_node)

def test_env_assignment_lineno_default_if_missing():
    """Test: Env assignment's lineno is 1 if original statement has no lineno."""
    code = """
async def test_async():
    await foo()
"""
    func_node = get_func_node(code, "test_async")
    call_pos = [CodePosition(3, 4)]
    # Patch call node but do NOT set stmt.lineno
    for node in ast.walk(func_node):
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            node.value.lineno = 3
            node.value.col_offset = 4
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="pytest",
        call_positions=call_pos,
    )
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output
    env_assigns = get_env_assignments(new_node)


def test_many_async_calls_scalability():
    """Test: Instrumenter can handle many async calls efficiently."""
    # Generate code with 500 awaits of foo()
    code_lines = ["async def test_many():"] + [
        f"    await foo()" for _ in range(500)
    ]
    code = "\n".join(code_lines)
    func_node = get_func_node(code, "test_many")
    positions = [CodePosition(i+2, 4) for i in range(500)]
    # Patch call nodes
    i = 0
    for node in ast.walk(func_node):
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            node.value.lineno = positions[i].lineno
            node.value.col_offset = positions[i].col_offset
            i += 1
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="pytest",
        call_positions=positions,
    )
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output
    env_assigns = get_env_assignments(new_node)
    # Values should be "0" to "499"
    values = [assign.value.value for assign in env_assigns]

def test_large_function_body_with_mixed_calls():
    """Test: Large function with mixed target and non-target awaits, only target instrumented."""
    code_lines = ["async def test_mixed():"] + [
        "    await foo()" if i % 3 == 0 else "    await bar()" for i in range(600)
    ]
    code = "\n".join(code_lines)
    func_node = get_func_node(code, "test_mixed")
    positions = [CodePosition(i+2, 4) for i in range(0, 600, 3)]
    # Patch call nodes
    foo_idx = 0
    for node in ast.walk(func_node):
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            # Only patch foo() calls
            if foo_idx < len(positions) and node.value.lineno is None:
                if foo_idx * 3 + 2 <= 601:  # Avoid index error
                    node.value.lineno = positions[foo_idx].lineno
                    node.value.col_offset = positions[foo_idx].col_offset
                    foo_idx += 1
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="pytest",
        call_positions=positions,
    )
    codeflash_output = instrumenter._process_test_function(func_node); new_node = codeflash_output
    env_assigns = get_env_assignments(new_node)
    values = [assign.value.value for assign in env_assigns]

def test_env_assignment_counter_is_per_function():
    """Test: Counter for env assignment is per function, not global."""
    code1 = """
async def test_one():
    await foo()
    await foo()
"""
    code2 = """
async def test_two():
    await foo()
    await foo()
"""
    func_node1 = get_func_node(code1, "test_one")
    func_node2 = get_func_node(code2, "test_two")
    positions1 = [CodePosition(3, 4), CodePosition(4, 4)]
    positions2 = [CodePosition(3, 4), CodePosition(4, 4)]
    # Patch call nodes
    i = 0
    for node in ast.walk(func_node1):
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            node.value.lineno = positions1[i].lineno
            node.value.col_offset = positions1[i].col_offset
            i += 1
    i = 0
    for node in ast.walk(func_node2):
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            node.value.lineno = positions2[i].lineno
            node.value.col_offset = positions2[i].col_offset
            i += 1
    instrumenter = AsyncCallInstrumenter(
        FunctionToOptimize("foo"),
        module_path="dummy.py",
        test_framework="pytest",
        call_positions=positions1 + positions2,
    )
    codeflash_output = instrumenter._process_test_function(func_node1); new_node1 = codeflash_output
    codeflash_output = instrumenter._process_test_function(func_node2); new_node2 = codeflash_output
    env_assigns1 = get_env_assignments(new_node1)
    env_assigns2 = get_env_assignments(new_node2)
    # Each function should have its own counter starting at 0
    values1 = [assign.value.value for assign in env_assigns1]
    values2 = [assign.value.value for assign in env_assigns2]
# 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
import sys
import types
# Dummy imports for the required classes and functions
from types import SimpleNamespace

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.instrument_existing_tests import \
    AsyncCallInstrumenter


# Minimal stubs for required classes and enums
class CodePosition:
    def __init__(self, lineno, col_offset):
        self.lineno = lineno
        self.col_offset = col_offset

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

# Helper function to parse and return an AST node for a test function
def make_test_func_ast(
    func_src: str,
    is_async: bool = False,
    decorator_list=None,
):
    tree = ast.parse(func_src)
    node = tree.body[0]
    # patch decorator_list if needed
    if decorator_list is not None:
        node.decorator_list = decorator_list
    return node

# Helper to simulate a parent class node
class DummyParent:
    def __init__(self, type):
        self.type = type

# unit tests

# BASIC TEST CASES

def test_timeout_decorator_added_for_unittest_sync_function():
    """Test that timeout_decorator is added for unittest framework on sync test functions."""
    func = FunctionToOptimize("target_func")
    instr = AsyncCallInstrumenter(func, "dummy_path", "unittest", [])
    src = "def test_func(): pass"
    node = make_test_func_ast(src)
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output # 11.5μs -> 8.62μs (33.6% faster)

def test_timeout_decorator_not_added_if_already_present():
    """Test that timeout_decorator is NOT added if already present."""
    func = FunctionToOptimize("target_func")
    instr = AsyncCallInstrumenter(func, "dummy_path", "unittest", [])
    src = "def test_func(): pass"
    # Simulate decorator already present
    dec = ast.Call(func=ast.Name(id="timeout_decorator.timeout", ctx=ast.Load()), args=[ast.Constant(value=15)], keywords=[])
    node = make_test_func_ast(src, decorator_list=[dec])
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output # 8.35μs -> 4.17μs (99.9% faster)
    # Should still have only one timeout_decorator
    timeout_decs = [d for d in new_node.decorator_list if isinstance(d, ast.Call) and isinstance(d.func, ast.Name) and d.func.id == "timeout_decorator.timeout"]

def test_no_timeout_decorator_for_pytest():
    """Test that timeout_decorator is NOT added for pytest framework."""
    func = FunctionToOptimize("target_func")
    instr = AsyncCallInstrumenter(func, "dummy_path", "pytest", [])
    src = "def test_func(): pass"
    node = make_test_func_ast(src)
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output # 5.79μs -> 3.41μs (69.7% faster)

def test_env_assignment_added_on_await_target_call():
    """Test that env assignment is added before await target_func call at target position."""
    func = FunctionToOptimize("target_func")
    call_pos = [CodePosition(lineno=2, col_offset=4)]
    instr = AsyncCallInstrumenter(func, "dummy_path", "pytest", call_pos)
    src = (
        "async def test_func():\n"
        "    await target_func()\n"
        "    await other_func()\n"
    )
    node = make_test_func_ast(src, is_async=True)
    # Patch lineno/col_offset for the target call
    await_stmt = node.body[0]
    await_stmt.value.lineno = 2
    await_stmt.value.col_offset = 4
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output
    assign = new_node.body[0]

def test_no_env_assignment_for_non_target_call():
    """Test that env assignment is NOT added for await of non-target function."""
    func = FunctionToOptimize("target_func")
    call_pos = [CodePosition(lineno=2, col_offset=4)]
    instr = AsyncCallInstrumenter(func, "dummy_path", "pytest", call_pos)
    src = (
        "async def test_func():\n"
        "    await other_func()\n"
    )
    node = make_test_func_ast(src, is_async=True)
    # Patch lineno/col_offset for the non-target call
    await_stmt = node.body[0]
    await_stmt.value.lineno = 2
    await_stmt.value.col_offset = 4
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output # 12.9μs -> 10.4μs (24.7% faster)

def test_multiple_env_assignments_for_multiple_target_calls():
    """Test that multiple env assignments are added for multiple target calls."""
    func = FunctionToOptimize("target_func")
    call_pos = [
        CodePosition(lineno=2, col_offset=4),
        CodePosition(lineno=3, col_offset=4),
    ]
    instr = AsyncCallInstrumenter(func, "dummy_path", "pytest", call_pos)
    src = (
        "async def test_func():\n"
        "    await target_func()\n"
        "    await target_func()\n"
    )
    node = make_test_func_ast(src, is_async=True)
    # Patch lineno/col_offset for both calls
    node.body[0].value.lineno = 2
    node.body[0].value.col_offset = 4
    node.body[1].value.lineno = 3
    node.body[1].value.col_offset = 4
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output

# EDGE TEST CASES

def test_no_decorator_list():
    """Test that function with no decorator_list attribute does not error."""
    func = FunctionToOptimize("target_func")
    instr = AsyncCallInstrumenter(func, "dummy_path", "unittest", [])
    src = "def test_func(): pass"
    node = make_test_func_ast(src)
    # Remove decorator_list attribute
    delattr(node, "decorator_list")
    # Should not raise
    try:
        instr._process_test_function(node)
    except Exception:
        pytest.fail("Should not raise if decorator_list missing")

def test_empty_body():
    """Test that function with empty body does not error and remains unchanged."""
    func = FunctionToOptimize("target_func")
    instr = AsyncCallInstrumenter(func, "dummy_path", "pytest", [])
    src = "def test_func():\n    pass"
    node = make_test_func_ast(src)
    node.body = []
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output # 1.71μs -> 1.94μs (11.9% slower)

def test_non_await_statements():
    """Test that non-await statements do not trigger env assignment."""
    func = FunctionToOptimize("target_func")
    call_pos = [CodePosition(lineno=2, col_offset=4)]
    instr = AsyncCallInstrumenter(func, "dummy_path", "pytest", call_pos)
    src = (
        "def test_func():\n"
        "    x = target_func()\n"
    )
    node = make_test_func_ast(src)
    # Patch lineno/col_offset for the call
    assign_stmt = node.body[0]
    assign_stmt.value.lineno = 2
    assign_stmt.value.col_offset = 4
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output # 12.6μs -> 10.3μs (21.9% faster)

def test_target_call_with_attribute():
    """Test that await of obj.target_func() is instrumented if target_func is the function name."""
    func = FunctionToOptimize("target_func")
    call_pos = [CodePosition(lineno=2, col_offset=4)]
    instr = AsyncCallInstrumenter(func, "dummy_path", "pytest", call_pos)
    src = (
        "async def test_func():\n"
        "    await obj.target_func()\n"
    )
    node = make_test_func_ast(src, is_async=True)
    await_stmt = node.body[0]
    await_stmt.value.lineno = 2
    await_stmt.value.col_offset = 4
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output

def test_call_node_missing_lineno_col_offset():
    """Test that call node missing lineno/col_offset does not get instrumented."""
    func = FunctionToOptimize("target_func")
    call_pos = [CodePosition(lineno=2, col_offset=4)]
    instr = AsyncCallInstrumenter(func, "dummy_path", "pytest", call_pos)
    src = (
        "async def test_func():\n"
        "    await target_func()\n"
    )
    node = make_test_func_ast(src, is_async=True)
    # Remove lineno/col_offset from call node
    await_stmt = node.body[0]
    if hasattr(await_stmt.value, "lineno"):
        delattr(await_stmt.value, "lineno")
    if hasattr(await_stmt.value, "col_offset"):
        delattr(await_stmt.value, "col_offset")
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output


def test_large_number_of_target_calls():
    """Test that env assignments are correctly added for many target calls (scalability)."""
    func = FunctionToOptimize("target_func")
    N = 500
    call_pos = [CodePosition(lineno=i+2, col_offset=4) for i in range(N)]
    instr = AsyncCallInstrumenter(func, "dummy_path", "pytest", call_pos)
    # Build a function with N await target_func() statements
    src_lines = ["async def test_func():"]
    for i in range(N):
        src_lines.append(f"    await target_func()  # {i}")
    src = "\n".join(src_lines)
    node = make_test_func_ast(src, is_async=True)
    # Patch lineno/col_offset for each call
    for i, stmt in enumerate(node.body):
        stmt.value.lineno = i+2
        stmt.value.col_offset = 4
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output
    # Should have N env assignments, each before the corresponding await
    assign_indices = [i for i, stmt in enumerate(new_node.body) if isinstance(stmt, ast.Assign)]
    for i in range(N):
        # Each env assignment should have correct value
        assign = new_node.body[2*i]

def test_large_number_of_non_target_calls():
    """Test that env assignments are NOT added for many non-target calls."""
    func = FunctionToOptimize("target_func")
    N = 500
    call_pos = [CodePosition(lineno=1000, col_offset=4)]  # No call matches this
    instr = AsyncCallInstrumenter(func, "dummy_path", "pytest", call_pos)
    src_lines = ["async def test_func():"]
    for i in range(N):
        src_lines.append(f"    await other_func()  # {i}")
    src = "\n".join(src_lines)
    node = make_test_func_ast(src, is_async=True)
    for i, stmt in enumerate(node.body):
        stmt.value.lineno = i+2
        stmt.value.col_offset = 4
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output # 1.98ms -> 1.53ms (29.6% faster)

def test_large_number_of_calls_mixed_target_and_non_target():
    """Test that env assignments are only added for target calls among many mixed calls."""
    func = FunctionToOptimize("target_func")
    N = 100
    call_pos = [CodePosition(lineno=i+2, col_offset=4) for i in range(N)]
    instr = AsyncCallInstrumenter(func, "dummy_path", "pytest", call_pos)
    src_lines = ["async def test_func():"]
    # Interleave target_func and other_func
    for i in range(N):
        src_lines.append(f"    await target_func()  # {i}")
        src_lines.append(f"    await other_func()  # {i}")
    src = "\n".join(src_lines)
    node = make_test_func_ast(src, is_async=True)
    # Patch lineno/col_offset for target_func calls only
    for i in range(N):
        node.body[2*i].value.lineno = i+2
        node.body[2*i].value.col_offset = 4
        node.body[2*i+1].value.lineno = 1000  # Non-target call, won't match
        node.body[2*i+1].value.col_offset = 4
    codeflash_output = instr._process_test_function(node); new_node = codeflash_output
    # Should have N env assignments, each before target_func call only
    assign_indices = [i for i, stmt in enumerate(new_node.body) if isinstance(stmt, ast.Assign)]
    for i in range(N):
        assign = new_node.body[2*i]
    # Non-target calls should not have env assignments
    for i in range(N):
        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-pr769-2025-09-26T23.14.56 and push.