⚡️ Speed up method TreeSitterAnalyzer.has_return_statement by 81% in PR #1561 (add/support_react)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1561

If you approve this dependent PR, these changes will be merged into the original PR branch add/support_react.

This PR will be automatically closed if the original PR is merged.

---

📄 81% (0.81x) speedup for TreeSitterAnalyzer.has_return_statement in codeflash/languages/javascript/treesitter_utils.py

⏱️ Runtime : 965 microseconds → 534 microseconds (best of 12 runs)

📝 Explanation and details

This optimization achieves an 80% speedup (965μs → 534μs) by replacing recursive tree traversal with an iterative stack-based approach and removing unnecessary operations.

Key Optimizations

1. Iterative Stack-Based Traversal (Primary Speedup)

The original _node_has_return used recursive calls with Python's call stack, which is expensive due to:

• Function call overhead (frame creation/destruction)
• Parameter passing on each recursive call
• Generator expressions with any() creating iterator overhead

The optimized version uses an explicit stack (stack = [node]) to traverse the AST iteratively. This eliminates:

• ~2000+ recursive function calls in typical runs (line profiler shows 2037 hits on the recursive version)
• Generator allocation overhead from any(self._node_has_return(child) for child in node.children)

2. Removed Unused source.encode("utf8") Call

The original code encoded the source string to bytes but never used source_bytes. This operation cost ~47μs per call (0.6% of total time) and was completely unnecessary.

3. Performance Characteristics by Test Case

• Large bodies (1000+ nodes): ~195% faster — iterative approach shines with deep/wide trees by avoiding stack frame overhead
• Simple cases: 9-34% faster — reduced overhead even for shallow trees
• Trade-off cases: 15-25% slower on trivial 2-3 node trees — stack setup overhead marginally exceeds recursive call cost for extremely small inputs

The optimization is particularly effective for real-world JavaScript/TypeScript code which often contains large function bodies with many statements, where the 195% speedup on large bodies demonstrates the practical value. The minor regression on trivial 2-3 node cases is negligible since production code rarely has such tiny functions, and the overall 80% speedup confirms this optimization benefits typical workloads.

The iterative approach also provides more predictable performance and avoids potential stack overflow issues with extremely deep nesting, making it more robust for production use.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 45 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import pytest  # used for our unit tests
from codeflash.languages.javascript.treesitter_utils import TreeSitterAnalyzer

# Minimal helper classes to construct node-like trees that mimic the attributes/methods
# the analyzer expects from tree-sitter Node objects. These helpers are intentionally
# tiny and only expose the specific API used by TreeSitterAnalyzer:
#   - .type (string)
#   - .children (list of child nodes)
#   - .child_by_field_name(name) -> returns a child node for the given field name
#
# NOTE: The test harness for this kata allows constructing small, focused objects like
# these for unit tests. These helpers do not try to emulate tree-sitter fully — only
# the behavior needed by the analyzer.
class FakeNode:
    def __init__(self, type_, children=None, body=None):
        # node type string (e.g., "function_declaration", "return_statement", etc.)
        self.type = type_
        # list of child nodes for generic traversal
        self.children = list(children) if children else []
        # helper reference used by child_by_field_name("body")
        self._body = body

    def child_by_field_name(self, name):
        # The analyzer only asks for the "body" field; return the prepared body node.
        if name == "body":
            return self._body
        return None

# Minimal FunctionNode-like object expected by TreeSitterAnalyzer.has_return_statement.
# The real project likely defines a similar structure; for testing we provide an object
# with the required attributes:
#   - .node (the FakeNode representing the function AST node)
#   - .is_generator (bool)
#   - .is_arrow (bool)
class FakeFunctionNode:
    def __init__(self, node: FakeNode, is_generator: bool = False, is_arrow: bool = False):
        self.node = node
        self.is_generator = is_generator
        self.is_arrow = is_arrow

def test_generator_function_always_reports_return():
    # Generators implicitly return an iterator; the method should return True
    # without inspecting the node tree at all.
    analyzer = TreeSitterAnalyzer("javascript")  # language argument is accepted as string
    fn = FakeFunctionNode(node=None, is_generator=True, is_arrow=False)
    codeflash_output = analyzer.has_return_statement(fn, "function* gen() { /* no explicit return */ }") # 711ns -> 541ns (31.4% faster)

