Optimize ensure_common_java_imports

Primary benefit — runtime improved from 64.3 ms to 10.7 ms (≈501% speedup). The optimized version was accepted because it materially reduces execution time, especially on the large and repeated-call workloads exercised in the tests.

What changed (specific optimizations)
- Substring pre-check before regex: the code first checks "if class_name not in code" using the fast C-level substring search, and only runs re.search(rf"\b{class_name}\b", ...) when the substring is present. This avoids dozens of expensive regex runs for names that aren’t even present.
- Batch import insertion: missing imports are collected into a list and inserted with a single _add_imports call that does one splitlines/join and one insertion of a block, instead of calling _add_import repeatedly.
- Minor local-variable rename (code = test_code) to avoid repeated attribute lookups and to make the flow clearer.

Why this yields the speedup
- Regex cost dominated the original function: the line profiler shows the re.search line consumed ~82% of the original function time. re.search with a \b word-boundary involves regex engine work and is substantially slower than a plain substring search.
- The substring test ("in") is implemented in C and is much cheaper; it quickly filters out most class names so the regex runs only for likely candidates. In the optimized profile the heavy regex line's relative cost dropped dramatically.
- Repeatedly calling _add_import caused repeated splitlines/join on the entire source for each added import (O(n) per insertion => O(k*n) when adding k imports). The new _add_imports builds the insertion block and performs a single split/join, giving roughly O(n + k) instead of O(k * n) behavior for that part.
- The profiler confirms these effects: total time for ensure_common_java_imports dropped from 0.105s to 0.027s; the number and cost of expensive operations (regex and split/join) fell accordingly.

Behavioral and compatibility notes
- Behavior is preserved: imports are still only added when needed, wildcard-package checks remain, and word-boundary checks are still performed (the regex is executed when necessary).
- No regressions in correctness were introduced by the changes; tests show identical behavior and faster runs.
- Memory and complexity trade-off: storing a small list of import statements is negligible; batching reduces overall CPU and memory churn.

When this helps most
- Large inputs or code with many references, and repeated calls (hot paths) benefit the most. The large-scale test shows a dramatic improvement (e.g., the 1000-iteration idempotent test went from ~31.5 ms to ~3.39 ms in one recorded case).
- Even small test cases became noticeably faster (microsecond-level improvements across the test suite), and empty inputs are sped up because cheap substring checks short-circuit further work.

Summary
- The optimization targets the two main costs: (1) many unnecessary regex searches and (2) repeated O(n) string recompositions when inserting imports. Replacing frequent regex invocations with a cheap substring pre-check and batching insertions cuts both CPU and memory work, producing the measured 501% runtime improvement without changing behavior.

Author: codeflash-ai[bot]

codeflash/languages/java/instrumentation.py

@@ -1277,16 +1277,25 @@ def remove_instrumentation(source: str) -> str:
 
 
 def ensure_common_java_imports(test_code: str) -> str:
+    imports_to_add: list[str] = []
+    code = test_code
+    # Fast path: avoid compiling regexes when class name substring isn't present
     for class_name, import_stmt in _COMMON_JAVA_IMPORTS.items():
-        if not re.search(rf"\b{class_name}\b", test_code):
+        if class_name not in code:
             continue
-        if import_stmt in test_code:
+        if import_stmt in code:
             continue
         package = import_stmt.split()[1].rsplit(".", 1)[0]
-        if f"import {package}.*;" in test_code:
+        if f"import {package}.*;" in code:
             continue
-        test_code = _add_import(test_code, import_stmt)
-    return test_code
+        # Only now do the (relatively expensive) regex check to ensure whole-word match
+        if not re.search(rf"\b{class_name}\b", code):
+            continue
+        imports_to_add.append(import_stmt)
+
+    if imports_to_add:
+        code = _add_imports(code, imports_to_add)
+    return code
 
 
 def instrument_generated_java_test(
@@ -1384,3 +1393,54 @@ def _add_import(source: str, import_statement: str) -> str:
 
     lines.insert(insert_idx, import_statement + "\n")
     return "".join(lines)
+
+
+
+def _add_imports(source: str, import_statements: list[str]) -> str:
+    """Add multiple import statements to the source.
+
+    This helper batches insertion of multiple imports at once to avoid repeated
+    split/join operations that would be performed by inserting each import individually.
+    """
+    lines = source.splitlines(keepends=True)
+    insert_idx = 0
+
+    # Find the last import or package statement
+    for i, line in enumerate(lines):
+        stripped = line.strip()
+        if stripped.startswith(("import ", "package ")):
+            insert_idx = i + 1
+        elif stripped and not stripped.startswith("//") and not stripped.startswith("/*"):
+            # First non-import, non-comment line
+            if insert_idx == 0:
+                insert_idx = i
+            break
+
+    block = "".join(stmt + "\n" for stmt in import_statements)
+    lines.insert(insert_idx, block)
+    return "".join(lines)
+
+
+def _add_imports(source: str, import_statements: list[str]) -> str:
+    """Add multiple import statements to the source.
+
+    This helper batches insertion of multiple imports at once to avoid repeated
+    split/join operations that would be performed by inserting each import individually.
+    """
+    lines = source.splitlines(keepends=True)
+    insert_idx = 0
+
+    # Find the last import or package statement
+    for i, line in enumerate(lines):
+        stripped = line.strip()
+        if stripped.startswith(("import ", "package ")):
+            insert_idx = i + 1
+        elif stripped and not stripped.startswith("//") and not stripped.startswith("/*"):
+            # First non-import, non-comment line
+            if insert_idx == 0:
+                insert_idx = i
+            break
+
+    block = "".join(stmt + "\n" for stmt in import_statements)
+    lines.insert(insert_idx, block)
+    return "".join(lines)