⚡️ Speed up method JavaAssertTransformer._find_balanced_parens by 35% 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.

---

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

⏱️ Runtime : 2.67 milliseconds → 1.98 milliseconds (best of 157 runs)

📝 Explanation and details

The optimized code achieves a 35% runtime improvement (from 2.67ms to 1.98ms) through two targeted optimizations that reduce overhead in the main parsing loop:

1. Cache len(code) as end variable

• Original: The loop condition while pos < len(code) and depth > 0 calls len(code) on every iteration (16,855 times in the profiler)
• Optimized: Pre-compute end = len(code) once and use while pos < end and depth > 0
• Impact: Reduces loop condition overhead from 16.1% to 14.1% of total time
• Why it's faster: Eliminates 16,854 redundant function calls to len(), replacing them with a simple integer comparison

2. Track prev_char incrementally instead of indexing backwards

• Original: prev_char = code[pos - 1] if pos > 0 else "" on every iteration (15.1% overhead)
• Optimized: Initialize prev_char = code[open_paren_pos] before the loop, then update prev_char = char at the end of each iteration (7.9% overhead)
• Impact: This is the primary optimization, nearly halving the cost of tracking the previous character
• Why it's faster: Eliminates 15,792 conditional expressions and backwards indexing operations, replacing them with simple variable assignments

Both optimizations target the hottest path in the code—the main parsing loop that executes thousands of times per invocation. The line profiler shows the while loop and character access operations consume 40-50% of total runtime, making these micro-optimizations highly effective.

Test results show consistent improvements across all scenarios:

• Simple cases: 5-15% faster
• Complex nested cases: 20-35% faster
• Large inputs (1000-depth nesting, 1000 repeated calls): 35-55% faster

The optimizations are particularly effective for complex Java code with deep nesting or long parenthesized expressions, which is exactly where this parser would be used most frequently in test transformation workflows.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 1140 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_not_open_paren_returns_none():
    # Create a real transformer instance with a dummy function name.
    t = JavaAssertTransformer("dummy")
    # If the provided position does not point to '(' we expect (None, -1).
    codeflash_output = t._find_balanced_parens("no paren here", 0); result = codeflash_output # 601ns -> 611ns (1.64% slower)

    # Also test position beyond the end-of-string returns the same sentinel.
    codeflash_output = t._find_balanced_parens("()", 10); result2 = codeflash_output # 310ns -> 290ns (6.90% faster)

def test_simple_and_nested_parentheses():
    t = JavaAssertTransformer("dummy")
    # Simple parentheses: "(abc)" should yield "abc" and position after closing paren.
    code = "(abc)"
    content, pos = t._find_balanced_parens(code, 0) # 2.23μs -> 2.02μs (10.4% faster)

    # Nested parentheses: "(a(b)c)" -> interior "a(b)c"
    code2 = "(a(b)c)"
    content2, pos2 = t._find_balanced_parens(code2, 0) # 1.88μs -> 1.51μs (24.5% faster)

def test_parentheses_ignored_inside_double_quoted_strings():
    t = JavaAssertTransformer("dummy")
    # Parentheses inside string literals should be ignored when balancing.
    code = '("a(b)c", x)'  # the '(' inside "a(b)c" should not affect depth.
    content, pos = t._find_balanced_parens(code, 0) # 3.37μs -> 2.94μs (14.3% faster)

def test_parentheses_ignored_inside_strings_with_escaped_quotes():
    t = JavaAssertTransformer("dummy")
    # This Java source has a string literal containing escaped double-quotes and parentheses.
    # Java source example desired: ("a\"(b)\"", x)
    # To have backslashes in the actual Python string, we use \\\" sequences.
    code = '("a\\\"(b)\\\"", x)'
    content, pos = t._find_balanced_parens(code, 0) # 3.85μs -> 3.21μs (20.0% faster)