def test_arrow_function_with_expression_body_has_implicit_return():
    # Arrow functions whose body is an expression (not a statement_block) implicitly return.
    # Create a function node whose body is a non-statement node (e.g., an identifier).
    body_node = FakeNode(type_="identifier", children=[])  # expression body
    arrow_node = FakeNode(type_="arrow_function", body=body_node)
    fn = FakeFunctionNode(node=arrow_node, is_generator=False, is_arrow=True)

    analyzer = TreeSitterAnalyzer("javascript")
    # Source string is irrelevant for this branch, but supply a realistic one.
    codeflash_output = analyzer.has_return_statement(fn, "const f = x => x + 1;") # 1.70μs -> 1.40μs (21.5% faster)

def test_arrow_function_with_block_body_relies_on_node_traversal_for_return():
    # Arrow function where the body is a statement_block should be inspected for
    # explicit return statements.
    # Build body: statement_block -> children include a return_statement node.
    return_node = FakeNode(type_="return_statement", children=[])
    body_block = FakeNode(type_="statement_block", children=[return_node])
    arrow_node = FakeNode(type_="arrow_function", body=body_block)
    fn = FakeFunctionNode(node=arrow_node, is_generator=False, is_arrow=True)

    analyzer = TreeSitterAnalyzer("javascript")
    codeflash_output = analyzer.has_return_statement(fn, "const f = () => { return 42; }") # 2.40μs -> 3.17μs (24.0% slower)

def test_arrow_function_with_block_body_without_return_is_false():
    # Arrow function with a statement block but no return inside should be False.
    body_block = FakeNode(type_="statement_block", children=[FakeNode(type_="expression_statement")])
    arrow_node = FakeNode(type_="arrow_function", body=body_block)
    fn = FakeFunctionNode(node=arrow_node, is_generator=False, is_arrow=True)

    analyzer = TreeSitterAnalyzer("javascript")
    codeflash_output = analyzer.has_return_statement(fn, "const f = () => { 1 + 2; }") # 3.21μs -> 2.73μs (17.2% faster)

def test_regular_function_with_top_level_return_is_true():
    # A normal (non-arrow, non-generator) function whose body contains a top-level
    # return statement should be detected.
    return_node = FakeNode(type_="return_statement")
    body_block = FakeNode(type_="statement_block", children=[return_node])
    func_node = FakeNode(type_="function_declaration", body=body_block)
    fn = FakeFunctionNode(node=func_node, is_generator=False, is_arrow=False)

    analyzer = TreeSitterAnalyzer("javascript")
    codeflash_output = analyzer.has_return_statement(fn, "function f(){ return 1; }") # 1.68μs -> 1.97μs (14.7% slower)

def test_regular_function_without_return_is_false():
    # No return at any level inside the function should yield False.
    body_block = FakeNode(type_="statement_block", children=[FakeNode(type_="expression_statement")])
    func_node = FakeNode(type_="function_declaration", body=body_block)
    fn = FakeFunctionNode(node=func_node, is_generator=False, is_arrow=False)

    analyzer = TreeSitterAnalyzer("javascript")
    codeflash_output = analyzer.has_return_statement(fn, "function f(){ 1 + 2; }") # 2.52μs -> 2.29μs (9.63% faster)

def test_nested_function_return_counts_as_return_for_outer_function():
    # According to the current implementation, if the outer function's traversal reaches
    # a nested function node which itself contains a return statement, that return will
    # be discovered during recursion and yield True.
    #
    # Build:
    # outer function body -> contains a nested function declaration node
    # nested function body -> contains a return_statement
    nested_return = FakeNode(type_="return_statement")
    nested_body = FakeNode(type_="statement_block", children=[nested_return])
    nested_func = FakeNode(type_="function_declaration", body=nested_body)

    # Outer body has the nested function as a child
    outer_body = FakeNode(type_="statement_block", children=[nested_func])
    outer_func = FakeNode(type_="function_declaration", body=outer_body)
    fn = FakeFunctionNode(node=outer_func, is_generator=False, is_arrow=False)

    analyzer = TreeSitterAnalyzer("javascript")
    # Current implementation will find the nested return during recursion -> True
    codeflash_output = analyzer.has_return_statement(fn, "function outer(){ function inner(){ return 5; } }") # 1.91μs -> 2.54μs (24.5% slower)

