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

---

📄 34% (0.34x) speedup for _extract_context_subscriptions in codeflash/languages/javascript/frameworks/react/context.py

⏱️ Runtime : 2.59 milliseconds → 1.93 milliseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 33% runtime speedup through two key changes:

1. Pre-compiled regex at module scope - Moving re.compile() from inside the function to module-level initialization (_CONTEXT_RE) eliminates redundant pattern compilation on every function call. The line profiler shows this overhead was consuming 15.1% of the original runtime (602μs across 54 calls). This is particularly impactful when the function is called repeatedly.

2. Using findall() instead of list comprehension with finditer() - Replacing [match.group(1) for match in context_re.finditer(component_source)] with _CONTEXT_RE.findall(component_source) leverages regex's built-in C implementation for direct group extraction. The original approach required:

• Iterating over match objects in Python
• Calling .group(1) on each match object
• Building a list through comprehension

This Python-level iteration consumed 83.4% of the original runtime. The findall() method performs the entire operation in optimized C code, returning the captured groups directly as a list.

Performance improvements across test cases:

• Empty/small inputs see dramatic gains (493-505% faster): The elimination of import and compilation overhead has outsized impact when actual matching is minimal
• Single matches improve 148-207% faster: Pre-compilation and direct extraction both contribute significantly
• Multiple matches improve 83-186% faster: Reduced per-match overhead from avoiding Python iteration
• Large-scale tests (1000 contexts) improve 39-71% faster: Benefits compound with volume, though relative gains decrease as actual matching work dominates

The optimization is universally beneficial across all test patterns - from single calls to large-scale parsing - making it valuable regardless of where this React context extraction function is called in the codebase.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 54 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.frameworks.react.context import \
    _extract_context_subscriptions

def test_simple_single_match():
    # Single straightforward useContext call with a bare identifier should be captured.
    src = "const value = useContext(MyContext);"
    # Expect the identifier 'MyContext' to be returned in a single-element list.
    codeflash_output = _extract_context_subscriptions(src) # 5.93μs -> 2.40μs (148% faster)

def test_multiple_matches_and_whitespace():
    # Multiple useContext calls with varying whitespace and placement should all be captured.
    src = """
    const a = useContext(FirstContext);
    function f() {
        const b = useContext(SecondContext)  // trailing comment
        const c = useContext  (   Third_Context_3   )
    }
    """
    # Expect the three identifiers in the order they appear.
    codeflash_output = _extract_context_subscriptions(src) # 7.93μs -> 4.05μs (96.1% faster)

def test_no_matches_for_empty_and_unrelated_text():
    # Empty source should produce an empty list.
    codeflash_output = _extract_context_subscriptions("") # 3.46μs -> 571ns (505% faster)
    # Source without any useContext occurrences should produce an empty list.
    codeflash_output = _extract_context_subscriptions("const x = 42; function foo() { return x; }") # 2.70μs -> 1.47μs (83.0% faster)
    # A token that contains 'useContext' as a suffix of a longer identifier should NOT match.
    # For example, 'myuseContext(Alpha)' should not match because there is no word boundary before 'useContext'.
    codeflash_output = _extract_context_subscriptions("const x = myuseContext(Alpha);") # 1.70μs -> 681ns (150% faster)

def test_quoted_argument_and_comment_block_block_matching():
    # When the argument is quoted (string literal), the regex doesn't capture it (quotes are not \w).
    src_quoted = "const ctxt = useContext('QuotedContext');"
    codeflash_output = _extract_context_subscriptions(src_quoted) # 5.11μs -> 2.13μs (139% faster)
    # A /* comment */ between '(' and the identifier prevents the \s* from matching the comment,
    # so it should not capture the identifier.
    src_commented = "const ctxt = useContext(/* comment */ CommentedContext);"
    codeflash_output = _extract_context_subscriptions(src_commented) # 2.31μs -> 1.19μs (94.0% faster)

