⚡️ Speed up function detect_unused_helper_functions by 15% in PR #1166 (skyvern-grace)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1166

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

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

---

📄 15% (0.15x) speedup for detect_unused_helper_functions in codeflash/context/unused_definition_remover.py

⏱️ Runtime : 4.80 milliseconds → 4.19 milliseconds (best of 5 runs)

📝 Explanation and details

This optimization achieves a 14% runtime improvement (4.80ms → 4.19ms) through several targeted micro-optimizations that reduce overhead in hot code paths:

Key Performance Improvements

1. Eliminated Redundant Dictionary Lookups via Caching

In CodeStringsMarkdown properties (flat, file_to_path), the original code called self._cache.get("key") twice per invocation. The optimized version caches the result in a local variable:

# Before: two lookups
if self._cache.get("flat") is not None:
    return self._cache["flat"]

# After: one lookup
cached = self._cache.get("flat")
if cached is not None:
    return cached

This eliminates redundant hash table lookups in frequently accessed properties.

2. Replaced dict.setdefault() for Atomic List Operations

In _analyze_imports_in_optimized_code, the original code used an if-check followed by assignment for the helpers dictionary:

# Before: check + assign (two operations)
if func_name in file_entry:
    file_entry[func_name].append(helper)
else:
    file_entry[func_name] = [helper]

# After: single atomic operation
helpers_by_file_and_func[module_name].setdefault(func_name, []).append(helper)

The setdefault() approach reduces the operation to a single dictionary call, eliminating the membership test.

3. Hoisted as_posix() Calls Outside String Formatting

In the markdown property, path conversion was moved outside the f-string:

# Before: as_posix() called inside f-string
f"```python:{code_string.file_path.as_posix()}\n..."

# After: precomputed in conditional branch
if code_string.file_path:
    file_path_str = code_string.file_path.as_posix()
    result.append(f"```python:{file_path_str}\n...")

This avoids repeated method calls during string formatting.

4. Optimized Set Membership Tests with Early Exit

The most impactful change replaced set.intersection() with short-circuit boolean checks:

# Before: creates intermediate set via intersection
is_called = bool(possible_call_names.intersection(called_function_names))

# After: early-exit on first match
if (helper_qualified_name in called_function_names or 
    helper_simple_name in called_function_names or 
    helper_fully_qualified_name in called_function_names):
    is_called = True

With ~200 helpers in large-scale tests, this avoids creating temporary sets for every comparison, showing 50% speedup in the large helper test (1.31ms → 868μs).

5. Minimized Repeated Attribute Access

Variables like entrypoint_file_path, attr_name, and value_id are now cached before use, reducing attribute lookups in the AST traversal loop.

Impact Based on Test Results

• Small workloads (10-50 helpers): 10-16% speedup from reduced dict lookups
• Large workloads (200 helpers): 50% speedup due to eliminated set operations in the helper-checking loop
• Edge cases (syntax errors, missing functions): Minimal overhead, consistent 2-3% improvement

This optimization is particularly valuable when detect_unused_helper_functions is called repeatedly during code analysis pipelines, as the cumulative effect of these micro-optimizations scales with the number of helper functions and code blocks analyzed.

✅ Correctness verification report:

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

