⚡️ Speed up function extract_dependent_function by 197% in PR #1457 (fix-coverage-qualified-name)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1457

If you approve this dependent PR, these changes will be merged into the original PR branch fix-coverage-qualified-name.

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

---

📄 197% (1.97x) speedup for extract_dependent_function in codeflash/code_utils/coverage_utils.py

⏱️ Runtime : 28.5 milliseconds → 9.57 milliseconds (best of 13 runs)

📝 Explanation and details

The optimized code achieves a 197% speedup (28.5ms → 9.57ms) through three strategic optimizations that dramatically reduce expensive AST parsing operations:

Key Optimizations

1. Early String Filtering (74% time reduction in parsing)
The optimization adds a lightweight heuristic check if "def" not in code_string.code before calling ast.parse(). Since function definitions require the def keyword, strings without it can be skipped entirely. In the profiler results, this reduced AST parsing from 32.5ms (80.5% of original runtime) to 9.9ms (74.2% of optimized runtime). The test results show dramatic improvements for large-scale scenarios:

• test_large_scale_many_code_strings_single_dependent_function: 6839% faster (4.45ms → 64.1μs)
• test_large_scale_with_preexisting_objects_and_many_irrelevant_entries: 4193% faster (2.26ms → 52.7μs)

2. Hoisted Main Function Name Computation
Moving bare_main calculation outside the loop (from line 13 to line 10) eliminates redundant string operations that were executed once per code string. This simple reordering saves repeated rsplit() calls.

3. Early Exit on Multiple Dependencies
The optimization checks if len(dependent_functions) > 1: return False immediately after adding each function name, rather than waiting until all code strings are processed. This allows the function to short-circuit as soon as it detects the failure condition, avoiding unnecessary AST parsing of remaining code strings.

Why This Matters

Based on the function references, extract_dependent_function is called during test generation workflows where it processes potentially hundreds or thousands of code strings. The optimization is particularly effective when:

• Most code strings don't contain function definitions (common in test contexts with imports, variables, etc.)
• Multiple dependent functions exist (early exit prevents wasted parsing)
• Code bases have many test-related code strings that aren't function definitions

The optimizations preserve exact behavior while intelligently avoiding expensive operations, making the code significantly more efficient in real-world usage patterns where the function processes large volumes of code strings.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from types import \
    SimpleNamespace  # lightweight object to hold attributes for our fake context

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.coverage_utils import (build_fully_qualified_name,
                                                 extract_dependent_function)

# helper to construct a minimal code context expected by extract_dependent_function
def make_context(code_strings, preexisting_objects):
    """
    Build a minimal stand-in for CodeOptimizationContext that has the attributes
    used by extract_dependent_function:
      - testgen_context.code_strings: iterable of objects with a .code string attribute
      - preexisting_objects: iterable of tuples (obj_name, parents) where parents is iterable
        of objects with .type and .name attributes
    We use types.SimpleNamespace to avoid creating any custom classes.
    """
    # Wrap each code string in a SimpleNamespace with attribute `code`
    code_string_objs = [SimpleNamespace(code=s) for s in code_strings]
    testgen_context = SimpleNamespace(code_strings=code_string_objs)
    # preexisting_objects is passed through as given; typical shape: [(name, (parent1, parent2, ...)), ...]
    return SimpleNamespace(testgen_context=testgen_context, preexisting_objects=preexisting_objects)

def test_returns_false_when_no_dependent_functions_present():
    # Only the main function is present in code strings; expecting False as no other dependent function
    code = "def main():\n    return 1\n"
    ctx = make_context([code], preexisting_objects=[])
    # Passing a qualified main name should still match and be discarded
    codeflash_output = extract_dependent_function("module.main", ctx); result = codeflash_output # 26.4μs -> 26.6μs (0.831% slower)

def test_extracts_single_dependent_function_and_qualifies_with_class_parent():
    # Code has one helper function besides main; preexisting_objects provides a ClassDef parent
    main_code = "def main():\n    return helper()\n"
    helper_code = "def helper():\n    return 42\n"
    parents = (SimpleNamespace(type="ClassDef", name="MyClass"),)
    ctx = make_context([main_code, helper_code], preexisting_objects=[("helper", parents)])
    # main passed as bare name
    codeflash_output = extract_dependent_function("main", ctx); result = codeflash_output # 39.5μs -> 39.9μs (0.855% slower)

def test_extracts_single_dependent_function_returns_bare_name_when_no_parents():
    # If there are no preexisting_objects for the dependent function, return bare function name
    main_code = "def main():\n    return x\n"
    dep_code = "def lone_helper():\n    return 'ok'\n"
    ctx = make_context([main_code, dep_code], preexisting_objects=[])
    codeflash_output = extract_dependent_function("main", ctx); result = codeflash_output # 33.9μs -> 32.8μs (3.45% faster)

def test_returns_false_when_multiple_dependent_functions_present():
    # If more than one non-main top-level function exists, function should return False
    main_code = "def main():\n    pass\n"
    a = "def a():\n    pass\n"
    b = "def b():\n    pass\n"
    ctx = make_context([main_code, a, b], preexisting_objects=[])
    codeflash_output = extract_dependent_function("main", ctx); result = codeflash_output # 31.9μs -> 31.7μs (0.599% faster)

