⚡️ Speed up function create_wrapper_function by 14% in PR #617 (alpha-async)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #617

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

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

---

📄 14% (0.14x) speedup for create_wrapper_function in codeflash/code_utils/instrument_existing_tests.py

⏱️ Runtime : 1.93 milliseconds → 1.70 milliseconds (best of 58 runs)

📝 Explanation and details

The optimized code achieves a 13% speedup by eliminating redundant AST node creation through aggressive caching and reuse.

What specific optimizations were applied:

1. Pre-constructed AST node cache: Instead of creating new ast.Name and ast.Attribute objects repeatedly (the original created hundreds of identical nodes), the optimized version creates each unique node once and reuses it. For example, n_test_module_name_L = name("test_module_name", LOAD) is created once and reused everywhere.

2. Constant node reuse: Common constants like ast.Constant(value=":") are created once as c_colon and reused throughout, eliminating duplicate allocations.

3. Factory functions for complex structures: JoinedStr constructions are moved into reusable factory functions like joinedstr_test_id(), reducing code duplication and enabling better optimization.

Why this leads to speedup:

• Reduced object allocation overhead: The original code was creating ~200+ identical AST nodes per call. The optimized version creates each unique node exactly once.
• Better memory locality: Reusing the same objects improves CPU cache efficiency.
• Fewer function calls: Instead of calling ast.Name() hundreds of times, it's called once per unique identifier.

Test case performance patterns:

The optimization shows consistent 15-20% improvements across most test cases, with particularly strong gains in tests that exercise the full AST construction pipeline (like test_assigns_test_id showing 19% improvement). The speedup is most pronounced for tests that trigger the complex JoinedStr constructions, which benefited most from the caching strategy.

This optimization is especially effective for workloads that call create_wrapper_function repeatedly, as the AST node reuse compound the benefits over multiple invocations.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_async_instrumentation.py::test_async_wrapper_preserves_return_value	73.6μs	70.1μs	✅4.97%
test_async_instrumentation.py::test_mixed_sync_async_instrumentation	74.3μs	69.5μs	✅6.88%
test_async_instrumentation.py::test_wrapper_function_includes_async_check	73.3μs	69.2μs	✅5.93%

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import ast
import sys
# Minimal stubs for enums, as required for the function
from enum import Enum

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


class TestingMode(Enum):
    BEHAVIOR = 1
    PERFORMANCE = 2

class VerificationType(Enum):
    FUNCTION_CALL = "function_call"
from codeflash.code_utils.instrument_existing_tests import \
    create_wrapper_function

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

# Helper function to extract argument names from ast.arguments
def get_arg_names(ast_args):
    return [a.arg for a in ast_args.args]

# Helper function to find a node of a given type in a body
def find_node(body, node_type):
    for node in body:
        if isinstance(node, node_type):
            return node
    return None

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

def test_functiondef_name_and_type():
    """Test that the returned object is an ast.FunctionDef with the correct name."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 86.6μs -> 72.2μs (19.9% faster)

def test_functiondef_arguments_behavior_mode():
    """Test that the arguments are correct in BEHAVIOR mode."""
    codeflash_output = create_wrapper_function(TestingMode.BEHAVIOR); fn = codeflash_output # 69.9μs -> 66.9μs (4.49% faster)
    arg_names = get_arg_names(fn.args)
    # Check required arguments
    expected = [
        "wrapped", "test_module_name", "test_class_name", "test_name",
        "function_name", "line_id", "loop_index", "codeflash_cur", "codeflash_con"
    ]

def test_functiondef_arguments_performance_mode():
    """Test that the arguments are correct in PERFORMANCE mode (no DB args)."""
    codeflash_output = create_wrapper_function(TestingMode.PERFORMANCE); fn = codeflash_output # 70.2μs -> 66.5μs (5.59% faster)
    arg_names = get_arg_names(fn.args)
    expected = [
        "wrapped", "test_module_name", "test_class_name", "test_name",
        "function_name", "line_id", "loop_index"
    ]

def test_assigns_test_id():
    """Test that the first statement assigns to 'test_id' as a JoinedStr."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 83.4μs -> 70.1μs (19.0% faster)
    assign = fn.body[0]
    # Check that all required FormattedValue and Constant nodes are present
    expected_names = [
        "test_module_name", "test_class_name", "test_name", "line_id", "loop_index"
    ]
    found_names = [v.value.id for v in assign.value.values if isinstance(v, ast.FormattedValue)]
    for name in expected_names:
        pass