⚙️ Click to see Existing Unit Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_unused_helper_revert.py::test_async_class_methods	227μs	221μs	2.61%✅
test_unused_helper_revert.py::test_async_entrypoint_with_async_helpers	170μs	160μs	6.57%✅
test_unused_helper_revert.py::test_async_generators_and_coroutines	311μs	292μs	6.50%✅
test_unused_helper_revert.py::test_class_method_calls_external_helper_functions	175μs	165μs	6.10%✅
test_unused_helper_revert.py::test_class_method_entrypoint_with_helper_methods	200μs	199μs	0.418%✅
test_unused_helper_revert.py::test_detect_unused_helper_functions	172μs	167μs	2.96%✅
test_unused_helper_revert.py::test_detect_unused_in_multi_file_project	157μs	147μs	6.90%✅
test_unused_helper_revert.py::test_mixed_sync_and_async_helpers	273μs	262μs	4.11%✅
test_unused_helper_revert.py::test_module_dot_function_import_style	152μs	149μs	1.48%✅
test_unused_helper_revert.py::test_multi_file_import_styles	221μs	206μs	7.22%✅
test_unused_helper_revert.py::test_nested_class_method_optimization	191μs	186μs	2.65%✅
test_unused_helper_revert.py::test_no_unused_helpers_no_revert	176μs	175μs	0.611%✅
test_unused_helper_revert.py::test_recursive_helper_function_not_detected_as_unused	137μs	130μs	5.00%✅
test_unused_helper_revert.py::test_static_method_and_class_method	229μs	216μs	6.12%✅
test_unused_helper_revert.py::test_sync_entrypoint_with_async_helpers	186μs	181μs	2.67%✅

🌀 Click to see Generated Regression Tests

# imports
from dataclasses import dataclass
from pathlib import Path

from codeflash.context.unused_definition_remover import detect_unused_helper_functions
from codeflash.discovery.functions_to_optimize import FunctionToOptimize  # noqa: redefinition
from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, FunctionSource  # noqa: redefinition


# A minimal "Jedi Name"-like object used by FunctionSource
@dataclass(frozen=True)
class JediName:
    type: str  # e.g., "function" or "class"


# Real-like FunctionSource dataclass matching the original shape used by the function.
@dataclass(frozen=True)
class FunctionSource:
    file_path: Path
    qualified_name: str
    fully_qualified_name: str
    only_function_name: str
    source_code: str
    jedi_definition: JediName

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, FunctionSource):
            return False
        return (
            self.file_path == other.file_path
            and self.qualified_name == other.qualified_name
            and self.fully_qualified_name == other.fully_qualified_name
            and self.only_function_name == other.only_function_name
            and self.source_code == other.source_code
        )

    def __hash__(self) -> int:
        return hash(
            (self.file_path, self.qualified_name, self.fully_qualified_name, self.only_function_name, self.source_code)
        )


# Minimal CodeString container used by CodeStringsMarkdown
@dataclass
class CodeString:
    code: str
    file_path: Path | None = None


# Minimal CodeStringsMarkdown with the minimal features the function relies on:
# - it's a distinct type for isinstance check
# - has attribute `code_strings` (list) where each element has `.code`
@dataclass
class CodeStringsMarkdown:
    code_strings: list[CodeString]


# Minimal CodeOptimizationContext matching the fields the function uses.
@dataclass
class CodeOptimizationContext:
    testgen_context: CodeStringsMarkdown | None
    read_writable_code: CodeStringsMarkdown | None
    read_only_context_code: str
    hashing_code_context: str
    hashing_code_context_hash: str
    helper_functions: list[FunctionSource]
    preexisting_objects: set[tuple[str, tuple]]  # not used in our tests; kept for compatibility


# Provide a minimal FunctionToOptimize used by the function.
@dataclass
class FunctionParent:
    name: str


@dataclass
class FunctionToOptimize:
    function_name: str
    file_path: Path
    parents: list[FunctionParent]


# ---------------------------------------------------------------------
# Unit tests for detect_unused_helper_functions
# These tests cover Basic, Edge, and Large Scale scenarios.
# ---------------------------------------------------------------------


# Helper factory to create FunctionSource objects quickly.
def make_helper(name: str, file_stem: str = "helpers", qual_prefix: str | None = None, jedi_type: str = "function"):
    """Create a FunctionSource with predictable naming:
    - only_function_name = name
    - file_path = Path(f"{file_stem}.py")
    - qualified_name = f"{qual_prefix or file_stem}.{name}"
    - fully_qualified_name = f"{file_stem}.{name}" (mirrors qualified)
    """
    file_path = Path(f"{file_stem}.py")
    qualified = f"{qual_prefix or file_stem}.{name}"
    fully_qualified = f"{file_stem}.{name}"
    src = f"def {name}(): pass"
    jedi = JediName(type=jedi_type)
    return FunctionSource(
        file_path=file_path,
        qualified_name=qualified,
        fully_qualified_name=fully_qualified,
        only_function_name=name,
        source_code=src,
        jedi_definition=jedi,
    )


