⚡️ Speed up method JavaAssertTransformer._generate_replacement by 85% in PR #1655 (feat/add/void/func)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1655

If you approve this dependent PR, these changes will be merged into the original PR branch feat/add/void/func.

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

---

📄 85% (0.85x) speedup for JavaAssertTransformer._generate_replacement in codeflash/languages/java/remove_asserts.py

⏱️ Runtime : 491 microseconds → 265 microseconds (best of 127 runs)

📝 Explanation and details

Runtime improvement (primary): The optimized version cuts the _generate_replacement runtime from ~491μs to ~265μs (≈85% speedup).

What changed:

• The loop-internal mutation self.invocation_counter += 1 was batched into a local_counter variable that is incremented inside the loop, and the instance attribute is updated exactly once at the end (self.invocation_counter = local_counter).
• The code also uses that local_counter when constructing variable names (f"_cf_result{local_counter}") so every attribute read/write inside the hot loop is removed.

Why this speeds things up:

• Attribute access and assignment on self require dictionary lookups and bytecode operations; doing that on each iteration is noticeably more expensive than operating on a local variable. Moving counter increments to a local integer eliminates O(k) attribute writes/loads inside the loop and replaces them with fast local variable ops.
• This reduces Python bytecode executed per iteration and cuts down on dict churn and attribute lookup overhead. The profiler shows the loop was the hot spot; eliminating repeated attribute increments directly reduces time spent in that hotspot.
• The large-scale test (1000 target calls) demonstrates the win: the optimized code avoids 1000 attribute writes and thus gets the biggest wall‑clock improvement. Small cases see modest or negligible changes, large cases see major gains.

Behavioral/compatibility notes:

• Semantics are preserved: the final invocation_counter value and the generated variable names remain the same as before. The only change is when the instance attribute is updated (after finishing the loop) — which is safe because the counter is an internal, per-instance sequence used only by this transformer.
• Exception-handling branch and other logic were left intact, so behavior for assertThrows/assertDoesNotThrow is unchanged.

When this matters (based on tests):

• Best for workloads where assertions contain many target calls (hot path in loop-heavy transformations). The annotated tests show the greatest speedup on the large-scale 1000-call case; microbenchmarks for single or few calls show small improvements (or negligible differences) as expected.

Summary:

• Primary optimization: avoid repeated attribute mutations/lookup in a hot loop by using a local counter and committing once. This simple change reduces per-iteration overhead and yields the significant runtime improvement you see in the profile and tests without altering observable behavior in normal single-threaded use.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

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

def test_basic_single_target_call_non_void():
    # Create a transformer for a non-void function (default is_void=False).
    transformer = JavaAssertTransformer(function_name="f")

    # Build an AssertionMatch representing a simple assertion with one target call.
    # leading_whitespace is a standard indentation string.
    assertion = AssertionMatch(
        is_exception_assertion=False,
        target_calls=[TargetCall(full_call="obj.f(1)")],
        leading_whitespace="    ",  # four spaces indentation
        lambda_body=None,
        assigned_var_name=None,
        assigned_var_type=None,
        assertion_method=None,
        exception_class=None,
    )

    # Call the method under test.
    codeflash_output = transformer._generate_replacement(assertion); replacement = codeflash_output

def test_multiple_target_calls_preserve_leading_newlines_for_first():
    # New transformer to ensure fresh invocation_counter.
    transformer = JavaAssertTransformer(function_name="g")

    # leading_whitespace contains two newlines followed by indentation.
    lw = "\n\n    "
    assertion = AssertionMatch(
        is_exception_assertion=False,
        target_calls=[
            TargetCall(full_call="a()"),
            TargetCall(full_call="b()"),
        ],
        leading_whitespace=lw,
        lambda_body=None,
        assigned_var_name=None,
        assigned_var_type=None,
        assertion_method=None,
        exception_class=None,
    )

    codeflash_output = transformer._generate_replacement(assertion); replacement = codeflash_output

    # The first line should preserve the leading newlines; the second should
    # use base_indent (leading newlines stripped) so it does not insert extra blank lines.
    expected = "\n\n    Object _cf_result1 = a();\n    Object _cf_result2 = b();"

def test_is_void_true_emits_call_statements_only():
    # When the transformer is configured for a void function, it should not
    # create Object variables but instead emit the calls followed by semicolons.
    transformer = JavaAssertTransformer(function_name="h", is_void=True)

    assertion = AssertionMatch(
        is_exception_assertion=False,
        target_calls=[TargetCall(full_call="obj.doStuff()")],
        leading_whitespace="  ",  # two spaces
        lambda_body=None,
        assigned_var_name=None,
        assigned_var_type=None,
        assertion_method=None,
        exception_class=None,
    )

    codeflash_output = transformer._generate_replacement(assertion); replacement = codeflash_output