def test_assigns_test_stdout_tag():
    """Test that 'test_stdout_tag' is assigned as a JoinedStr and used in print."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 82.6μs -> 70.1μs (17.8% faster)
    # Find assignment to test_stdout_tag
    assign = None
    for node in fn.body:
        if isinstance(node, ast.Assign):
            if node.targets[0].id == "test_stdout_tag":
                assign = node
                break
    # Check that the print uses test_stdout_tag
    found = False
    for node in fn.body:
        if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
            call = node.value
            if (isinstance(call.func, ast.Name) and call.func.id == "print" and
                any(isinstance(arg, ast.JoinedStr) and
                    any(isinstance(v, ast.FormattedValue) and
                        getattr(v.value, "id", None) == "test_stdout_tag"
                        for v in arg.values)
                    for arg in call.args)):
                found = True

def test_assigns_codeflash_test_index():
    """Test that 'codeflash_test_index' is assigned as a subscript of codeflash_wrap.index."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 83.0μs -> 69.5μs (19.3% faster)
    assign = None
    for node in fn.body:
        if isinstance(node, ast.Assign):
            if node.targets[0].id == "codeflash_test_index":
                assign = node
                break
    sub = assign.value

def test_assigns_invocation_id():
    """Test that 'invocation_id' is assigned as a JoinedStr with line_id and codeflash_test_index."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 82.7μs -> 69.5μs (19.1% faster)
    assign = None
    for node in fn.body:
        if isinstance(node, ast.Assign):
            if node.targets[0].id == "invocation_id":
                assign = node
                break
    # Should use line_id and codeflash_test_index
    ids = [v.value.id for v in assign.value.values if isinstance(v, ast.FormattedValue)]

def test_assigns_exception_to_none():
    """Test that 'exception' is assigned to None before try block."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 82.2μs -> 69.1μs (18.9% faster)
    found = False
    for node in fn.body:
        if isinstance(node, ast.Assign):
            if node.targets[0].id == "exception":
                found = True

def test_try_block_and_except_handler():
    """Test that there is a Try block with ExceptHandler for Exception as e."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 81.6μs -> 68.9μs (18.5% faster)
    try_node = find_node(fn.body, ast.Try)
    handler = try_node.handlers[0]

def test_final_return_statement():
    """Test that the last statement is a Return of return_value."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 82.2μs -> 68.9μs (19.3% faster)
    # Last non-empty statement should be Return
    for node in reversed(fn.body):
        if isinstance(node, ast.Return):
            break
    else:
        pytest.fail("No Return statement found in function body.")

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

def test_if_codeflash_wrap_index_created():
    """Test that there is an If statement that creates codeflash_wrap.index if not present."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 82.6μs -> 68.3μs (20.9% faster)
    # Find the If node
    if_node = fn.body[1]
    # The test should be 'not hasattr(codeflash_wrap, "index")'
    test = if_node.test
    args = test.operand.args
    # The body should assign codeflash_wrap.index = {}
    assign = if_node.body[0]

def test_if_test_id_in_index_branching():
    """Test that the If statement checks if test_id in codeflash_wrap.index and does correct branch."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 82.1μs -> 68.4μs (20.2% faster)
    # Find the If node (third statement)
    if_node = fn.body[2]
    # The test should be 'test_id in codeflash_wrap.index'
    test = if_node.test
    # The comparator is codeflash_wrap.index
    comp = test.comparators[0]

def test_behavior_mode_has_db_statements():
    """Test that in BEHAVIOR mode, the function includes DB execute and commit."""
    codeflash_output = create_wrapper_function(TestingMode.BEHAVIOR); fn = codeflash_output # 69.2μs -> 65.2μs (6.02% faster)
    # Look for ast.Expr nodes with ast.Call to codeflash_cur.execute and codeflash_con.commit
    found_execute = False
    found_commit = False
    for node in fn.body:
        if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
            call = node.value
            if (isinstance(call.func, ast.Attribute) and
                call.func.attr == "execute" and
                getattr(call.func.value, "id", None) == "codeflash_cur"):
                found_execute = True
            if (isinstance(call.func, ast.Attribute) and
                call.func.attr == "commit" and
                getattr(call.func.value, "id", None) == "codeflash_con"):
                found_commit = True

def test_performance_mode_has_no_db_statements():
    """Test that in PERFORMANCE mode, the function does NOT include DB execute/commit."""
    codeflash_output = create_wrapper_function(TestingMode.PERFORMANCE); fn = codeflash_output # 68.9μs -> 65.0μs (5.92% faster)
    for node in fn.body:
        if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
            call = node.value
            if (isinstance(call.func, ast.Attribute) and
                call.func.attr in ("execute", "commit") and
                getattr(call.func.value, "id", None) in ("codeflash_cur", "codeflash_con")):
                pytest.fail("Found DB statement in PERFORMANCE mode")

