⚡️ Speed up function extract_react_context by 15% in PR #1561 (add/support_react)

codeflash/languages/javascript/frameworks/detector.py

@@ -10,7 +10,10 @@
 import logging
 from dataclasses import dataclass, field
 from functools import lru_cache
-from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pathlib import Path
 
 logger = logging.getLogger(__name__)
 

codeflash/languages/javascript/frameworks/react/context.py

@@ -7,17 +7,27 @@
 from __future__ import annotations
 
 import logging
+import re
 from dataclasses import dataclass, field
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from tree_sitter import Node
 
     from codeflash.languages.javascript.frameworks.react.analyzer import OptimizationOpportunity
     from codeflash.languages.javascript.frameworks.react.discovery import ReactComponentInfo
     from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer
 
+_BUILTIN_COMPONENTS = frozenset(("React.Fragment", "Fragment", "Suspense", "React.Suspense"))
+
+HOOK_PATTERN = re.compile(r"\b(use[A-Z]\w*)\s*\(")
+
+JSX_COMPONENT_RE = re.compile(r"<([A-Z][a-zA-Z0-9.]*)")
+
+CONTEXT_RE = re.compile(r"\buseContext\s*\(\s*(\w+)")
+
 logger = logging.getLogger(__name__)
 
 
@@ -108,42 +118,51 @@ def extract_react_context(
 
 def _extract_hook_usages(component_source: str) -> list[HookUsage]:
     """Parse hook calls and their dependency arrays from component source."""
-    import re
-
     hooks: list[HookUsage] = []
-    # Match useXxx( patterns
-    hook_pattern = re.compile(r"\b(use[A-Z]\w*)\s*\(")
 
-    for match in hook_pattern.finditer(component_source):
+    # Use precompiled HOOK_PATTERN
+    for match in HOOK_PATTERN.finditer(component_source):
         hook_name = match.group(1)
-        # Try to determine if there's a dependency array
-        # Look for ], [ pattern after the hook call (simplified heuristic)
-        rest_of_line = component_source[match.end() :]
         has_deps = False
         dep_count = 0
 
         # Simple heuristic: count brackets to find dependency array
         bracket_depth = 1
-        for i, char in enumerate(rest_of_line):
-            if char == "(":
+        pos = match.end()
+        source_len = len(component_source)
+
+        while pos < source_len:
+            # Find next '(' or ')' from current position
+            next_open = component_source.find("(", pos)
+            next_close = component_source.find(")", pos)
+
+            # Determine which paren comes next
+            if next_close == -1 and next_open == -1:
+                # No more parentheses; abort
+                break
+            if next_open != -1 and (next_open < next_close or next_close == -1):
+                # Found an opening paren before the next closing paren
                 bracket_depth += 1
-            elif char == ")":
-                bracket_depth -= 1
-                if bracket_depth == 0:
-                    # Check if the last argument before closing paren is an array
-                    preceding = rest_of_line[:i].rstrip()
-                    if preceding.endswith("]"):
-                        has_deps = True
-                        # Count items in the array (rough: count commas + 1 for non-empty)
-                        array_start = preceding.rfind("[")
-                        if array_start >= 0:
-                            array_content = preceding[array_start + 1 : -1].strip()
-                            if array_content:
-                                dep_count = array_content.count(",") + 1
-                            else:
-                                dep_count = 0  # empty deps []
-                                has_deps = True
-                    break
+                pos = next_open + 1
+                continue
+            # Otherwise, we found a closing paren next
+            pos = next_close + 1
+            bracket_depth -= 1
+            if bracket_depth == 0:
+                # Check if the last argument before closing paren is an array
+                preceding = component_source[match.end() : next_close].rstrip()
+                if preceding.endswith("]"):
+                    has_deps = True
+                    array_start = preceding.rfind("[")
+                    if array_start >= 0:
+                        # Extract content inside the brackets (exclude the closing bracket)
+                        array_content = preceding[array_start + 1 : -1].strip()
+                        if array_content:
+                            dep_count = array_content.count(",") + 1
+                        else:
+                            dep_count = 0  # empty deps []
+                            has_deps = True
+                break
 
         hooks.append(HookUsage(name=hook_name, has_dependency_array=has_deps, dependency_count=dep_count))
 
@@ -152,25 +171,16 @@ def _extract_hook_usages(component_source: str) -> list[HookUsage]:
 
 def _extract_child_components(component_source: str, analyzer: TreeSitterAnalyzer, full_source: str) -> list[str]:
     """Find child component names rendered in JSX."""
-    import re
-
-    # Match JSX tags that start with uppercase (React components)
-    jsx_component_re = re.compile(r"<([A-Z][a-zA-Z0-9.]*)")
-    children = set()
-    for match in jsx_component_re.finditer(component_source):
-        name = match.group(1)
-        # Skip React built-ins like React.Fragment
-        if name not in ("React.Fragment", "Fragment", "Suspense", "React.Suspense"):
-            children.add(name)
+    children = set(JSX_COMPONENT_RE.findall(component_source))
+    # Skip React built-ins like React.Fragment
+    if children:
+        children.difference_update(_BUILTIN_COMPONENTS)
     return sorted(children)
 
 
 def _extract_context_subscriptions(component_source: str) -> list[str]:
     """Find React context subscriptions via useContext calls."""
-    import re
-
-    context_re = re.compile(r"\buseContext\s*\(\s*(\w+)")
-    return [match.group(1) for match in context_re.finditer(component_source)]
+    return [match.group(1) for match in CONTEXT_RE.finditer(component_source)]
 
 
 def _find_type_definition(type_name: str, source: str, analyzer: TreeSitterAnalyzer) -> str | None:

codeflash/languages/javascript/frameworks/react/discovery.py

@@ -10,10 +10,11 @@
 import re
 from dataclasses import dataclass
 from enum import Enum
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from tree_sitter import Node
 
     from codeflash.languages.javascript.treesitter import FunctionNode, TreeSitterAnalyzer
@@ -191,11 +192,7 @@ def _node_contains_jsx(node: Node) -> bool:
             if _node_contains_jsx(child):
                 return True
 
-    for child in node.children:
-        if _node_contains_jsx(child):
-            return True
-
-    return False
+    return any(_node_contains_jsx(child) for child in node.children)
 
 
 def _extract_hooks_used(function_source: str) -> list[str]:

codeflash/languages/javascript/frameworks/react/profiler.py

@@ -12,10 +12,11 @@
 
 import logging
 import re
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from tree_sitter import Node
 
     from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer
@@ -76,9 +77,7 @@ def instrument_component_with_profiler(source: str, component_name: str, analyze
     result = _insert_after_imports(result, counter_code, analyzer)
 
     # Ensure React is imported
-    result = _ensure_react_import(result)
-
-    return result
+    return _ensure_react_import(result)
 
 
 def instrument_all_components_for_tracing(source: str, file_path: Path, analyzer: TreeSitterAnalyzer) -> str:
@@ -161,11 +160,12 @@ def walk(node: Node) -> None:
 
 def _contains_jsx(node: Node) -> bool:
     """Check if a tree-sitter node contains JSX elements."""
-    if node.type in ("jsx_element", "jsx_self_closing_element", "jsx_fragment"):
-        return True
-    for child in node.children:
-        if _contains_jsx(child):
+    stack = [node]
+    while stack:
+        node = stack.pop()
+        if node.type in ("jsx_element", "jsx_self_closing_element", "jsx_fragment"):
             return True
+        stack.extend(node.children)
     return False
 
 

codeflash/languages/javascript/treesitter_utils.py

@@ -1580,9 +1580,9 @@ def get_analyzer_for_file(file_path: Path) -> TreeSitterAnalyzer:
     """
     suffix = file_path.suffix.lower()
 
-    if suffix in (".ts",):
+    if suffix == ".ts":
         return TreeSitterAnalyzer(TreeSitterLanguage.TYPESCRIPT)
-    if suffix in (".tsx",):
+    if suffix == ".tsx":
         return TreeSitterAnalyzer(TreeSitterLanguage.TSX)
     # Default to JavaScript for .js, .jsx, .mjs, .cjs
     return TreeSitterAnalyzer(TreeSitterLanguage.JAVASCRIPT)