def test_property_access_partial_capture_and_preceding_underscore_blocks():
    # If a property access is used, e.g., useContext(Some.Context),
    # the regex captures only the leading \w+ before the dot -> 'Some'.
    src_prop = "const p = useContext(Some.Context);"
    codeflash_output = _extract_context_subscriptions(src_prop) # 5.61μs -> 2.18μs (157% faster)
    # If useContext is preceded by an underscore (word character), there is no word boundary
    # before 'useContext' and it should not match.
    src_preceded = "const x = a_useContext(MyCtx);"
    codeflash_output = _extract_context_subscriptions(src_preceded) # 2.02μs -> 771ns (163% faster)

def test_identifier_character_variations_and_newlines():
    # Identifiers with digits and underscores are allowed by \w and should match.
    src_ids = "a=useContext(_ctx123); b=useContext(ctx_2);"
    codeflash_output = _extract_context_subscriptions(src_ids) # 6.23μs -> 2.46μs (154% faster)
    # Newlines and other whitespace between '(' and the identifier are allowed (\s*) and should match.
    src_newline = "const v = useContext(\n    NewlineCtx\n);"
    codeflash_output = _extract_context_subscriptions(src_newline) # 2.38μs -> 872ns (172% faster)

def test_large_scale_many_matches():
    # Construct a source string with 1000 sequential useContext calls with unique identifiers.
    n = 1000  # use the upper-bound size requested
    # Build lines like: const v0 = useContext(Context0); ... const v999 = useContext(Context999);
    lines = [f"const v{i} = useContext(Context{i});" for i in range(n)]
    big_src = "\n".join(lines)
    # Run the extraction and verify length and first/last few entries to keep assertions efficient.
    codeflash_output = _extract_context_subscriptions(big_src); result = codeflash_output # 441μs -> 316μs (39.4% faster)

def test_large_scale_mixed_valid_and_ignored_calls():
    # Build a large source mixing valid and intentionally unmatchable useContext invocations.
    n = 500
    parts = []
    for i in range(n):
        # valid call
        parts.append(f"let a{i} = useContext(ValidCtx{i});")
        # quoted argument (should be ignored)
        parts.append(f"let b{i} = useContext('Ignored{i}');")
        # comment between parentheses (should be ignored)
        parts.append(f"let c{i} = useContext(/*c*/IgnoredByComment{i});")
        # property access -> captures only leading part before dot
        parts.append(f"let d{i} = useContext(Parent{i}.Child);")
    big_src = "\n".join(parts)
    codeflash_output = _extract_context_subscriptions(big_src); res = codeflash_output # 996μs -> 844μs (18.0% faster)
    # Expect exactly n valid full matches + n partial matches from property access (Parent{i})
    # The quoted and commented ones should not contribute matches.
    expected_count = n + n  # ValidCtx + Parent parts
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import re

# imports
import pytest
from codeflash.languages.javascript.frameworks.react.context import \
    _extract_context_subscriptions

def test_single_context_subscription():
    """Test extraction of a single useContext call."""
    source = "const theme = useContext(ThemeContext);"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.96μs -> 2.26μs (163% faster)

def test_multiple_context_subscriptions():
    """Test extraction of multiple useContext calls."""
    source = """
    const theme = useContext(ThemeContext);
    const user = useContext(UserContext);
    const settings = useContext(SettingsContext);
    """
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 7.13μs -> 3.44μs (108% faster)

def test_usecontext_with_spaces():
    """Test useContext with varying amounts of whitespace."""
    source = "useContext  (  MyContext  )"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.16μs -> 1.80μs (186% faster)

def test_usecontext_with_newlines_and_tabs():
    """Test useContext with newlines and tabs in arguments."""
    source = "const ctx = useContext(\n\tAppContext\n);"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.37μs -> 1.98μs (171% faster)