def test_no_target_calls_returns_empty_string():
    transformer = JavaAssertTransformer(function_name="x")

    # No target calls => nothing to capture; should return empty string.
    assertion = AssertionMatch(
        is_exception_assertion=False,
        target_calls=[],
        leading_whitespace="",
        lambda_body=None,
        assigned_var_name=None,
        assigned_var_type=None,
        assertion_method=None,
        exception_class=None,
    )

    codeflash_output = transformer._generate_replacement(assertion); replacement = codeflash_output

def test_exception_lambda_body_assert_does_not_throw_with_assignment_without_semicolon():
    # Test assertDoesNotThrow behavior when assigned to a variable and lambda body
    # does not include a semicolon: the code path should synthesize "Type name = body;"
    transformer = JavaAssertTransformer(function_name="y")

    assertion = AssertionMatch(
        is_exception_assertion=True,
        target_calls=[],
        leading_whitespace="  ",
        lambda_body="someCall()",  # no trailing semicolon intentionally
        assigned_var_name="ex",
        assigned_var_type="IllegalArgumentException",
        assertion_method="assertDoesNotThrow",
        exception_class=None,
    )

    codeflash_output = transformer._generate_replacement(assertion); replacement = codeflash_output

    # Since there's no semicolon in the stripped lambda_body, the method should
    # return a variable assignment with a semicolon appended.
    expected = "  IllegalArgumentException ex = someCall();"

def test_exception_assert_throws_with_assignment_uses_exception_class_and_catches():
    # Test assertThrows when assigned to a variable: it should generate a
    # null-initialized variable and try/catch blocks that assign caught exception.
    transformer = JavaAssertTransformer(function_name="z")

    assertion = AssertionMatch(
        is_exception_assertion=True,
        target_calls=[],
        leading_whitespace="    ",
        lambda_body="code();",
        assigned_var_name="ex",
        assigned_var_type="IllegalArgumentException",
        assertion_method="assertThrows",
        exception_class="IllegalArgumentException",
    )

    codeflash_output = transformer._generate_replacement(assertion); replacement = codeflash_output

    # The invocation_counter starts at 0 and is incremented to 1 inside the method.
    expected = (
        "    IllegalArgumentException ex = null;\n"
        "    try { code(); } catch (IllegalArgumentException _cf_caught1) { ex = _cf_caught1; } "
        "catch (Exception _cf_ignored1) {}"
    )

def test_exception_with_no_lambda_but_target_calls_uses_first_call_in_try():
    transformer = JavaAssertTransformer(function_name="m")

    assertion = AssertionMatch(
        is_exception_assertion=True,
        target_calls=[TargetCall(full_call="calculator.divide(1, 0)")],
        leading_whitespace=" ",
        lambda_body=None,
        assigned_var_name=None,
        assigned_var_type=None,
        assertion_method="assertThrows",
        exception_class=None,
    )

    codeflash_output = transformer._generate_replacement(assertion); replacement = codeflash_output

    # Should wrap the call in try/catch with invocation counter 1.
    expected = " try { calculator.divide(1, 0); } catch (Exception _cf_ignored1) {}"

def test_exception_fallback_no_callable_results_in_commented_out_assertion():
    transformer = JavaAssertTransformer(function_name="n")

    assertion = AssertionMatch(
        is_exception_assertion=True,
        target_calls=[],
        leading_whitespace="\t",  # a tab
        lambda_body=None,
        assigned_var_name=None,
        assigned_var_type=None,
        assertion_method="assertThrows",
        exception_class=None,
    )

    codeflash_output = transformer._generate_replacement(assertion); replacement = codeflash_output

    # No callable found -> fallback comment with the original leading whitespace.
    expected = "\t// Removed assertThrows: could not extract callable"

def test_base_indent_strips_both_newline_and_carriage_return_for_subsequent_calls():
    transformer = JavaAssertTransformer(function_name="o")

    # Use CRLF style leading whitespace; base_indent must strip both \r and \n.
    lw = "\r\n\t"
    assertion = AssertionMatch(
        is_exception_assertion=False,
        target_calls=[TargetCall(full_call="x()"), TargetCall(full_call="y()")],
        leading_whitespace=lw,
        lambda_body=None,
        assigned_var_name=None,
        assigned_var_type=None,
        assertion_method=None,
        exception_class=None,
    )

    codeflash_output = transformer._generate_replacement(assertion); replacement = codeflash_output

    # First line preserves the CRLF; second line uses only the tab (base_indent).
    expected = "\r\n\tObject _cf_result1 = x();\n\tObject _cf_result2 = y();"

