⚡️ Speed up function create_trace_replay_test_code by 36% in PR #586 (benchmark-fixture-fix)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #586

If you approve this dependent PR, these changes will be merged into the original PR branch benchmark-fixture-fix.

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

---

📄 36% (0.36x) speedup for create_trace_replay_test_code in codeflash/benchmarking/replay_test.py

⏱️ Runtime : 10.5 milliseconds → 7.71 milliseconds (best of 325 runs)

📝 Explanation and details

Here is an optimized version, rewritten for maximum runtime performance.

• Major hotspots (from the profiler) are:
  • get_function_alias and get_unique_test_name – called thousands of times, simple string ops.
  • Many calls to get_function_alias duplicate input in many places.
  • Many repeated lookups (func.get("key")).
  • textwrap.indent and especially textwrap.dedent called many times with the same strings.
• Optimization strategies:
  • Cache results of get_function_alias and get_unique_test_name per argument tuple (using functools.lru_cache).
  • Pre-dedent test templates once at function scope, not in loop.
  • Minimize string concatenation, and loop variable lookups.
  • Use list building + ''.join() for string results where appropriate.
  • Avoid repeated str.format/f-string reinterpretation in tight loops: compose test bodies fully with known variables, not dynamic format.
  • Inline/copy critical logic to avoid redundant function calls per-loop where possible (without changing output).
  • Inline simple get-alias pattern for unique test name.
  • Batch key access from func via local variables, and give fast locals in all loops.

---

---

Key optimizations explained:

• @lru_cache for get_function_alias and get_unique_test_name: avoids recomputation for repeated arguments, which are very frequent.
• Templates are dedented once at module load, dramatically cutting cost of repeated dedentation per test case.
• All func.get("key") inside tight loops are replaced by direct func["key"] to avoid re-parsing and make local lookups faster.
• String concatenations are gathered in lists and joined at the end for efficiency (instead of +=).
• All template string field subs are consolidated to minimize calls to slow formatting/f-strings inside loops.
• Minimizes repeated lookups and attribute access by using local variables for everything repeatedly accessed in loops.

This should result in very significant runtime reduction for large numbers of test generation, as all major hotspots are addressed. Output remains identical to the original.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import re
import textwrap
from types import SimpleNamespace
from typing import Any

# imports
import pytest
from codeflash.benchmarking.replay_test import create_trace_replay_test_code

# -----------------
# unit tests below
# -----------------

def _funcprops(is_classmethod=False, is_staticmethod=False):
    # Helper to create a function_properties object
    return SimpleNamespace(is_classmethod=is_classmethod, is_staticmethod=is_staticmethod)

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

def test_basic_single_function_pytest():
    # Test a single function, no class, with basic names and file path
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "baz_bench",
            "function_properties": _funcprops(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=10,
    ); code = codeflash_output # 91.5μs -> 17.9μs (411% faster)

def test_basic_single_function_unittest():
    # Test a single function, no class, with unittest framework
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "baz_bench",
            "function_properties": _funcprops(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="unittest",
        max_run_count=5,
    ); code = codeflash_output # 91.5μs -> 17.1μs (436% faster)

def test_basic_multiple_functions():
    # Test with two functions from different modules
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "baz_bench",
            "function_properties": _funcprops(),
        },
        {
            "module_name": "spam.eggs",
            "function_name": "qux",
            "class_name": None,
            "file_path": "/tmp/spam/eggs.py",
            "benchmark_function_name": "qux_bench",
            "function_properties": _funcprops(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=3,
    ); code = codeflash_output # 95.4μs -> 23.8μs (302% faster)

def test_basic_class_method():
    # Test a class instance method (not static or classmethod)
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": "MyClass",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "baz_bench",
            "function_properties": _funcprops(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=7,
    ); code = codeflash_output # 94.0μs -> 21.0μs (349% faster)

def test_basic_class_classmethod():
    # Test a classmethod
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": "MyClass",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "baz_bench",
            "function_properties": _funcprops(is_classmethod=True),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=11,
    ); code = codeflash_output # 91.2μs -> 18.6μs (391% faster)

def test_basic_class_staticmethod():
    # Test a staticmethod
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": "MyClass",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "baz_bench",
            "function_properties": _funcprops(is_staticmethod=True),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=12,
    ); code = codeflash_output # 89.7μs -> 17.3μs (420% faster)

