Optimize extract_dependent_function

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.

Author: codeflash-ai[bot]

codeflash/code_utils/coverage_utils.py

@@ -13,16 +13,29 @@
 def extract_dependent_function(main_function: str, code_context: CodeOptimizationContext) -> str | Literal[False]:
     """Extract the single dependent function from the code context excluding the main function."""
     dependent_functions = set()
-    for code_string in code_context.testgen_context.code_strings:
-        ast_tree = ast.parse(code_string.code)
-        dependent_functions.update(
-            {node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))}
-        )
 
     # Compare using bare name since AST extracts bare function names
     bare_main = main_function.rsplit(".", 1)[-1] if "." in main_function else main_function
-    if bare_main in dependent_functions:
-        dependent_functions.discard(bare_main)
+
+    for code_string in code_context.testgen_context.code_strings:
+        # Quick heuristic: skip parsing entirely if there is no 'def' token,
+        # since no function definitions can be present without it.
+        if "def" not in code_string.code:
+            continue
+
+        ast_tree = ast.parse(code_string.code)
+        # Add function names directly, skipping the bare main name.
+        for node in ast_tree.body:
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                name = node.name
+                if name == bare_main:
+                    continue
+                dependent_functions.add(name)
+                # If more than one dependent function (other than the main) is found,
+                # we can return False early since the final result cannot be a single name.
+                if len(dependent_functions) > 1:
+                    return False
+
 
     if not dependent_functions:
         return False