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

codeflash/code_utils/static_analysis.py

@@ -116,31 +116,41 @@ def analyze_imported_modules(
 def get_first_top_level_object_def_ast(
     object_name: str, object_type: type[ObjectDefT], node: ast.AST
 ) -> ObjectDefT | None:
-    for child in ast.iter_child_nodes(node):
-        if isinstance(child, object_type) and child.name == object_name:
-            return child
-        if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
-            continue
-        if descendant := get_first_top_level_object_def_ast(object_name, object_type, child):
-            return descendant
+    # Use local bindings for attribute/func lookups and type tuples for speed in tight loops
+    iter_child_nodes = ast.iter_child_nodes
+    fn_type_tuple = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)
+
+    for child in iter_child_nodes(node):
+        # Fast path: check type and name directly
+        if isinstance(child, object_type):
+            if getattr(child, "name", None) == object_name:
+                return child
+        # Avoid recursing into function/class defs as per original logic
+        elif not isinstance(child, fn_type_tuple):
+            descendant = get_first_top_level_object_def_ast(object_name, object_type, child)
+            if descendant is not None:
+                return descendant
     return None
 
 
 def get_first_top_level_function_or_method_ast(
     function_name: str, parents: list[FunctionParent], node: ast.AST
 ) -> ast.FunctionDef | ast.AsyncFunctionDef | None:
+    # Fast path: no parents, search from the root node
     if not parents:
         result = get_first_top_level_object_def_ast(function_name, ast.FunctionDef, node)
         if result is not None:
             return result
         return get_first_top_level_object_def_ast(function_name, ast.AsyncFunctionDef, node)
-    if parents[0].type == "ClassDef" and (
-        class_node := get_first_top_level_object_def_ast(parents[0].name, ast.ClassDef, node)
-    ):
-        result = get_first_top_level_object_def_ast(function_name, ast.FunctionDef, class_node)
-        if result is not None:
-            return result
-        return get_first_top_level_object_def_ast(function_name, ast.AsyncFunctionDef, class_node)
+    # Only ClassDef parent handled, search in class scope only once
+    parent0 = parents[0]
+    if parent0.type == "ClassDef":
+        class_node = get_first_top_level_object_def_ast(parent0.name, ast.ClassDef, node)
+        if class_node is not None:
+            result = get_first_top_level_object_def_ast(function_name, ast.FunctionDef, class_node)
+            if result is not None:
+                return result
+            return get_first_top_level_object_def_ast(function_name, ast.AsyncFunctionDef, class_node)
     return None