From e1544e0ad4e9396055b072456886ef44159ceeca Mon Sep 17 00:00:00 2001
From: abhay-codes07 <abhaysingh0293@gmail.com>
Date: Fri, 11 Sep 2026 23:29:22 +0530
Subject: [PATCH] perf(resolution): walk the Python tree iteratively, not via
 recursive yield-from
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_walk_python_tree drove the symbol-resolution passes through a recursive
generator: one suspended frame per tree level, every node re-propagated up
the whole chain — ~25M frame resumptions on a 364-file corpus for ~2.8M
nodes. An explicit stack yields the identical preorder (node identity and
order verified) in ~30% less time, and no longer risks RecursionError on a
deeply nested tree. Same rewrite _walk_js_tree already carries.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JJbLfztSxm2tH5cJwBbe9q
---
 graphify/extractors/resolution.py        | 17 +++++-
 tests/test_walk_python_tree_iterative.py | 71 ++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_walk_python_tree_iterative.py

diff --git a/graphify/extractors/resolution.py b/graphify/extractors/resolution.py
index 983b81be2..0cd097647 100644
--- a/graphify/extractors/resolution.py
+++ b/graphify/extractors/resolution.py
@@ -1949,9 +1949,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)
+    """Preorder walk of a tree-sitter tree, iteratively.
+
+    The recursive ``yield from`` form built one suspended generator frame per
+    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.
+    """
+    stack = [node]
+    while stack:
+        current = stack.pop()
+        yield current
+        stack.extend(reversed(current.children))
 
 def _python_import_from_module(node, source: bytes) -> tuple[int, str] | None:
     level = 0
diff --git a/tests/test_walk_python_tree_iterative.py b/tests/test_walk_python_tree_iterative.py
new file mode 100644
index 000000000..426b379ba
--- /dev/null
+++ b/tests/test_walk_python_tree_iterative.py
@@ -0,0 +1,71 @@
+"""_walk_python_tree yields identical preorder, iteratively (#perf).
+
+The recursive ``yield from`` form built one generator frame per tree level and
+re-propagated every node up the chain — ~25M frame resumptions on a 364-file
+corpus for ~2.8M nodes. The iterative rewrite must yield the exact same nodes
+in the exact same preorder; only the frame overhead goes away.
+"""
+
+import pytest
+
+tspython = pytest.importorskip("tree_sitter_python")
+from tree_sitter import Language, Parser  # noqa: E402
+
+from graphify.extractors.resolution import _walk_python_tree
+
+
+def _reference_preorder(node):
+    """The pre-optimization recursive walk, verbatim."""
+    yield node
+    for child in node.children:
+        yield from _reference_preorder(child)
+
+
+def _parse(src: str):
+    parser = Parser(Language(tspython.language()))
+    return parser.parse(src.encode()).root_node
+
+
+SOURCES = [
+    "x = 1\n",
+    (
+        "import os\n"
+        "from a.b import c, d\n\n"
+        "class Foo(Base):\n"
+        "    attr: int = 0\n"
+        "    def method(self, x):\n"
+        "        def inner():\n"
+        "            return [i for i in range(x) if i % 2]\n"
+        "        return inner()\n\n"
+        "async def bar():\n"
+        "    async with ctx() as c:\n"
+        "        await c.run(lambda z: z + 1)\n"
+    ),
+    "",  # empty module
+]
+
+
+@pytest.mark.parametrize("src", SOURCES)
+def test_matches_recursive_preorder(src):
+    root = _parse(src)
+    got = list(_walk_python_tree(root))
+    expected = list(_reference_preorder(root))
+    # Identity, not just equality: same node objects, same order.
+    assert [id(n) for n in got] == [id(n) for n in expected]
+
+
+def test_visits_every_node_once():
+    root = _parse(SOURCES[1])
+    got = list(_walk_python_tree(root))
+    assert len(got) == len(set(id(n) for n in got))
+    assert got[0] is root  # preorder: root first
+
+
+def test_deeply_nested_does_not_recurse():
+    """A pathologically deep tree that would overflow the recursion limit for
+    the old form walks fine iteratively."""
+    depth = 2000
+    src = "x = " + "(" * depth + "1" + ")" * depth + "\n"
+    root = _parse(src)
+    count = sum(1 for _ in _walk_python_tree(root))
+    assert count > depth