⚡ Bolt: [performance improvement] Avoid eager list materialization in AST traversal

.jules/bolt.md

@@ -0,0 +1,3 @@
+## 2024-08-20 - Avoid eagerly materializing generators in AST traversal
+**Learning:** In AST traversal functions, accepting `list[ast.AST]` and passing materialized generators (e.g. `list(ast.iter_child_nodes(node))`) forces unnecessary list allocations and iterations, which causes a performance bottleneck and wastes memory, especially on deep or large ASTs.
+**Action:** Change the signature of AST traversal functions to accept `Iterable[ast.AST]` and pass the generators directly (e.g. `ast.iter_child_nodes(node)`) without eagerly materializing them.

src/wardline/install/block.py

@@ -31,6 +31,7 @@
 
 _BLOCK_VERSION = "1"
 
+
 def _compose_body(grant_suffix: str = "", grant_sentence: str = "") -> str:
     return (
         "This project uses **wardline** as its trust-boundary gate. Before handing "
@@ -81,6 +82,7 @@ def _pack_guidance(project_root: Path) -> tuple[str, str]:
     )
     return suffix, sentence
 
+
 _OWN_NS = "wardline"
 _END_MARKER = f"<!-- /{_OWN_NS}:instructions -->"
 _WRITER_MARKER = f"<!-- {_OWN_NS}:last-writer:wardline install -->"

src/wardline/mcp/server.py

@@ -5103,9 +5103,7 @@ def _grants_merged_arguments(self, arguments: dict[str, Any]) -> dict[str, Any]:
             if caller_packs is None or (
                 isinstance(caller_packs, list) and all(isinstance(p, str) for p in caller_packs)
             ):
-                merged["trust_packs"] = list(
-                    dict.fromkeys([*(caller_packs or []), *self._default_trusted_packs])
-                )
+                merged["trust_packs"] = list(dict.fromkeys([*(caller_packs or []), *self._default_trusted_packs]))
         if self._default_trust_local_packs:
             caller_local = merged.get("trust_local_packs")
             # Identity checks, not equality: 0 == False, and masking a caller's

src/wardline/scanner/taint/variable_level.py

@@ -26,6 +26,7 @@
 
 import ast
 import contextvars
+from collections.abc import Iterable
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
@@ -2519,7 +2520,7 @@ def compute_return_taint(
     """
     returns: list[tuple[TaintState, str | None, ast.expr]] = []
     _collect_return_paths(
-        list(func_node.body),
+        func_node.body,
         function_taint,
         taint_map,
         var_taints,
@@ -2568,7 +2569,7 @@ def compute_return_callee(
     """
     returns: list[tuple[TaintState, str | None, ast.expr]] = []
     _collect_return_paths(
-        list(func_node.body),
+        func_node.body,
         function_taint,
         taint_map,
         var_taints,
@@ -2589,14 +2590,16 @@ 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],
+    # OPTIMIZATION: Accept Iterable[ast.AST] to allow direct generator iteration
+    # without eager list materialization, saving memory on deep/large ASTs.
+    nodes: Iterable[ast.AST],
     name: str,
     worst: TaintState,
     function_taint: TaintState,
@@ -2629,9 +2632,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
@@ -2648,7 +2649,9 @@ def _return_callee(node: ast.expr) -> str | None:
 
 
 def _collect_return_paths(
-    nodes: list[ast.AST],
+    # OPTIMIZATION: Accept Iterable[ast.AST] to allow direct generator iteration
+    # without eager list materialization, saving memory on deep/large ASTs.
+    nodes: Iterable[ast.AST],
     function_taint: TaintState,
     taint_map: dict[str, TaintState],
     var_taints: dict[str, TaintState],
@@ -2686,7 +2689,7 @@ def _collect_return_paths(
                     _CURRENT_VAR_TYPES.reset(token_types)
             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,

tests/unit/install/test_doctor_pack_grants.py

@@ -66,9 +66,7 @@ def test_project_mcp_check_accepts_grant_flags(tmp_path: Path, monkeypatch: pyte
     assert check.ok, check.message
 
 
-def test_project_mcp_check_names_divergence_not_missing(
-    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
-) -> None:
+def test_project_mcp_check_names_divergence_not_missing(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
     # A present-but-noncanonical entry is a different failure than an absent one;
     # "missing wardline server" for a visibly present entry sent the operator
     # chasing the wrong problem.

tests/unit/install/test_mcp_json.py

@@ -549,9 +549,7 @@ def test_repair_preserves_trust_pack_grant_flags(tmp_path: Path, monkeypatch: py
         "--allow-custom-packs",
     ]
     (tmp_path / ".mcp.json").write_text(
-        json.dumps(
-            {"mcpServers": {"wardline": {"type": "stdio", "command": "/bin/wardline", "args": list(args)}}}
-        ),
+        json.dumps({"mcpServers": {"wardline": {"type": "stdio", "command": "/bin/wardline", "args": list(args)}}}),
         encoding="utf-8",
     )
     assert merge_mcp_entry(tmp_path) == "unchanged"
@@ -581,9 +579,7 @@ def test_repair_preserves_repeated_trust_pack_grants(tmp_path: Path, monkeypatch
     assert merge_mcp_entry(tmp_path) == "unchanged"
 
 
-def test_repair_drops_dangling_trust_pack_but_keeps_bare_grant(
-    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
-) -> None:
+def test_repair_drops_dangling_trust_pack_but_keeps_bare_grant(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
     # A --trust-pack with a missing or flag-shaped value is malformed: it must be
     # dropped cleanly, and must never swallow the following --allow-custom-packs.
     monkeypatch.setattr("wardline.install.mcp_json._find_wardline_command", lambda: "/bin/wardline")

tests/unit/mcp/test_server_trust_grants.py

@@ -24,9 +24,7 @@
 def _pack_project(tmp_path: Path) -> Path:
     proj = tmp_path / "proj"
     (proj / "scripts").mkdir(parents=True)
-    (proj / "scripts" / "grantpack.py").write_text(
-        'config = {"exclude": ["skipped_by_pack.py"]}\n', encoding="utf-8"
-    )
+    (proj / "scripts" / "grantpack.py").write_text('config = {"exclude": ["skipped_by_pack.py"]}\n', encoding="utf-8")
     (proj / "weft.toml").write_text(f'[wardline]\npacks = ["{PACK_NAME}"]\n', encoding="utf-8")
     (proj / "kept.py").write_text("def kept():\n    return 1\n", encoding="utf-8")
     (proj / "skipped_by_pack.py").write_text("def skipped():\n    return 1\n", encoding="utf-8")