def test_parentheses_ignored_inside_char_literals():
    t = JavaAssertTransformer("dummy")
    # Parenthesis inside a character literal should be ignored: Java: ('(', x)
    code = "('(', x)"
    content, pos = t._find_balanced_parens(code, 0) # 2.85μs -> 2.56μs (11.4% faster)

def test_unbalanced_parentheses_returns_none():
    t = JavaAssertTransformer("dummy")
    # Missing closing parenthesis -> should return the sentinel (None, -1).
    code = "(unclosed"
    content, pos = t._find_balanced_parens(code, 0) # 2.44μs -> 2.04μs (19.6% faster)

    # Extra closing without matching open (starting at a later '(' should still be invalid if not balanced)
    code2 = "prefix ) (a)"
    # open_paren_pos points at the '(' before 'a' which is balanced, so it's valid:
    idx = code2.index("(")
    content2, pos2 = t._find_balanced_parens(code2, idx) # 1.26μs -> 1.19μs (5.78% faster)

def test_open_paren_not_at_position_returns_none():
    t = JavaAssertTransformer("dummy")
    code = "(a)"
    # ask for position 1 which is 'a' not '('
    codeflash_output = t._find_balanced_parens(code, 1) # 601ns -> 631ns (4.75% slower)

def test_large_nested_parentheses_depth_1000():
    t = JavaAssertTransformer("dummy")
    # Build a deeply nested parentheses string: 1000 opens, then 'x', then 1000 closes.
    depth = 1000
    code = "(" * depth + "x" + ")" * depth
    # Call on the very first '('
    content, pos = t._find_balanced_parens(code, 0) # 323μs -> 211μs (52.9% faster)
    # Interior should be (depth-1) opens + 'x' + (depth-1) closes
    expected_interior = "(" * (depth - 1) + "x" + ")" * (depth - 1)

def test_repeated_invocations_stability_and_performance():
    t = JavaAssertTransformer("dummy")
    # Run the same small extraction 1000 times to ensure stability and reasonable performance.
    small_code = "(alpha, beta)"
    for _ in range(1000):
        content, pos = t._find_balanced_parens(small_code, 0) # 1.98ms -> 1.49ms (33.3% faster)
# 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.remove_asserts import JavaAssertTransformer