def test_basic_class_init():
    # Test __init__ method
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "__init__",
            "class_name": "MyClass",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "init_bench",
            "function_properties": _funcprops(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=2,
    ); code = codeflash_output # 92.7μs -> 19.3μs (379% faster)

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

def test_edge_empty_functions_data():
    # No functions_data should still produce valid code
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=[],
        test_framework="pytest",
        max_run_count=8,
    ); code = codeflash_output # 76.5μs -> 4.17μs (1734% faster)

def test_edge_special_characters_in_names():
    # Test that special characters in benchmark_function_name are sanitized
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "bench mark!@#",
            "function_properties": _funcprops(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=1,
    ); code = codeflash_output # 88.7μs -> 17.3μs (414% faster)

def test_edge_duplicate_function_names():
    # Test two functions with same name but different modules/classes
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "bench1",
            "function_properties": _funcprops(),
        },
        {
            "module_name": "spam.eggs",
            "function_name": "baz",
            "class_name": "Eggs",
            "file_path": "/tmp/spam/eggs.py",
            "benchmark_function_name": "bench2",
            "function_properties": _funcprops(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=4,
    ); code = codeflash_output # 100μs -> 28.3μs (255% faster)

def test_edge_long_module_and_function_names():
    # Test with very long module and function names
    long_mod = "a" * 100
    long_func = "f" * 80
    functions_data = [
        {
            "module_name": long_mod,
            "function_name": long_func,
            "class_name": None,
            "file_path": "/tmp/long.py",
            "benchmark_function_name": "bench_long",
            "function_properties": _funcprops(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=1,
    ); code = codeflash_output # 89.5μs -> 17.2μs (420% faster)

def test_edge_invalid_test_framework():
    # Should raise AssertionError for invalid test_framework
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "baz_bench",
            "function_properties": _funcprops(),
        }
    ]
    with pytest.raises(AssertionError):
        create_trace_replay_test_code(
            trace_file="trace.sqlite",
            functions_data=functions_data,
            test_framework="nose",  # invalid
            max_run_count=1,
        ) # 43.4μs -> 42.8μs (1.36% faster)

def test_edge_file_path_with_backslashes():
    # Should preserve backslashes in file_path
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": r"C:\foo\bar.py",
            "benchmark_function_name": "baz_bench",
            "function_properties": _funcprops(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=1,
    ); code = codeflash_output # 88.9μs -> 16.6μs (436% faster)

def test_edge_function_name_init_is_excluded_from_functions():
    # __init__ should not be in the functions list
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "__init__",
            "class_name": "MyClass",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "init_bench",
            "function_properties": _funcprops(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=1,
    ); code = codeflash_output # 92.7μs -> 20.3μs (357% faster)

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

def test_large_many_functions():
    # Test with 100 functions, all different names
    functions_data = [
        {
            "module_name": f"mod{i}",
            "function_name": f"func{i}",
            "class_name": None,
            "file_path": f"/tmp/mod{i}.py",
            "benchmark_function_name": f"bench{i}",
            "function_properties": _funcprops(),
        }
        for i in range(100)
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=100,
    ); code = codeflash_output # 561μs -> 464μs (21.0% faster)
    # All imports and test functions should be present
    for i in range(100):
        pass
    # Functions list should contain all function names, sorted
    names = [f"func{i}" for i in range(100)]
    sorted_names = sorted(names)

def test_large_many_methods_and_classes():
    # Test with 50 classes, each with a method and a staticmethod
    functions_data = []
    for i in range(50):
        functions_data.append({
            "module_name": f"pkg.mod{i}",
            "function_name": f"meth{i}",
            "class_name": f"Class{i}",
            "file_path": f"/tmp/pkg/mod{i}.py",
            "benchmark_function_name": f"bench_meth{i}",
            "function_properties": _funcprops(),
        })
        functions_data.append({
            "module_name": f"pkg.mod{i}",
            "function_name": f"stat{i}",
            "class_name": f"Class{i}",
            "file_path": f"/tmp/pkg/mod{i}.py",
            "benchmark_function_name": f"bench_stat{i}",
            "function_properties": _funcprops(is_staticmethod=True),
        })
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=50,
    ); code = codeflash_output # 817μs -> 727μs (12.4% faster)
    # All class imports present
    for i in range(50):
        pass
    # All test functions present for methods and staticmethods
    for i in range(50):
        pass

