⚡️ Speed up function _node_contains_jsx by 685% 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.

---

📄 685% (6.85x) speedup for _node_contains_jsx in codeflash/languages/javascript/frameworks/react/discovery.py

⏱️ Runtime : 1.03 milliseconds → 131 microseconds (best of 16 runs)

📝 Explanation and details

The optimized code achieves a 684% speedup (from 1.03ms to 131μs) by replacing recursive calls with an iterative stack-based approach and using a frozenset for type checking.

Key Optimizations

1. Eliminated Recursive Overhead
The original code uses recursion with multiple call sites, creating significant overhead from:

• Function call stack frames for each node traversal
• Generator expression overhead in any(_node_contains_jsx(child) for child in node.children)
• Repeated function entry/exit costs

The optimized version uses an explicit stack with a while loop, eliminating all recursion overhead. This is particularly impactful for tree traversal where every node would trigger a new function call in the original version.

2. Frozenset Lookup for Type Checking
Moving JSX type checks into a frozenset (_JSX_TYPES) provides O(1) average-case lookup instead of tuple membership testing, which is O(n) linear scan. While the performance difference is small for 5 types, frozensets are optimized for membership testing and signal intent clearly.

3. Simplified Control Flow
The original code had special-case handling for return_statement nodes before the generic child traversal. The optimized version treats all nodes uniformly—checking the type first, then processing children—reducing branching and making the code path more predictable.

Performance Impact by Test Case

The optimization excels on breadth-heavy trees:

• test_large_breadth_many_children_with_one_jsx_at_end: 7541% faster (324μs → 4.25μs)
• test_large_breadth_many_children_no_jsx: 292% faster (338μs → 86.5μs)

For these cases, the original recursive approach with generator expressions incurs massive overhead when traversing 1000 siblings, while the iterative approach efficiently processes them with minimal per-node cost.

For simple cases with few nodes, the speedup is more modest (47-99% faster) due to lower recursion depth, though there is a slight regression (~1-25% slower) for the absolute simplest cases with a single node, likely due to the stack initialization overhead. This is an acceptable trade-off given the dramatic improvements on realistic tree structures.

The optimization preserves all behavior including the TypeError when children is None, ensuring backward compatibility.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from argparse import \
    Namespace  # a lightweight real class we can instantiate to mimic nodes

# imports
import pytest  # used for our unit tests
from codeflash.languages.javascript.frameworks.react.discovery import \
    _node_contains_jsx
# function to test
from tree_sitter import Node

def test_direct_jsx_types_recognized():
    # For each JSX-related node type, a node with that type should immediately return True.
    jsx_types = [
        "jsx_element",
        "jsx_self_closing_element",
        "jsx_fragment",
        "jsx_expression",
        "jsx_opening_element",
    ]
    for t in jsx_types:
        # Construct a real argparse.Namespace instance with the required attributes.
        node = Namespace(type=t, children=[])  # no children needed; the type itself should trigger True
        codeflash_output = _node_contains_jsx(node) # 1.76μs -> 2.35μs (25.1% slower)

def test_non_jsx_node_returns_false_when_no_children():
    # A node with a non-JSX type and no children should return False.
    node = Namespace(type="program", children=[])
    codeflash_output = _node_contains_jsx(node) # 1.62μs -> 1.10μs (47.3% faster)

def test_nested_child_with_jsx_detected():
    # A parent node that does not itself indicate JSX should return True if any descendant is JSX.
    child = Namespace(type="jsx_fragment", children=[])  # a descendant that is JSX
    parent = Namespace(type="program", children=[child])  # parent should detect descendant JSX
    codeflash_output = _node_contains_jsx(parent) # 2.21μs -> 1.13μs (95.6% faster)

def test_return_statement_checks_children_recursively_true():
    # A return_statement should search its children and return True if any child contains JSX.
    inner = Namespace(type="jsx_expression", children=[])  # JSX inside the returned expression
    returned = Namespace(type="return_statement", children=[Namespace(type="expression_statement", children=[inner])])
    codeflash_output = _node_contains_jsx(returned) # 2.50μs -> 1.32μs (89.3% faster)

def test_return_statement_checks_children_recursively_false():
    # A return_statement with children that do not contain JSX should return False.
    returned = Namespace(type="return_statement", children=[Namespace(type="expression_statement", children=[])])
    codeflash_output = _node_contains_jsx(returned) # 2.56μs -> 1.28μs (99.3% faster)

def test_none_children_raises_type_error():
    # If children is None (not iterable), Python should raise a TypeError when trying to iterate.
    node = Namespace(type="program", children=None)
    with pytest.raises(TypeError):
        _node_contains_jsx(node) # 3.50μs -> 3.54μs (1.13% slower)

def test_large_breadth_many_children_with_one_jsx_at_end():
    # Create a root with 1000 children where only the last child is JSX.
    many_children = [Namespace(type="expression", children=[]) for _ in range(999)]
    many_children.append(Namespace(type="jsx_element", children=[]))  # the single JSX node at the end
    root = Namespace(type="program", children=many_children)
    # The algorithm should iterate through siblings and eventually return True.
    codeflash_output = _node_contains_jsx(root) # 324μs -> 4.25μs (7541% faster)

def test_large_breadth_many_children_no_jsx():
    # A root with 1000 non-JSX children should return False.
    many_children = [Namespace(type="expression", children=[]) for _ in range(1000)]
    root = Namespace(type="program", children=many_children)
    codeflash_output = _node_contains_jsx(root) # 338μs -> 86.5μs (292% faster)

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

[claude]

PR Review Summary

Prek Checks

✅ Passed — ruff format required a minor fix in discovery.py (frozenset formatting). Auto-fixed and pushed in commit a947e43a.

Mypy

✅ No issues found in codeflash/languages/javascript/frameworks/react/discovery.py.

Code Review

✅ No critical issues found. This is a clean optimization of _node_contains_jsx:

• Replaces recursive tree traversal with iterative stack-based DFS
• Uses frozenset for O(1) JSX type membership checks
• Behavior is preserved (correctness verified by codeflash regression tests)

Note: _function_returns_jsx (line 171) has an unused source_bytes variable — this is a pre-existing issue from the parent branch, not introduced by this PR.

Test Coverage

File	PR Coverage	Main Coverage	Change
react/discovery.py	94%	🆕 new file	—
react/analyzer.py	100%	🆕 new file	—
react/benchmarking.py	100%	🆕 new file	—
react/context.py	99%	🆕 new file	—
react/profiler.py	13%	🆕 new file	⚠️
react/testgen.py	88%	🆕 new file	—
frameworks/detector.py	100%	🆕 new file	—
api/aiservice.py	20%	20%	—
languages/base.py	98%	98%	—
javascript/parse.py	51%	49%	+2%
javascript/support.py	70%	71%	-1%
models/function_types.py	100%	100%	—
result/critic.py	73%	70%	+3%
result/explanation.py	45%	46%	-1%

Summary: New React framework files have good coverage (88-100%) except profiler.py (13%) which lacks test coverage. No coverage regressions from this optimization PR specifically — the _node_contains_jsx function itself is covered by tests.

---

Last updated: 2026-02-20