def test_usecontext_in_function_body():
    """Test useContext calls within a function body."""
    source = """
    function MyComponent() {
        const auth = useContext(AuthContext);
        return <div>{auth.user}</div>;
    }
    """
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 6.67μs -> 3.16μs (111% faster)

def test_usecontext_with_underscore_context_name():
    """Test useContext with context name containing underscores."""
    source = "useContext(My_Theme_Context)"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.19μs -> 1.72μs (201% faster)

def test_usecontext_with_numeric_context_name():
    """Test useContext with context name containing numbers."""
    source = "useContext(Context123)"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.10μs -> 1.66μs (207% faster)

def test_usecontext_in_jsx():
    """Test useContext within JSX element."""
    source = """
    function App() {
        const theme = useContext(ThemeContext);
        return <Component theme={theme} />;
    }
    """
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 6.69μs -> 3.15μs (113% faster)

def test_usecontext_in_hook():
    """Test useContext as part of a hook definition."""
    source = """
    export function useTheme() {
        return useContext(ThemeContext);
    }
    """
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 6.13μs -> 2.50μs (145% faster)

def test_duplicate_context_subscriptions():
    """Test that duplicate context subscriptions are both returned."""
    source = """
    const theme1 = useContext(ThemeContext);
    const theme2 = useContext(ThemeContext);
    """
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 6.31μs -> 2.71μs (132% faster)

def test_usecontext_order_preserved():
    """Test that the order of context subscriptions is preserved."""
    source = """
    useContext(FirstContext);
    useContext(SecondContext);
    useContext(ThirdContext);
    """
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 6.39μs -> 2.56μs (150% faster)

def test_empty_string():
    """Test with empty source string."""
    source = ""
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 3.33μs -> 561ns (493% faster)

def test_no_usecontext_calls():
    """Test source with no useContext calls."""
    source = "const x = useState(0);"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 4.23μs -> 1.41μs (199% faster)

def test_usecontext_with_no_arguments():
    """Test useContext with no arguments (invalid but should not match)."""
    source = "const x = useContext();"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 4.67μs -> 1.80μs (159% faster)

def test_usecontext_not_word_boundary():
    """Test that UseContext without word boundary is not matched."""
    source = "const myuseContext = useContext(RealContext);"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.54μs -> 2.16μs (156% faster)

def test_usecontext_with_comment():
    """Test useContext calls with comments."""
    source = """
    // useContext(CommentedContext)
    const theme = useContext(RealContext);
    """
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 6.43μs -> 2.62μs (145% faster)

def test_usecontext_in_string():
    """Test useContext within a string literal."""
    source = 'const x = "useContext(StringContext)";'
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.38μs -> 2.01μs (167% faster)

def test_single_char_context_name():
    """Test context name with single character."""
    source = "useContext(A)"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.16μs -> 1.71μs (201% faster)

def test_very_long_context_name():
    """Test with very long context name."""
    long_name = "A" * 1000
    source = f"useContext({long_name})"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 7.77μs -> 4.24μs (83.5% faster)

def test_context_name_starting_with_number_invalid():
    """Test that context names starting with numbers are not matched."""
    source = "useContext(123Context)"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.11μs -> 1.74μs (193% faster)

def test_usecontext_with_special_chars_nearby():
    """Test useContext surrounded by special characters."""
    source = "(useContext(MyContext))[0]"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.06μs -> 1.83μs (176% faster)

def test_multiple_usecontext_same_line():
    """Test multiple useContext calls on the same line."""
    source = "useContext(A), useContext(B), useContext(C)"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 6.34μs -> 2.48μs (156% faster)

def test_nested_parentheses():
    """Test useContext with extra parentheses around argument."""
    source = "useContext((MyContext))"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 4.76μs -> 1.79μs (165% faster)

def test_usecontext_with_trailing_comma():
    """Test useContext with trailing comma in arguments."""
    source = "useContext(MyContext,)"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.31μs -> 1.84μs (188% faster)

