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

.jules/bolt.md

@@ -0,0 +1,3 @@
+## 2023-10-27 - Explicit stack AST traversal over recursive generators
+**Learning:** In highly recursive AST hot-paths (`_walk_own`, `_iter_l2_body_nodes`), replacing `yield from` recursion with an explicit stack eliminates generator overhead and prevents deep callstacks while preserving short-circuit capability. It is essential to push nodes in reverse order and use generator comprehensions inside `stack.extend()` to avoid eager evaluation and high memory overhead.
+**Action:** When working on deep AST traversals, always favor an explicit stack with `yield` and inline generator comprehensions over nested `yield from` recursion.

src/wardline/scanner/analyzer.py

@@ -453,16 +453,31 @@ 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)
+            # Perf opt: Stack-based AST traversal instead of deep recursive yield from.
+            # Child nodes pushed in reverse to preserve expected AST iteration order, avoiding eager list overhead.
+            stack: list[ast.AST] = []
+            stack.extend(reversed(node.body))
+
+            while stack:
+                current = stack.pop()
+                yield current
 
-            for stmt in node.body:
-                yield stmt
-                yield from walk(stmt)
+                for field in reversed(current._fields):
+                    try:
+                        value = getattr(current, field)
+                    except AttributeError:
+                        continue
+                    if isinstance(value, list):
+                        stack.extend(
+                            n
+                            for n in reversed(value)
+                            if isinstance(n, ast.AST)
+                            and not isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda))
+                        )
+                    elif isinstance(value, ast.AST) and not isinstance(
+                        value, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)
+                    ):
+                        stack.append(value)
 
         def _assignment_targets(node: ast.AST) -> list[ast.expr]:
             if isinstance(node, ast.Assign):

src/wardline/scanner/ast_primitives.py

@@ -105,38 +105,59 @@ def iter_calls_in_function_body(
     values, base classes, metaclass keywords) are still attributed to ``node``.
     """
 
-    def walk_node(current: ast.AST) -> Iterator[ast.Call]:
+    # Perf opt: Using explicit stack over recursive "yield from" to eliminate generator nesting
+    # overhead in hot path traversal. Elements are pushed in reverse to maintain left-to-right processing order.
+    stack: list[ast.AST] = []
+    stack.extend(reversed(node.body))
+
+    while stack:
+        current = stack.pop()
+
         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_default in reversed(current.args.kw_defaults):
+                if kw_default is not None:
+                    stack.append(kw_default)
+            for default in reversed(current.args.defaults):
+                if default is not None:
+                    stack.append(default)
+            for decorator in reversed(current.decorator_list):
+                if decorator is not None:
+                    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):
+                if keyword.value is not None:
+                    stack.append(keyword.value)
+            for base in reversed(current.bases):
+                if base is not None:
+                    stack.append(base)
+            for decorator in reversed(current.decorator_list):
+                if decorator is not None:
+                    stack.append(decorator)
+            continue
+
         if isinstance(current, ast.Lambda):
-            yield from _walk_argument_defaults(current.args)
-            return
+            for kw_default in reversed(current.args.kw_defaults):
+                if kw_default is not None:
+                    stack.append(kw_default)
+            for default in reversed(current.args.defaults):
+                if default is not None:
+                    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)
+        for field in reversed(current._fields):
+            try:
+                value = getattr(current, field)
+            except AttributeError:
+                continue
+            if isinstance(value, list):
+                stack.extend(n for n in reversed(value) if isinstance(n, ast.AST))
+            elif isinstance(value, ast.AST):
+                stack.append(value)
 
 
 def resolve_self_method_fqn(

src/wardline/scanner/rules/_ast_helpers.py

@@ -67,14 +67,38 @@ 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):
+    # Perf opt: Replaced recursive yield from with explicit stack to avoid
+    # deep callstack overhead and generator creation during hot-path AST traversal.
+    # We must push elements in reverse order to preserve original top-down, left-to-right
+    # traversal, and we use generator comprehension to avoid eager large list instantiation.
+    stack: list[ast.AST] = []
+    for field in reversed(node._fields):
+        try:
+            value = getattr(node, field)
+        except AttributeError:
+            continue
+        if isinstance(value, list):
+            stack.extend(n for n in reversed(value) if isinstance(n, ast.AST))
+        elif isinstance(value, ast.AST):
+            stack.append(value)
+
+    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
+            for field in reversed(current._fields):
+                try:
+                    value = getattr(current, field)
+                except AttributeError:
+                    continue
+                if isinstance(value, list):
+                    stack.extend(n for n in reversed(value) if isinstance(n, ast.AST))
+                elif isinstance(value, ast.AST):
+                    stack.append(value)
 
 
 def _reachable_statements_in_block(
@@ -639,9 +663,33 @@ 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
+    # Perf opt: Explicit stack replaces yield from recursion for hot-path AST traversal.
+    # Preserves lazy evaluation and short-circuiting capabilities.
+    # Child nodes are pushed to the stack in reverse order to keep exact traversal sequence.
+    # We use generator comprehension inside extend() to prevent eager subtree evaluation.
+    stack: list[ast.AST] = []
+    for field in reversed(node._fields):
+        try:
+            value = getattr(node, field)
+        except AttributeError:
+            continue
+        if isinstance(value, list):
+            stack.extend(n for n in reversed(value) if isinstance(n, ast.AST))
+        elif isinstance(value, ast.AST):
+            stack.append(value)
+
+    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
+            for field in reversed(current._fields):
+                try:
+                    value = getattr(current, field)
+                except AttributeError:
+                    continue
+                if isinstance(value, list):
+                    stack.extend(n for n in reversed(value) if isinstance(n, ast.AST))
+                elif isinstance(value, ast.AST):
+                    stack.append(value)