diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 00000000..0f619032 --- /dev/null +++ b/.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. diff --git a/src/wardline/scanner/analyzer.py b/src/wardline/scanner/analyzer.py index b3660ad3..a192f7ae 100644 --- a/src/wardline/scanner/analyzer.py +++ b/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): diff --git a/src/wardline/scanner/ast_primitives.py b/src/wardline/scanner/ast_primitives.py index 70f565b3..1ed7922d 100644 --- a/src/wardline/scanner/ast_primitives.py +++ b/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) diff --git a/src/wardline/scanner/rules/_ast_helpers.py b/src/wardline/scanner/rules/_ast_helpers.py index 7c3b52ff..714cf710 100644 --- a/src/wardline/scanner/rules/_ast_helpers.py +++ b/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))