def test_large_scale_many_target_calls_unique_variable_names_and_performance():
    # Large scale test: 1000 target calls to verify scalability and uniqueness.
    transformer = JavaAssertTransformer(function_name="big")

    n = 1000
    # Create 1000 TargetCall instances with distinct calls.
    calls = [TargetCall(full_call=f"fn_call({i})") for i in range(n)]

    assertion = AssertionMatch(
        is_exception_assertion=False,
        target_calls=calls,
        leading_whitespace="",  # no leading whitespace to simplify assertions
        lambda_body=None,
        assigned_var_name=None,
        assigned_var_type=None,
        assertion_method=None,
        exception_class=None,
    )

    codeflash_output = transformer._generate_replacement(assertion); replacement = codeflash_output

    # Ensure we got n lines back.
    lines = replacement.splitlines()

    # Ensure that each generated variable name is unique and in sequence.
    # The produced pattern is 'Object _cf_result{i} = fn_call(i);'
    for i, line in enumerate(lines, start=1):
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests
from codeflash.languages.java.remove_asserts import JavaAssertTransformer

# We create small, simple helper classes to act as assertion-like and call-like objects.
# These are thin duck-typed helpers that provide only the attributes accessed by
# JavaAssertTransformer._generate_replacement. Using simple objects keeps tests focused
# on the transformer's behavior (the transformer uses attribute access only).
class _CallLike:
    def __init__(self, full_call: str):
        # The transformer reads `.full_call` for each target call.
        self.full_call = full_call

class _AssertionLike:
    def __init__(
        self,
        *,
        is_exception_assertion: bool = False,
        target_calls=None,
        leading_whitespace: str = "",
        lambda_body: str | None = None,
        assigned_var_name: str | None = None,
        assigned_var_type: str | None = None,
        assertion_method: str | None = None,
        exception_class: str | None = None,
    ):
        # Attributes the transformer expects on its `assertion` parameter.
        self.is_exception_assertion = is_exception_assertion
        self.target_calls = target_calls or []
        self.leading_whitespace = leading_whitespace
        self.lambda_body = lambda_body
        self.assigned_var_name = assigned_var_name
        self.assigned_var_type = assigned_var_type
        self.assertion_method = assertion_method
        self.exception_class = exception_class

def test_single_target_call_non_void_generates_object_assignment():
    # Create transformer for a function; default invocation_counter == 0.
    t = JavaAssertTransformer(function_name="compute")
    # Single target call; not an exception assertion.
    call = _CallLike("obj.compute(42)")
    a = _AssertionLike(
        is_exception_assertion=False,
        target_calls=[call],
        leading_whitespace="    ",  # four-space indent preserved for first replacement
    )

    codeflash_output = t._generate_replacement(a); out = codeflash_output # 2.32μs -> 2.42μs (3.77% slower)

def test_multiple_target_calls_non_void_preserves_leading_newline_behavior():
    # Leading whitespace includes a leading newline; first line must keep it, subsequent lines use base_indent.
    t = JavaAssertTransformer(function_name="f")
    calls = [_CallLike("a()"), _CallLike("b()")]
    a = _AssertionLike(is_exception_assertion=False, target_calls=calls, leading_whitespace="\n    ")
    codeflash_output = t._generate_replacement(a); out = codeflash_output # 3.00μs -> 2.85μs (5.25% faster)
    # Expect two lines. First retains the exact leading_whitespace, second uses base_indent (leading newline stripped).
    expected = "\n    Object _cf_result1 = a();\n    Object _cf_result2 = b();"

def test_void_calls_generate_statements_without_assignment():
    # When transformer.is_void is True we should produce bare calls with semicolons (no Object declarations).
    t = JavaAssertTransformer(function_name="v", is_void=True)
    calls = [_CallLike("sink.write(x)"), _CallLike("sink.flush()")]
    a = _AssertionLike(is_exception_assertion=False, target_calls=calls, leading_whitespace="  ")
    codeflash_output = t._generate_replacement(a); out = codeflash_output # 2.33μs -> 2.29μs (1.79% faster)
    expected = "  sink.write(x);\n  sink.flush();"

def test_no_target_calls_returns_empty_string():
    t = JavaAssertTransformer(function_name="none")
    a = _AssertionLike(is_exception_assertion=False, target_calls=[], leading_whitespace=" ")
    codeflash_output = t._generate_replacement(a); out = codeflash_output # 581ns -> 532ns (9.21% faster)

def test_exception_assertion_with_lambda_body_generates_try_catch_and_adds_semicolon_if_missing():
    t = JavaAssertTransformer(function_name="ex")
    # Simulate an assertThrows with a lambda body missing a trailing semicolon.
    a = _AssertionLike(
        is_exception_assertion=True,
        target_calls=[],
        leading_whitespace="  ",
        lambda_body="calculator.divide(1, 0)",
    )
    codeflash_output = t._generate_replacement(a); out = codeflash_output # 2.02μs -> 2.00μs (0.948% faster)

