⚡️ Speed up method JavaAssertTransformer._infer_return_type by 127% in PR #1663 (fix/java-maven-test-execution-bugs)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1663

If you approve this dependent PR, these changes will be merged into the original PR branch fix/java-maven-test-execution-bugs.

This PR will be automatically closed if the original PR is merged.

---

📄 127% (1.27x) speedup for JavaAssertTransformer._infer_return_type in codeflash/languages/java/remove_asserts.py

⏱️ Runtime : 11.9 milliseconds → 5.23 milliseconds (best of 230 runs)

📝 Explanation and details

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.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 2060 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

from dataclasses import dataclass

# imports
import pytest  # used for our unit tests
from codeflash.languages.java.remove_asserts import (JUNIT5_VALUE_ASSERTIONS,
                                                     JavaAssertTransformer)

# NOTE:
# The code under test expects an "AssertionMatch"-like object with at least two attributes:
# - assertion_method: the assertion method name (e.g. "assertEquals")
# - original_text: the full assertion call text (e.g. 'assertEquals(1, x);')
# The repository does not expose such a class in the provided snippets, and the implementation
# only accesses attributes dynamically (no isinstance checks). For clarity and explicitness
# we define a minimal real dataclass here to hold those attributes. This is a real class
# (not a mock) and provides concrete instances for testing.
@dataclass
class AssertionMatch:
    assertion_method: str
    original_text: str

# Create a transformer instance to use in tests.
# Use a realistic function_name as required by the constructor.
_transformer = JavaAssertTransformer(function_name="dummyFunction")

def test_assert_true_and_false_return_boolean():
    # assertTrue -> "boolean"
    a_true = AssertionMatch(assertion_method="assertTrue", original_text="assertTrue(result);")
    codeflash_output = _transformer._infer_return_type(a_true) # 561ns -> 581ns (3.44% slower)
    # assertFalse -> "boolean"
    a_false = AssertionMatch(assertion_method="assertFalse", original_text="assertFalse(flag);")
    codeflash_output = _transformer._infer_return_type(a_false) # 281ns -> 320ns (12.2% slower)

def test_assert_null_and_notnull_return_object():
    # assertNull -> "Object"
    a_null = AssertionMatch(assertion_method="assertNull", original_text="assertNull(obj);")
    codeflash_output = _transformer._infer_return_type(a_null) # 631ns -> 601ns (4.99% faster)
    # assertNotNull -> "Object"
    a_notnull = AssertionMatch(assertion_method="assertNotNull", original_text="assertNotNull(obj);")
    codeflash_output = _transformer._infer_return_type(a_notnull) # 361ns -> 371ns (2.70% slower)

def test_simple_integer_literal_in_assert_equals():
    # assertEquals with a plain integer literal expected value -> "int"
    am = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals(42, calc());")
    codeflash_output = _transformer._infer_return_type(am) # 11.2μs -> 8.55μs (31.5% faster)

def test_negative_integer_and_long_literals():
    # negative int
    am_neg = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals(-7, f());")
    codeflash_output = _transformer._infer_return_type(am_neg) # 9.56μs -> 7.91μs (20.9% faster)
    # long with trailing L
    am_long = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals(1234567890123L, f());")
    codeflash_output = _transformer._infer_return_type(am_long) # 10.9μs -> 9.30μs (17.2% faster)

def test_float_and_double_literals_are_distinguished():
    # float with suffix f -> "float"
    am_float = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals(1.23f, f());")
    codeflash_output = _transformer._infer_return_type(am_float) # 8.87μs -> 7.33μs (20.9% faster)
    # float without decimal but with f suffix -> "float"
    am_float2 = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals(1f, f());")
    codeflash_output = _transformer._infer_return_type(am_float2) # 4.85μs -> 3.44μs (41.1% faster)
    # double with decimal (no f) -> "double"
    am_double = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals(3.1415, f());")
    codeflash_output = _transformer._infer_return_type(am_double) # 5.70μs -> 4.62μs (23.4% faster)
    # explicit double with d -> "double"
    am_double_d = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals(2d, f());")
    codeflash_output = _transformer._infer_return_type(am_double_d) # 3.95μs -> 3.00μs (31.8% faster)

def test_char_and_string_and_null_and_boolean_literals():
    # char literal -> "char"
    am_char = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals('a', x());")
    codeflash_output = _transformer._infer_return_type(am_char) # 9.40μs -> 7.80μs (20.4% faster)
    # escaped char -> "char" (e.g. newline escaped)
    am_escaped_char = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals('\\n', x());")
    codeflash_output = _transformer._infer_return_type(am_escaped_char) # 5.82μs -> 4.36μs (33.6% faster)
    # string literal -> "String"
    am_str = AssertionMatch(assertion_method="assertEquals", original_text='assertEquals("hello", f());')
    codeflash_output = _transformer._infer_return_type(am_str) # 5.42μs -> 4.25μs (27.6% faster)
    # boolean literal expected inside assertEquals -> "boolean"
    am_bool = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals(true, f());")
    codeflash_output = _transformer._infer_return_type(am_bool) # 3.60μs -> 2.50μs (43.7% faster)
    # null literal -> "Object"
    am_null = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals(null, f());")
    codeflash_output = _transformer._infer_return_type(am_null) # 3.32μs -> 2.22μs (49.1% faster)

