⚡️ Speed up function find_target_node by 18% in PR #763 (fix/correctly-find-funtion-node-when-reverting-helpers)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #763

If you approve this dependent PR, these changes will be merged into the original PR branch fix/correctly-find-funtion-node-when-reverting-helpers.

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

---

📄 18% (0.18x) speedup for find_target_node in codeflash/context/unused_definition_remover.py

⏱️ Runtime : 892 microseconds → 753 microseconds (best of 101 runs)

📝 Explanation and details

The optimized version eliminates recursive function calls by replacing the recursive _find helper with an iterative approach. This provides significant performance benefits:

Key Optimizations:

1. Removed Recursion Overhead: The original code used a recursive helper function _find that created new stack frames for each parent traversal. The optimized version uses a simple iterative loop that traverses parents sequentially without function call overhead.

2. Eliminated Function Creation: The original code defined the _find function on every call to find_target_node. The optimized version removes this repeated function definition entirely.

3. Early Exit with for-else: The optimized code uses Python's for-else construct to immediately return None when a parent class isn't found, avoiding unnecessary continued searching.

4. Reduced Attribute Access: By caching function_to_optimize.function_name in a local variable target_name and reusing body variables, the code reduces repeated attribute lookups.

Performance Impact by Test Case:

• Simple cases (top-level functions, basic class methods): 23-62% faster due to eliminated recursion overhead
• Nested class scenarios: 45-84% faster, with deeper nesting showing greater improvements as recursion elimination has more impact
• Large-scale tests: 12-22% faster, showing consistent benefits even with many nodes to traverse
• Edge cases (empty modules, non-existent classes): 52-76% faster due to more efficient early termination

The optimization is particularly effective for deeply nested class hierarchies where the original recursive approach created multiple stack frames, while the iterative version maintains constant memory usage regardless of nesting depth.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import ast
from typing import Optional

# imports
import pytest  # used for our unit tests
from codeflash.context.unused_definition_remover import find_target_node


# Minimal stubs for FunctionToOptimize and FunctionParent for testing
class FunctionParent:
    def __init__(self, name: str):
        self.name = name

class FunctionToOptimize:
    def __init__(self, function_name: str, parents: list[FunctionParent]):
        self.function_name = function_name
        self.parents = parents
from codeflash.context.unused_definition_remover import find_target_node

# unit tests

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