def test_whitespace_only_string():
    """Test with whitespace-only source."""
    source = "   \n\t\n   "
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 3.27μs -> 541ns (504% faster)

def test_context_subscriptions_case_sensitive():
    """Test that context names are case-sensitive."""
    source = "useContext(mycontext)"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.25μs -> 1.88μs (179% faster)
    # Different from "MyContext"
    source2 = "useContext(MyContext)"
    codeflash_output = _extract_context_subscriptions(source2); result2 = codeflash_output # 2.29μs -> 741ns (210% faster)

def test_usecontext_followed_by_method_call():
    """Test useContext result being used immediately."""
    source = "useContext(ConfigContext).getValue()"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.19μs -> 1.95μs (166% faster)

def test_usecontext_in_conditional():
    """Test useContext within conditional logic."""
    source = """
    if (condition) {
        const theme = useContext(ThemeContext);
    }
    """
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.80μs -> 2.67μs (117% faster)

def test_usecontext_in_loop():
    """Test useContext within a loop (unusual pattern but valid syntax)."""
    source = """
    for (let i = 0; i < 3; i++) {
        const ctx = useContext(Context);
    }
    """
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 6.34μs -> 2.85μs (122% faster)

def test_usecontext_with_mixed_whitespace_types():
    """Test useContext with mixed whitespace (spaces, tabs, newlines)."""
    source = "useContext  \t  (  \n  MixedContext  \n  )"
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 5.23μs -> 1.80μs (190% faster)

def test_many_context_subscriptions():
    """Test extraction of 100 context subscriptions."""
    # Generate source code with 100 useContext calls
    contexts = [f"Context{i}" for i in range(100)]
    source = ";\n".join([f"useContext({ctx})" for ctx in contexts])
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 37.8μs -> 22.0μs (71.5% faster)

def test_very_large_source_file():
    """Test with a large source file containing many lines."""
    # Create a large source file with 500 lines
    lines = []
    context_count = 0
    expected_contexts = []
    for i in range(500):
        if i % 5 == 0:  # Insert useContext every 5 lines
            ctx_name = f"Context{context_count}"
            lines.append(f"const ctx{context_count} = useContext({ctx_name});")
            expected_contexts.append(ctx_name)
            context_count += 1
        else:
            lines.append(f"// Line {i}")
    source = "\n".join(lines)
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 98.5μs -> 82.4μs (19.6% faster)

def test_deeply_nested_component_tree():
    """Test with source representing deeply nested components."""
    # Simulate multiple nested components, each using different contexts
    source = ""
    expected = []
    for i in range(50):
        source += f"""
        function Component{i}() {{
            const ctx{i} = useContext(Context{i});
            return <div>{i}</div>;
        }}
        """
        expected.append(f"Context{i}")
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 77.6μs -> 67.7μs (14.7% faster)

def test_large_number_same_context():
    """Test with many subscriptions to the same context."""
    # Create source with 100 subscriptions to the same context
    source = ";\n".join(["useContext(SharedContext)" for _ in range(100)])
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 38.7μs -> 23.1μs (67.4% faster)

def test_alternating_different_contexts():
    """Test with alternating different context subscriptions."""
    source = ""
    expected = []
    for i in range(200):
        ctx = f"Context{i % 10}"  # 10 different contexts, alternating
        source += f"useContext({ctx});\n"
        expected.append(ctx)
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 68.4μs -> 40.4μs (69.3% faster)

def test_performance_with_long_context_names():
    """Test performance with very long context names."""
    long_name = "VeryLongContextNameWithManyCharacters" * 10
    source = f"useContext({long_name})" * 50
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 67.0μs -> 57.5μs (16.5% faster)

