diff --git a/graphify/extract.py b/graphify/extract.py index 1822c556f..0e69ee34b 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -79,6 +79,7 @@ _contained_in_package, _decldef_class_stem, _disambiguate_colliding_node_ids, + _SOURCE_KEY_CACHE, _find_workspace_root, _go_import_path_for_file, _is_type_like_definition, @@ -5319,6 +5320,9 @@ def extract( # Workspace package manifests/globs can change during watch or repeated extraction. _WORKSPACE_PACKAGE_CACHE.clear() _XAML_CSHARP_CLASS_CACHE.clear() + # Per-run memo for _source_key (resolution.py) — bounds it to one extract() + # call's corpus instead of growing unboundedly across a long-lived watch process. + _SOURCE_KEY_CACHE.clear() # Infer a common root for cache keys (use first diverging segment, not sum of all matches) try: diff --git a/graphify/extractors/resolution.py b/graphify/extractors/resolution.py index 8af9aac35..5165a7dcc 100644 --- a/graphify/extractors/resolution.py +++ b/graphify/extractors/resolution.py @@ -637,14 +637,29 @@ def _blank(s: str) -> str: out.append(_blank(src[pos:])) return "".join(out), lang +# (source_file, root) -> resolved key. _disambiguate_colliding_node_ids calls +# this once per node/edge/raw_call, and the same source_file repeats for every +# item extracted from one file, so the uncached path was re-doing a +# filesystem resolve() (stat + realpath) for every edge and raw_call in the +# batch. Pure function of its arguments for the lifetime of one extract() run +# (the corpus doesn't move mid-run), so memoizing is safe. +_SOURCE_KEY_CACHE: dict[tuple[str, Path], str] = {} + + def _source_key(source_file: str, root: Path) -> str: if not source_file: return "" + cache_key = (source_file, root) + cached = _SOURCE_KEY_CACHE.get(cache_key) + if cached is not None: + return cached source_path = Path(source_file) try: - return str(source_path.resolve().relative_to(root)) + result = str(source_path.resolve().relative_to(root)) except Exception: - return str(source_path) + result = str(source_path) + _SOURCE_KEY_CACHE[cache_key] = result + return result def _node_disambiguation_source_key(node: dict, root: Path) -> str: source_file = str(node.get("source_file", "")) @@ -1668,9 +1683,20 @@ def _parse_python_tree(path: Path): return None def _walk_python_tree(node): - yield node - for child in node.children: - yield from _walk_python_tree(child) + """Pre-order DFS over the tree (node, then its children left to right). + + Iterative with an explicit stack rather than ``yield from`` recursion: each + ``yield from`` frame re-forwards every descendant's value through the whole + call chain, so a recursive generator costs O(depth) per yielded node on a + tree this large. The stack walk yields each node exactly once, directly. + """ + stack = [node] + while stack: + n = stack.pop() + yield n + children = n.children + for i in range(len(children) - 1, -1, -1): + stack.append(children[i]) def _python_import_from_module(node, source: bytes) -> tuple[int, str] | None: level = 0 diff --git a/tests/test_hermes_verified_optimizations.py b/tests/test_hermes_verified_optimizations.py new file mode 100644 index 000000000..896651409 --- /dev/null +++ b/tests/test_hermes_verified_optimizations.py @@ -0,0 +1,71 @@ +"""Regression tests for the performance changes verified by Hermes QA. + +Covers: +- _walk_python_tree: iterative pre-order DFS must yield the identical node set + and order as the previous recursive generator (tree shape preserved). +- _source_key: memoization must not change the resolved key, and repeat calls + must hit the cache (same result). +""" +from pathlib import Path + +import graphify.extractors.resolution as res +from graphify.extractors.resolution import _walk_python_tree + + +def _simple_node(name="n", children=()): + class _N: + def __init__(self, name, children): + self.name = name + self.children = tuple(children) + + def __repr__(self): + return f"<{self.name}>" + + return _N(name, children) + + +def _walk_recursive(node): + """Reference: the pre-change recursive generator implementation.""" + yield node + for child in node.children: + yield from _walk_recursive(child) + + +def test_walk_python_tree_pre_order_identical_to_recursive(): + # Build a 3-level tree with a wide middle row. + leaves = [_simple_node(f"l{i}") for i in range(4)] + mid = [ + _simple_node("m0", (leaves[0], leaves[1])), + _simple_node("m1", (leaves[2],)), + _simple_node("m2", (leaves[3],)), + ] + root = _simple_node("root", mid) + + iterative = list(_walk_python_tree(root)) + recursive = list(_walk_recursive(root)) + + # Same total node count and same membership. + assert len(iterative) == len(recursive) == 8 + ids_iter = [n.name for n in iterative] + ids_rec = [n.name for n in recursive] + # Pre-order on the given tree: root, m0, l0, l1, m1, l2, m2, l3. + assert ids_iter == ["root", "m0", "l0", "l1", "m1", "l2", "m2", "l3"] + assert ids_iter == ids_rec + + +def test_source_key_memoized_and_stable(tmp_path): + """reset cache, two distinct source files resolve once; repeat calls identical.""" + res._SOURCE_KEY_CACHE.clear() + (tmp_path / "a.py").write_text("x = 1") + (tmp_path / "sub").mkdir() + (tmp_path / "sub" / "b.py").write_text("y = 2") + + k_a1 = res._source_key(str(tmp_path / "a.py"), tmp_path) + k_a2 = res._source_key(str(tmp_path / "a.py"), tmp_path) + k_b = res._source_key(str(tmp_path / "sub" / "b.py"), tmp_path) + + assert k_a1 == k_a2 == "a.py" + assert k_b == "sub/b.py" + # Memoization populated the cache. + assert res._SOURCE_KEY_CACHE.get((str(tmp_path / "a.py"), tmp_path)) == k_a1 + res._SOURCE_KEY_CACHE.clear() \ No newline at end of file