def test_cast_expression_in_expected_becomes_cast_type():
    # cast like (byte)0 should infer "byte"
    am_cast = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals((byte)0, f());")
    codeflash_output = _transformer._infer_return_type(am_cast) # 13.1μs -> 10.8μs (21.2% faster)
    # another primitive cast
    am_short = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals((short)1, f());")
    codeflash_output = _transformer._infer_return_type(am_short) # 8.21μs -> 6.57μs (24.8% faster)

def test_no_parenthesis_returns_object():
    # If there's no "(" the method should fall back to "Object"
    am_no_paren = AssertionMatch(assertion_method="assertEquals", original_text="assertEquals")
    codeflash_output = _transformer._infer_return_type(am_no_paren) # 1.29μs -> 1.30μs (0.768% slower)

def test_assert_equals_with_message_first_returns_message_type():
    # Overloads that put a message first will cause this implementation to treat the message as expected.
    # If the message is a string literal, return "String".
    am_msg_first = AssertionMatch(assertion_method="assertEquals", original_text='assertEquals("oops", 5, actual);')
    # Implementation picks the first arg -> a string literal -> "String"
    codeflash_output = _transformer._infer_return_type(am_msg_first) # 11.5μs -> 7.93μs (44.4% faster)

def test_expected_with_inner_commas_and_braces_not_split_at_top_level():
    # The expected value contains braces and commas; top-level split must ignore those inner commas.
    # e.g. new int[]{1, 2, 3} should be kept as a single top-level arg.
    original = "assertEquals(new int[]{1, 2, 3}, compute());"
    am = AssertionMatch(assertion_method="assertEquals", original_text=original)
    # It's not a literal the inference recognizes, so we expect "Object".
    codeflash_output = _transformer._infer_return_type(am) # 15.9μs -> 12.3μs (29.2% faster)

def test_string_argument_with_commas_is_not_split():
    # If the expected is a string containing commas, the splitter must treat it as one arg.
    original = 'assertEquals("a, b, c", value);'
    am = AssertionMatch(assertion_method="assertEquals", original_text=original)
    codeflash_output = _transformer._infer_return_type(am) # 10.5μs -> 8.12μs (29.5% faster)

def test_trailing_whitespace_and_semicolon_variants():
    # Variations in trailing characters like ending ")" vs ");" vs " ); " should be handled.
    samples = [
        "assertEquals(10, x)",
        "assertEquals(10, x);",
        "assertEquals(10, x );",
        "assertEquals( 10 ,x);",
    ]
    for s in samples:
        am = AssertionMatch(assertion_method="assertEquals", original_text=s)
        codeflash_output = _transformer._infer_return_type(am) # 22.8μs -> 19.1μs (19.0% faster)

def test_many_inferences_loop_stress():
    # Build 1000 assertion texts alternating among types to check consistent behavior at scale.
    types_and_literals = [
        ("assertEquals", "1"),         # int
        ("assertEquals", "2L"),        # long
        ("assertEquals", "3.0"),       # double
        ("assertEquals", "4.0f"),      # float
        ("assertEquals", "'z'"),       # char
        ("assertEquals", '"s"'),       # string
        ("assertTrue", "dummy()"),     # boolean via method name
        ("assertNull", "dummy()"),     # Object via method name
    ]
    results_expected = {
        "1": "int",
        "2L": "long",
        "3.0": "double",
        "4.0f": "float",
        "'z'": "char",
        '"s"': "String",
    }

    # Run 1000 iterations, cycling through the patterns.
    for i in range(1000):
        method, lit = types_and_literals[i % len(types_and_literals)]
        # For methods that are JUNIT5 value assertions, build an assertEquals call using the literal.
        if method == "assertEquals":
            text = f"assertEquals({lit}, f());"
            am = AssertionMatch(assertion_method=method, original_text=text)
            # If literal is in our expected map, assert the expected type; else fallback Object.
            if lit in results_expected:
                codeflash_output = _transformer._infer_return_type(am)
            else:
                codeflash_output = _transformer._infer_return_type(am)
        else:
            # For other assertion method names, just rely on declared behavior
            am = AssertionMatch(assertion_method=method, original_text=f"{method}(x);")
            if method in ("assertTrue", "assertFalse"):
                codeflash_output = _transformer._infer_return_type(am)
            elif method in ("assertNull", "assertNotNull"):
                codeflash_output = _transformer._infer_return_type(am)

