From 3ec75ee68d73c6866479b086bf91e25e0edcee69 Mon Sep 17 00:00:00 2001
From: "codeflash-ai[bot]"
 <148906541+codeflash-ai[bot]@users.noreply.github.com>
Date: Fri, 20 Feb 2026 17:27:30 +0000
Subject: [PATCH] Optimize TreeSitterAnalyzer._node_has_return
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This optimization achieves a **14% runtime improvement** (395μs → 347μs) by eliminating repeated allocations and attribute lookups in a tree traversal algorithm.

## Key Optimizations

**1. Module-level frozenset for function types**
The original code recreated a tuple `("function_declaration", "function_expression", "arrow_function", "method_definition")` on every call. The optimized version moves this to a module-level `frozenset` (`_FUNC_LIKE_TYPES`), eliminating ~5.2% of the original function's overhead (5218ns per call). The frozenset also provides O(1) membership testing that's optimized at the C level.

**2. Localized `stack.extend` method**
By caching `stack.extend` as a local variable (`stack_extend = stack.extend`), the code avoids repeated attribute lookups on the list object. In tight loops with many iterations (the while loop executes 2521 times in the profiler), this removes thousands of attribute resolution steps. Line profiler shows this particularly benefits the two `stack_extend(reversed(children))` calls, which now execute ~10% faster (357305ns → 322655ns for the main case).

## Performance Characteristics

The optimization excels on workloads with:
- **Large/deep trees**: The `test_many_siblings_without_return_performance` (1000 nodes) shows **28.1% speedup** (86.1μs → 67.3μs)
- **Deep nesting**: The `test_deep_nesting_with_return_at_bottom` (1000 levels) shows **12.4% speedup** (198μs → 176μs)
- **Many function nodes**: The `test_many_functions_with_internal_returns_ignored_if_no_body` (500 functions) shows **10.1% speedup** (98.0μs → 89.0μs)

Small trees see modest overhead (1-8% slower) due to the additional local variable assignment, but this is outweighed by gains on realistic code analysis workloads where ASTs are typically large and deeply nested.

The optimizations preserve all original behavior—same traversal order, same return detection logic—making this a safe drop-in replacement that scales better with tree size.
---
 .../languages/javascript/treesitter_utils.py     | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/codeflash/languages/javascript/treesitter_utils.py b/codeflash/languages/javascript/treesitter_utils.py
index 3b9910343..c648391e4 100644
--- a/codeflash/languages/javascript/treesitter_utils.py
+++ b/codeflash/languages/javascript/treesitter_utils.py
@@ -18,6 +18,10 @@
 
     from tree_sitter import Node, Tree
 
+_FUNC_LIKE_TYPES: frozenset[str] = frozenset(
+    ("function_declaration", "function_expression", "arrow_function", "method_definition")
+)
+
 _FUNCTION_BODY_TYPES = frozenset(
     {
         "function_declaration",
@@ -1242,9 +1246,9 @@ def has_return_statement(self, function_node: FunctionNode, source: str) -> bool
 
     def _node_has_return(self, node: Node) -> bool:
         """Recursively check if a node contains a return statement."""
-        # Use an explicit stack to avoid recursion overhead while preserving traversal order.
-        func_types = ("function_declaration", "function_expression", "arrow_function", "method_definition")
         stack = [node]
+        # Localize extend for minor speedup and avoid repeated attribute lookups on the list object
+        stack_extend = stack.extend
         while stack:
             current = stack.pop()
             # Direct return statement check
@@ -1252,20 +1256,22 @@ def _node_has_return(self, node: Node) -> bool:
                 return True
 
             # If this node is a function-like node, only traverse its body children
-            if current.type in func_types:
+            if current.type in _FUNC_LIKE_TYPES:
                 body_node = current.child_by_field_name("body")
                 if body_node:
                     # Push children in reverse so they are processed in original order
                     children = body_node.children
                     if children:
-                        stack.extend(reversed(children))
+                        # reversed() returns an iterator; extend consumes it without creating an intermediate list
+                        stack_extend(reversed(children))
+                # Do not traverse other parts of the function node
                 # Do not traverse other parts of the function node
                 continue
 
             # General case: traverse all children
             children = current.children
             if children:
-                stack.extend(reversed(children))
+                stack_extend(reversed(children))
 
         return False