⚡️ Speed up function _add_behavior_instrumentation by 22% in PR #1199 (omni-java)

codeflash/languages/java/instrumentation.py

@@ -16,6 +16,7 @@
 
 import logging
 import re
+from functools import lru_cache
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -257,6 +258,10 @@ def _add_behavior_instrumentation(source: str, class_name: str, func_name: str)
     i = 0
     iteration_counter = 0
 
+
+    # Pre-compile the regex pattern once
+    method_call_pattern = _get_method_call_pattern(func_name)
+
     while i < len(lines):
         line = lines[i]
         stripped = line.strip()
@@ -299,11 +304,11 @@ def _add_behavior_instrumentation(source: str, class_name: str, func_name: str)
 
             while i < len(lines) and brace_depth > 0:
                 body_line = lines[i]
-                for ch in body_line:
-                    if ch == "{":
-                        brace_depth += 1
-                    elif ch == "}":
-                        brace_depth -= 1
+                # Count braces more efficiently using string methods
+                open_count = body_line.count('{')
+                close_count = body_line.count('}')
+                brace_depth += open_count - close_count
+
 
                 if brace_depth > 0:
                     body_lines.append(body_line)
@@ -318,17 +323,6 @@ def _add_behavior_instrumentation(source: str, class_name: str, func_name: str)
             call_counter = 0
             wrapped_body_lines = []
 
-            # Use regex to find method calls with the target function
-            # Pattern matches: receiver.funcName(args) where receiver can be:
-            # - identifier (counter, calc, etc.)
-            # - new ClassName()
-            # - new ClassName(args)
-            # - this
-            method_call_pattern = re.compile(
-                rf"((?:new\s+\w+\s*\([^)]*\)|[a-zA-Z_]\w*))\s*\.\s*({re.escape(func_name)})\s*\(([^)]*)\)",
-                re.MULTILINE
-            )
-
             for body_line in body_lines:
                 # Check if this line contains a call to the target function
                 if func_name in body_line and "(" in body_line:
@@ -726,3 +720,22 @@ def _add_import(source: str, import_statement: str) -> str:
 
     lines.insert(insert_idx, import_statement + "\n")
     return "".join(lines)
+
+
+
+@lru_cache(maxsize=128)
+def _get_method_call_pattern(func_name: str):
+    """Cache compiled regex patterns for method call matching."""
+    return re.compile(
+        rf"((?:new\s+\w+\s*\([^)]*\)|[a-zA-Z_]\w*))\s*\.\s*({re.escape(func_name)})\s*\(([^)]*)\)",
+        re.MULTILINE
+    )
+
+
+@lru_cache(maxsize=128)
+def _get_method_call_pattern(func_name: str):
+    """Cache compiled regex patterns for method call matching."""
+    return re.compile(
+        rf"((?:new\s+\w+\s*\([^)]*\)|[a-zA-Z_]\w*))\s*\.\s*({re.escape(func_name)})\s*\(([^)]*)\)",
+        re.MULTILINE
+    )