def test_split_with_nested_generics_and_parentheses():
    # Complex expected arg containing nested generics and parentheses should be recognized as one arg.
    original = "assertEquals(Collections.<String, List<Integer>>emptyMap(), calc());"
    am = AssertionMatch(assertion_method="assertEquals", original_text=original)
    # Not a literal we recognize -> "Object"
    codeflash_output = _transformer._infer_return_type(am) # 22.1μs -> 18.8μs (17.7% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from types import \
    SimpleNamespace  # lightweight real class used to hold attributes

# imports
import pytest  # used for our unit tests
from codeflash.languages.java.remove_asserts import (JUNIT5_VALUE_ASSERTIONS,
                                                     JavaAssertTransformer)

# function to test
# We will test JavaAssertTransformer._infer_return_type (and related internals it uses)
# by creating a real JavaAssertTransformer instance and passing minimal assertion-like
# objects (SimpleNamespace) that expose the attributes the function expects:
# - assertion_method (str)
# - original_text (str)
#
# NOTE: SimpleNamespace is a real class from the standard library and provides
# attribute storage. It is used here only to provide a concrete object with the
# required attributes for _infer_return_type to inspect.

def make_transformer() -> JavaAssertTransformer:
    # Construct a transformer with a sample function name. Analyzer is optional.
    return JavaAssertTransformer(function_name="foo")

def test_assert_true_and_false_return_boolean():
    # Create transformer instance
    t = make_transformer()
    # assertTrue should always produce "boolean"
    a_true = SimpleNamespace(assertion_method="assertTrue", original_text="assertTrue(x)")
    codeflash_output = t._infer_return_type(a_true) # 682ns -> 581ns (17.4% faster)
    # assertFalse should always produce "boolean"
    a_false = SimpleNamespace(assertion_method="assertFalse", original_text="assertFalse(x)")
    codeflash_output = t._infer_return_type(a_false) # 340ns -> 291ns (16.8% faster)

def test_assert_null_and_not_null_return_object():
    # assertNull and assertNotNull always map to Object (reference type)
    t = make_transformer()
    a_null = SimpleNamespace(assertion_method="assertNull", original_text="assertNull(x)")
    codeflash_output = t._infer_return_type(a_null) # 721ns -> 681ns (5.87% faster)
    a_notnull = SimpleNamespace(assertion_method="assertNotNull", original_text="assertNotNull(x)")
    codeflash_output = t._infer_return_type(a_notnull) # 380ns -> 371ns (2.43% faster)

def test_junit5_value_assertions_with_simple_literals():
    # For JUnit5 value assertions, expected literal determines returned type.
    # Build transformer and iterate over a few literal cases.
    t = make_transformer()
    method = "assertEquals"
    # integer literal
    a_int = SimpleNamespace(assertion_method=method, original_text="assertEquals(42, actual)")
    codeflash_output = t._infer_return_type(a_int) # 10.9μs -> 8.42μs (28.9% faster)
    # negative integer
    a_neg = SimpleNamespace(assertion_method=method, original_text="assertEquals(-5, actual)")
    codeflash_output = t._infer_return_type(a_neg) # 6.61μs -> 4.36μs (51.7% faster)
    # long literal (trailing L)
    a_long = SimpleNamespace(assertion_method=method, original_text="assertEquals(123456789L, actual)")
    codeflash_output = t._infer_return_type(a_long) # 8.92μs -> 6.96μs (28.0% faster)
    # float literal (trailing f)
    a_float = SimpleNamespace(assertion_method=method, original_text="assertEquals(3.14f, actual)")
    codeflash_output = t._infer_return_type(a_float) # 5.17μs -> 3.31μs (56.4% faster)
    # double literal (decimal without f, or trailing d)
    a_double1 = SimpleNamespace(assertion_method=method, original_text="assertEquals(2.71828, actual)")
    codeflash_output = t._infer_return_type(a_double1) # 6.06μs -> 4.12μs (47.2% faster)
    a_double2 = SimpleNamespace(assertion_method=method, original_text="assertEquals(1d, actual)")
    codeflash_output = t._infer_return_type(a_double2) # 4.51μs -> 2.73μs (65.5% faster)
    # char literal
    a_char = SimpleNamespace(assertion_method=method, original_text="assertEquals('x', actual)")
    codeflash_output = t._infer_return_type(a_char) # 5.27μs -> 3.43μs (53.8% faster)
    # string literal
    a_str = SimpleNamespace(assertion_method=method, original_text='assertEquals("hello", actual)')
    codeflash_output = t._infer_return_type(a_str) # 5.65μs -> 3.76μs (50.4% faster)
    # boolean literal
    a_bool = SimpleNamespace(assertion_method=method, original_text="assertEquals(true, actual)")
    codeflash_output = t._infer_return_type(a_bool) # 4.31μs -> 2.40μs (79.2% faster)
    # null literal
    a_null = SimpleNamespace(assertion_method=method, original_text="assertEquals(null, actual)")
    codeflash_output = t._infer_return_type(a_null) # 4.14μs -> 2.24μs (84.4% faster)

