⚡ Bolt: [performance improvement] Optimize hot-path AST traversals with explicit stacks

.jules/bolt.md

@@ -0,0 +1,3 @@
+## 2025-02-23 - Explicit Stack AST Traversal
+**Learning:** Recursive `yield from` in hot-path AST traversals introduces significant overhead. Precomputing reversed children via `node._fields` and using an explicit stack preserves traversal order and lazy evaluation while being faster.
+**Action:** Use an explicit stack combined with `yield` and `reversed()` children list instead of `yield from` recursion for hot-path AST traversal.

src/wardline/scanner/ast_primitives.py

@@ -93,6 +93,25 @@ def build_import_alias_map(
     return alias_map
 
 
+def _reversed_children(node: ast.AST) -> list[ast.AST]:
+    """Return a list of AST children in reverse order, extracting directly from _fields
+    to avoid ast.iter_child_nodes overhead in hot paths."""
+    children: list[ast.AST] = []
+    for field in node._fields:
+        try:
+            value = getattr(node, field)
+        except AttributeError:
+            continue
+        if isinstance(value, list):
+            for item in value:
+                if isinstance(item, ast.AST):
+                    children.append(item)
+        elif isinstance(value, ast.AST):
+            children.append(value)
+    children.reverse()
+    return children
+
+
 def iter_calls_in_function_body(
     node: ast.FunctionDef | ast.AsyncFunctionDef,
 ) -> Iterator[ast.Call]:
@@ -104,39 +123,39 @@ def iter_calls_in_function_body(
     Header expressions that execute in the enclosing scope (decorators, default
     values, base classes, metaclass keywords) are still attributed to ``node``.
     """
+    stack: list[ast.AST] = list(reversed(node.body))
+
+    while stack:
+        current = stack.pop()
 
-    def walk_node(current: ast.AST) -> Iterator[ast.Call]:
         if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef)):
-            for decorator in current.decorator_list:
-                yield from walk_node(decorator)
-            yield from _walk_argument_defaults(current.args)
-            return
+            args = current.args
+            for kw_default in reversed(args.kw_defaults):
+                if kw_default is not None:
+                    stack.append(kw_default)
+            stack.extend(reversed(args.defaults))
+            stack.extend(reversed(current.decorator_list))
+            continue
+
         if isinstance(current, ast.ClassDef):
-            for decorator in current.decorator_list:
-                yield from walk_node(decorator)
-            for base in current.bases:
-                yield from walk_node(base)
-            for keyword in current.keywords:
-                yield from walk_node(keyword.value)
-            return
+            for keyword in reversed(current.keywords):
+                stack.append(keyword.value)
+            stack.extend(reversed(current.bases))
+            stack.extend(reversed(current.decorator_list))
+            continue
+
         if isinstance(current, ast.Lambda):
-            yield from _walk_argument_defaults(current.args)
-            return
+            args = current.args
+            for kw_default in reversed(args.kw_defaults):
+                if kw_default is not None:
+                    stack.append(kw_default)
+            stack.extend(reversed(args.defaults))
+            continue
+
         if isinstance(current, ast.Call):
             yield current
-        for child in ast.iter_child_nodes(current):
-            yield from walk_node(child)
-
-    def _walk_argument_defaults(args: ast.arguments) -> Iterator[ast.Call]:
-        for default in args.defaults:
-            yield from walk_node(default)
-        for kw_default in args.kw_defaults:
-            if kw_default is None:
-                continue
-            yield from walk_node(kw_default)
 
-    for stmt in node.body:
-        yield from walk_node(stmt)
+        stack.extend(_reversed_children(current))
 
 
 def resolve_self_method_fqn(

src/wardline/scanner/rules/_ast_helpers.py

@@ -33,15 +33,36 @@
 _RAISING_CONVERSION_NAMES: frozenset[str] = frozenset({"int", "float", "complex", "Decimal", "Fraction", "UUID"})
 
 
+def _reversed_children(node: ast.AST) -> list[ast.AST]:
+    """Return a list of AST children in reverse order, extracting directly from _fields
+    to avoid ast.iter_child_nodes overhead in hot paths."""
+    children: list[ast.AST] = []
+    for field in node._fields:
+        try:
+            value = getattr(node, field)
+        except AttributeError:
+            continue
+        if isinstance(value, list):
+            for item in value:
+                if isinstance(item, ast.AST):
+                    children.append(item)
+        elif isinstance(value, ast.AST):
+            children.append(value)
+    children.reverse()
+    return children
+
+
 def _own_statements(node: ast.AST) -> Iterator[ast.stmt]:
     """Yield every statement in *node*'s own scope, not descending into nested
     def/class bodies. Includes the bodies of if/for/while/try/with at any depth."""
-    for child in ast.iter_child_nodes(node):
-        if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+    stack = _reversed_children(node)
+    while stack:
+        current = stack.pop()
+        if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
             continue
-        if isinstance(child, ast.stmt):
-            yield child
-        yield from _own_statements(child)
+        if isinstance(current, ast.stmt):
+            yield current
+        stack.extend(_reversed_children(current))
 
 
 def _own_reachable_statements(
@@ -67,14 +88,16 @@ def _own_nodes_in_reachable_stmt(stmt: ast.stmt) -> Iterator[ast.AST]:
 
 
 def _walk_own_non_stmt_children(node: ast.AST) -> Iterator[ast.AST]:
-    for child in ast.iter_child_nodes(node):
-        if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
-            yield child
-        elif isinstance(child, ast.stmt):
+    stack = _reversed_children(node)
+    while stack:
+        current = stack.pop()
+        if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
+            yield current
+        elif isinstance(current, ast.stmt):
             continue
         else:
-            yield child
-            yield from _walk_own_non_stmt_children(child)
+            yield current
+            stack.extend(_reversed_children(current))
 
 
 def _reachable_statements_in_block(
@@ -639,9 +662,11 @@ def own_nodes(node: ast.AST) -> Iterator[ast.AST]:
 
 
 def _walk_own(node: ast.AST) -> Iterator[ast.AST]:
-    for child in ast.iter_child_nodes(node):
-        if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
-            yield child
+    stack = _reversed_children(node)
+    while stack:
+        current = stack.pop()
+        if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
+            yield current
         else:
-            yield child
-            yield from _walk_own(child)
+            yield current
+            stack.extend(_reversed_children(current))