def test_large_max_run_count():
    # Test that max_run_count is used for all test bodies
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "baz_bench",
            "function_properties": _funcprops(),
        },
        {
            "module_name": "spam.eggs",
            "function_name": "qux",
            "class_name": None,
            "file_path": "/tmp/spam/eggs.py",
            "benchmark_function_name": "qux_bench",
            "function_properties": _funcprops(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=999,
    ); code = codeflash_output # 97.1μs -> 24.5μs (297% faster)

def test_large_benchmark_name_variation():
    # Test that benchmark names with spaces and symbols are sanitized for many functions
    functions_data = [
        {
            "module_name": f"mod{i}",
            "function_name": f"func{i}",
            "class_name": None,
            "file_path": f"/tmp/mod{i}.py",
            "benchmark_function_name": f"bench {i}!@#",
            "function_properties": _funcprops(),
        }
        for i in range(20)
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=20,
    ); code = codeflash_output # 190μs -> 111μs (70.7% faster)
    # All test function names should be sanitized
    for i in range(20):
        pass

def test_large_mixed_functions_and_methods():
    # Mix of functions, methods, classmethods, staticmethods
    functions_data = []
    for i in range(10):
        functions_data.append({
            "module_name": f"mod{i}",
            "function_name": f"func{i}",
            "class_name": None,
            "file_path": f"/tmp/mod{i}.py",
            "benchmark_function_name": f"bench_func{i}",
            "function_properties": _funcprops(),
        })
        functions_data.append({
            "module_name": f"mod{i}",
            "function_name": f"meth{i}",
            "class_name": f"Class{i}",
            "file_path": f"/tmp/mod{i}.py",
            "benchmark_function_name": f"bench_meth{i}",
            "function_properties": _funcprops(),
        })
        functions_data.append({
            "module_name": f"mod{i}",
            "function_name": f"clsmeth{i}",
            "class_name": f"Class{i}",
            "file_path": f"/tmp/mod{i}.py",
            "benchmark_function_name": f"bench_clsmeth{i}",
            "function_properties": _funcprops(is_classmethod=True),
        })
        functions_data.append({
            "module_name": f"mod{i}",
            "function_name": f"statmeth{i}",
            "class_name": f"Class{i}",
            "file_path": f"/tmp/mod{i}.py",
            "benchmark_function_name": f"bench_statmeth{i}",
            "function_properties": _funcprops(is_staticmethod=True),
        })
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=10,
    ); code = codeflash_output # 348μs -> 267μs (30.3% faster)
    # All test function names should be present and correct for each type
    for i in range(10):
        pass
    # All class imports present
    for i in range(10):
        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 re
import textwrap
from types import SimpleNamespace
from typing import Any

# imports
import pytest  # used for our unit tests
from codeflash.benchmarking.replay_test import create_trace_replay_test_code

# --- Unit tests for create_trace_replay_test_code ---

# Helper: mock function_properties object
def make_func_props(is_classmethod=False, is_staticmethod=False):
    return SimpleNamespace(is_classmethod=is_classmethod, is_staticmethod=is_staticmethod)

# --- BASIC TEST CASES ---

def test_basic_single_function_pytest():
    """Test a single top-level function with minimal fields and pytest framework."""
    functions_data = [{
        "module_name": "foo.bar",
        "function_name": "baz",
        "class_name": None,
        "file_path": "/tmp/foo/bar.py",
        "benchmark_function_name": "foo_baz_bench",
        "function_properties": make_func_props(),
    }]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data, test_framework="pytest", max_run_count=10); code = codeflash_output # 90.2μs -> 16.5μs (446% faster)

def test_basic_single_function_unittest():
    """Test a single top-level function with unittest framework."""
    functions_data = [{
        "module_name": "foo.bar",
        "function_name": "baz",
        "class_name": None,
        "file_path": "/tmp/foo/bar.py",
        "benchmark_function_name": "foo_baz_bench",
        "function_properties": make_func_props(),
    }]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data, test_framework="unittest", max_run_count=5); code = codeflash_output # 89.9μs -> 16.2μs (455% faster)

