diff --git a/codeflash/languages/java/instrumentation.py b/codeflash/languages/java/instrumentation.py index 89408ee63..3c4495fa1 100644 --- a/codeflash/languages/java/instrumentation.py +++ b/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 + )