def test_non_function_node_with_return_statement_detected():
    # _node_has_return is recursive and will find return statements in generic node trees.
    # Simulate calling has_return_statement with a function-like wrapper whose node is not
    # a function type but still contains a return statement somewhere in its children.
    # For the purposes of has_return_statement this scenario shouldn't normally happen,
    # but we ensure the underlying recursion behaves consistently.
    return_node = FakeNode(type_="return_statement")
    arbitrary_block = FakeNode(type_="some_block", children=[FakeNode(type_="expression"), return_node])
    fn = FakeFunctionNode(node=arbitrary_block, is_generator=False, is_arrow=False)

    analyzer = TreeSitterAnalyzer("javascript")
    codeflash_output = analyzer.has_return_statement(fn, "/* arbitrary code with return */") # 3.46μs -> 2.57μs (34.3% faster)

def test_large_body_with_return_near_end_performs_and_detects():
    # Build a large statement_block with 1000 children, the last of which is a return.
    many_children = [FakeNode(type_="expression_statement") for _ in range(999)]
    many_children.append(FakeNode(type_="return_statement"))  # return at the end
    big_body = FakeNode(type_="statement_block", children=many_children)
    big_func = FakeNode(type_="function_declaration", body=big_body)
    fn = FakeFunctionNode(node=big_func, is_generator=False, is_arrow=False)

    analyzer = TreeSitterAnalyzer("javascript")
    # This should be True and should complete promptly even for 1000 children.
    codeflash_output = analyzer.has_return_statement(fn, "function big(){ /* large body */ }") # 331μs -> 112μs (195% faster)

def test_many_sibling_functions_without_return_are_false_and_scale():
    # Create a function body that contains 1000 expression statements (no return).
    many_children = [FakeNode(type_="expression_statement") for _ in range(1000)]
    big_body = FakeNode(type_="statement_block", children=many_children)
    big_func = FakeNode(type_="function_declaration", body=big_body)
    fn = FakeFunctionNode(node=big_func, is_generator=False, is_arrow=False)

    analyzer = TreeSitterAnalyzer("javascript")
    # Ensure scanning 1000 children without any return results in False.
    codeflash_output = analyzer.has_return_statement(fn, "function big(){ /* lots of statements */ }") # 333μs -> 112μs (196% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr1561-2026-02-20T17.17.41 and push.

[codeflash-ai]

⚡️ Codeflash found optimizations for this PR

📄 14% (0.14x) speedup for TreeSitterAnalyzer._node_has_return in codeflash/languages/javascript/treesitter_utils.py

⏱️ Runtime : 395 microseconds → 347 microseconds (best of 13 runs)

A dependent PR with the suggested changes has been created. Please review:

• ⚡️ Speed up method TreeSitterAnalyzer._node_has_return by 14% in PR #1615 (codeflash/optimize-pr1561-2026-02-20T17.17.41) #1616

If you approve, it will be merged into this PR (branch codeflash/optimize-pr1561-2026-02-20T17.17.41).

[claude]

PR Review Summary

Prek Checks

Fixed 1 issue:

• D202 (blank-line-after-function) in codeflash/languages/javascript/treesitter_utils.py — auto-fixed and committed (4023b73)

mypy: No type errors found.

Code Review

No critical issues found. The optimization is straightforward and correct:

1. Removed unused source_bytes = source.encode("utf8") — dead code that was never referenced.
2. Converted recursive _node_has_return to iterative stack-based traversal — standard optimization that preserves the original traversal semantics. The iterative version visits the same nodes in the same order as the recursive version.

The logic for handling function-type nodes (traversing only their body children) is preserved identically in the iterative version.

Test Coverage

File	Stmts	Miss	Cover
codeflash/languages/javascript/treesitter_utils.py	—	—	0%*

* treesitter_utils.py is a new file from the parent add/support_react branch. No existing unit tests directly import this module (tests use treesitter.py which has a separate TreeSitterAnalyzer class). Coverage for this file should be addressed in the parent PR. The codeflash bot verified correctness via 45 generated regression tests (all passing).

Overall project coverage: 79% (unchanged by this PR since the modified file has no test coverage on either branch).

---

Last updated: 2026-02-20