def test_basic_multiple_functions():
    """Test multiple top-level functions from different modules."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo_baz_bench",
            "function_properties": make_func_props(),
        },
        {
            "module_name": "alpha.beta",
            "function_name": "gamma",
            "class_name": None,
            "file_path": "/tmp/alpha/beta.py",
            "benchmark_function_name": "alpha_gamma_bench",
            "function_properties": make_func_props(),
        }
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 96.1μs -> 23.8μs (303% faster)

def test_basic_class_method():
    """Test a class method."""
    functions_data = [{
        "module_name": "foo.bar",
        "function_name": "baz",
        "class_name": "Qux",
        "file_path": "/tmp/foo/bar.py",
        "benchmark_function_name": "qux_baz_bench",
        "function_properties": make_func_props(is_classmethod=True),
    }]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 90.1μs -> 18.7μs (381% faster)

def test_basic_static_method():
    """Test a static method."""
    functions_data = [{
        "module_name": "foo.bar",
        "function_name": "baz",
        "class_name": "Qux",
        "file_path": "/tmp/foo/bar.py",
        "benchmark_function_name": "qux_baz_bench",
        "function_properties": make_func_props(is_staticmethod=True),
    }]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 90.6μs -> 17.4μs (420% faster)

def test_basic_instance_method():
    """Test a regular instance method."""
    functions_data = [{
        "module_name": "foo.bar",
        "function_name": "baz",
        "class_name": "Qux",
        "file_path": "/tmp/foo/bar.py",
        "benchmark_function_name": "qux_baz_bench",
        "function_properties": make_func_props(),
    }]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 93.2μs -> 19.8μs (370% faster)

# --- EDGE TEST CASES ---

def test_edge_function_name_with_special_chars():
    """Test function/benchmark names with special characters."""
    functions_data = [{
        "module_name": "foo.bar",
        "function_name": "b@z$",
        "class_name": None,
        "file_path": "/tmp/foo/bar.py",
        "benchmark_function_name": "bench:mark!@#",
        "function_properties": make_func_props(),
    }]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 89.2μs -> 16.6μs (437% faster)

def test_edge_class_name_with_special_chars():
    """Test class names with special characters."""
    functions_data = [{
        "module_name": "foo.bar",
        "function_name": "baz",
        "class_name": "Q@x!",
        "file_path": "/tmp/foo/bar.py",
        "benchmark_function_name": "bench",
        "function_properties": make_func_props(),
    }]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 93.2μs -> 19.8μs (370% faster)

def test_edge_empty_functions_data():
    """Test empty functions_data list."""
    codeflash_output = create_trace_replay_test_code("trace.sqlite", [], test_framework="pytest"); code = codeflash_output # 76.2μs -> 4.47μs (1604% faster)

def test_edge_function_name_dupes():
    """Test functions with the same name but different modules/classes/benchmarks."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "bench1",
            "function_properties": make_func_props(),
        },
        {
            "module_name": "foo.baz",
            "function_name": "baz",
            "class_name": "Qux",
            "file_path": "/tmp/foo/baz.py",
            "benchmark_function_name": "bench2",
            "function_properties": make_func_props(),
        },
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 99.8μs -> 27.8μs (259% faster)

def test_edge_init_method():
    """Test __init__ method handling."""
    functions_data = [{
        "module_name": "foo.bar",
        "function_name": "__init__",
        "class_name": "Qux",
        "file_path": "/tmp/foo/bar.py",
        "benchmark_function_name": "qux_init_bench",
        "function_properties": make_func_props(),
    }]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 90.7μs -> 18.9μs (381% faster)

def test_edge_assert_on_invalid_framework():
    """Test that invalid test_framework raises AssertionError."""
    with pytest.raises(AssertionError):
        create_trace_replay_test_code("trace.sqlite", [], test_framework="nose") # 43.5μs -> 43.0μs (1.21% faster)

def test_edge_max_run_count_zero():
    """Test with max_run_count set to zero."""
    functions_data = [{
        "module_name": "foo.bar",
        "function_name": "baz",
        "class_name": None,
        "file_path": "/tmp/foo/bar.py",
        "benchmark_function_name": "foo_baz_bench",
        "function_properties": make_func_props(),
    }]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data, max_run_count=0); code = codeflash_output # 90.4μs -> 16.5μs (447% faster)

