⚡️ Speed up function transform_java_assertions by 32% 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.

---

📄 32% (0.32x) speedup for transform_java_assertions in codeflash/languages/java/remove_asserts.py

⏱️ Runtime : 576 microseconds → 437 microseconds (best of 250 runs)

📝 Explanation and details

Reordered the imports to match the original module ordering (moved the parser import before the dataclasses import). This reverts a non-performance-related formatting change introduced in the candidate optimization while keeping the two meaningful optimizations that provided the speedup: replacing the intermediate replacements list with on-the-fly generation and replacing max() with a direct comparison. The change is minimal and preserves the optimized behavior and performance.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

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

def _patch_transformer_methods(detect_framework_impl, find_assertions_impl):
    """
    Helper to patch JavaAssertTransformer._detect_framework and
    JavaAssertTransformer._find_assertions, returning a restore function.
    We patch attributes on the real class and return a callable to restore them.
    """
    cls = remove_asserts.JavaAssertTransformer
    orig_detect = getattr(cls, "_detect_framework")
    orig_find = getattr(cls, "_find_assertions")

    # Apply patches
    setattr(cls, "_detect_framework", detect_framework_impl)
    setattr(cls, "_find_assertions", find_assertions_impl)

    # Return restore callable
    def _restore():
        setattr(cls, "_detect_framework", orig_detect)
        setattr(cls, "_find_assertions", orig_find)

    return _restore

def test_empty_string_returns_same():
    # Arrange: patch to avoid any interaction with the real analyzer or AST code.
    restore = _patch_transformer_methods(lambda self, s: "junit5", lambda self, s: [])
    try:
        # Act: transforming an empty source should return the exact same empty string.
        src = ""
        codeflash_output = remove_asserts.transform_java_assertions(src, function_name="someFunc"); out = codeflash_output
    finally:
        # Always restore patched methods so other tests are unaffected.
        restore()

def test_whitespace_only_returns_same():
    # Leading/trailing whitespace only should be returned unchanged.
    restore = _patch_transformer_methods(lambda self, s: "junit5", lambda self, s: [])
    try:
        src = "   \n\t  \r\n"
        codeflash_output = remove_asserts.transform_java_assertions(src, function_name="someFunc"); out = codeflash_output
    finally:
        restore()

def test_no_assertions_returns_identical_source():
    # A small Java snippet with no assertion-like constructs should be preserved.
    restore = _patch_transformer_methods(lambda self, s: "junit5", lambda self, s: [])
    try:
        src = (
            "package com.example;\n\n"
            "import java.util.List;\n\n"
            "public class TestClass {\n"
            "    public void helper() {\n"
            "        System.out.println(\"no asserts here\");\n"
            "    }\n"
            "}\n"
        )
        codeflash_output = remove_asserts.transform_java_assertions(src, function_name="helper"); out = codeflash_output
    finally:
        restore()

def test_source_with_only_comments_and_special_characters():
    # Ensure special characters and comment-only source are untouched.
    restore = _patch_transformer_methods(lambda self, s: "junit5", lambda self, s: [])
    try:
        src = (
            "// This is a comment with braces {} and parentheses () and quotes \"'\"\n"
            "/* Multi-line\n"
            "   comment with ( ) { } and a ; semicolon\n"
            "*/\n"
        )
        codeflash_output = remove_asserts.transform_java_assertions(src, function_name="doesNothing"); out = codeflash_output
    finally:
        restore()

def test_large_scale_no_assertions_performance_like_behavior():
    # Construct a large source (1000 lines) that contains no assertions.
    # The function should return the exact same content and handle the size.
    restore = _patch_transformer_methods(lambda self, s: "junit5", lambda self, s: [])
    try:
        # Build 1000 lines of innocuous Java-like code
        lines = [f"// line {i} - harmless" for i in range(1000)]
        src = "\n".join(lines) + "\n"
        codeflash_output = remove_asserts.transform_java_assertions(src, function_name="largeTest"); out = codeflash_output
    finally:
        restore()

def test_looks_like_assertions_but_none_detected_results_in_identity():
    # Simulate the analyzer detecting a framework (to exercise _detect_framework path)
    # but simulate that no assertions were discovered by returning an empty list.
    restore = _patch_transformer_methods(
        lambda self, s: "assertj",  # pretend an assertion library was detected
        lambda self, s: [],  # but no actual assertions found
    )
    try:
        src = (
            "import org.assertj.core.api.Assertions;\n"
            "public class MyTest {\n"
            "    public void testSomething() {\n"
            "        // contains the word assert but not an assertion call\n"
            "        String x = \"assert\";\n"
            "    }\n"
            "}\n"
        )
        codeflash_output = remove_asserts.transform_java_assertions(src, function_name="testSomething"); out = codeflash_output
    finally:
        restore()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from typing import List

