⚡️ Speed up method JavaAssertTransformer._find_balanced_parens by 62% in PR #1199 (omni-java)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1199

If you approve this dependent PR, these changes will be merged into the original PR branch omni-java.

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

---

📄 62% (0.62x) speedup for JavaAssertTransformer._find_balanced_parens in codeflash/languages/java/remove_asserts.py

⏱️ Runtime : 1.65 milliseconds → 1.02 milliseconds (best of 138 runs)

📝 Explanation and details

The optimized code achieves a 62% runtime improvement by eliminating redundant operations in the tight character-scanning loop of _find_balanced_parens.

Key optimizations:

1. Cached len(code) computation: The original code called len(code) on every loop iteration in the while condition. The optimized version computes code_len = len(code) once before the loop, eliminating ~10,000 redundant length calculations for typical workloads.

2. Eliminated repeated indexing for prev_char: The original code computed code[pos - 1] if pos > 0 else "" on every iteration. The optimized version maintains prev_char as a variable that's updated at the end of each iteration with the current char. This removes one index operation and one conditional check per iteration.

Performance impact by test case type:

• Simple cases (2-5 chars): 7-26% faster - modest gains since loop iterations are minimal
• Nested/complex cases (nested parens, strings with special chars): 20-43% faster - substantial gains as these execute more loop iterations where the per-iteration savings compound
• Large-scale cases (100+ nesting levels, 500+ arguments, 1000+ char strings): 37-83% faster - dramatic speedups because they execute thousands of loop iterations where eliminating two operations per iteration (length check + index lookup) provides massive cumulative benefit

The optimizations are particularly effective for Java code parsing scenarios with deeply nested method calls or large argument lists - common in test frameworks and assertion libraries. Since this is a JavaAssertTransformer used for test code transformation, these patterns are likely to occur frequently in practice, making the optimization highly valuable for real-world usage.

The changes maintain identical behavior including edge case handling for escaped characters, unbalanced parentheses, and invalid positions.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 122 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 JavaAssertTransformer

def test_basic_balanced_parens_simple():
    # Basic: simple single-level parentheses
    transformer = JavaAssertTransformer("assert")  # create instance
    code = "foo(bar)"  # sample code
    open_pos = code.find("(")  # find opening paren position
    content, after_pos = transformer._find_balanced_parens(code, open_pos) # 2.31μs -> 1.96μs (17.9% faster)

def test_nested_parentheses_outer_and_inner():
    # Basic: nested parentheses, ensure outer returns the whole inner structure
    transformer = JavaAssertTransformer("assert")
    code = "a(b(c,d),e)"  # nested structure
    open_pos_outer = code.find("(")  # position of outer '('
    content_outer, pos_outer = transformer._find_balanced_parens(code, open_pos_outer) # 3.11μs -> 2.69μs (15.7% faster)

    # Now test the inner parentheses by finding the second '('
    open_pos_inner = code.find("(", open_pos_outer + 1)
    content_inner, pos_inner = transformer._find_balanced_parens(code, open_pos_inner) # 1.37μs -> 1.17μs (17.1% faster)

def test_parentheses_ignore_double_quoted_strings():
    # Edge: parentheses inside double-quoted strings must be ignored when balancing
    transformer = JavaAssertTransformer("assert")
    # The string contains a closing paren inside quotes which should not end the outer group.
    code = 'method("this has ) inside", other)'  # double-quote contains ')'
    open_pos = code.find("(")
    content, after_pos = transformer._find_balanced_parens(code, open_pos) # 5.42μs -> 4.18μs (29.7% faster)

    # Also test that escaped quotes inside the string do not prematurely terminate the string
    code2 = 'call("escaped \\" quote and )", x)'
    open_pos2 = code2.find("(")
    content2, after_pos2 = transformer._find_balanced_parens(code2, open_pos2) # 4.41μs -> 3.08μs (43.3% faster)