class TestFindBalancedParens:
    """Comprehensive unit tests for _find_balanced_parens method."""

    def test_simple_empty_parentheses(self):
        """Test finding content in simple empty parentheses."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("()", 0) # 1.60μs -> 1.56μs (2.56% faster)

    def test_simple_single_argument(self):
        """Test finding content with a single simple argument."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(42)", 0) # 2.18μs -> 2.06μs (5.81% faster)

    def test_simple_multiple_arguments(self):
        """Test finding content with multiple arguments separated by commas."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(1, 2, 3)", 0) # 2.79μs -> 2.50μs (11.6% faster)

    def test_nested_parentheses(self):
        """Test finding content with nested parentheses."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(foo(bar))", 0) # 3.17μs -> 2.75μs (14.9% faster)

    def test_deeply_nested_parentheses(self):
        """Test finding content with deeply nested parentheses."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(a(b(c(d))))", 0) # 3.39μs -> 2.88μs (17.8% faster)

    def test_string_with_parenthesis_inside(self):
        """Test that parentheses inside strings are ignored."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('("hello(world)")', 0) # 3.81μs -> 3.13μs (21.8% faster)

    def test_string_with_double_quote(self):
        """Test handling of strings delimited by double quotes."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('(message)', 0) # 2.58μs -> 2.38μs (8.84% faster)

    def test_character_literal_with_parenthesis(self):
        """Test that parentheses inside character literals are ignored."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("('(')", 0) # 2.29μs -> 2.13μs (7.50% faster)

    def test_character_literal_with_single_quote(self):
        """Test handling of character literals delimited by single quotes."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("('a')", 0) # 2.23μs -> 2.09μs (6.69% faster)

    def test_mixed_strings_and_literals(self):
        """Test handling of mixed string and character literals."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('(")", \'(\')', 0) # 3.37μs -> 2.89μs (16.6% faster)

    def test_escaped_quote_in_string(self):
        """Test that escaped quotes don't terminate strings."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('("say \\"hello\\"")', 0) # 3.91μs -> 3.24μs (20.8% faster)

    def test_escaped_single_quote_in_char(self):
        """Test that escaped quotes don't terminate character literals."""
        transformer = JavaAssertTransformer("test")
        # Note: In the actual implementation, backslash before a character terminates it
        content, pos = transformer._find_balanced_parens("('\\'')", 0) # 2.50μs -> 2.24μs (11.6% faster)

    def test_invalid_starting_position_not_paren(self):
        """Test that function returns None when starting position is not an opening paren."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("abc", 0) # 571ns -> 612ns (6.70% slower)

    def test_invalid_starting_position_past_end(self):
        """Test that function returns None when starting position is past the end."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(test)", 10) # 470ns -> 481ns (2.29% slower)

    def test_invalid_unbalanced_parens_missing_close(self):
        """Test that function returns None when closing paren is missing."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(unclosed", 0) # 2.42μs -> 2.10μs (15.2% faster)

    def test_invalid_unbalanced_parens_extra_close(self):
        """Test handling when there's an extra closing paren (we only match the first balanced pair)."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(test))", 0) # 2.29μs -> 2.00μs (14.6% faster)

    def test_position_in_middle_of_string(self):
        """Test starting position in the middle of a string."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("prefix(content)", 6) # 2.65μs -> 2.27μs (16.3% faster)

    def test_empty_string_input(self):
        """Test with empty string input."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("", 0) # 461ns -> 481ns (4.16% slower)

    def test_only_opening_paren(self):
        """Test with only an opening parenthesis."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(", 0) # 811ns -> 912ns (11.1% slower)

    def test_method_call_like_expression(self):
        """Test with a method call-like expression."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(obj.method(arg1, arg2))", 0) # 4.99μs -> 4.05μs (23.3% faster)

    def test_complex_expression_with_operators(self):
        """Test with complex expression containing various operators."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(a + b * c - d / e)", 0) # 4.14μs -> 3.32μs (24.8% faster)

    def test_lambda_expression(self):
        """Test with Java lambda expression inside parentheses."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(x -> x * 2)", 0) # 3.14μs -> 2.58μs (21.8% faster)

    def test_generic_type_brackets(self):
        """Test with generic type parameters (angle brackets are not parens)."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(List<String> list)", 0) # 4.13μs -> 3.35μs (23.4% faster)

    def test_whitespace_preservation(self):
        """Test that whitespace inside parentheses is preserved."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(  a  ,  b  )", 0) # 3.31μs -> 2.77μs (19.6% faster)

    def test_newline_inside_parens(self):
        """Test with newlines inside parentheses."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(a\nb)", 0) # 2.15μs -> 1.94μs (10.8% faster)

    def test_tab_inside_parens(self):
        """Test with tab characters inside parentheses."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(a\tb)", 0) # 2.13μs -> 1.99μs (7.07% faster)

    def test_multiple_string_literals(self):
        """Test with multiple string literals in sequence."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('("hello", "world")', 0) # 4.27μs -> 3.53μs (21.0% faster)

    def test_string_with_escaped_backslash_then_quote(self):
        """Test string with escaped backslash followed by quote."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('("test\\\\")', 0) # 2.54μs -> 1.99μs (27.7% faster)

    def test_adjacent_char_literals(self):
        """Test with adjacent character literals."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("('a')('b')", 0) # 2.28μs -> 2.12μs (7.53% faster)

    def test_large_balanced_parens(self):
        """Test with a large number of balanced parentheses."""
        transformer = JavaAssertTransformer("test")
        # Create string like ((((((((((a))))))))))
        code = "(" * 50 + "a" + ")" * 50
        content, pos = transformer._find_balanced_parens(code, 0) # 16.1μs -> 11.9μs (35.8% faster)

    def test_alternating_nested_parens(self):
        """Test with alternating patterns of nested parentheses."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("((a)(b)(c))", 0) # 3.19μs -> 2.81μs (13.2% faster)

    def test_realistic_java_assertion(self):
        """Test with a realistic Java assertion argument."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens(
            '(assertEquals("expected", actual.getValue()))', 0
        ) # 8.21μs -> 6.14μs (33.6% faster)

    def test_string_containing_escaped_quote_sequence(self):
        """Test string containing multiple escaped quotes."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('("a\\"b\\"c")', 0) # 3.17μs -> 2.71μs (16.6% faster)

    def test_char_literal_with_backslash_n(self):
        """Test character literal with escape sequence."""
        transformer = JavaAssertTransformer("test")
        # '\n' represents a newline character literal
        content, pos = transformer._find_balanced_parens("('\\n')", 0) # 2.44μs -> 2.20μs (10.9% faster)

    def test_position_at_exact_closing_paren_location(self):
        """Test starting position exactly at the closing paren (should fail)."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(test)", 5) # 571ns -> 611ns (6.55% slower)

    def test_complex_realistic_junit_assertion(self):
        """Test with realistic JUnit assertion pattern."""
        transformer = JavaAssertTransformer("test")
        code = '(assertTrue(obj.method(arg1, arg2) > 0 && flag))'
        content, pos = transformer._find_balanced_parens(code, 0) # 8.40μs -> 6.33μs (32.7% faster)

    def test_string_with_unmatched_single_quote(self):
        """Test string containing an unmatched single quote."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('("don\'t")', 0) # 3.02μs -> 2.56μs (17.9% faster)

    def test_very_long_content(self):
        """Test with very long content inside parentheses."""
        transformer = JavaAssertTransformer("test")
        long_content = "x" * 1000
        code = "(" + long_content + ")"
        content, pos = transformer._find_balanced_parens(code, 0) # 152μs -> 99.1μs (54.3% faster)

    def test_multiple_levels_with_strings(self):
        """Test nested parentheses with strings at various levels."""
        transformer = JavaAssertTransformer("test")
        code = '(outer("inner(fake)", real(param)))'
        content, pos = transformer._find_balanced_parens(code, 0) # 6.73μs -> 5.21μs (29.2% faster)

    def test_opening_paren_at_end_of_string(self):
        """Test when opening paren is the last character."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("test(", 4) # 802ns -> 881ns (8.97% slower)

    def test_consecutive_method_calls(self):
        """Test parsing consecutive method call parentheses."""
        transformer = JavaAssertTransformer("test")
        code = "(obj.method1(arg))"
        content, pos = transformer._find_balanced_parens(code, 0) # 4.28μs -> 3.48μs (23.1% faster)

    def test_char_literal_at_boundary(self):
        """Test character literal immediately before closing paren."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(x, 'y')", 0) # 2.91μs -> 2.50μs (16.0% faster)

    def test_empty_string_literal(self):
        """Test with empty string literal."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('("")', 0) # 2.15μs -> 2.03μs (5.90% faster)

    def test_empty_char_literal_like(self):
        """Test sequence that looks like empty char but isn't valid Java."""
        transformer = JavaAssertTransformer("test")
        # This tests the parser behavior, not Java validity
        content, pos = transformer._find_balanced_parens("('')", 0) # 2.05μs -> 1.90μs (7.83% faster)

    def test_string_with_digit_and_paren(self):
        """Test string containing digits and parentheses."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('(test(123))', 0) # 3.24μs -> 2.77μs (16.6% faster)

    def test_offset_start_position_in_valid_parens(self):
        """Test starting from non-zero position with valid opening paren."""
        transformer = JavaAssertTransformer("test")
        code = "junk(content)"
        content, pos = transformer._find_balanced_parens(code, 4) # 2.69μs -> 2.37μs (14.0% faster)

    def test_only_closing_paren(self):
        """Test with only a closing parenthesis."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens(")", 0) # 571ns -> 610ns (6.39% slower)

    def test_multiple_parens_choose_first(self):
        """Test that starting at first paren captures correct balanced pair."""
        transformer = JavaAssertTransformer("test")
        code = "(first)(second)"
        content, pos = transformer._find_balanced_parens(code, 0) # 2.46μs -> 2.16μs (13.9% faster)

    def test_special_chars_in_content(self):
        """Test with various special characters inside parentheses."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(!@#$%^&*)", 0) # 2.85μs -> 2.50μs (14.0% faster)

    def test_hyphen_and_arithmetic_operators(self):
        """Test with arithmetic operators and negative numbers."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(a - b + c * -5)", 0) # 3.65μs -> 3.04μs (20.1% faster)

    def test_equals_and_comparison_operators(self):
        """Test with comparison and equality operators."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(a == b && c != d)", 0) # 4.01μs -> 3.24μs (23.8% faster)

    def test_bitwise_operators(self):
        """Test with bitwise operators."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(a & b | c ^ d)", 0) # 3.52μs -> 2.92μs (20.6% faster)

    def test_ternary_operator(self):
        """Test with ternary conditional operator."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(a ? b : c)", 0) # 2.97μs -> 2.49μs (18.9% faster)

    def test_cast_expression(self):
        """Test with type cast expressions."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("((String) obj)", 0) # 3.62μs -> 2.98μs (21.2% faster)

    def test_array_access(self):
        """Test with array access (note: square brackets, not parens)."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(arr[i])", 0) # 2.58μs -> 2.31μs (12.1% faster)

    def test_dot_operator(self):
        """Test with dot operator for member access."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(obj.field.method())", 0) # 4.41μs -> 3.60μs (22.5% faster)

    def test_string_with_multiple_escape_sequences(self):
        """Test string with multiple different escape sequences."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('("a\\tb\\nc\\\\d")', 0) # 3.51μs -> 2.98μs (17.8% faster)

    def test_alternate_quote_marks_in_strings(self):
        """Test string containing the other type of quote."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens('("He said \'hello\'")', 0) # 4.23μs -> 3.41μs (24.1% faster)

    def test_numeric_literals(self):
        """Test with various numeric literals."""
        transformer = JavaAssertTransformer("test")
        content, pos = transformer._find_balanced_parens("(42, 3.14, 0xFF, 1_000_000L)", 0) # 5.27μs -> 4.15μs (27.0% faster)

    def test_realistic_assertion_with_message(self):
        """Test realistic assertion with message parameter."""
        transformer = JavaAssertTransformer("test")
        code = '(result, "Expected X but got Y")'
        content, pos = transformer._find_balanced_parens(code, 0) # 5.79μs -> 4.55μs (27.3% 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-pr1199-2026-02-20T12.22.18 and push.

[claude]

PR Review Summary

Prek Checks

✅ All checks pass — ruff check and ruff format both passed with no issues.

Mypy

✅ No type errors found in codeflash/languages/java/remove_asserts.py.

Code Review

✅ No critical issues found. The optimization is behavior-preserving:

• end = len(code) caching: Safe — code is an immutable string parameter unchanged during the loop.
• Incremental prev_char tracking: Correctly initialized to code[open_paren_pos] (equals code[pos - 1] since pos = open_paren_pos + 1), then updated to char after each iteration. The original if pos > 0 else "" guard was always true since pos >= 1.

Test Coverage

File	Stmts	Miss	Coverage
codeflash/languages/java/remove_asserts.py	447	55	88%

• ✅ All changed executable lines (800, 808, 810, 832) are covered by tests
• Line 807 is a comment (not executable)
• 157 tests pass covering this file
• No coverage regression (file is modified, not new)

Optimization PRs Status

No codeflash optimization PRs were merged — all 26 open PRs have CI failures (common: code/snyk rate limit, js-esm-async-optimization, tracer-replay, unit-tests (windows-latest, 3.13)). PR #1602 itself still has pending checks.

---

Last updated: 2026-02-20