From c7795495e4b5dfddfeef01e039fb7a2c647c15b3 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Mon, 29 Jun 2026 16:41:18 +0000 Subject: [PATCH] perf: optimize hot-path AST traversals with explicit stacks Replaced recursive `yield from` with explicit stack traversals in hot-paths `_walk_own`, `_walk_own_non_stmt_children`, `walk_node`, and `_iter_l2_body_nodes`. This eliminates generator nesting overhead and preserves short-circuit capability, significantly improving iteration speed. Co-authored-by: tachyon-beep <544926+tachyon-beep@users.noreply.github.com> --- .jules/bolt.md | 3 + src/wardline/scanner/analyzer.py | 33 +++++++--- src/wardline/scanner/ast_primitives.py | 73 ++++++++++++++-------- src/wardline/scanner/rules/_ast_helpers.py | 70 +++++++++++++++++---- 4 files changed, 133 insertions(+), 46 deletions(-) create mode 100644 .jules/bolt.md diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 00000000..85e5a3df --- /dev/null +++ b/.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. diff --git a/src/wardline/scanner/analyzer.py b/src/wardline/scanner/analyzer.py index b3660ad3..95d8920c 100644 --- a/src/wardline/scanner/analyzer.py +++ b/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): diff --git a/src/wardline/scanner/ast_primitives.py b/src/wardline/scanner/ast_primitives.py index 70f565b3..361904d0 100644 --- a/src/wardline/scanner/ast_primitives.py +++ b/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( diff --git a/src/wardline/scanner/rules/_ast_helpers.py b/src/wardline/scanner/rules/_ast_helpers.py index 7c3b52ff..bf2a0968 100644 --- a/src/wardline/scanner/rules/_ast_helpers.py +++ b/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)