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

---

📄 30% (0.30x) speedup for _wrap_return_with_profiler in codeflash/languages/javascript/frameworks/react/profiler.py

⏱️ Runtime : 25.7 microseconds → 19.7 microseconds (best of 250 runs)

📝 Explanation and details

I combined the two child-iteration passes in _wrap_return_with_profiler into a single loop (while still preserving the original parenthesis-handling semantics), and replaced the generator-based any(...) in _contains_jsx with an explicit loop to avoid generator overhead. I also delay decoding bytes until after determining the final slice to avoid redundant encodings/decodings. These changes reduce CPU work and memory churn while preserving exact behavior.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from types import \
    SimpleNamespace  # lightweight real object to mimic node attributes

# imports
import pytest  # used for our unit tests
from codeflash.languages.javascript.frameworks.react.profiler import \
    _wrap_return_with_profiler

# We will create SimpleNamespace instances to mimic the minimal parts of a tree-sitter Node
# that _wrap_return_with_profiler uses at runtime: .type, .start_byte, .end_byte, .children.
# This keeps tests deterministic and focused on the function's string-manipulation logic.

def make_child(typ: str, start: int, end: int, children=None):
    """Helper to create a minimal node-like object with required attributes."""
    if children is None:
        children = []
    return SimpleNamespace(type=typ, start_byte=start, end_byte=end, children=children)

def test_basic_wrap_simple_jsx():
    # Basic case: return <div>Hi</div>;
    source = "function Comp() { return <div>Hi</div>; }"
    src_bytes = source.encode("utf-8")

    # Find byte indices for JSX element
    jsx_start = source.index("<div")
    jsx_end = source.index("</div>") + len("</div>")

    # Semicolon position (to ensure function picks correct jsx_end if a ';' child exists)
    semicolon_pos = source.index(";")

    # Construct a return_node with a 'return' token child, a jsx element child, and a ';' child.
    return_token = make_child("return", source.index("return"), source.index("return") + len("return"))
    jsx_child = make_child("jsx_element", jsx_start, jsx_end)
    semicolon_child = make_child(";", semicolon_pos, semicolon_pos + 1)
    return_node = make_child("return_statement", start=source.index("return"), end=semicolon_pos + 1,
                             children=[return_token, jsx_child, semicolon_child])

    # Call the function under test
    codeflash_output = _wrap_return_with_profiler(source, return_node, profiler_id="p1", safe_name="Comp_safe"); wrapped = codeflash_output # 3.53μs -> 3.24μs (8.99% faster)

def test_wrap_preserves_parentheses_inside_return():
    # Case: return ( <span>OK</span> );
    source = "function A(){return (\n  <span>OK</span>\n);}"
    # Find positions
    paren_open = source.index("(")
    paren_close = source.index(")")  # first closing paren
    jsx_start = source.index("<span")
    jsx_end = source.index("</span>") + len("</span>")

    # Build a parenthesized_expression node that spans the parens
    parenthesized = make_child("parenthesized_expression", paren_open, paren_close + 1,
                               children=[make_child("jsx_element", jsx_start, jsx_end)])
    return_token = make_child("return", source.index("return"), source.index("return") + len("return"))
    return_node = make_child("return_statement", start=source.index("return"), end=paren_close + 2,
                             children=[return_token, parenthesized])

    # Wrap it
    codeflash_output = _wrap_return_with_profiler(source, return_node, profiler_id="pid", safe_name="A_safe"); wrapped = codeflash_output # 5.36μs -> 3.58μs (49.9% faster)

def test_no_jsx_returns_source_unchanged():
    # If there is no JSX in the return statement, the source should be returned unchanged.
    source = "function B(){ return 42; }"
    # Create a 'return' node with only a numeric literal child (no jsx)
    literal_start = source.index("42")
    literal_end = literal_start + len("42")
    literal_child = make_child("number", literal_start, literal_end)
    return_token = make_child("return", source.index("return"), source.index("return") + len("return"))
    return_node = make_child("return_statement", start=source.index("return"), end=literal_end + 1,
                             children=[return_token, literal_child])

    codeflash_output = _wrap_return_with_profiler(source, return_node, profiler_id="x", safe_name="B_safe"); wrapped = codeflash_output # 2.41μs -> 1.90μs (26.8% faster)

