diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 00000000..e973c48f --- /dev/null +++ b/.jules/bolt.md @@ -0,0 +1,3 @@ +## 2024-05-24 - Avoid eager list materialization in AST traversal +**Learning:** In highly recursive AST traversal functions like `_assignment_callee` and `_collect_return_paths`, wrapping `ast.iter_child_nodes()` or `func_node.body` in `list()` eagerly materializes the generator. This forces unnecessary memory allocations for intermediate list structures, hurting performance across large Python codebases. +**Action:** Accept `Iterable[ast.AST]` in traversal signatures instead of `list[ast.AST]` and pass the generators directly. \ No newline at end of file diff --git a/src/wardline/scanner/taint/variable_level.py b/src/wardline/scanner/taint/variable_level.py index a51a6989..54ad12e7 100644 --- a/src/wardline/scanner/taint/variable_level.py +++ b/src/wardline/scanner/taint/variable_level.py @@ -33,7 +33,7 @@ from wardline.core.taints import _PROVENANCE_CLASH, RAW_ZONE, TRUST_RANK, TaintState, combine if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Iterable, Iterator # Serialisation sinks — calls that cross the representation boundary. Their # output sheds validation provenance (raw bytes/str), so → UNKNOWN_RAW. This is @@ -2414,7 +2414,7 @@ def compute_return_taint( (the function's anchored *body* taint, pinned to its declaration). """ returns: list[tuple[TaintState, str | None, ast.expr]] = [] - _collect_return_paths(list(func_node.body), function_taint, taint_map, var_taints, returns, return_snapshots) + _collect_return_paths(func_node.body, function_taint, taint_map, var_taints, returns, return_snapshots) if not returns: return None result = returns[0][0] @@ -2454,7 +2454,7 @@ def compute_return_callee( beyond one hop stay ``None`` (the N-hop walk lives in the Loomweave stored-fact path). """ returns: list[tuple[TaintState, str | None, ast.expr]] = [] - _collect_return_paths(list(func_node.body), function_taint, taint_map, var_taints, returns, return_snapshots) + _collect_return_paths(func_node.body, function_taint, taint_map, var_taints, returns, return_snapshots) if not returns: return None worst = returns[0][0] @@ -2468,14 +2468,14 @@ def compute_return_callee( # a direct call. Provenance only — never changes a fire/no-fire decision. for taint, callee, node in returns: if taint == worst and callee is None and isinstance(node, ast.Name): - indirect = _assignment_callee(list(func_node.body), node.id, worst, function_taint, taint_map, var_taints) + indirect = _assignment_callee(func_node.body, node.id, worst, function_taint, taint_map, var_taints) if indirect is not None: return indirect return None def _assignment_callee( - nodes: list[ast.AST], + nodes: Iterable[ast.AST], name: str, worst: TaintState, function_taint: TaintState, @@ -2508,9 +2508,7 @@ def _assignment_callee( and _resolve_expr(node.value, function_taint, taint_map, var_taints) == worst ): result = callee - nested = _assignment_callee( - list(ast.iter_child_nodes(node)), name, worst, function_taint, taint_map, var_taints - ) + nested = _assignment_callee(ast.iter_child_nodes(node), name, worst, function_taint, taint_map, var_taints) if nested is not None: result = nested return result @@ -2527,7 +2525,7 @@ def _return_callee(node: ast.expr) -> str | None: def _collect_return_paths( - nodes: list[ast.AST], + nodes: Iterable[ast.AST], function_taint: TaintState, taint_map: dict[str, TaintState], var_taints: dict[str, TaintState], @@ -2558,7 +2556,7 @@ def _collect_return_paths( taint = _resolve_expr(node.value, function_taint, taint_map, dict(snapshot or var_taints)) out.append((taint, _return_callee(node.value), node.value)) _collect_return_paths( - list(ast.iter_child_nodes(node)), + ast.iter_child_nodes(node), function_taint, taint_map, var_taints,