# imports
import pytest  # used for our unit tests
# import the function and class under test
from codeflash.languages.java.remove_asserts import (JavaAssertTransformer,
                                                     transform_java_assertions)

# NOTE:
# The real JavaAssertTransformer depends on a JavaAnalyzer that uses tree-sitter.
# To keep tests deterministic and avoid heavy external dependencies, we temporarily
# monkeypatch a few methods on JavaAssertTransformer in each test to control behavior.
# We avoid using unittest.mock.patch (disallowed by the problem constraints) and use
# pytest's monkeypatch fixture which assigns attributes directly and restores them after.
#
# Small helper class used purely inside tests to represent minimal assertion-like
# objects consumed by transform() — we keep these minimal and deterministic.
class _AssertionLike:
    def __init__(self, start_pos: int, end_pos: int, leading_whitespace: str = "", index: int = 0):
        # transform() only requires start_pos and end_pos for ordering/filtering.
        # _generate_replacement (monkeypatched in tests) may also use other attributes.
        self.start_pos = start_pos
        self.end_pos = end_pos
        self.leading_whitespace = leading_whitespace
        self.index = index
        # property used by original code; tests that need it can set it on the instance.
        self.is_exception_assertion = False
        self.target_calls = None

def test_empty_string_returns_same():
    # If the source is an empty string, transform_java_assertions should return it unchanged.
    src = ""
    codeflash_output = transform_java_assertions(src, function_name="foo"); out = codeflash_output # 4.71μs -> 4.81μs (2.08% slower)

def test_whitespace_only_returns_same():
    # If the source contains only whitespace (no code), it should be returned unchanged.
    src = "   \n\t  "
    codeflash_output = transform_java_assertions(src, function_name="foo"); out = codeflash_output # 4.45μs -> 4.33μs (2.77% faster)

def test_no_assertions_returns_source(monkeypatch):
    # When no assertions are found, transform should return the original source unchanged.
    src = "class A { void test() { int x = 1; } }"

    # Prevent analyzer usage by forcing framework detection to a no-op return.
    monkeypatch.setattr(JavaAssertTransformer, "_detect_framework", lambda self, s: "junit5")
    # Make _find_assertions return an empty list (no assertions found).
    monkeypatch.setattr(JavaAssertTransformer, "_find_assertions", lambda self, s: [])

    codeflash_output = transform_java_assertions(src, function_name="foo"); out = codeflash_output # 5.55μs -> 5.33μs (4.15% faster)

def test_single_replacement_replaces_correct_range(monkeypatch):
    # Test that a single assertion region is replaced with the generated replacement.
    # Build a source string with a distinct "ASSERT_BLOCK" region we expect to be replaced.
    prefix = "class C { void t() { "
    assert_block = "/*ASSERT*/System.out.println(1);/*END*/"
    suffix = " } }"
    src = prefix + assert_block + suffix

    # Compute start/end indices for the region to replace (cover the assert_block).
    start = src.index(assert_block)
    end = start + len(assert_block)

    # Monkeypatch framework detection and found assertions to control flow.
    monkeypatch.setattr(JavaAssertTransformer, "_detect_framework", lambda self, s: "junit5")
    # Return a single assertion object that covers the ASSERT_BLOCK region.
    monkeypatch.setattr(JavaAssertTransformer, "_find_assertions", lambda self, s: [_AssertionLike(start, end, leading_whitespace="    ")])

    # Replace generation should produce a deterministic string.
    monkeypatch.setattr(JavaAssertTransformer, "_generate_replacement", lambda self, a: a.leading_whitespace + "REPLACED_CALL();")

    codeflash_output = transform_java_assertions(src, function_name="foo"); out = codeflash_output # 10.0μs -> 9.03μs (11.1% faster)
    # Expect the assert_block to be removed and replaced by our generated string.
    expected = prefix + "    REPLACED_CALL();" + suffix

def test_multiple_replacements_and_ordering(monkeypatch):
    # Ensure multiple assertion regions are replaced and ordering is determined by start_pos.
    parts = ["p0", "[A]", "p1", "[B]", "p2", "[C]", "p3"]
    src = "|".join(parts)  # create a separable source string

    # find indices for the markers [A], [B], [C]
    markers = ["[A]", "[B]", "[C]"]
    assertions: List[_AssertionLike] = []
    for idx, m in enumerate(markers):
        st = src.index(m)
        assertions.append(_AssertionLike(st, st + len(m), leading_whitespace="", index=idx))

    # Shuffle the returned assertions to confirm transform sorts them by start_pos internally.
    monkeypatch.setattr(JavaAssertTransformer, "_detect_framework", lambda self, s: "junit5")
    monkeypatch.setattr(JavaAssertTransformer, "_find_assertions", lambda self, s: [assertions[2], assertions[0], assertions[1]])

    # Generate unique replacement strings based on the index attribute.
    def gen_repl(self, assertion):
        return f"<REPL{assertion.index}>"
    monkeypatch.setattr(JavaAssertTransformer, "_generate_replacement", gen_repl)

    codeflash_output = transform_java_assertions(src, function_name="foo"); out = codeflash_output # 10.7μs -> 9.20μs (15.8% faster)
    # Build expected by replacing markers in source order.
    expected = src.replace("[A]", "<REPL0>").replace("[B]", "<REPL1>").replace("[C]", "<REPL2>")

