⚡️ Speed up function get_first_top_level_function_or_method_ast by 17% in PR #678 (standalone-fto-async)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #678

If you approve this dependent PR, these changes will be merged into the original PR branch standalone-fto-async.

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

---

📄 17% (0.17x) speedup for get_first_top_level_function_or_method_ast in codeflash/code_utils/static_analysis.py

⏱️ Runtime : 390 microseconds → 334 microseconds (best of 59 runs)

📝 Explanation and details

The optimized code achieves a 16% speedup through several targeted micro-optimizations that reduce overhead in the tight loops that traverse AST nodes:

Key Optimizations:

1. Local Variable Bindings: Assigns ast.iter_child_nodes and the type tuple to local variables, eliminating repeated attribute lookups during iteration. The profiler shows this reduces the per-hit cost of the main loop from 1560.1ns to 1530.6ns.

2. Restructured Type Checking: Splits the combined isinstance(child, object_type) and child.name == object_name check into separate conditions. This allows early exit after the type check fails and uses getattr(child, "name", None) for safer attribute access, reducing the attribute lookup overhead shown in the profiler (from 403.8ns to 308ns per hit).

3. Optimized Control Flow: Changes the nested if-statements to elif structure, reducing redundant type checks. The isinstance(child, fn_type_tuple) check now only runs when needed, improving branch prediction.

4. Direct Parent Access: Caches parents[0] as parent0 to avoid repeated list indexing, though this has minimal impact on the overall performance.

Performance Impact by Test Type:

• Large-scale tests (500+ functions/classes): Benefit most from reduced per-node overhead in deep traversals
• Basic cases: See consistent but smaller improvements due to fewer nodes processed
• Edge cases: Minimal impact since they often involve early returns or empty searches

The optimizations are most effective for codebases with complex AST structures where the functions traverse many nodes, making the micro-optimizations compound significantly.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import ast
from typing import TypeVar

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.static_analysis import \
    get_first_top_level_function_or_method_ast
from codeflash.models.models import FunctionParent

# unit tests

# Helper to build FunctionParent objects
def make_class_parent(name: str):
    return FunctionParent(name=name, type="ClassDef")

def make_func_parent(name: str):
    return FunctionParent(name=name, type="FunctionDef")

# Basic Test Cases









def foo(): pass
def foo(): pass
"""
    tree = ast.parse(src)
    codeflash_output = get_first_top_level_function_or_method_ast("foo", [], tree); result = codeflash_output



def outer():
    def inner(): pass
"""
    tree = ast.parse(src)
    codeflash_output = get_first_top_level_function_or_method_ast("inner", [], tree); result = codeflash_output




def foo(): pass
"""
    tree = ast.parse(src)
    codeflash_output = get_first_top_level_function_or_method_ast("foo", [], tree); result = codeflash_output



def foo():
    """Docstring"""
    pass
'''
    tree = ast.parse(src)
    codeflash_output = get_first_top_level_function_or_method_ast("foo", [], tree); result = codeflash_output







def outer():
    def foo(): pass
"""
    tree = ast.parse(src)
    codeflash_output = get_first_top_level_function_or_method_ast("foo", [], tree); result = codeflash_output


def foo(): pass
class foo: pass
"""
    tree = ast.parse(src)
    codeflash_output = get_first_top_level_function_or_method_ast("foo", [], tree); result_func = codeflash_output
    codeflash_output = get_first_top_level_function_or_method_ast(
        "foo", [make_class_parent("foo")], tree
    ); result_method = codeflash_output


def foo(): pass
"""
    tree = ast.parse(src)
    codeflash_output = get_first_top_level_function_or_method_ast("foo", [], tree); result_func = codeflash_output
    codeflash_output = get_first_top_level_function_or_method_ast(
        "foo", [make_class_parent("foo")], tree
    ); result_method = codeflash_output