def test_edge_long_module_and_function_names():
    """Test with very long module and function names."""
    long_mod = "a" * 100
    long_func = "b" * 100
    functions_data = [{
        "module_name": long_mod,
        "function_name": long_func,
        "class_name": None,
        "file_path": "/tmp/foo/bar.py",
        "benchmark_function_name": "bench",
        "function_properties": make_func_props(),
    }]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 89.8μs -> 17.0μs (428% faster)

# --- LARGE SCALE TEST CASES ---

def test_large_scale_many_functions():
    """Test with a large number of functions (up to 1000)."""
    n = 1000
    functions_data = [
        {
            "module_name": f"mod{i}",
            "function_name": f"func{i}",
            "class_name": None,
            "file_path": f"/tmp/mod{i}.py",
            "benchmark_function_name": f"bench{i}",
            "function_properties": make_func_props(),
        }
        for i in range(n)
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 4.84ms -> 4.58ms (5.63% faster)
    # Check that all function imports and test functions are present
    for i in (0, n//2, n-1):
        pass

def test_large_scale_mixed_methods():
    """Test with a mix of functions, instance, static, and class methods."""
    n = 100
    functions_data = []
    for i in range(n):
        if i % 4 == 0:
            # top-level function
            functions_data.append({
                "module_name": f"mod{i}",
                "function_name": f"func{i}",
                "class_name": None,
                "file_path": f"/tmp/mod{i}.py",
                "benchmark_function_name": f"bench{i}",
                "function_properties": make_func_props(),
            })
        elif i % 4 == 1:
            # instance method
            functions_data.append({
                "module_name": f"mod{i}",
                "function_name": f"func{i}",
                "class_name": f"Cls{i}",
                "file_path": f"/tmp/mod{i}.py",
                "benchmark_function_name": f"bench{i}",
                "function_properties": make_func_props(),
            })
        elif i % 4 == 2:
            # class method
            functions_data.append({
                "module_name": f"mod{i}",
                "function_name": f"func{i}",
                "class_name": f"Cls{i}",
                "file_path": f"/tmp/mod{i}.py",
                "benchmark_function_name": f"bench{i}",
                "function_properties": make_func_props(is_classmethod=True),
            })
        else:
            # static method
            functions_data.append({
                "module_name": f"mod{i}",
                "function_name": f"func{i}",
                "class_name": f"Cls{i}",
                "file_path": f"/tmp/mod{i}.py",
                "benchmark_function_name": f"bench{i}",
                "function_properties": make_func_props(is_staticmethod=True),
            })
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 730μs -> 632μs (15.4% faster)

def test_large_scale_long_names_and_special_chars():
    """Test with long and special-character-laden names in large batch."""
    n = 50
    functions_data = [
        {
            "module_name": f"mod{str(i)}.submod{str(i)}",
            "function_name": f"f@nc_{i}$",
            "class_name": f"Cl@s_{i}!" if i % 2 == 0 else None,
            "file_path": f"/tmp/mod{i}.py",
            "benchmark_function_name": f"bench#{i}!",
            "function_properties": make_func_props(is_classmethod=(i % 3 == 0), is_staticmethod=(i % 5 == 0)),
        }
        for i in range(n)
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 429μs -> 346μs (24.0% faster)
    # Check that special chars are handled in test function names
    for i in range(0, n, 10):
        mod = f"mod{str(i)}.submod{str(i)}"
        func = f"f@nc_{i}$"
        bench = f"bench#{i}!"
        class_name = f"Cl@s_{i}!" if i % 2 == 0 else None
        # Import present
        if class_name:
            pass
        else:
            pass
        # Test function name should have special chars replaced
        clean_bench = re.sub(r"[^a-zA-Z0-9_]+", "_", bench).strip("_")
        if class_name:
            test_func = f"def test_{mod.replace('.', '_')}_{class_name}_{func}_{clean_bench}():"
        else:
            test_func = f"def test_{mod.replace('.', '_')}_{func}_{clean_bench}():"
# 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-pr586-2025-07-28T06.14.39 and push.