def test_mixed_realistic_component_code():
    """Test with realistic React component code containing various patterns."""
    source = """
    import React, { useContext, useState, useEffect } from 'react';
    
    function ComplexComponent() {
        const theme = useContext(ThemeContext);
        const [count, setCount] = useState(0);
        const user = useContext(UserContext);
        const auth = useContext(AuthContext);
        
        useEffect(() => {
            const config = useContext(ConfigContext);
            console.log(config);
        }, []);
        
        function nestedFunction() {
            const locale = useContext(LocaleContext);
            return locale;
        }
        
        return (
            <div>
                {/* useContext(CommentedContext) */}
                <ThemeProvider value={theme}>
                    <UserDisplay user={user} />
                </ThemeProvider>
            </div>
        );
    }
    
    export default ComplexComponent;
    """
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 16.3μs -> 11.6μs (41.1% faster)
    # Should extract all useContext calls (including the one in comment and useEffect)
    expected = ["ThemeContext", "UserContext", "AuthContext", "ConfigContext", "LocaleContext"]

def test_1000_random_contexts():
    """Test with 1000 random context subscriptions."""
    source_parts = []
    expected = []
    for i in range(1000):
        ctx = f"RandomContext{i}"
        source_parts.append(f"useContext({ctx})")
        expected.append(ctx)
    source = ";\n".join(source_parts)
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 376μs -> 248μs (51.2% faster)

def test_performance_complex_regex_matching():
    """Test performance with complex source requiring extensive regex matching."""
    # Create source with many patterns similar to useContext but not matching
    source = ""
    expected = []
    for i in range(100):
        source += f"""
        MyUseContext(Fake{i});
        useContext  (Real{i});
        notUseContext(Fake{i});
        use_context(Fake{i});
        """
        expected.append(f"Real{i}")
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 145μs -> 129μs (12.5% faster)

def test_all_printable_ascii_in_context_names():
    """Test with various alphanumeric patterns in context names."""
    contexts = []
    source_parts = []
    # Test contexts with different valid identifier patterns
    valid_names = [
        "SimpleContext",
        "Context_With_Underscores",
        "Context123",
        "_PrivateContext",
        "UPPERCASE_CONTEXT",
        "camelCaseContext",
        "PascalCaseContext",
    ]
    for name in valid_names:
        contexts.append(name)
        source_parts.append(f"useContext({name})")
    source = ";\n".join(source_parts)
    codeflash_output = _extract_context_subscriptions(source); result = codeflash_output # 8.32μs -> 3.91μs (113% 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-20T03.27.09 and push.

[claude]

PR Review Summary

Prek Checks

• I001 (unsorted-imports): Auto-fixed in context.py — import re moved to correct alphabetical position. Committed and pushed.
• Remaining (9 unfixable): All require --unsafe-fixes and are pre-existing in the parent branch (add/support_react), not introduced by this PR:
  • TC003 (Path import in type-checking block) in context.py, detector.py, discovery.py, profiler.py
  • SIM110 (use any()) in discovery.py, profiler.py
  • RET504 (unnecessary assignment) in profiler.py
  • FURB171 (single-item container) in treesitter_utils.py
• mypy: No issues found in context.py.

Code Review

No critical issues found. The optimization is correct and behavior-preserving:

1. Pre-compiles regex _CONTEXT_RE at module scope instead of inside the function — avoids redundant re.compile() on every call.
2. Replaces [match.group(1) for match in re.finditer(...)] with re.findall(...) — findall() returns captured groups directly when the pattern has one group, producing identical results with less Python-level iteration.

Test Coverage

File	Stmts	Miss	Coverage	Status
codeflash/languages/javascript/frameworks/react/context.py	—	—	0%	New file (parent branch)

• context.py is a new file introduced by the parent branch add/support_react and has 0% test coverage. No tests import or exercise this module.
• The optimization PR itself does not change behavior, so the lack of coverage is a pre-existing gap from the parent branch, not a regression introduced here.
• 8 unrelated test failures in tests/test_tracer.py (pre-existing, not related to this PR).

---

Last updated: 2026-02-20