def test_parentheses_ignore_single_char_literals():
    # Edge: parentheses inside character literals must be ignored
    transformer = JavaAssertTransformer("assert")
    code = "foo('(', x)"  # char literal contains '(' that should not increase depth
    open_pos = code.find("(")
    content, after_pos = transformer._find_balanced_parens(code, open_pos) # 2.89μs -> 2.50μs (15.7% faster)

    # Also test a closing parenthesis inside a char literal
    code2 = "bar(')', y)"
    open_pos2 = code2.find("(")
    content2, after_pos2 = transformer._find_balanced_parens(code2, open_pos2) # 1.85μs -> 1.50μs (23.3% faster)

def test_unbalanced_parentheses_returns_none_minus_one():
    # Edge: return (None, -1) when parentheses are unbalanced or open_paren_pos invalid
    transformer = JavaAssertTransformer("assert")
    # Unbalanced: missing closing paren
    code = "(abc"
    content, pos = transformer._find_balanced_parens(code, 0) # 1.81μs -> 1.49μs (21.5% faster)

    # Invalid open_paren_pos: position not at '('
    code2 = "no_paren_here"
    content2, pos2 = transformer._find_balanced_parens(code2, 0) # 380ns -> 390ns (2.56% slower)

    # open_paren_pos at end of string (out of range for '(')
    code3 = "("
    content3, pos3 = transformer._find_balanced_parens(code3, 0) # 401ns -> 461ns (13.0% slower)

    # open_paren_pos beyond string length
    content4, pos4 = transformer._find_balanced_parens("abc", 10) # 230ns -> 220ns (4.55% faster)

def test_opening_at_inner_position():
    # Basic: if the provided open_paren_pos is for an inner '(', it should return the inner content
    transformer = JavaAssertTransformer("assert")
    code = "outer(a(inner), b)"
    # Find inner '(' (the one before "inner")
    outer_open = code.find("(")
    inner_open = code.find("(", outer_open + 1)
    # inner_open should be index of '(' before 'inner'
    content_inner, pos_inner = transformer._find_balanced_parens(code, inner_open) # 2.51μs -> 2.20μs (14.1% faster)

def test_escaped_characters_and_backslashes_in_strings_and_chars():
    # Edge: complex mix of escapes - ensure that backslashes prevent toggling
    transformer = JavaAssertTransformer("assert")
    # A double quote is escaped inside the string; the parser uses prev_char != "\\" check
    code = 'f("an escaped \\" quote ) and more", z)'
    open_pos = code.find("(")
    content, after_pos = transformer._find_balanced_parens(code, open_pos) # 6.39μs -> 4.86μs (31.5% faster)

    # Character literal with escaped backslash before single-quote should not toggle incorrectly:
    code2 = r"g('\\'', p)"  # contains a backslash-escaped single-quote sequence in a char literal context (raw string for readability)
    # This example includes two characters inside the quoted part; we ensure the function doesn't error out.
    open_pos2 = code2.find("(")
    content2, after_pos2 = transformer._find_balanced_parens(code2, open_pos2) # 2.38μs -> 1.91μs (24.7% faster)

def test_large_scale_many_nested_parentheses():
    # Large Scale: many nested parentheses to test scalability (but keep total length < 1000)
    transformer = JavaAssertTransformer("assert")
    depth = 300  # nested depth, kept below 1000 to satisfy constraints
    # Build a deeply nested string: f(((...x...)))
    inner = "x"
    nested = "(" * depth + inner + ")" * depth
    code = "f(" + nested + ")"
    # open_paren_pos is the first '(' after the 'f'
    open_pos = code.find("(")
    content, after_pos = transformer._find_balanced_parens(code, open_pos) # 99.9μs -> 62.6μs (59.7% faster)

