From 243deb39e74479c0bedd97e2e81bf0e1d574c06d Mon Sep 17 00:00:00 2001 From: Pini Vaknin <37843584+pishuv@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:56:07 +0300 Subject: [PATCH] perf(resolution): walk each JS/TS tree once when collecting symbol facts _collect_js_symbol_resolution_facts walked every parsed tree in full four times, and each pass kept only a few node types. The first pass now walks once and keeps the nodes of the six types the whole-tree passes filter on, in walk order. Later passes iterate that list. Each fact list is filled by a single pass, so fact order does not change. _walk_js_tree now uses a TreeCursor instead of reading .children at every node. A cursor is bounded by its start node, so the walks that start inside a tree (function bodies, and the import and export nodes passed to _js_named_specifiers) stay inside that subtree. Overlaps #3353, which also indexes JS nodes once and frees each tree after use. Co-Authored-By: Claude Opus 5.5 --- graphify/extractors/resolution.py | 56 +++++++++++++++++++---------- tests/test_js_walk_order.py | 59 +++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 19 deletions(-) create mode 100644 tests/test_js_walk_order.py diff --git a/graphify/extractors/resolution.py b/graphify/extractors/resolution.py index a447d3bd8..cb4eecda0 100644 --- a/graphify/extractors/resolution.py +++ b/graphify/extractors/resolution.py @@ -1702,14 +1702,17 @@ def _parse_js_tree(path: Path): return None def _walk_js_tree(node): - # Iterative DFS avoids Python's O(depth) generator-chain overhead. - # Recursive yield-from creates one generator frame per level — at 26+ - # levels deep each leaf's value had to propagate through 26 frames. - stack = [node] - while stack: - n = stack.pop() - yield n - stack.extend(reversed(n.children)) + # Pre-order walk with a TreeCursor, which skips building a ``.children`` + # list at every node. A cursor is bounded by the node it starts from, so + # goto_parent() fails at ``node`` and a subtree walk never leaves it. + cursor = node.walk() + while True: + yield cursor.node + if cursor.goto_first_child(): + continue + while not cursor.goto_next_sibling(): + if not cursor.goto_parent(): + return def _js_module_specifier(node, source: bytes) -> str | None: source_node = node.child_by_field_name("source") @@ -2073,6 +2076,18 @@ def _ts_walk_class_members(class_node, source: bytes, path: Path, class_nid: str _SymbolUseFact(path, class_nid, name, "references", ctx, m_line) ) +# Node types the whole-tree passes of _collect_js_symbol_resolution_facts +# filter on. +_JS_FACT_NODE_TYPES = frozenset({ + "export_statement", + "import_statement", + "lexical_declaration", + "class_declaration", + "abstract_class_declaration", + "interface_declaration", +}) + + def _collect_js_symbol_resolution_facts(paths: list[Path], facts: _SymbolResolutionFacts) -> None: js_paths = [ path for path in paths @@ -2081,7 +2096,7 @@ def _collect_js_symbol_resolution_facts(paths: list[Path], facts: _SymbolResolut if not js_paths: return - trees: dict[Path, tuple[bytes, object]] = {} + trees: dict[Path, tuple[bytes, object, list]] = {} for path in js_paths: resolved_path = _resolve_cached(path) @@ -2089,9 +2104,13 @@ def _collect_js_symbol_resolution_facts(paths: list[Path], facts: _SymbolResolut if parsed is None: continue source, root_node = parsed - trees[resolved_path] = parsed + indexed = [ + node for node in _walk_js_tree(root_node) + if node.type in _JS_FACT_NODE_TYPES + ] + trees[resolved_path] = (source, root_node, indexed) - for node in _walk_js_tree(root_node): + for node in indexed: if node.type == "export_statement": for name in _js_exported_declaration_names(node, source): facts.declarations.append( @@ -2129,7 +2148,7 @@ def _collect_js_symbol_resolution_facts(paths: list[Path], facts: _SymbolResolut ) ) - for node in _walk_js_tree(root_node): + for node in indexed: for alias, target in _js_lexical_aliases(node, source): facts.aliases.append( _SymbolAliasFact(path, alias, target, node.start_point[0] + 1) @@ -2140,9 +2159,9 @@ def _collect_js_symbol_resolution_facts(paths: list[Path], facts: _SymbolResolut parsed = trees.get(resolved_path) if parsed is None: continue - source, root_node = parsed + source, _root_node, indexed = parsed - for node in _walk_js_tree(root_node): + for node in indexed: if node.type != "export_statement": continue @@ -2236,7 +2255,7 @@ def _collect_js_symbol_resolution_facts(paths: list[Path], facts: _SymbolResolut parsed = trees.get(resolved_path) if parsed is None: continue - source, root_node = parsed + source, root_node, _indexed = parsed for source_id, body in _js_top_level_function_bodies(path, root_node, source): for node in _walk_js_tree(body): imported_name = _js_call_identifier(node, source) @@ -2258,9 +2277,9 @@ def _collect_js_symbol_resolution_facts(paths: list[Path], facts: _SymbolResolut parsed = trees.get(resolved_path) if parsed is None: continue - source, root_node = parsed + source, _root_node, indexed = parsed stem = _file_stem(path) - for node in _walk_js_tree(root_node): + for node in indexed: if node.type not in ( "class_declaration", "abstract_class_declaration", @@ -2310,8 +2329,7 @@ def _walk_python_tree(node): ancestor and re-propagated every node up the whole chain — ~25M frame resumptions on a 364-file corpus for ~2.8M actual nodes. An explicit stack yields each node exactly once in the identical preorder (children pushed - reversed so the first child pops first). Same rewrite, same reasoning as - ``_walk_js_tree`` above. + reversed so the first child pops first). """ stack = [node] while stack: diff --git a/tests/test_js_walk_order.py b/tests/test_js_walk_order.py new file mode 100644 index 000000000..6c04015c9 --- /dev/null +++ b/tests/test_js_walk_order.py @@ -0,0 +1,59 @@ +"""_walk_js_tree: pre-order over the whole tree, bounded by its start node. + +The JS/TS symbol-fact passes rely on walk order for fact order, and the +call-use pass walks function bodies as subtrees. Both must hold for the +TreeCursor walk exactly as they did for the ``.children`` stack walk. +""" + +from graphify.extractors.resolution import _parse_js_tree, _walk_js_tree + +SOURCE = """\ +import { a, b as c } from "./dep.js"; +export const x = a, y = c; +export class K extends Base { + m() { return helper(a, () => b()); } +} +function outer() { + const inner = () => { call1(); if (x) { call2(); } }; + return inner; +} +export { outer as default }; +""" + + +def _span(node): + return (node.type, node.start_byte, node.end_byte) + + +def _reference_preorder(node): + out = [node] + for child in node.children: + out.extend(_reference_preorder(child)) + return out + + +def test_walk_is_children_preorder(tmp_path): + path = tmp_path / "m.ts" + path.write_text(SOURCE, encoding="utf-8") + parsed = _parse_js_tree(path) + assert parsed is not None + _source, root = parsed + assert [_span(n) for n in _walk_js_tree(root)] == [ + _span(n) for n in _reference_preorder(root) + ] + + +def test_walk_from_any_node_stays_inside_its_subtree(tmp_path): + """Start a walk at every node, including ones with later siblings, as + _js_named_specifiers does with an import_statement or export_clause.""" + path = tmp_path / "m.ts" + path.write_text(SOURCE, encoding="utf-8") + parsed = _parse_js_tree(path) + assert parsed is not None + _source, root = parsed + starts = _reference_preorder(root) + assert any(n.next_sibling is not None and n.child_count for n in starts) + for start in starts: + assert [_span(n) for n in _walk_js_tree(start)] == [ + _span(n) for n in _reference_preorder(start) + ]