def make_context(helpers: list[FunctionSource]):
    """Create a CodeOptimizationContext with minimal required fields."""
    # We won't use most fields; set placeholders.
    return CodeOptimizationContext(
        testgen_context=None,
        read_writable_code=None,
        read_only_context_code="",
        hashing_code_context="",
        hashing_code_context_hash="",
        helper_functions=helpers,
        preexisting_objects=set(),
    )


def test_detects_no_unused_when_helper_called_directly():
    # Basic scenario: helper is in same file and called by name in the target function.
    target_path = Path("module.py")
    # Create a helper in the same file
    helper = make_helper("helper", file_stem="module", qual_prefix="module")
    ctx = make_context([helper])

    # optimized code contains the target function and a direct call to helper()
    optimized_code = """
def target():
    helper()
"""
    # define the FunctionToOptimize that points to top-level 'target'
    fto = FunctionToOptimize(function_name="target", file_path=target_path, parents=[])

    codeflash_output = detect_unused_helper_functions(fto, ctx, optimized_code)
    unused = codeflash_output  # 79.1μs -> 70.9μs (11.6% faster)


def test_detects_no_unused_when_imported_as_alias_and_called():
    # Helper is defined in a different file 'utils.py' and imported with alias
    helper = make_helper("util_fn", file_stem="utils", qual_prefix="utils")
    ctx = make_context([helper])

    optimized_code = """
from utils import util_fn as ufn
def target():
    ufn()
"""
    fto = FunctionToOptimize(function_name="target", file_path=Path("main.py"), parents=[])

    codeflash_output = detect_unused_helper_functions(fto, ctx, optimized_code)
    unused = codeflash_output  # 92.3μs -> 79.5μs (16.0% faster)


def test_returns_empty_when_entrypoint_not_found():
    # If the target function is missing from the optimized code, return []
    helper = make_helper("something", file_stem="mod")
    ctx = make_context([helper])

    optimized_code = """
# No function named 'missing_target' here
def other():
    pass
"""
    fto = FunctionToOptimize(function_name="missing_target", file_path=Path("mod.py"), parents=[])
    codeflash_output = detect_unused_helper_functions(fto, ctx, optimized_code)
    unused = codeflash_output  # 24.8μs -> 24.2μs (2.40% faster)


def test_handles_codestrings_markdown_with_multiple_blocks():
    # CodeStringsMarkdown: returns combined results of analyzing each block separately.
    helper1 = make_helper("a", file_stem="utils")
    helper2 = make_helper("b", file_stem="utils")
    ctx = make_context([helper1, helper2])

    # First block calls `a`, second block does not call `b`
    block1 = CodeString(code="from utils import a\n\ndef target():\n    a()")
    block2 = CodeString(code="def target():\n    pass")  # 'b' not called here
    csm = CodeStringsMarkdown(code_strings=[block1, block2])

    fto = FunctionToOptimize(function_name="target", file_path=Path("main.py"), parents=[])
    codeflash_output = detect_unused_helper_functions(fto, ctx, csm)
    unused = codeflash_output  # 11.1μs -> 10.9μs (2.21% faster)