def test_non_value_assertions_default_to_object():
    # Methods not in JUNIT5_VALUE_ASSERTIONS should default to Object
    t = make_transformer()
    # Using a fluent assertion method name; _infer_return_type should fall back to Object
    a = SimpleNamespace(assertion_method="assertThat", original_text='assertThat(x).isEqualTo(y)')
    codeflash_output = t._infer_return_type(a) # 921ns -> 852ns (8.10% faster)
    # Also verify an unknown method name falls back to Object
    a2 = SimpleNamespace(assertion_method="someOtherAssertion", original_text='someOtherAssertion(...)')
    codeflash_output = t._infer_return_type(a2) # 391ns -> 411ns (4.87% slower)

def test_malformed_assertion_text_returns_object():
    # If the original_text does not contain parentheses, the inference should default to Object
    t = make_transformer()
    a = SimpleNamespace(assertion_method="assertEquals", original_text="assertEquals")  # no '('
    codeflash_output = t._infer_return_type(a) # 1.50μs -> 1.40μs (7.13% faster)

def test_message_first_overload_treated_as_expected():
    # Some overloads place a message as the first argument; code treats the first arg as expected.
    # So if the message is a string literal, the inferred type should be String.
    t = make_transformer()
    # Simulate message-first overload: ("message", expected, actual)
    a = SimpleNamespace(assertion_method="assertEquals", original_text='assertEquals("oops", 5, actual)')
    # Because the transformer grabs the first argument, it will think expected is the message => String
    codeflash_output = t._infer_return_type(a) # 11.9μs -> 8.24μs (44.5% faster)

def test_cast_expression_and_escaped_char_literal():
    # Cast expression like (byte)0 should return the cast type 'byte'
    t = make_transformer()
    a_cast = SimpleNamespace(assertion_method="assertEquals", original_text="assertEquals((byte)0, actual)")
    codeflash_output = t._infer_return_type(a_cast) # 13.0μs -> 10.6μs (22.8% faster)
    # Escaped char literal such as '\n' must be passed with backslash preserved in the string;
    # in source code it looks like: '\n' -> we represent it as "'\\n'" so the literal contains backslash.
    a_escaped = SimpleNamespace(assertion_method="assertEquals", original_text="assertEquals('\\n', actual)")
    codeflash_output = t._infer_return_type(a_escaped) # 7.06μs -> 4.93μs (43.3% faster)

def test_top_level_arg_splitting_respects_strings_and_generics():
    # Ensure splitting logic handles commas embedded in strings and nested generics/parentheses.
    t = make_transformer()
    # Expected is a string containing a comma; it should be preserved as the first argument.
    original = 'assertEquals("a,b,c", someMethod(Collections.<String, Integer>emptyList()), actual)'
    a = SimpleNamespace(assertion_method="assertEquals", original_text=original)
    # The first argument is a quoted string -> String
    codeflash_output = t._infer_return_type(a) # 23.1μs -> 8.02μs (188% faster)
    # Also test that generics and parentheses do not confuse splitting when the first arg is a numeric literal
    original2 = 'assertEquals(0, foo(bar(1,2), baz<Inner>(x, y)), actual)'
    a2 = SimpleNamespace(assertion_method="assertEquals", original_text=original2)
    codeflash_output = t._infer_return_type(a2) # 13.4μs -> 3.99μs (235% faster)

def test_large_number_of_args_split_and_infer_first_argument():
    # Construct a large assertion with many comma-separated args to exercise _split_top_level_args
    t = make_transformer()
    # Make 1000 arguments where the first is '0' and the rest are nested calls; ensures splitting scales.
    many = ", ".join(f"methodCall({i})" for i in range(1, 1000))
    original = f"assertEquals(0, {many})"
    a = SimpleNamespace(assertion_method="assertEquals", original_text=original)
    # The first argument is '0' -> int
    codeflash_output = t._infer_return_type(a) # 3.97ms -> 9.19μs (43097% faster)

def test_repeated_inference_over_many_iterations():
    # Verify deterministic behavior under repeated calls (1000 iterations).
    t = make_transformer()
    a = SimpleNamespace(assertion_method="assertEquals", original_text="assertEquals(123L, actual)")
    # Call repeatedly and ensure the result is stable and correct each time.
    for _ in range(1000):
        codeflash_output = t._infer_return_type(a) # 4.93ms -> 3.07ms (60.7% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr1663-2026-02-25T20.29.24 and push.