def test_top_level_function_found():
    """Test finding a top-level function."""
    code = "def foo(): pass"
    tree = ast.parse(code)
    fto = FunctionToOptimize("foo", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.87μs -> 1.34μs (39.5% faster)

def test_top_level_function_not_found():
    """Test searching for a non-existent top-level function."""
    code = "def foo(): pass"
    tree = ast.parse(code)
    fto = FunctionToOptimize("bar", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.76μs -> 1.29μs (36.5% faster)

def test_class_method_found():
    """Test finding a method inside a class."""
    code = "class A:\n    def foo(self): pass"
    tree = ast.parse(code)
    fto = FunctionToOptimize("foo", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.65μs -> 1.87μs (41.1% faster)

def test_class_method_not_found():
    """Test searching for a non-existent method in a class."""
    code = "class A:\n    def foo(self): pass"
    tree = ast.parse(code)
    fto = FunctionToOptimize("bar", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.50μs -> 1.65μs (51.5% faster)

def test_async_function_found():
    """Test finding an async function."""
    code = "async def foo(): pass"
    tree = ast.parse(code)
    fto = FunctionToOptimize("foo", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.76μs -> 1.43μs (23.1% faster)

def test_method_in_nested_class_found():
    """Test finding a method inside a nested class."""
    code = "class A:\n    class B:\n        def foo(self): pass"
    tree = ast.parse(code)
    fto = FunctionToOptimize("foo", [FunctionParent("A"), FunctionParent("B")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 3.08μs -> 1.96μs (56.7% faster)

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

def test_empty_module():
    """Test with an empty module."""
    code = ""
    tree = ast.parse(code)
    fto = FunctionToOptimize("foo", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.16μs -> 721ns (61.3% faster)

def test_function_with_same_name_in_different_classes():
    """Test finding the correct method when multiple classes have a method with the same name."""
    code = """
class A:
    def foo(self): pass
class B:
    def foo(self): pass
"""
    tree = ast.parse(code)
    fto_a = FunctionToOptimize("foo", [FunctionParent("A")])
    fto_b = FunctionToOptimize("foo", [FunctionParent("B")])
    codeflash_output = find_target_node(tree, fto_a); node_a = codeflash_output # 2.61μs -> 1.70μs (53.5% faster)
    codeflash_output = find_target_node(tree, fto_b); node_b = codeflash_output # 1.63μs -> 1.06μs (53.7% faster)

def test_function_with_same_name_top_level_and_class():
    """Test that the correct node is found when function exists both at top-level and inside a class."""
    code = """
def foo(): pass
class A:
    def foo(self): pass
"""
    tree = ast.parse(code)
    fto_top = FunctionToOptimize("foo", [])
    fto_class = FunctionToOptimize("foo", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto_top); node_top = codeflash_output # 1.67μs -> 1.29μs (29.5% faster)
    codeflash_output = find_target_node(tree, fto_class); node_class = codeflash_output # 2.00μs -> 1.26μs (58.7% faster)

def test_class_not_found():
    """Test searching for a method in a non-existent class."""
    code = "def foo(): pass"
    tree = ast.parse(code)
    fto = FunctionToOptimize("foo", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.46μs -> 1.01μs (44.4% faster)

def test_deeply_nested_class_not_found():
    """Test searching for a method in a deeply nested, non-existent class."""
    code = "class A:\n    def bar(self): pass"
    tree = ast.parse(code)
    fto = FunctionToOptimize("foo", [FunctionParent("A"), FunctionParent("B")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.29μs -> 1.51μs (51.0% faster)

def test_function_in_class_with_other_top_level_items():
    """Test finding a method in a class when there are other top-level items."""
    code = """
def top(): pass
class A:
    def foo(self): pass
def another(): pass
"""
    tree = ast.parse(code)
    fto = FunctionToOptimize("foo", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.75μs -> 1.88μs (45.7% faster)

def test_function_in_class_with_inner_class_and_method():
    """Test finding a method in an inner class when outer class also has a method of same name."""
    code = """
class A:
    def foo(self): pass
    class B:
        def foo(self): pass
"""
    tree = ast.parse(code)
    fto_outer = FunctionToOptimize("foo", [FunctionParent("A")])
    fto_inner = FunctionToOptimize("foo", [FunctionParent("A"), FunctionParent("B")])
    codeflash_output = find_target_node(tree, fto_outer); node_outer = codeflash_output # 2.48μs -> 1.69μs (46.2% faster)
    codeflash_output = find_target_node(tree, fto_inner); node_inner = codeflash_output # 1.95μs -> 1.34μs (45.6% faster)

def test_non_function_nodes_ignored():
    """Test that non-function nodes with matching names are ignored."""
    code = """
foo = 42
class foo: pass
"""
    tree = ast.parse(code)
    fto = FunctionToOptimize("foo", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.03μs -> 1.68μs (20.9% faster)

def test_function_in_class_with_decorator():
    """Test finding a decorated method in a class."""
    code = """
class A:
    @staticmethod
    def foo(): pass
"""
    tree = ast.parse(code)
    fto = FunctionToOptimize("foo", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.54μs -> 1.63μs (55.8% faster)

def test_function_with_unicode_name():
    """Test finding a function with a unicode name."""
    code = "def f\u00f6\u00f6(): pass"
    tree = ast.parse(code)
    fto = FunctionToOptimize("föö", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.73μs -> 1.29μs (34.1% faster)

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

def test_many_top_level_functions():
    """Test with many top-level functions, ensuring correct one is found."""
    code = "\n".join(f"def func{i}(): pass" for i in range(100))
    tree = ast.parse(code)
    fto = FunctionToOptimize("func42", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 8.31μs -> 6.59μs (26.1% faster)

def test_many_classes_with_methods():
    """Test with many classes, each with a method, ensuring correct one is found."""
    code = "\n".join(
        f"class C{i}:\n    def foo(self): pass"
        for i in range(100)
    )
    tree = ast.parse(code)
    fto = FunctionToOptimize("foo", [FunctionParent("C77")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 12.3μs -> 9.99μs (22.9% faster)

def test_deeply_nested_classes_large():
    """Test with deeply nested classes (up to 10 levels), ensuring correct method is found."""
    code = ""
    for i in range(10):
        code += "    " * i + f"class C{i}:\n"
    code += "    " * 10 + "def foo(self): pass\n"
    tree = ast.parse(code)
    parents = [FunctionParent(f"C{i}") for i in range(10)]
    fto = FunctionToOptimize("foo", parents)
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 6.42μs -> 3.51μs (83.1% faster)

def test_large_number_of_methods_in_class():
    """Test with a class containing many methods, ensuring correct one is found."""
    code = "class A:\n" + "\n".join(
        f"    def func{i}(self): pass" for i in range(100)
    )
    tree = ast.parse(code)
    fto = FunctionToOptimize("func55", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 10.8μs -> 8.73μs (23.3% faster)

def test_large_number_of_nested_classes_and_methods():
    """Test with many nested classes, each with a method, ensuring correct method in deepest class is found."""
    code = ""
    for i in range(5):
        code += "    " * i + f"class C{i}:\n"
    code += "    " * 5 + "def foo(self): pass\n"
    tree = ast.parse(code)
    parents = [FunctionParent(f"C{i}") for i in range(5)]
    fto = FunctionToOptimize("foo", parents)
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 4.24μs -> 2.52μs (67.8% faster)

def test_performance_with_max_elements():
    """Test performance with maximum allowed elements (999 top-level functions)."""
    code = "\n".join(f"def func{i}(): pass" for i in range(999))
    tree = ast.parse(code)
    fto = FunctionToOptimize("func998", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 122μs -> 108μs (12.7% 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 typing import Optional

# imports
import pytest  # used for our unit tests
from codeflash.context.unused_definition_remover import find_target_node


# Stubs for required classes
class FunctionParent:
    def __init__(self, name: str):
        self.name = name

class FunctionToOptimize:
    def __init__(self, function_name: str, parents: list[FunctionParent]):
        self.function_name = function_name
        self.parents = parents
from codeflash.context.unused_definition_remover import find_target_node

# unit tests

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

def test_basic_top_level_function_found():
    # Test finding a top-level function
    src = """
def foo():
    pass

def bar():
    pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("foo", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.91μs -> 1.41μs (35.5% faster)

def test_basic_top_level_function_not_found():
    # Test when function does not exist
    src = """
def foo():
    pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("bar", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.72μs -> 1.25μs (37.6% faster)

def test_basic_class_method_found():
    # Test finding a method inside a class
    src = """
class A:
    def foo(self):
        pass

    def bar(self):
        pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("foo", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.75μs -> 1.91μs (43.9% faster)

def test_basic_class_method_not_found():
    # Test when method does not exist in class
    src = """
class A:
    def foo(self):
        pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("bar", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.51μs -> 1.55μs (61.9% faster)

def test_basic_async_function_found():
    # Test finding a top-level async function
    src = """
async def foo():
    pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("foo", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.74μs -> 1.32μs (31.8% faster)

def test_basic_nested_class_method_found():
    # Test finding a method inside nested classes
    src = """
class Outer:
    class Inner:
        def foo(self):
            pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("foo", [FunctionParent("Outer"), FunctionParent("Inner")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.90μs -> 1.87μs (54.5% faster)

def test_basic_nested_class_method_not_found():
    # Test when method does not exist in nested classes
    src = """
class Outer:
    class Inner:
        def bar(self):
            pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("foo", [FunctionParent("Outer"), FunctionParent("Inner")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.98μs -> 1.81μs (64.1% faster)

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

def test_edge_empty_module():
    # Test with empty module
    src = ""
    tree = ast.parse(src)
    fto = FunctionToOptimize("foo", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.16μs -> 661ns (75.9% faster)

def test_edge_class_with_no_body():
    # Test class with no body
    src = """
class A:
    pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("foo", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.47μs -> 1.61μs (53.5% faster)

def test_edge_function_with_same_name_in_multiple_classes():
    # Test function with same name in different classes
    src = """
class A:
    def foo(self):
        pass

class B:
    def foo(self):
        pass
"""
    tree = ast.parse(src)
    fto_a = FunctionToOptimize("foo", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto_a); node_a = codeflash_output # 2.41μs -> 1.61μs (49.7% faster)

    fto_b = FunctionToOptimize("foo", [FunctionParent("B")])
    codeflash_output = find_target_node(tree, fto_b); node_b = codeflash_output # 1.61μs -> 1.05μs (53.3% faster)

def test_edge_function_with_same_name_top_level_and_class():
    # Test function with same name at top level and inside class
    src = """
def foo():
    pass

class A:
    def foo(self):
        pass
"""
    tree = ast.parse(src)
    fto_top = FunctionToOptimize("foo", [])
    codeflash_output = find_target_node(tree, fto_top); node_top = codeflash_output # 1.65μs -> 1.17μs (41.0% faster)

    fto_class = FunctionToOptimize("foo", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto_class); node_class = codeflash_output # 2.15μs -> 1.53μs (40.5% faster)

def test_edge_class_with_same_name_nested():
    # Test nested classes with same name
    src = """
class A:
    class A:
        def foo(self):
            pass
"""
    tree = ast.parse(src)
    # Should find the inner A.foo
    fto = FunctionToOptimize("foo", [FunctionParent("A"), FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.88μs -> 1.86μs (54.4% faster)

def test_edge_function_with_decorator():
    # Test function with decorator
    src = """
def foo():
    pass

@staticmethod
def bar():
    pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("bar", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.83μs -> 1.40μs (30.7% faster)

def test_edge_class_method_with_decorator():
    # Test class method with decorator
    src = """
class A:
    @classmethod
    def foo(cls):
        pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("foo", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.40μs -> 1.66μs (44.6% faster)

def test_edge_nonexistent_parent_class():
    # Test with non-existent parent class
    src = """
def foo():
    pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("foo", [FunctionParent("A")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.48μs -> 972ns (52.6% faster)

def test_edge_nonexistent_nested_parent_class():
    # Test with non-existent nested parent class
    src = """
class A:
    def foo(self):
        pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("foo", [FunctionParent("A"), FunctionParent("B")])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 2.25μs -> 1.38μs (63.0% faster)

def test_edge_function_with_non_ascii_name():
    # Test function with non-ascii name
    src = """
def föö():
    pass
"""
    tree = ast.parse(src)
    fto = FunctionToOptimize("föö", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 1.78μs -> 1.25μs (42.4% faster)

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

def test_large_many_top_level_functions():
    # Test with many top-level functions
    src = "\n".join([f"def func_{i}():\n    pass" for i in range(1000)])
    tree = ast.parse(src)
    # Pick a few to test
    for i in [0, 499, 999]:
        fto = FunctionToOptimize(f"func_{i}", [])
        codeflash_output = find_target_node(tree, fto); node = codeflash_output # 180μs -> 154μs (17.1% faster)

def test_large_many_classes_and_methods():
    # Test with many classes, each with a method
    src = "\n".join([f"class C{i}:\n    def foo(self):\n        pass" for i in range(1000)])
    tree = ast.parse(src)
    # Pick a few to test
    for i in [0, 499, 999]:
        fto = FunctionToOptimize("foo", [FunctionParent(f"C{i}")])
        codeflash_output = find_target_node(tree, fto); node = codeflash_output # 146μs -> 123μs (18.7% faster)

def test_large_deeply_nested_classes():
    # Test with deeply nested classes
    depth = 10
    src = ""
    indent = ""
    for i in range(depth):
        src += f"{indent}class C{i}:\n"
        indent += "    "
    src += f"{indent}def foo(self):\n{indent}    pass\n"
    tree = ast.parse(src)
    parents = [FunctionParent(f"C{i}") for i in range(depth)]
    fto = FunctionToOptimize("foo", parents)
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 6.34μs -> 3.49μs (81.9% faster)

def test_large_many_methods_in_one_class():
    # Test with one class and many methods
    src = "class A:\n" + "\n".join([f"    def func_{i}(self):\n        pass" for i in range(1000)])
    tree = ast.parse(src)
    # Pick a few to test
    for i in [0, 499, 999]:
        fto = FunctionToOptimize(f"func_{i}", [FunctionParent("A")])
        codeflash_output = find_target_node(tree, fto); node = codeflash_output # 189μs -> 167μs (13.0% faster)

def test_large_no_matching_function():
    # Test with large module, no matching function
    src = "\n".join([f"def func_{i}():\n    pass" for i in range(1000)])
    tree = ast.parse(src)
    fto = FunctionToOptimize("not_found", [])
    codeflash_output = find_target_node(tree, fto); node = codeflash_output # 123μs -> 109μs (13.1% faster)
# 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-pr763-2025-09-25T14.28.58 and push.