def test_main_name_can_be_given_as_qualified_or_bare_equally():
    # Ensure both 'mod.main' and 'main' discard the same bare name
    main_code = "def main():\n    pass\n"
    helper_code = "def helper():\n    return True\n"
    ctx = make_context([main_code, helper_code], preexisting_objects=[])
    # qualified main
    codeflash_output = extract_dependent_function("pkg.module.main", ctx); res1 = codeflash_output # 30.1μs -> 29.5μs (2.10% faster)
    # bare main
    codeflash_output = extract_dependent_function("main", ctx); res2 = codeflash_output # 16.6μs -> 16.5μs (0.851% faster)

def test_async_functions_are_detected_as_dependent():
    # AsyncFunctionDef nodes should also be gathered as dependent functions
    main_code = "def main():\n    pass\n"
    async_code = "async def async_helper():\n    return 1\n"
    ctx = make_context([main_code, async_code], preexisting_objects=[])
    codeflash_output = extract_dependent_function("main", ctx); result = codeflash_output # 31.0μs -> 30.3μs (2.11% faster)

def test_nested_function_definitions_are_ignored():
    # A nested function inside main should not be considered top-level and thus ignored
    # Only the outer main is top-level here
    code = (
        "def main():\n"
        "    def nested():\n"
        "        return 10\n"
        "    return nested()\n"
    )
    ctx = make_context([code], preexisting_objects=[])
    codeflash_output = extract_dependent_function("main", ctx); result = codeflash_output # 29.6μs -> 29.5μs (0.267% faster)

def test_duplicate_function_definitions_across_code_strings_count_once():
    # Duplicate top-level definitions of the same function name across multiple code strings
    # should be deduplicated by set semantics and still return the single dependent function
    main_code = "def main():\n    pass\n"
    helper_code1 = "def helper():\n    return 1\n"
    helper_code2 = "def helper():\n    return 2\n"
    ctx = make_context([main_code, helper_code1, helper_code2], preexisting_objects=[])
    codeflash_output = extract_dependent_function("main", ctx); result = codeflash_output # 37.3μs -> 36.9μs (1.20% faster)

def test_syntax_error_in_code_string_propagates():
    # If a code string contains invalid Python, ast.parse will raise SyntaxError and it should propagate
    bad_code = "def oops(:\n    pass\n"  # invalid syntax
    ctx = make_context([bad_code], preexisting_objects=[])
    with pytest.raises(SyntaxError):
        extract_dependent_function("main", ctx) # 37.0μs -> 35.4μs (4.64% faster)

def test_build_fully_qualified_name_ignores_non_class_parents_and_applies_multiple_class_parents():
    # If parents include non-ClassDef entries they should be ignored; ClassDef parents
    # present multiple times should be prefixed in order.
    # For illustration, parents are given in order Outer, Inner; resulting prefixing yields Inner.Outer.name
    parents = (
        SimpleNamespace(type="Module", name="mod"),  # ignored
        SimpleNamespace(type="ClassDef", name="Outer"),
        SimpleNamespace(type="ClassDef", name="Inner"),
    )
    ctx = make_context([], preexisting_objects=[("fn", parents)])
    # Call build_fully_qualified_name directly to assert prefixing order
    full = build_fully_qualified_name("fn", ctx)

def test_large_scale_many_code_strings_single_dependent_function():
    # Create many (1000) code strings that do not define functions and one that defines the dependent function.
    non_func_code = "x = 1\n"  # harmless non-function code
    many = [non_func_code] * 1000  # repeated non-function code
    dependent_code = "def lone_dep():\n    return 'ok'\n"
    # Put the dependent function at the end to ensure the loop parses many items before finding it
    code_list = many + [dependent_code]
    ctx = make_context(code_list, preexisting_objects=[])
    codeflash_output = extract_dependent_function("main", ctx); result = codeflash_output # 4.45ms -> 64.1μs (6839% faster)

def test_large_scale_with_preexisting_objects_and_many_irrelevant_entries():
    # Many preexisting_objects entries and many code strings but exactly one dependent function
    code_list = ["a = 1\n"] * 500  # irrelevant code
    code_list.append("def only_one():\n    return 5\n")
    # Create many preexisting object entries that do not match the function name
    fake_parents = [(f"name{i}", ()) for i in range(300)]
    # Add a matching preexisting_objects entry with nested ClassDef parents to test qualification logic
    matching_parents = (SimpleNamespace(type="ClassDef", name="Top"), SimpleNamespace(type="ClassDef", name="Inner"))
    preexisting_objects = fake_parents + [("only_one", matching_parents)]
    ctx = make_context(code_list, preexisting_objects=preexisting_objects)
    codeflash_output = extract_dependent_function("main", ctx); result = codeflash_output # 2.26ms -> 52.7μs (4193% faster)
# 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 unittest.mock import MagicMock, Mock