def foo(): pass
"""
    tree = ast.parse(src)
    codeflash_output = get_first_top_level_function_or_method_ast("foo", [], tree); result = codeflash_output




def foo(): pass
"""
    tree = ast.parse(src)
    # Simulate parent of type FunctionDef (not supported by implementation)
    parents = [make_func_parent("foo")]
    codeflash_output = get_first_top_level_function_or_method_ast("foo", parents, tree); result = codeflash_output

# Large Scale Test Cases






#------------------------------------------------
from __future__ import annotations

import ast
from typing import TypeVar

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.static_analysis import \
    get_first_top_level_function_or_method_ast
from codeflash.models.models import FunctionParent

# unit tests

# Helper: parse source and call the function
def get_func_ast(source, func_name, parents=None):
    tree = ast.parse(source)
    return get_first_top_level_function_or_method_ast(
        func_name,
        parents or [],
        tree
    )

# Helper: create FunctionParent for class
def class_parent(name):
    return [FunctionParent(name=name, type="ClassDef")]

# 1. Basic Test Cases

def test_basic_top_level_function():
    """Should find a top-level function definition."""
    src = "def foo(): pass"
    node = get_func_ast(src, "foo")

def test_basic_top_level_async_function():
    """Should find a top-level async function definition."""
    src = "async def bar(): pass"
    node = get_func_ast(src, "bar")

def test_basic_class_method():
    """Should find a method inside a class."""
    src = "class C:\n    def baz(self): pass"
    node = get_func_ast(src, "baz", class_parent("C"))

def test_basic_class_async_method():
    """Should find an async method inside a class."""
    src = "class D:\n    async def qux(self): pass"
    node = get_func_ast(src, "qux", class_parent("D"))

def test_basic_multiple_functions():
    """Should find the first matching function if there are multiple."""
    src = "def foo(): pass\ndef foo(): pass"
    node = get_func_ast(src, "foo")

def test_basic_multiple_classes():
    """Should find method in the first matching class."""
    src = ("class A:\n    def m(self): pass\n"
           "class A:\n    def m(self): pass")
    node = get_func_ast(src, "m", class_parent("A"))

# 2. Edge Test Cases

def test_edge_no_matching_function():
    """Should return None if function does not exist."""
    src = "def foo(): pass"
    node = get_func_ast(src, "bar")

def test_edge_no_matching_class():
    """Should return None if class does not exist."""
    src = "class C:\n    def foo(self): pass"
    node = get_func_ast(src, "foo", class_parent("D"))

def test_edge_function_nested_in_function():
    """Should NOT find a function nested inside another function."""
    src = "def outer():\n    def inner(): pass"
    node = get_func_ast(src, "inner")

def test_edge_method_nested_in_method():
    """Should NOT find a method nested inside another method."""
    src = "class C:\n    def outer(self):\n        def inner(self): pass"
    node = get_func_ast(src, "inner", class_parent("C"))

def test_edge_function_in_class_not_found():
    """Should NOT find a method in a class if no parents given."""
    src = "class C:\n    def foo(self): pass"
    node = get_func_ast(src, "foo")

def test_edge_function_name_collision():
    """Should find top-level function even if class method of same name exists."""
    src = ("def foo(): pass\n"
           "class C:\n    def foo(self): pass")
    node = get_func_ast(src, "foo")

def test_edge_method_name_collision():
    """Should find class method even if top-level function of same name exists."""
    src = ("def foo(): pass\n"
           "class C:\n    def foo(self): pass")
    node = get_func_ast(src, "foo", class_parent("C"))

def test_edge_async_and_sync_collision():
    """Should prefer sync over async if both exist at top-level."""
    src = "def foo(): pass\nasync def foo(): pass"
    node = get_func_ast(src, "foo")

def test_edge_async_and_sync_method_collision():
    """Should prefer sync over async if both exist in class."""
    src = ("class C:\n    def foo(self): pass\n    async def foo(self): pass")
    node = get_func_ast(src, "foo", class_parent("C"))

def test_edge_empty_source():
    """Should return None for empty source."""
    src = ""
    node = get_func_ast(src, "foo")