def test_detects_method_call_via_self_and_class_qualified_name():
    # If entrypoint is a method and calls self.helper(), the detection should consider ClassName.helper as called.
    class_name = "MyClass"
    # Helper defined as part of the same class (qualified name)
    helper = FunctionSource(
        file_path=Path("module.py"),
        qualified_name=f"{class_name}.helper",
        fully_qualified_name=f"{class_name}.helper",
        only_function_name="helper",
        source_code="def helper(self): pass",
        jedi_definition=JediName(type="function"),
    )
    ctx = make_context([helper])

    # Define a class with target method calling self.helper()
    optimized_code = f"""
class {class_name}:
    def target(self):
        self.helper()
"""
    # FunctionToOptimize parents indicates it's inside MyClass
    fto = FunctionToOptimize(
        function_name="target", file_path=Path("module.py"), parents=[FunctionParent(name=class_name)]
    )
    codeflash_output = detect_unused_helper_functions(fto, ctx, optimized_code)
    unused = codeflash_output  # 91.7μs -> 83.3μs (10.0% faster)


def test_detects_import_module_and_usage_via_module_function():
    # Import module as alias and call module.helper()
    helper = make_helper("work", file_stem="utils", qual_prefix="utils")
    ctx = make_context([helper])

    optimized_code = """
import utils as u
def target():
    u.work()
"""
    fto = FunctionToOptimize(function_name="target", file_path=Path("main.py"), parents=[])
    codeflash_output = detect_unused_helper_functions(fto, ctx, optimized_code)
    unused = codeflash_output  # 91.0μs -> 79.9μs (14.0% faster)


def test_considers_module_dot_function_calls_without_imports():
    # Even if there's no 'import' statement, a call like utils.fn() adds 'utils.fn' to called names.
    helper = make_helper("external", file_stem="utils", qual_prefix="utils")
    ctx = make_context([helper])

    optimized_code = """
def target():
    utils.external()
"""
    fto = FunctionToOptimize(function_name="target", file_path=Path("main.py"), parents=[])
    codeflash_output = detect_unused_helper_functions(fto, ctx, optimized_code)
    unused = codeflash_output  # 77.7μs -> 66.9μs (16.1% faster)


def test_returns_empty_on_syntax_error_in_optimized_code():
    # If ast.parse raises (SyntaxError), the function catches and returns []
    helper = make_helper("x", file_stem="m")
    ctx = make_context([helper])

    optimized_code = "def target(:\n    pass"  # invalid syntax
    fto = FunctionToOptimize(function_name="target", file_path=Path("m.py"), parents=[])
    codeflash_output = detect_unused_helper_functions(fto, ctx, optimized_code)
    unused = codeflash_output  # 39.1μs -> 38.2μs (2.38% faster)


def test_large_scale_many_helpers_with_some_called():
    # Create a larger set of helpers (200) and call a subset (10).
    total_helpers = 200
    called_count = 10
    helpers = []
    # Helpers live in the same file "bigmod.py" except one that is cross-file to test module naming.
    for i in range(total_helpers):
        name = f"helper_{i}"
        # make half in same file, half in another file to exercise module-name logic
        file_stem = "bigmod" if (i % 2 == 0) else "othermod"
        # prefix qualified name to be unique
        qual_prefix = f"{file_stem}"
        helpers.append(make_helper(name, file_stem=file_stem, qual_prefix=qual_prefix))

    ctx = make_context(helpers)

    # Build optimized code that calls only helper_0 ... helper_9
    calls = "\n    ".join(f"{'othermod.' if i % 2 else ''}helper_{i}()" for i in range(called_count))
    optimized_code = f"""
def target():
    {calls}
"""
    fto = FunctionToOptimize(function_name="target", file_path=Path("bigmod.py"), parents=[])

    codeflash_output = detect_unused_helper_functions(fto, ctx, optimized_code)
    unused = codeflash_output  # 1.31ms -> 868μs (50.2% faster)
    # Ensure none of the first called_count helpers are in the unused list
    called_names = {f"helper_{i}" for i in range(called_count)}
    # Map the unused elements to their simple names for quick check
    unused_simple_names = {h.only_function_name for h in unused}


# 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-pr1166-2026-01-24T15.53.33 and push.