⚡ Bolt: Optimize AST Traversal Hot Paths using Explicit Stack

.jules/bolt.md

@@ -0,0 +1,3 @@
+## 2025-02-18 - Optimized AST traversal using explicit stack and manual `_fields` iteration
+**Learning:** Python's recursive generator syntax (`yield from`) incurs noticeable overhead in deep or wide ASTs, making it slower and more memory-intensive on hot paths than an explicit stack loop. Additionally, `ast.iter_child_nodes()` overhead can be bypassed by manually iterating over a node's `_fields`. To correctly reverse the traversal order (so stack `pop()` processes items depth-first left-to-right) and avoid eager whole-tree list materialization, the children for a *single node* should be extracted into a list and reversed using `stack.extend(reversed(children))`. Also, `isinstance(val, ast.AST)` is necessary when iterating over `_fields` since some fields hold literal strings or booleans, which lack `_fields` themselves.
+**Action:** Use an explicit stack (with `yield`, not `yield from`) for hot-path AST traversals. Manually iterate over `_fields`, safely handle `AttributeError`s for missing fields, collect a single node's children into a list, and reverse that list before extending the stack.

src/wardline/scanner/analyzer.py

@@ -453,16 +453,32 @@ def _bind_call_site_arguments_to_parameters(
             return result
 
         def _iter_l2_body_nodes(node: ast.FunctionDef | ast.AsyncFunctionDef) -> Iterator[ast.AST]:
-            def walk(current: ast.AST) -> Iterator[ast.AST]:
-                for child in ast.iter_child_nodes(current):
-                    if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
+            stack: list[ast.AST] = []
+
+            def push_children(current: ast.AST) -> None:
+                children = []
+                for field in current._fields:
+                    try:
+                        val = getattr(current, field)
+                    except AttributeError:
                         continue
-                    yield child
-                    yield from walk(child)
+                    if isinstance(val, list):
+                        for item in val:
+                            if isinstance(item, ast.AST):
+                                children.append(item)
+                    elif isinstance(val, ast.AST):
+                        children.append(val)
+                stack.extend(reversed(children))
 
             for stmt in node.body:
                 yield stmt
-                yield from walk(stmt)
+                push_children(stmt)
+                while stack:
+                    current = stack.pop()
+                    if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
+                        continue
+                    yield current
+                    push_children(current)
 
         def _assignment_targets(node: ast.AST) -> list[ast.expr]:
             if isinstance(node, ast.Assign):

src/wardline/scanner/ast_primitives.py

@@ -104,39 +104,57 @@ 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))
+
+    def push_children(current: ast.AST) -> None:
+        children = []
+        for field in current._fields:
+            try:
+                val = getattr(current, field)
+            except AttributeError:
+                continue
+            if isinstance(val, list):
+                for item in val:
+                    if isinstance(item, ast.AST):
+                        children.append(item)
+            elif isinstance(val, ast.AST):
+                children.append(val)
+        stack.extend(reversed(children))
+
+    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
+            for kw in reversed(current.args.kw_defaults):
+                if kw is not None:
+                    stack.append(kw)
+            for default in reversed(current.args.defaults):
+                stack.append(default)
+            for decorator in reversed(current.decorator_list):
+                stack.append(decorator)
+            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)
+            for base in reversed(current.bases):
+                stack.append(base)
+            for decorator in reversed(current.decorator_list):
+                stack.append(decorator)
+            continue
+
         if isinstance(current, ast.Lambda):
-            yield from _walk_argument_defaults(current.args)
-            return
+            for kw in reversed(current.args.kw_defaults):
+                if kw is not None:
+                    stack.append(kw)
+            for default in reversed(current.args.defaults):
+                stack.append(default)
+            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)
+        push_children(current)
 
 
 def resolve_self_method_fqn(

src/wardline/scanner/rules/_ast_helpers.py

@@ -36,12 +36,32 @@
 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):
+    stack: list[ast.AST] = []
+
+    def push_children(current: ast.AST) -> None:
+        children = []
+        for field in current._fields:
+            try:
+                val = getattr(current, field)
+            except AttributeError:
+                continue
+            if isinstance(val, list):
+                for item in val:
+                    if isinstance(item, ast.AST):
+                        children.append(item)
+            elif isinstance(val, ast.AST):
+                children.append(val)
+        stack.extend(reversed(children))
+
+    push_children(node)
+
+    while stack:
+        child = stack.pop()
         if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
             continue
         if isinstance(child, ast.stmt):
             yield child
-        yield from _own_statements(child)
+        push_children(child)
 
 
 def _own_reachable_statements(
@@ -67,14 +87,34 @@ 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):
+    stack: list[ast.AST] = []
+
+    def push_children(current: ast.AST) -> None:
+        children = []
+        for field in current._fields:
+            try:
+                val = getattr(current, field)
+            except AttributeError:
+                continue
+            if isinstance(val, list):
+                for item in val:
+                    if isinstance(item, ast.AST):
+                        children.append(item)
+            elif isinstance(val, ast.AST):
+                children.append(val)
+        stack.extend(reversed(children))
+
+    push_children(node)
+
+    while stack:
+        child = stack.pop()
         if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
             yield child
         elif isinstance(child, ast.stmt):
             continue
         else:
             yield child
-            yield from _walk_own_non_stmt_children(child)
+            push_children(child)
 
 
 def _reachable_statements_in_block(
@@ -635,13 +675,34 @@ def handler_substitutes_on_failure(handler: ast.ExceptHandler, returned_names: f
 def own_nodes(node: ast.AST) -> Iterator[ast.AST]:
     """Yield *node* itself and all descendant nodes in its own scope (skipping nested scopes)."""
     yield node
-    yield from _walk_own(node)
 
-
-def _walk_own(node: ast.AST) -> Iterator[ast.AST]:
-    for child in ast.iter_child_nodes(node):
+    stack: list[ast.AST] = []
+
+    def push_children(current: ast.AST) -> None:
+        children = []
+        for field in current._fields:
+            try:
+                val = getattr(current, field)
+            except AttributeError:
+                continue
+            if isinstance(val, list):
+                for item in val:
+                    if isinstance(item, ast.AST):
+                        children.append(item)
+            elif isinstance(val, ast.AST):
+                children.append(val)
+        stack.extend(reversed(children))
+
+    push_children(node)
+
+    while stack:
+        child = stack.pop()
         if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
             yield child
         else:
             yield child
-            yield from _walk_own(child)
+            push_children(child)
+
+
+def _walk_own(node: ast.AST) -> Iterator[ast.AST]:  # pragma: no cover
+    raise NotImplementedError("Replaced by inline own_nodes logic")