def test_edge_class_with_no_methods():
    """Should return None if class exists but method does not."""
    src = "class C:\n    pass"
    node = get_func_ast(src, "foo", class_parent("C"))

def test_edge_nested_class():
    """Should NOT find method in nested class."""
    src = ("class Outer:\n"
           "    class Inner:\n"
           "        def foo(self): pass")
    node = get_func_ast(src, "foo", class_parent("Inner"))

def test_edge_function_in_if_block():
    """Should find top-level function defined in an if block."""
    src = ("if True:\n"
           "    def foo(): pass")
    node = get_func_ast(src, "foo")

def test_edge_class_in_if_block():
    """Should find top-level class and its method defined in an if block."""
    src = ("if True:\n"
           "    class C:\n"
           "        def foo(self): pass")
    node = get_func_ast(src, "foo", class_parent("C"))

def test_edge_function_with_decorator():
    """Should find function with decorator."""
    src = ("@staticmethod\n"
           "def foo(): pass")
    node = get_func_ast(src, "foo")

def test_edge_method_with_decorator():
    """Should find method with decorator inside class."""
    src = ("class C:\n"
           "    @staticmethod\n"
           "    def foo(): pass")
    node = get_func_ast(src, "foo", class_parent("C"))

def test_edge_function_with_args():
    """Should find function with arguments."""
    src = "def foo(a, b): pass"
    node = get_func_ast(src, "foo")

# 3. Large Scale Test Cases

def test_large_many_functions():
    """Should find the first of many top-level functions."""
    src = "\n".join(f"def func{i}(): pass" for i in range(500))
    node = get_func_ast(src, "func0")

def test_large_many_classes_and_methods():
    """Should find method in the first of many classes."""
    src = "\n".join(f"class C{i}:\n    def foo(self): pass" for i in range(300))
    node = get_func_ast(src, "foo", class_parent("C0"))

def test_large_many_methods_same_name():
    """Should find method in the correct class among many classes with same method name."""
    src = "\n".join(f"class C{i}:\n    def foo(self): pass" for i in range(100))
    node = get_func_ast(src, "foo", class_parent("C50"))
    # Should be in class C50
    # Find the line number for C50's method
    expected_lineno = 2 + 2 * 50

def test_large_functions_in_if_blocks():
    """Should find function defined in an if block among many functions."""
    src = "\n".join(f"if True:\n    def func{i}(): pass" for i in range(200))
    node = get_func_ast(src, "func199")
    # Should be the last function
    expected_lineno = 1 + 2 * 199

def test_large_class_with_many_methods():
    """Should find a method among many in a single class."""
    methods = "\n".join(f"    def m{i}(self): pass" for i in range(500))
    src = f"class Big:\n{methods}"
    node = get_func_ast(src, "m499", class_parent("Big"))
    # Should be the last method
    expected_lineno = 2 + 499

def test_large_many_classes_with_same_method_name():
    """Should find method in the first class when all have same method name."""
    src = "\n".join(f"class C{i}:\n    def foo(self): pass" for i in range(100))
    node = get_func_ast(src, "foo", class_parent("C0"))

def test_large_many_functions_name_collision():
    """Should find the first function when many have the same name."""
    src = "\n".join("def foo(): pass" for _ in range(500))
    node = get_func_ast(src, "foo")

def test_large_class_and_top_level_function_same_name():
    """Should find top-level function even if class method of same name exists in large file."""
    src = ("def foo(): pass\n" +
           "\n".join(f"class C{i}:\n    def foo(self): pass" for i in range(500)))
    node = get_func_ast(src, "foo")

def test_large_class_and_top_level_function_same_name_method():
    """Should find class method even if top-level function exists in large file."""
    src = ("def foo(): pass\n" +
           "\n".join(f"class C{i}:\n    def foo(self): pass" for i in range(500)))
    node = get_func_ast(src, "foo", class_parent("C499"))
    # Should be in class C499
    expected_lineno = 2 + 2 * 499
# 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-pr678-2025-09-02T17.37.12 and push.