Optimize JavaAssertTransformer._infer_return_type

Runtime improvement (primary): the optimized version cuts the measured wall-clock time from ~11.9 ms to ~5.23 ms (≈127% speedup). Most of the previous time was spent parsing the entire argument list for JUnit value assertions; the profiler shows _split_top_level_args accounted for the dominant portion of runtime.

What changed (specific optimizations):
- Introduced _extract_first_arg that scans args_str once and stops as soon as the first top-level comma is encountered instead of calling _split_top_level_args to produce the full list.
- The new routine keeps parsing state inline (depth, in_string, escape handling) and builds only the first-argument string (one small list buffer) rather than accumulating all arguments into a list of substrings.
- Early-trimming and early-return avoid unnecessary work when the first argument is empty or when there are no commas.

Why this is faster (mechanics):
- Less work: in common cases we only need the first top-level argument to infer the expected type. Splitting all top-level arguments does O(n) work and allocates O(m) substrings for the entire argument list; extracting only the first arg is usually much cheaper (O(k) where k is length up to first top-level comma).
- Fewer allocations: avoids creating many intermediate strings and list entries, which reduces Python object overhead and GC pressure.
- Better branch locality: the loop exits earlier in the typical case (simple literals), so average time per call drops significantly — this shows up strongly in the large-loop and many-arg tests.

Behavioral impact and trade-offs:
- Semantics are preserved for the intended use: the function only needs the first argument to infer the return type, so replacing a full-split with a single-arg extractor keeps correctness for all existing tests.
- Microbenchmarks for very trivial cases (e.g., assertTrue/assertFalse) show tiny per-call regressions (a few tens of ns) in some test samples; this is a reasonable trade-off for the substantial end-to-end runtime improvement, especially since the optimized code targets the hot path (value-assertion type inference) where gains are largest.

When this helps most:
- Calls with long argument lists or many nested/comma-containing constructs (nested generics, long sequences of arguments) — see the huge improvements in tests like large number of args and nested generics.
- Hot loops and repeated inference (many_inferences_loop_stress, repeated_inference) — fewer allocations and earlier exits compound into large throughput gains.

In short: the optimization reduces unnecessary parsing and allocations by only extracting what is required (the first top-level argument), which directly reduced CPU time and memory churn and produced the measured ~2x runtime improvement while keeping behavior for the intended use-cases.

Author: codeflash-ai[bot]

codeflash/languages/java/remove_asserts.py

@@ -941,15 +941,15 @@ def _infer_type_from_assertion_args(self, original_text: str, method: str) -> st
         elif args_str.endswith(")"):
             args_str = args_str[:-1]
 
-        # Split top-level args (respecting parens, strings, generics)
-        args = self._split_top_level_args(args_str)
-        if not args:
+        # Fast-path: only extract the first top-level argument instead of splitting all arguments.
+        first_arg = self._extract_first_arg(args_str)
+        if not first_arg:
             return "Object"
 
         # assertEquals has (expected, actual) or (expected, actual, message/delta)
         # Some overloads have (message, expected, actual) in JUnit 4 but JUnit 5 uses (expected, actual[, message])
         # Try the first argument as the expected value
-        expected = args[0].strip()
+        expected = first_arg.strip()
 
         return self._type_from_literal(expected)
 
@@ -1108,6 +1108,58 @@ def _generate_exception_replacement(self, assertion: AssertionMatch) -> str:
         # Fallback: comment out the assertion
         return f"{ws}// Removed assertThrows: could not extract callable"
 
+    def _extract_first_arg(self, args_str: str) -> str | None:
+        """Extract the first top-level argument from args_str.
+
+        This is a lightweight alternative to splitting all top-level arguments;
+        it stops at the first top-level comma, respects nested delimiters and strings,
+        and avoids constructing the full argument list for better performance.
+        """
+        n = len(args_str)
+        i = 0
+
+        # skip leading whitespace
+        while i < n and args_str[i].isspace():
+            i += 1
+        if i >= n:
+            return None
+
+        depth = 0
+        in_string = False
+        string_char = ""
+        cur: list[str] = []
+
+        while i < n:
+            ch = args_str[i]
+
+            if in_string:
+                cur.append(ch)
+                if ch == "\\" and i + 1 < n:
+                    i += 1
+                    cur.append(args_str[i])
+                elif ch == string_char:
+                    in_string = False
+            elif ch in ('"', "'"):
+                in_string = True
+                string_char = ch
+                cur.append(ch)
+            elif ch in ("(", "<", "[", "{"):
+                depth += 1
+                cur.append(ch)
+            elif ch in (")", ">", "]", "}"):
+                depth -= 1
+                cur.append(ch)
+            elif ch == "," and depth == 0:
+                break
+            else:
+                cur.append(ch)
+            i += 1
+
+        # Trim trailing whitespace from the extracted argument
+        if not cur:
+            return None
+        return "".join(cur).rstrip()
+
 
 def transform_java_assertions(source: str, function_name: str, qualified_name: str | None = None) -> str:
     """Transform Java test code by removing assertions and capturing function calls.