def test_nested_assertions_filtered(monkeypatch):
    # If one assertion is nested within another (its span lies within the other),
    # it should be filtered out (only the outer assertion should be replaced).
    base = "START--"
    inner = "INNER_ASSERT"
    outer = "OUTER_ASSERT_BEGIN[" + inner + "]OUTER_ASSERT_END"
    tail = "--END"
    src = base + outer + tail

    # Determine spans. Outer covers entire outer string, inner covers the inner substring.
    outer_start = src.index("OUTER_ASSERT_BEGIN")
    outer_end = outer_start + len("OUTER_ASSERT_BEGIN[" + inner + "]OUTER_ASSERT_END")
    inner_start = src.index(inner)
    inner_end = inner_start + len(inner)

    # Return both assertions but unordered to exercise sorting + nesting filter.
    monkeypatch.setattr(JavaAssertTransformer, "_detect_framework", lambda self, s: "junit5")
    monkeypatch.setattr(JavaAssertTransformer, "_find_assertions", lambda self, s: [_AssertionLike(inner_start, inner_end, index=1), _AssertionLike(outer_start, outer_end, index=0)])

    # Generate replacements that show which assertion was used.
    def gen_repl(self, a):
        # Use index to indicate inner vs outer; if inner were used we'd see <I>, outer => <O>
        return "<I>" if a.index == 1 else "<O>"
    monkeypatch.setattr(JavaAssertTransformer, "_generate_replacement", gen_repl)

    codeflash_output = transform_java_assertions(src, function_name="foo"); out = codeflash_output # 10.3μs -> 9.38μs (10.4% faster)
    # Only the outer replacement should appear (inner is nested and filtered out)
    expected = base + "<O>" + tail

def test_empty_replacement_removes_assertion_body(monkeypatch):
    # When _generate_replacement returns an empty string, the assertion region should be removed.
    src = "X[TO_REMOVE]Y"
    start = src.index("[TO_REMOVE]")
    end = start + len("[TO_REMOVE]")

    monkeypatch.setattr(JavaAssertTransformer, "_detect_framework", lambda self, s: "junit5")
    monkeypatch.setattr(JavaAssertTransformer, "_find_assertions", lambda self, s: [_AssertionLike(start, end)])
    # _generate_replacement returns empty string to signal deletion of the assertion content.
    monkeypatch.setattr(JavaAssertTransformer, "_generate_replacement", lambda self, a: "")

    codeflash_output = transform_java_assertions(src, function_name="foo"); out = codeflash_output # 9.18μs -> 8.10μs (13.2% faster)

def test_large_scale_many_replacements(monkeypatch):
    # Stress test: create a source with 1000 small assertion markers and ensure all are replaced.
    N = 1000  # scale up to 1000 as requested
    pieces = []
    markers = []
    for i in range(N):
        pieces.append(f"seg{i}")
        marker = f"[A{i}]"
        pieces.append(marker)
        markers.append(marker)
    pieces.append("end")
    src = "|".join(pieces)

    # Build assertion objects for every marker with correct spans.
    assertions: List[_AssertionLike] = []
    for i, marker in enumerate(markers):
        st = src.index(marker)
        assertions.append(_AssertionLike(st, st + len(marker), index=i))

    # Return the assertions in reverse order to ensure transform sorts by start_pos.
    monkeypatch.setattr(JavaAssertTransformer, "_detect_framework", lambda self, s: "junit5")
    monkeypatch.setattr(JavaAssertTransformer, "_find_assertions", lambda self, s: list(reversed(assertions)))

    # Generate replacement that includes the marker index so we can assert correctness.
    def gen_repl(self, a):
        return f"<R{a.index}>"
    monkeypatch.setattr(JavaAssertTransformer, "_generate_replacement", gen_repl)

    codeflash_output = transform_java_assertions(src, function_name="foo"); out = codeflash_output # 487μs -> 354μs (37.6% faster)

    # Verify all expected replacements appear and no original markers remain.
    for i in range(N):
        pass
# 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.54.26 and push.