⚡ Bolt: Optimize hot-path AST traversals via explicit stacks

.jules/bolt.md

@@ -0,0 +1,3 @@
+## 2024-07-02 - AST Traversal with explicit stack
+**Learning:** For hot-path AST traversal, use an explicit stack combined with `yield` instead of `yield from` recursion to preserve lazy evaluation and short-circuiting capabilities. Avoid eager list-appending (`list.append()`) to prevent computing the whole subtree when early matches exist. Reverse child nodes via `reversed()` before extending the stack to preserve traversal order. Critically, retain `isinstance(node, ast.AST)` checks to prevent `AttributeError`s when traversing `_fields`, as not all field values are AST nodes (some are strings or literals).
+**Action:** Replace `yield from` with explicit stack traversal and `yield` for hot-path AST traversals.

src/wardline/scanner/analyzer.py

@@ -453,16 +453,33 @@ 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)):
-                        continue
-                    yield child
-                    yield from walk(child)
-
             for stmt in node.body:
                 yield stmt
-                yield from walk(stmt)
+                stack: list[ast.AST] = [stmt]
+                while stack:
+                    current = stack.pop()
+                    if current is not stmt:
+                        yield current
+                    if hasattr(current, "_fields"):
+                        # PERF: Explicit stack traversal avoids yield from recursion overhead on hot paths.
+                        # Eagerly pushing reversed children ensures left-to-right depth-first order.
+                        children = []
+                        for field in current._fields:
+                            try:
+                                value = getattr(current, 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)
+                        if children:
+                            for child in reversed(children):
+                                if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
+                                    continue
+                                stack.append(child)
 
         def _assignment_targets(node: ast.AST) -> list[ast.expr]:
             if isinstance(node, ast.Assign):

src/wardline/scanner/ast_primitives.py

@@ -106,34 +106,59 @@ def iter_calls_in_function_body(
     """
 
     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
-        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
-        if isinstance(current, ast.Lambda):
-            yield from _walk_argument_defaults(current.args)
-            return
-        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:
+        stack = [current]
+        while stack:
+            node = stack.pop()
+
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                args = node.args
+                for kw_default in reversed(args.kw_defaults):
+                    if kw_default is not None:
+                        stack.append(kw_default)
+                for default in reversed(args.defaults):
+                    stack.append(default)
+                for decorator in reversed(node.decorator_list):
+                    stack.append(decorator)
                 continue
-            yield from walk_node(kw_default)
+
+            if isinstance(node, ast.ClassDef):
+                for keyword in reversed(node.keywords):
+                    stack.append(keyword.value)
+                for base in reversed(node.bases):
+                    stack.append(base)
+                for decorator in reversed(node.decorator_list):
+                    stack.append(decorator)
+                continue
+
+            if isinstance(node, ast.Lambda):
+                args = node.args
+                for kw_default in reversed(args.kw_defaults):
+                    if kw_default is not None:
+                        stack.append(kw_default)
+                for default in reversed(args.defaults):
+                    stack.append(default)
+                continue
+
+            if isinstance(node, ast.Call):
+                yield node
+
+            if hasattr(node, "_fields"):
+                # PERF: Explicit stack traversal avoids yield from recursion overhead on hot paths.
+                # Eagerly pushing reversed children ensures left-to-right depth-first order.
+                children = []
+                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)
+                if children:
+                    stack.extend(reversed(children))
 
     for stmt in node.body:
         yield from walk_node(stmt)

src/wardline/scanner/rules/_ast_helpers.py

@@ -67,14 +67,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):
-        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)
+    stack = [node]
+    while stack:
+        current = stack.pop()
+        if current is not node:
+            if isinstance(current, ast.stmt) and not isinstance(
+                current, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)
+            ):
+                continue
+            yield current
+            if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
+                continue
+        if hasattr(current, "_fields"):
+            # PERF: Explicit stack traversal avoids yield from recursion overhead on hot paths.
+            # Eagerly pushing reversed children ensures left-to-right depth-first order.
+            children = []
+            for field in current._fields:
+                try:
+                    value = getattr(current, 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)
+            if children:
+                stack.extend(reversed(children))
 
 
 def _reachable_statements_in_block(
@@ -639,9 +659,27 @@ 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
-        else:
-            yield child
-            yield from _walk_own(child)
+    stack = [node]
+    while stack:
+        current = stack.pop()
+        if current is not node:
+            yield current
+            if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
+                continue
+        if hasattr(current, "_fields"):
+            # PERF: Explicit stack traversal avoids yield from recursion overhead on hot paths.
+            # Eagerly pushing reversed children ensures left-to-right depth-first order.
+            children = []
+            for field in current._fields:
+                try:
+                    value = getattr(current, 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)
+            if children:
+                stack.extend(reversed(children))