Optimize _find_type_node

The optimized code achieves a **143% speedup** (from 169μs to 69.3μs) through three key performance improvements:

**1. Eliminated Dictionary Reconstruction Overhead (25.3% of original time)**
The original code rebuilt the `type_declarations` dictionary on every recursive call (329 times in the profiler). By hoisting it to a module-level constant `TYPE_DECLARATIONS`, this overhead is completely eliminated. The profiler shows ~340μs spent creating this dict repeatedly in the original version.

**2. Replaced Recursion with Iterative Stack-Based DFS**
The original code made 319 recursive calls (`_find_type_node(child, type_name, source_bytes)` at 638μs). Each recursive call incurs Python function call overhead including frame creation, argument passing, and local variable setup. The optimized version uses an explicit stack to traverse the tree iteratively, eliminating this overhead entirely. This is especially impactful in the deep tree test case, which shows **211% speedup** (146μs → 47μs).

**3. Direct Byte Comparison Instead of UTF-8 Decoding**
The original code decoded byte slices to strings 12 times (`source_bytes[...].decode("utf8")` at 13μs). The optimized version encodes `type_name` to bytes once at function entry (10.7μs for 10 calls), then performs direct byte-to-byte comparison without any decoding. This is particularly effective for multibyte UTF-8 names, as shown in the UTF-8 test case with **25.1% speedup** (2.24μs → 1.79μs).

**Performance Analysis by Test Case:**
- Simple cases show modest improvements (0-2μs) due to lower overhead
- Nested/deep traversals show dramatic gains (e.g., 211% on 300-depth tree) where recursion elimination matters most
- UTF-8 handling improves 25% by avoiding repeated decode operations
- A few edge cases show minor regression (1-6% slower) due to stack manipulation overhead, but these are dwarfed by gains in realistic workloads

The optimization preserves exact behavior including traversal order (reversed children maintain left-to-right DFS), return types, and edge case handling while delivering significant runtime improvements especially for deep syntax trees—a common scenario when parsing Java source code.

Author: codeflash-ai[bot]

codeflash/languages/java/context.py

@@ -19,6 +19,12 @@
 if TYPE_CHECKING:
     from tree_sitter import Node
 
+TYPE_DECLARATIONS = {
+    "class_declaration": "class",
+    "interface_declaration": "interface",
+    "enum_declaration": "enum",
+}
+
 logger = logging.getLogger(__name__)
 
 
@@ -253,23 +259,25 @@ def _find_type_node(node: Node, type_name: str, source_bytes: bytes) -> tuple[No
         Tuple of (node, type_kind) where type_kind is "class", "interface", or "enum".
 
     """
-    type_declarations = {
-        "class_declaration": "class",
-        "interface_declaration": "interface",
-        "enum_declaration": "enum",
-    }
-
-    if node.type in type_declarations:
-        name_node = node.child_by_field_name("name")
-        if name_node:
-            node_name = source_bytes[name_node.start_byte : name_node.end_byte].decode("utf8")
-            if node_name == type_name:
-                return node, type_declarations[node.type]
-
-    for child in node.children:
-        result, kind = _find_type_node(child, type_name, source_bytes)
-        if result:
-            return result, kind
+    # Encode the search name once to avoid repeated UTF-8 decodes of slices.
+    type_name_bytes = type_name.encode("utf8")
+
+    # Use an explicit stack for DFS to avoid recursion overhead.
+    stack: list[Node] = [node]
+
+    while stack:
+        current = stack.pop()
+        if current.type in TYPE_DECLARATIONS:
+            name_node = current.child_by_field_name("name")
+            if name_node:
+                # Compare bytes directly to avoid decoding the slice to str.
+                if source_bytes[name_node.start_byte : name_node.end_byte] == type_name_bytes:
+                    return current, TYPE_DECLARATIONS[current.type]
+
+        # Push children in reverse order so that the leftmost child is processed first,
+        # preserving the original recursive traversal order.
+        for child in reversed(current.children):
+            stack.append(child)
 
     return None, ""