def test_exception_assertion_assigned_assertDoesNotThrow_without_semicolon_produces_assignment_line():
    t = JavaAssertTransformer(function_name="ex2")
    # When assigned and assertion_method is assertDoesNotThrow and lambda_body has no semicolon,
    # the code should be converted into a direct assignment statement of the variable.
    a = _AssertionLike(
        is_exception_assertion=True,
        target_calls=[],
        leading_whitespace="\t",  # tab indent should be preserved
        lambda_body="createThing()",  # no trailing semicolon
        assigned_var_name="thing",
        assigned_var_type="Thing",
        assertion_method="assertDoesNotThrow",
    )
    codeflash_output = t._generate_replacement(a); out = codeflash_output # 2.15μs -> 2.12μs (1.41% faster)

def test_exception_assertion_assigned_assertThrows_uses_exception_class_and_catches_and_assigns():
    t = JavaAssertTransformer(function_name="ex3")
    # Assigned assertThrows case: should initialize var to null, try the code and on catch assign the caught exception.
    a = _AssertionLike(
        is_exception_assertion=True,
        target_calls=[],
        leading_whitespace="  ",
        lambda_body="doRisky();",
        assigned_var_name="ex",
        assigned_var_type="IllegalArgumentException",
        assertion_method="assertThrows",
        exception_class="IllegalArgumentException",
    )
    codeflash_output = t._generate_replacement(a); out = codeflash_output # 2.38μs -> 2.40μs (0.790% slower)
    # Build expected string: note base_indent is leading_whitespace stripped of leading newlines (same in this case)
    expected = (
        "  IllegalArgumentException ex = null;\n"
        "  try { doRisky(); } catch (IllegalArgumentException _cf_caught1) { ex = _cf_caught1; } catch (Exception _cf_ignored1) {}"
    )

def test_exception_assertion_without_lambda_falls_back_to_target_call_extraction():
    t = JavaAssertTransformer(function_name="ex4")
    # No lambda body, but there is a target call; should wrap the first target call in try/catch.
    call = _CallLike("subject.perform()")
    a = _AssertionLike(is_exception_assertion=True, target_calls=[call], leading_whitespace="")
    codeflash_output = t._generate_replacement(a); out = codeflash_output # 1.57μs -> 1.48μs (6.14% faster)

def test_leading_whitespace_with_multiple_newlines_handles_base_indent_correctly():
    t = JavaAssertTransformer(function_name="indent_test")
    calls = [_CallLike("one()"), _CallLike("two()"), _CallLike("three()")]
    # Leading whitespace has multiple leading newlines; base_indent should strip newlines.
    a = _AssertionLike(is_exception_assertion=False, target_calls=calls, leading_whitespace="\n\n    ")
    codeflash_output = t._generate_replacement(a); out = codeflash_output # 3.53μs -> 3.33μs (6.04% faster)
    # First line begins with the full leading whitespace; subsequent lines should start with "    " (no leading newlines)
    expected = "\n\n    Object _cf_result1 = one();\n    Object _cf_result2 = two();\n    Object _cf_result3 = three();"

def test_large_scale_many_target_calls_increments_counter_and_generates_many_assignments():
    t = JavaAssertTransformer(function_name="big")
    n = 1000  # large-scale: 1000 calls
    calls = [_CallLike(f"doWork({i})") for i in range(1, n + 1)]
    a = _AssertionLike(is_exception_assertion=False, target_calls=calls, leading_whitespace="")
    codeflash_output = t._generate_replacement(a); out = codeflash_output # 466μs -> 241μs (93.3% faster)
    # Ensure that the output contains exactly n lines and that the last declared variable uses the final counter.
    lines = out.splitlines()

def test_mixed_calls_and_exception_flags_do_not_leak_counter_values_between_transformers():
    # Create two transformers to verify internal invocation_counter is per-instance.
    t1 = JavaAssertTransformer(function_name="a")
    t2 = JavaAssertTransformer(function_name="b")

    # t1: two normal calls
    a1 = _AssertionLike(is_exception_assertion=False, target_calls=[_CallLike("x()"), _CallLike("y()")], leading_whitespace="")
    codeflash_output = t1._generate_replacement(a1); out1 = codeflash_output # 3.10μs -> 2.83μs (9.91% faster)

    # t2: one exception assert with lambda; should start counter at 0 for t2 and increment only for t2.
    a2 = _AssertionLike(is_exception_assertion=True, lambda_body="danger()", leading_whitespace="")
    codeflash_output = t2._generate_replacement(a2); out2 = codeflash_output # 1.55μs -> 1.50μs (3.33% 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-pr1655-2026-02-25T05.48.36 and push.