def test_behavior_mode_has_pickled_return_value():
    """Test that in BEHAVIOR mode, pickled_return_value is assigned using pickle.dumps."""
    codeflash_output = create_wrapper_function(TestingMode.BEHAVIOR); fn = codeflash_output # 69.0μs -> 65.0μs (6.06% faster)
    found = False
    for node in fn.body:
        if isinstance(node, ast.Assign):
            if hasattr(node.targets[0], "id") and node.targets[0].id == "pickled_return_value":
                for branch in [node.value.body, node.value.orelse]:
                    pass
                found = True

def test_performance_mode_has_no_pickled_return_value():
    """Test that in PERFORMANCE mode, pickled_return_value is NOT assigned."""
    codeflash_output = create_wrapper_function(TestingMode.PERFORMANCE); fn = codeflash_output # 68.9μs -> 65.9μs (4.59% faster)
    for node in fn.body:
        if isinstance(node, ast.Assign):
            if hasattr(node.targets[0], "id") and node.targets[0].id == "pickled_return_value":
                pytest.fail("Found pickled_return_value in PERFORMANCE mode")

def test_if_exception_is_raised():
    """Test that there is an If statement that raises exception if set."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 83.0μs -> 69.3μs (19.8% faster)
    # Last If node before Return
    if_nodes = [n for n in fn.body if isinstance(n, ast.If)]
    found = False
    for if_node in if_nodes:
        if isinstance(if_node.test, ast.Name) and if_node.test.id == "exception":
            found = True

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

def test_many_arguments_and_long_names():
    """Test that the function can handle long argument names and still builds correct AST."""
    # Simulate long test names and module/class names by checking that the JoinedStr handles them
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 82.5μs -> 69.6μs (18.5% faster)
    assign = fn.body[0]
    # Should still be a JoinedStr with all FormattedValues present
    formatted = [v for v in assign.value.values if isinstance(v, ast.FormattedValue)]

def test_large_body_structure():
    """Test that the function body contains all expected major blocks, even with large number of statements."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 82.0μs -> 68.2μs (20.4% faster)

def test_scalability_with_many_invocations():
    """Test that the function's AST structure allows for many unique test_id values (simulate 1000 invocations)."""
    codeflash_output = create_wrapper_function(); fn = codeflash_output # 82.4μs -> 68.6μs (20.1% faster)
    # The code uses codeflash_wrap.index as a dict with test_id as key
    # Simulate 1000 unique test_id values (no actual execution, just check AST supports it)
    # The AST must NOT hardcode any limit on number of test_id entries
    # Check that the relevant statements use test_id as a key
    for node in fn.body:
        if isinstance(node, ast.Assign):
            if hasattr(node.targets[0], "id") and node.targets[0].id == "codeflash_test_index":
                sub = node.value
    # If the code used a fixed-length array or similar, this would fail


def test_ast_isolated_behavior_and_performance_modes():
    """Test that BEHAVIOR and PERFORMANCE modes produce different ASTs for key statements."""
    codeflash_output = create_wrapper_function(TestingMode.BEHAVIOR); fn_behavior = codeflash_output # 73.0μs -> 69.5μs (4.99% faster)
    codeflash_output = create_wrapper_function(TestingMode.PERFORMANCE); fn_perf = codeflash_output # 65.1μs -> 58.9μs (10.5% faster)
# 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
from enum import Enum
from types import CodeType, FunctionType

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


# Minimal stubs for TestingMode and VerificationType
class TestingMode(Enum):
    BEHAVIOR = 1
    PERFORMANCE = 2

class VerificationType(Enum):
    FUNCTION_CALL = 1
from codeflash.code_utils.instrument_existing_tests import \
    create_wrapper_function

# ========== Unit Tests ==========

# Helper function to compile and retrieve the wrapper function from the AST
def get_wrapper_func(mode=TestingMode.BEHAVIOR):
    import gc
    import inspect
    import pickle
    import time
    import types

    ast_mod = ast.Module(
        body=[create_wrapper_function(mode)],
        type_ignores=[] if sys.version_info >= (3, 8) else [],
    )
    ast.fix_missing_locations(ast_mod)
    ns = {
        "pickle": pickle,
        "gc": gc,
        "inspect": inspect,
        "time": time,
        "codeflash_async_wrap_inner": lambda f, *a, **k: f(*a, **k),  # fallback for sync in tests
        "VerificationType": VerificationType,
        "TestingMode": TestingMode,
    }
    code = compile(ast_mod, "<test>", "exec")
    exec(code, ns)
    return ns["codeflash_wrap"]

# --- Basic Test Cases ---

To edit these changes git checkout codeflash/optimize-pr617-2025-08-06T00.05.08 and push.

[codeflash-ai]

This PR has been automatically closed because the original PR #617 by KRRT7 was closed.