# imports
import pytest
from codeflash.code_utils.coverage_utils import (build_fully_qualified_name,
                                                 extract_dependent_function)
from codeflash.models.function_types import FunctionParent
from codeflash.models.models import CodeOptimizationContext

# Helper function to create a CodeOptimizationContext with test data
def create_context(code_strings_list, preexisting_objects=None):
    """Helper to create a CodeOptimizationContext for testing."""
    # Create mock CodeString objects
    mock_code_strings = []
    for code in code_strings_list:
        mock_code_string = Mock()
        mock_code_string.code = code
        mock_code_strings.append(mock_code_string)
    
    # Create mock testgen_context
    mock_testgen_context = Mock()
    mock_testgen_context.code_strings = mock_code_strings
    
    # Create context with mocked dependencies
    context = Mock(spec=CodeOptimizationContext)
    context.testgen_context = mock_testgen_context
    context.preexisting_objects = preexisting_objects or set()
    
    return context

def test_build_fully_qualified_name_unqualified():
    """Test build_fully_qualified_name with unqualified function name."""
    # Given: an unqualified function name and empty preexisting_objects
    function_name = "simple_func"
    context = Mock(spec=CodeOptimizationContext)
    context.preexisting_objects = set()
    
    # When: build_fully_qualified_name is called
    result = build_fully_qualified_name(function_name, context)

def test_build_fully_qualified_name_already_qualified():
    """Test build_fully_qualified_name with already qualified name."""
    # Given: a fully qualified function name
    function_name = "MyClass.my_method"
    context = Mock(spec=CodeOptimizationContext)
    context.preexisting_objects = set()
    
    # When: build_fully_qualified_name is called
    result = build_fully_qualified_name(function_name, context)

def test_build_fully_qualified_with_class_parent():
    """Test build_fully_qualified_name with class parent."""
    # Given: function name and preexisting class parent
    function_name = "my_method"
    parent = Mock(spec=FunctionParent)
    parent.type = "ClassDef"
    parent.name = "MyClass"
    context = Mock(spec=CodeOptimizationContext)
    context.preexisting_objects = {("my_method", (parent,))}
    
    # When: build_fully_qualified_name is called
    result = build_fully_qualified_name(function_name, context)

def test_build_fully_qualified_with_multiple_parents():
    """Test build_fully_qualified_name with multiple class parents."""
    # Given: function with multiple parent classes
    function_name = "my_method"
    parent1 = Mock(spec=FunctionParent)
    parent1.type = "ClassDef"
    parent1.name = "OuterClass"
    parent2 = Mock(spec=FunctionParent)
    parent2.type = "ClassDef"
    parent2.name = "InnerClass"
    context = Mock(spec=CodeOptimizationContext)
    context.preexisting_objects = {("my_method", (parent1, parent2))}
    
    # When: build_fully_qualified_name is called
    result = build_fully_qualified_name(function_name, context)

def test_build_fully_qualified_non_class_parent_ignored():
    """Test that non-class parents are ignored in qualification."""
    # Given: function with non-ClassDef parent
    function_name = "my_func"
    parent = Mock(spec=FunctionParent)
    parent.type = "ModuleDef"  # Not ClassDef
    parent.name = "SomeModule"
    context = Mock(spec=CodeOptimizationContext)
    context.preexisting_objects = {("my_func", (parent,))}
    
    # When: build_fully_qualified_name is called
    result = build_fully_qualified_name(function_name, context)

def test_function_name_not_in_preexisting_objects():
    """Test when function name is not in preexisting_objects."""
    # Given: function name that doesn't exist in preexisting_objects
    function_name = "unknown_func"
    context = Mock(spec=CodeOptimizationContext)
    context.preexisting_objects = {("other_func", ())}
    
    # When: build_fully_qualified_name is called
    result = build_fully_qualified_name(function_name, context)

def test_empty_preexisting_objects():
    """Test build_fully_qualified_name with empty preexisting_objects."""
    # Given: empty preexisting_objects set
    function_name = "func"
    context = Mock(spec=CodeOptimizationContext)
    context.preexisting_objects = set()
    
    # When: build_fully_qualified_name is called
    result = build_fully_qualified_name(function_name, context)

def test_performance_with_large_preexisting_objects():
    """Test build_fully_qualified_name with large preexisting_objects set."""
    # Given: preexisting_objects with 1000 entries
    function_name = "my_func"
    preexisting = set()
    for i in range(1000):
        parent = Mock(spec=FunctionParent)
        if i == 500:
            # Place our target function in the middle
            parent.type = "ClassDef"
            parent.name = "TargetClass"
            preexisting.add(("my_func", (parent,)))
        else:
            parent.type = "ClassDef"
            parent.name = f"Class_{i}"
            preexisting.add((f"func_{i}", (parent,)))
    
    context = Mock(spec=CodeOptimizationContext)
    context.preexisting_objects = preexisting
    
    # When: build_fully_qualified_name is called
    result = build_fully_qualified_name(function_name, context)

To edit these changes git checkout codeflash/optimize-pr1457-2026-02-12T04.58.15 and push.