def test_handles_trailing_semicolon_child_correctly():
    # Ensure that if the function detects a ';' child it uses its start_byte to set jsx_end
    source = "function C(){ return <b>Bold</b> ; }"
    jsx_start = source.index("<b>")
    jsx_end = source.index("</b>") + len("</b>")
    semicolon_pos = source.index(";")

    return_token = make_child("return", source.index("return"), source.index("return") + len("return"))
    jsx_child = make_child("jsx_element", jsx_start, jsx_end)
    semicolon_child = make_child(";", semicolon_pos, semicolon_pos + 1)
    return_node = make_child("return_statement", start=source.index("return"), end=semicolon_pos + 2,
                             children=[return_token, jsx_child, semicolon_child])

    codeflash_output = _wrap_return_with_profiler(source, return_node, profiler_id="s1", safe_name="Csafe"); wrapped = codeflash_output # 3.45μs -> 3.15μs (9.26% faster)

def test_large_scale_wrapping_performance_and_correctness():
    # Large-scale test: create a large JSX payload inside a return and ensure wrapping occurs correctly.
    inner_text = "X" * 1000  # 1000 characters to simulate a large JSX body
    source = "function Large(){ return <div>" + inner_text + "</div>; }"
    jsx_start = source.index("<div>")
    jsx_end = source.index("</div>") + len("</div>")
    semicolon_pos = source.index(";")

    # Create nodes mimicking parser structure
    return_token = make_child("return", source.index("return"), source.index("return") + len("return"))
    jsx_child = make_child("jsx_element", jsx_start, jsx_end)
    semicolon_child = make_child(";", semicolon_pos, semicolon_pos + 1)
    return_node = make_child("return_statement", start=source.index("return"), end=semicolon_pos + 1,
                             children=[return_token, jsx_child, semicolon_child])

    # Wrap and ensure it runs quickly and produces the expected wrapper
    codeflash_output = _wrap_return_with_profiler(source, return_node, profiler_id="big", safe_name="Large_safe"); wrapped = codeflash_output # 4.37μs -> 3.79μs (15.3% faster)

def test_multiple_children_and_nested_jsx_detection():
    # More complex children list: some non-jsx children before the jsx, and nested jsx child deeper in structure.
    source = "function Nested(){ return /*comment*/ (\n  <section><div>Deep</div></section>\n); }"
    # Compute indices of nested JSX; outer section starts at:
    section_start = source.index("<section>")
    section_end = source.index("</section>") + len("</section>")
    div_start = source.index("<div>")
    div_end = source.index("</div>") + len("</div>")

    # Create a nested child: an intermediate wrapper (non-jsx) whose child contains jsx_element
    inner_jsx_child = make_child("jsx_element", div_start, div_end)
    wrapper_child = make_child("parent", section_start, section_end, children=[inner_jsx_child])

    # Parenthesized expression that contains the wrapper_child
    paren_open = source.index("(")
    paren_close = source.index(")")
    parenthesized = make_child("parenthesized_expression", paren_open, paren_close + 1, children=[wrapper_child])

    return_token = make_child("return", source.index("return"), source.index("return") + len("return"))
    comment_child = make_child("comment", source.index("/*comment*/"), source.index("/*comment*/") + len("/*comment*/"))
    return_node = make_child("return_statement", start=source.index("return"), end=paren_close + 2,
                             children=[return_token, comment_child, parenthesized])

    codeflash_output = _wrap_return_with_profiler(source, return_node, profiler_id="n1", safe_name="Nested_safe"); wrapped = codeflash_output # 6.56μs -> 4.09μs (60.5% 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-20T06.29.37 and push.

[claude]

PR Review Summary

Prek Checks

✅ Fixed — 2 issues resolved and pushed:

• SIM110: Replaced explicit for loop with any() in _contains_jsx (ruff lint)
• Formatting: ruff format auto-fixed

All checks now pass.

Mypy

✅ No type errors found in changed file.

Code Review

✅ No critical issues found. This is a clean optimization that combines two child-iteration passes in _wrap_return_with_profiler into a single loop, eliminating the redundant second traversal of return_node.children. The jsx_content decoding is also deferred until after the parenthesis check, avoiding redundant encode/decode cycles. Semantics are preserved.

Minor note (non-blocking): has_parens (line 197/201) is assigned but never used — this is pre-existing dead code from the base branch, not introduced by this PR.

Test Coverage

File	Stmts	Miss	Cover
codeflash/languages/javascript/frameworks/react/profiler.py	127	110	13%

• ⚠️ profiler.py has low coverage (13%). The existing test file (tests/react/test_profiler.py) only covers generate_render_counter_code and marker parsing, not the core instrumentation functions (_wrap_return_with_profiler, instrument_component_with_profiler, etc.).
• This is a pre-existing coverage gap from the base branch add/support_react, not introduced by this optimization PR.
• The optimization itself is correctness-preserving (single-loop refactor with identical control flow).
• Overall project coverage: 79% (unchanged by this PR).

---

Last updated: 2026-02-20