def test_balanced_with_multiple_arguments_and_whitespace():
    # Basic/Edge: multiple arguments, whitespace, and newlines between tokens
    transformer = JavaAssertTransformer("assert")
    code = "call(  first, \n  second( inner1, inner2 ),  third ) end"
    open_pos = code.find("(")
    content, after_pos = transformer._find_balanced_parens(code, open_pos) # 8.23μs -> 6.19μs (33.0% faster)
    # The content should exactly be the substring between the parentheses (preserve whitespace/newlines)
    expected = "  first, \n  second( inner1, inner2 ),  third "
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest
from codeflash.languages.java.parser import get_java_analyzer
from codeflash.languages.java.remove_asserts import JavaAssertTransformer

class TestFindBalancedParens:
    """Test suite for JavaAssertTransformer._find_balanced_parens method."""

    def setup_method(self):
        """Set up test fixtures before each test."""
        self.transformer = JavaAssertTransformer("testMethod")

    # ============================================================================
    # BASIC TEST CASES - Test fundamental functionality under normal conditions
    # ============================================================================

    def test_simple_empty_parentheses(self):
        """Test finding content in simple empty parentheses."""
        code = "()"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 1.44μs -> 1.51μs (4.63% slower)

    def test_simple_single_argument(self):
        """Test finding content with a single simple argument."""
        code = "(arg)"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 2.31μs -> 2.08μs (11.1% faster)

    def test_simple_multiple_arguments(self):
        """Test finding content with multiple comma-separated arguments."""
        code = "(arg1, arg2, arg3)"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 4.07μs -> 3.23μs (26.1% faster)

    def test_nested_parentheses(self):
        """Test finding content with nested parentheses."""
        code = "(func(nested))"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 3.72μs -> 3.04μs (22.1% faster)

    def test_deeply_nested_parentheses(self):
        """Test finding content with deeply nested parentheses."""
        code = "(a(b(c(d))))"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 3.46μs -> 2.85μs (21.5% faster)

    def test_parentheses_with_spaces(self):
        """Test finding content with whitespace."""
        code = "( arg1 , arg2 )"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 3.58μs -> 2.96μs (21.0% faster)

    def test_string_with_parentheses(self):
        """Test that parentheses inside double-quoted strings are ignored."""
        code = '("text(with)parens")'
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 4.37μs -> 3.36μs (30.2% faster)

    def test_char_with_parentheses(self):
        """Test that parentheses inside single-quoted character literals are ignored."""
        code = "('(')"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 2.44μs -> 2.15μs (13.5% faster)

    def test_offset_starting_position(self):
        """Test finding balanced parens when starting position is not 0."""
        code = "func(args)"
        content, end_pos = self.transformer._find_balanced_parens(code, 4) # 2.31μs -> 2.00μs (15.6% faster)

    def test_function_call_in_code(self):
        """Test finding parens in the middle of larger code."""
        code = "int x = method(param1, param2); other code"
        content, end_pos = self.transformer._find_balanced_parens(code, 13) # 672ns -> 612ns (9.80% faster)

    # ============================================================================
    # EDGE TEST CASES - Test extreme or unusual conditions
    # ============================================================================

    def test_invalid_start_position_negative(self):
        """Test that negative position returns None."""
        code = "(arg)"
        content, end_pos = self.transformer._find_balanced_parens(code, -1) # 671ns -> 681ns (1.47% slower)

    def test_invalid_start_position_beyond_length(self):
        """Test that position beyond code length returns None."""
        code = "(arg)"
        content, end_pos = self.transformer._find_balanced_parens(code, 10) # 451ns -> 501ns (9.98% slower)

    def test_start_position_not_on_paren(self):
        """Test that starting on non-paren character returns None."""
        code = "arg(param)"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 621ns -> 621ns (0.000% faster)

    def test_unbalanced_missing_closing_paren(self):
        """Test that unbalanced parens (missing closing) returns None."""
        code = "(arg"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 1.87μs -> 1.55μs (20.6% faster)

    def test_unbalanced_extra_closing_paren(self):
        """Test that code with extra closing paren after balanced set."""
        code = "(arg))"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 2.30μs -> 2.07μs (11.1% faster)

    def test_string_with_escaped_quote(self):
        """Test handling of escaped quotes in strings."""
        code = '("text\\"with\\"quotes")'
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 4.73μs -> 3.71μs (27.6% faster)

    def test_char_with_escaped_quote(self):
        """Test handling of escaped quotes in char literals."""
        code = "('\\''))"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 2.65μs -> 2.32μs (14.2% faster)

    def test_single_opening_paren_only(self):
        """Test with only opening parenthesis."""
        code = "("
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 842ns -> 931ns (9.56% slower)

    def test_multiple_consecutive_parens(self):
        """Test with multiple consecutive parentheses."""
        code = "()()()"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 1.58μs -> 1.47μs (7.47% faster)

    def test_string_with_unbalanced_paren(self):
        """Test that unbalanced parens in strings don't affect balance counting."""
        code = '("unpaired(paren", "another")'
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 5.70μs -> 4.41μs (29.3% faster)

    def test_empty_code_string(self):
        """Test with empty code string."""
        code = ""
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 491ns -> 521ns (5.76% slower)

    def test_code_with_only_closing_paren(self):
        """Test with only closing parenthesis."""
        code = ")"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 631ns -> 591ns (6.77% faster)

    def test_newlines_in_parens(self):
        """Test that newlines inside parens are preserved."""
        code = "(arg1,\narg2)"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 3.31μs -> 2.67μs (23.6% faster)

    def test_tabs_in_parens(self):
        """Test that tabs inside parens are preserved."""
        code = "(arg1,\targ2)"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 3.29μs -> 2.63μs (24.8% faster)

    def test_special_characters_in_parens(self):
        """Test with special characters inside parens."""
        code = "(arg@#$%^&*)"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 3.19μs -> 2.58μs (23.2% faster)

    def test_java_keywords_in_parens(self):
        """Test with Java keywords inside parens."""
        code = "(if else while for class int)"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 5.52μs -> 4.22μs (30.9% faster)

    def test_string_containing_backslash_before_quote(self):
        """Test string with backslash that doesn't escape next char."""
        code = '("path\\\\file")'
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 3.65μs -> 3.02μs (20.9% faster)

    def test_multiple_string_literals(self):
        """Test with multiple string literals inside parens."""
        code = '("string1", "string2", "string3")'
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 6.22μs -> 4.75μs (31.0% faster)

    def test_multiple_char_literals(self):
        """Test with multiple character literals inside parens."""
        code = "('a', 'b', 'c')"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 3.98μs -> 3.27μs (21.8% faster)

    def test_mixed_string_and_char_literals(self):
        """Test with both string and character literals."""
        code = '("string", \'c\', "another")'
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 5.61μs -> 4.40μs (27.6% faster)

    def test_string_followed_by_paren(self):
        """Test string ending with paren-like character."""
        code = '("text)")'
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 3.00μs -> 2.52μs (18.6% faster)

    def test_char_followed_by_paren(self):
        """Test character literal ending with paren-like character."""
        code = "(')')'"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 2.44μs -> 2.08μs (17.3% faster)

    def test_escaped_backslash_before_quote(self):
        """Test handling of escaped backslash followed by quote."""
        code = '("test\\\\")'
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 2.58μs -> 2.12μs (21.7% faster)

    def test_position_at_exact_string_length(self):
        """Test position at exactly the length of the code."""
        code = "text"
        content, end_pos = self.transformer._find_balanced_parens(code, 4) # 570ns -> 471ns (21.0% faster)

    # ============================================================================
    # LARGE SCALE TEST CASES - Test performance and scalability
    # ============================================================================

    def test_many_nested_parentheses(self):
        """Test with many levels of nesting (100 levels)."""
        # Build nested parens: ((((...))))
        code = "(" * 100 + ")" * 100
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 29.9μs -> 21.4μs (39.2% faster)

    def test_many_arguments(self):
        """Test with many comma-separated arguments (500)."""
        args = ", ".join([f"arg{i}" for i in range(500)])
        code = f"({args})"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 625μs -> 370μs (68.8% faster)

    def test_large_string_inside_parens(self):
        """Test with large string literal inside parens."""
        large_string = "x" * 1000
        code = f'("{large_string}")'
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 141μs -> 77.2μs (83.1% faster)

    def test_many_nested_function_calls(self):
        """Test with many nested function calls (50 levels)."""
        code = "func" + "(".join([""] * 50) + "arg" + ")".join([""] * 50)
        # Extract just the parentheses part
        start_paren = code.index("(")
        content, end_pos = self.transformer._find_balanced_parens(code, start_paren) # 16.1μs -> 11.7μs (37.7% faster)

    def test_alternating_nested_and_sequential(self):
        """Test with alternating nested and sequential parens."""
        # Pattern: (a(b)c(d(e)f)g)
        code = "(a(b)c(d(e)f)g)"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 3.87μs -> 3.22μs (20.2% faster)

    def test_very_long_code_with_late_paren(self):
        """Test with very long code but parens appear late."""
        # Create long code before the parenthesis
        long_prefix = "x = " * 200  # Creates 800 character prefix
        code = long_prefix + "(arg)"
        start_pos = len(long_prefix)
        content, end_pos = self.transformer._find_balanced_parens(code, start_pos) # 2.73μs -> 2.42μs (12.8% faster)

    def test_many_string_literals_large_scale(self):
        """Test with 200 string literals inside parens."""
        strings = ", ".join([f'"str{i}"' for i in range(200)])
        code = f"({strings})"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 288μs -> 167μs (72.1% faster)

    def test_many_char_literals_large_scale(self):
        """Test with 100 character literals inside parens."""
        chars = ", ".join([f"'{chr(65 + (i % 26))}'" for i in range(100)])
        code = f"({chars})"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 79.5μs -> 46.6μs (70.8% faster)

    def test_balanced_parens_with_large_mixed_content(self):
        """Test with large mixed content (strings, chars, parens, args)."""
        code = '(arg1, "string(with)parens", \'c\', func(nested), "another", x, y, z, 123, true, false, null)'
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 14.5μs -> 10.5μs (37.5% faster)

    def test_deeply_nested_with_multiple_strings(self):
        """Test deeply nested structure (50 levels) with strings at various levels."""
        # Build: (arg1, "str1", (arg2, "str2", (arg3, ...)))
        inner = "arg50"
        for i in range(49, 0, -1):
            inner = f'arg{i}, "str{i}", ({inner})'
        code = f"({inner})"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 140μs -> 84.1μs (66.9% faster)

    def test_sequential_paren_groups_after_first(self):
        """Test that we correctly identify closing paren of first group in sequence."""
        code = "(first)(second)(third)"
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 2.62μs -> 2.23μs (17.5% faster)
        # Verify second group
        content, end_pos = self.transformer._find_balanced_parens(code, 7) # 1.84μs -> 1.46μs (26.0% faster)

    def test_performance_many_unbalanced_nested(self):
        """Test performance with many nested levels without balance issues."""
        # Create 200 levels of nesting
        code = "(" * 200 + "content" + ")" * 200
        content, end_pos = self.transformer._find_balanced_parens(code, 0) # 67.2μs -> 42.6μs (57.7% faster)

    def test_large_code_with_multiple_parens_positions(self):
        """Test finding different parens groups in large code."""
        code = "func1(arg1) code func2(arg2) more code func3(arg3)"
        # Find first
        content1, pos1 = self.transformer._find_balanced_parens(code, 6) # 651ns -> 661ns (1.51% slower)
        # Find second
        content2, pos2 = self.transformer._find_balanced_parens(code, 29) # 301ns -> 291ns (3.44% faster)
        # Find third
        content3, pos3 = self.transformer._find_balanced_parens(code, 51) # 220ns -> 260ns (15.4% slower)
# 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-pr1199-2026-02-04T01.43.13 and push.