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

---

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

⏱️ Runtime : 2.56 milliseconds → 1.81 milliseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 41% runtime improvement by replacing character-by-character iteration with regex-based scanning to find special characters (', ", (, )).

Key Optimization

Original approach: Iterates through every character in the code string (26,253 iterations in profiler), checking each one against multiple conditions.

Optimized approach: Uses self._special_re.search(code, pos) to jump directly to the next special character (only 4,621 iterations in profiler), reducing iteration count by ~82%.

Why This Works

1. Reduces iteration overhead: In typical Java code, special characters are sparse. The regex engine (implemented in C) efficiently scans to the next occurrence, skipping irrelevant characters like alphanumerics, whitespace, and operators.

2. Per-character cost reduction: The profiler shows the original while pos < end and depth > 0: line alone consumed 15.6% of runtime with ~190ns per hit. The optimized version's m = self._special_re.search(code, pos) takes ~525ns per hit but executes 5.6x fewer times, resulting in net savings.

3. Elimination of escape tracking: The original tracked prev_char for every iteration. The optimized version checks code[i - 1] only when needed (at special character positions), avoiding 26,253 assignments.

Performance Characteristics

The optimization excels when processing:

• Large flat content (many arguments): 1051% faster on 1000 comma-separated elements because it skips over all the commas and identifiers
• Long strings with few special chars: 74.5% faster on large strings because it jumps past text content
• Mixed content: 13.5-53% faster on realistic mixed structures

Trade-off for deeply nested structures:

• Deep nesting (500 levels): 68% slower because regex overhead dominates when every character is a paren. This is acceptable since deeply nested structures are rare in practice.

The acceptance is justified by the significant runtime improvement on realistic code patterns where special characters represent a small fraction of total characters.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 94 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_simple_balanced_parens_basic():
    # Create a real instance of the transformer (use real constructor)
    t = JavaAssertTransformer("dummy")
    # Simple, balanced parentheses at the very start
    code = "(a, b)"
    # open paren at index 0
    content, pos = t._find_balanced_parens(code, 0) # 2.12μs -> 3.35μs (36.5% slower)

def test_nested_parens_basic():
    t = JavaAssertTransformer("dummy")
    # Parentheses that include a nested pair inside
    code = "prefix (outer(inner), x) suffix"
    open_pos = code.index("(")  # find the actual '(' position
    content, pos = t._find_balanced_parens(code, open_pos) # 3.40μs -> 4.59μs (26.0% slower)

def test_parentheses_with_strings_ignored_inside():
    t = JavaAssertTransformer("dummy")
    # A string literal inside contains a closing paren which should be ignored.
    code = 'call(") not end", other)'
    open_pos = code.index("(")
    content, pos = t._find_balanced_parens(code, open_pos) # 3.66μs -> 5.09μs (28.2% slower)

def test_char_literals_do_not_affect_parentheses_count():
    t = JavaAssertTransformer("dummy")
    # Char literals that contain parentheses should not be treated as paren tokens.
    # Two char literals containing '(' and ')' respectively are inside the parens.
    code = "f('(', ')')"
    open_pos = code.index("(")
    content, pos = t._find_balanced_parens(code, open_pos) # 2.71μs -> 5.79μs (53.3% slower)

def test_escaped_double_quote_within_string_does_not_end_string():
    t = JavaAssertTransformer("dummy")
    # The string contains an escaped double-quote and a closing paren; the escaped quote
    # should not prematurely terminate the string handling.
    code = 'g("a \\" ) end", b)'
    open_pos = code.index("(")
    content, pos = t._find_balanced_parens(code, open_pos) # 3.46μs -> 5.31μs (34.9% slower)

def test_unbalanced_parens_returns_none_and_minus_one():
    t = JavaAssertTransformer("dummy")
    # Missing closing parenthesis => should return (None, -1)
    code = "(incomplete"
    content, pos = t._find_balanced_parens(code, 0) # 2.20μs -> 1.51μs (45.7% faster)

def test_non_paren_position_returns_none_and_minus_one():
    t = JavaAssertTransformer("dummy")
    # Position points to a character that is not '(' => expected failure sentinel
    code = "no parentheses here"
    content, pos = t._find_balanced_parens(code, 0) # 632ns -> 632ns (0.000% faster)

def test_open_paren_pos_out_of_range_returns_none_and_minus_one():
    t = JavaAssertTransformer("dummy")
    code = "()"
    # Position equal to len(code) is out of bounds (no character there)
    content, pos = t._find_balanced_parens(code, len(code)) # 491ns -> 471ns (4.25% faster)

def test_large_number_of_comma_separated_elements_1000():
    t = JavaAssertTransformer("dummy")
    # Create a large argument list of 1000 elements to exercise loops and slicing performance.
    elements = ",".join(str(i) for i in range(1000))  # "0,1,2,...,999"
    code = "(" + elements + ")"
    content, pos = t._find_balanced_parens(code, 0) # 375μs -> 32.6μs (1051% faster)

def test_deeply_nested_parentheses_depth_500():
    t = JavaAssertTransformer("dummy")
    # Build a nested structure with depth 500 inside the outermost parentheses.
    # Outer '(' followed by 499 '(' then 'x' then 499 ')' then final ')'.
    inner_depth = 499
    nested = "(" * inner_depth + "x" + ")" * inner_depth
    code = "(" + nested + ")"
    # Starting at the very first '(' should find the corresponding closing paren at the end.
    content, pos = t._find_balanced_parens(code, 0) # 105μs -> 331μs (68.3% slower)
# 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 JavaAnalyzer
from codeflash.languages.java.remove_asserts import JavaAssertTransformer

class TestFindBalancedParensBasic:
    """Test basic functionality of _find_balanced_parens."""
    
    def test_simple_single_pair(self):
        """Test finding content within a simple pair of parentheses."""
        transformer = JavaAssertTransformer("testMethod")
        code = "method(hello)"
        # Position 6 is the opening parenthesis
        content, end_pos = transformer._find_balanced_parens(code, 6) # 2.24μs -> 3.12μs (28.0% slower)
    
    def test_empty_parentheses(self):
        """Test finding content within empty parentheses."""
        transformer = JavaAssertTransformer("testMethod")
        code = "method()"
        content, end_pos = transformer._find_balanced_parens(code, 6) # 1.46μs -> 2.88μs (49.3% slower)
    
    def test_nested_parentheses(self):
        """Test finding content with nested parentheses."""
        transformer = JavaAssertTransformer("testMethod")
        code = "method(func(inner))"
        content, end_pos = transformer._find_balanced_parens(code, 6) # 2.88μs -> 4.45μs (35.4% slower)
    
    def test_content_with_string_literal(self):
        """Test finding content when string contains parenthesis."""
        transformer = JavaAssertTransformer("testMethod")
        code = 'method("text(with)paren")'
        content, end_pos = transformer._find_balanced_parens(code, 6) # 3.33μs -> 5.42μs (38.6% slower)
    
    def test_multiple_arguments(self):
        """Test finding content with multiple comma-separated arguments."""
        transformer = JavaAssertTransformer("testMethod")
        code = "method(a, b, c)"
        content, end_pos = transformer._find_balanced_parens(code, 6) # 2.23μs -> 2.98μs (25.2% slower)
    
    def test_content_with_char_literal(self):
        """Test finding content when char literal contains parenthesis."""
        transformer = JavaAssertTransformer("testMethod")
        code = "method('(')"
        content, end_pos = transformer._find_balanced_parens(code, 6) # 2.06μs -> 4.63μs (55.4% slower)

class TestFindBalancedParensEdge:
    """Test edge cases of _find_balanced_parens."""
    
    def test_invalid_position_too_large(self):
        """Test with position beyond string length."""
        transformer = JavaAssertTransformer("testMethod")
        code = "method(hello)"
        content, end_pos = transformer._find_balanced_parens(code, 100) # 470ns -> 461ns (1.95% faster)
    
    def test_invalid_position_not_paren(self):
        """Test with position not pointing to opening parenthesis."""
        transformer = JavaAssertTransformer("testMethod")
        code = "method(hello)"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 641ns -> 642ns (0.156% slower)
    
    def test_unbalanced_missing_close(self):
        """Test with unbalanced parentheses (missing closing)."""
        transformer = JavaAssertTransformer("testMethod")
        code = "method(hello"
        content, end_pos = transformer._find_balanced_parens(code, 6) # 1.72μs -> 1.38μs (24.7% faster)
    
    def test_escaped_quote_in_string(self):
        """Test handling of escaped quotes in string literals."""
        transformer = JavaAssertTransformer("testMethod")
        code = r'method("text\"with(paren)")'
        content, end_pos = transformer._find_balanced_parens(code, 6) # 3.56μs -> 5.95μs (40.2% slower)
    
    def test_escaped_quote_in_char(self):
        """Test handling of escaped quotes in char literals."""
        transformer = JavaAssertTransformer("testMethod")
        code = r"method('\'')"
        content, end_pos = transformer._find_balanced_parens(code, 6) # 2.21μs -> 4.71μs (53.0% slower)
    
    def test_position_at_zero(self):
        """Test with position zero pointing to opening paren."""
        transformer = JavaAssertTransformer("testMethod")
        code = "(test)"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 1.86μs -> 2.98μs (37.6% slower)
    
    def test_deeply_nested_parens(self):
        """Test with deeply nested parentheses."""
        transformer = JavaAssertTransformer("testMethod")
        code = "(((((a)))))"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 2.83μs -> 6.34μs (55.5% slower)
    
    def test_string_with_escaped_backslash_before_quote(self):
        """Test string with escaped backslash followed by quote."""
        transformer = JavaAssertTransformer("testMethod")
        code = r'method("text\\(paren)")'
        content, end_pos = transformer._find_balanced_parens(code, 6) # 3.08μs -> 5.30μs (41.8% slower)
    
    def test_alternating_strings_and_parens(self):
        """Test content with alternating string literals and parentheses."""
        transformer = JavaAssertTransformer("testMethod")
        code = 'method("a", (b), "c")'
        content, end_pos = transformer._find_balanced_parens(code, 6) # 3.36μs -> 6.01μs (44.2% slower)
    
    def test_char_literal_with_backslash(self):
        """Test char literal containing escaped character."""
        transformer = JavaAssertTransformer("testMethod")
        code = r"method('\n')"
        content, end_pos = transformer._find_balanced_parens(code, 6) # 2.09μs -> 4.19μs (50.0% slower)
    
    def test_position_at_exact_end(self):
        """Test with position exactly at end of code."""
        transformer = JavaAssertTransformer("testMethod")
        code = "method()"
        content, end_pos = transformer._find_balanced_parens(code, len(code)) # 461ns -> 451ns (2.22% faster)
    
    def test_single_char_content(self):
        """Test with single character inside parentheses."""
        transformer = JavaAssertTransformer("testMethod")
        code = "(x)"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 1.72μs -> 3.09μs (44.2% slower)

class TestFindBalancedParensLargeScale:
    """Test performance and scalability of _find_balanced_parens."""
    
    def test_large_nested_depth(self):
        """Test with deeply nested parentheses (500 levels)."""
        transformer = JavaAssertTransformer("testMethod")
        # Build nested structure: (((((...))))
        code = "(" * 500 + "content" + ")" * 500
        content, end_pos = transformer._find_balanced_parens(code, 0) # 105μs -> 328μs (67.8% slower)
    
    def test_large_flat_arguments(self):
        """Test with many comma-separated arguments."""
        transformer = JavaAssertTransformer("testMethod")
        # Create 1000 arguments
        args = ", ".join([f"arg{i}" for i in range(1000)])
        code = f"({args})"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 762μs -> 63.8μs (1095% faster)
    
    def test_large_string_with_special_chars(self):
        """Test with large string literal containing many special characters."""
        transformer = JavaAssertTransformer("testMethod")
        # Create large string with parentheses, quotes, etc.
        large_string = '"' + "text(with)special\"chars" * 100 + '"'
        code = f"({large_string})"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 211μs -> 121μs (74.5% faster)
    
    def test_mixed_complexity_large(self):
        """Test with complex structure: mixed nesting, strings, and chars."""
        transformer = JavaAssertTransformer("testMethod")
        # Build: (arg1, "string()", func(x), 'c', ((nested)), ...)
        parts = []
        for i in range(100):
            if i % 4 == 0:
                parts.append(f'arg{i}')
            elif i % 4 == 1:
                parts.append(f'"string{i}()"')
            elif i % 4 == 2:
                parts.append(f"func{i}(x)")
            else:
                parts.append(f"'c'")
        args = ", ".join(parts)
        code = f"({args})"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 89.7μs -> 79.0μs (13.5% faster)
    
    def test_long_code_with_position_late(self):
        """Test with long code string and opening paren near the end."""
        transformer = JavaAssertTransformer("testMethod")
        prefix = "x" * 900
        suffix = "(content)"
        code = prefix + suffix
        paren_pos = len(prefix)
        content, end_pos = transformer._find_balanced_parens(code, paren_pos) # 2.67μs -> 3.46μs (22.6% slower)
    
    def test_many_string_literals_sequential(self):
        """Test with 500 sequential string literals."""
        transformer = JavaAssertTransformer("testMethod")
        strings = ", ".join([f'"string{i}"' for i in range(500)])
        code = f"({strings})"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 587μs -> 383μs (53.0% faster)
    
    def test_many_char_literals_sequential(self):
        """Test with 500 sequential char literals."""
        transformer = JavaAssertTransformer("testMethod")
        chars = ", ".join([f"'c'" for _ in range(500)])
        code = f"({chars})"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 229μs -> 325μs (29.7% slower)

class TestFindBalancedParensComprehensive:
    """Comprehensive tests covering corner cases and real-world scenarios."""
    
    def test_real_world_method_call_1(self):
        """Test realistic method call with multiple arguments."""
        transformer = JavaAssertTransformer("testMethod")
        code = 'assertEquals("expected", actual, "message")'
        # Find balanced parens starting at position 10 (opening paren after assertEquals)
        content, end_pos = transformer._find_balanced_parens(code, 10) # 632ns -> 621ns (1.77% faster)
    
    def test_real_world_method_call_2(self):
        """Test realistic nested method calls."""
        transformer = JavaAssertTransformer("testMethod")
        code = 'assertTrue(obj.method("test()", 42))'
        content, end_pos = transformer._find_balanced_parens(code, 10) # 4.40μs -> 6.38μs (31.1% slower)
    
    def test_real_world_lambda_expression(self):
        """Test lambda expression in parentheses."""
        transformer = JavaAssertTransformer("testMethod")
        code = '(x) -> x.getValue()'
        content, end_pos = transformer._find_balanced_parens(code, 0) # 1.76μs -> 3.01μs (41.4% slower)
    
    def test_empty_nested_structure(self):
        """Test empty nested parentheses."""
        transformer = JavaAssertTransformer("testMethod")
        code = '(()())'
        content, end_pos = transformer._find_balanced_parens(code, 0) # 2.15μs -> 5.03μs (57.2% slower)
    
    def test_string_only_content(self):
        """Test where entire content is a string."""
        transformer = JavaAssertTransformer("testMethod")
        code = '("hello world")'
        content, end_pos = transformer._find_balanced_parens(code, 0) # 2.94μs -> 4.43μs (33.5% slower)
    
    def test_char_literal_only_content(self):
        """Test where entire content is a char literal."""
        transformer = JavaAssertTransformer("testMethod")
        code = "('a')"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 2.05μs -> 4.24μs (51.5% slower)
    
    def test_multiple_parens_in_string(self):
        """Test string containing multiple parentheses."""
        transformer = JavaAssertTransformer("testMethod")
        code = '("test((()))test")'
        content, end_pos = transformer._find_balanced_parens(code, 0) # 3.18μs -> 6.34μs (49.9% slower)
    
    def test_multiple_quotes_in_content(self):
        """Test content with multiple quoted strings."""
        transformer = JavaAssertTransformer("testMethod")
        code = '("first", "second", "third")'
        content, end_pos = transformer._find_balanced_parens(code, 0) # 4.51μs -> 6.06μs (25.6% slower)
    
    def test_quote_after_escaped_backslash(self):
        """Test quote appearing after escaped backslash in string."""
        transformer = JavaAssertTransformer("testMethod")
        code = r'("\\\"test)")'  # String containing \\\" followed by test
        content, end_pos = transformer._find_balanced_parens(code, 0) # 2.83μs -> 5.26μs (46.3% slower)
    
    def test_position_one(self):
        """Test with position index 1 (second character, not paren)."""
        transformer = JavaAssertTransformer("testMethod")
        code = "a(test)"
        content, end_pos = transformer._find_balanced_parens(code, 1) # 2.04μs -> 2.88μs (28.9% slower)
    
    def test_complex_numeric_expressions(self):
        """Test parentheses with complex numeric expressions."""
        transformer = JavaAssertTransformer("testMethod")
        code = "(1 + 2 * 3 / 4 - 5)"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 3.29μs -> 2.87μs (14.7% faster)
    
    def test_generic_types_in_parens(self):
        """Test generics/angle brackets within parentheses."""
        transformer = JavaAssertTransformer("testMethod")
        code = "(List<String> list)"
        content, end_pos = transformer._find_balanced_parens(code, 0) # 3.31μs -> 3.06μs (8.25% 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-21T00.19.00 and push.

[claude]

PR Review Summary

Prek Checks

Status: ✅ Passing after fix

Fixed 1 issue:

• codeflash/languages/registry.py: Added # noqa: F401 to side-effect import (from codeflash.languages.java import support as _) to prevent ruff from incorrectly removing it. The other two language imports (python, javascript) didn't need the annotation.

mypy: ✅ No issues found in codeflash/languages/java/remove_asserts.py

Code Review

No critical issues found.

The optimization replaces character-by-character iteration in _find_balanced_parens with regex-based scanning (re.compile(r"""['"()]"""))) to jump directly to special characters. The approach is sound:

• Escape checking logic (code[i - 1] == "\\") is functionally equivalent to the original prev_char tracking
• String/char literal handling is preserved correctly
• Edge cases (unbalanced parens, out-of-bounds positions) still return (None, -1) sentinel values
• The _special_re regex is precompiled in __init__ (line 195), avoiding per-call overhead

Note: Both original and optimized code share a limitation with double-escaped characters (e.g., "\\\\" followed by a quote), but this is not a regression.

Test Coverage

File	Stmts	Miss	Coverage	Status
codeflash/languages/java/remove_asserts.py	449	55	88%	✅ New file, exceeds 75% threshold

This file is new (does not exist on main). The 88% coverage is well above the 75% minimum required for new files. The 94 generated regression tests provide good behavioral coverage of the optimized method.

---

Last updated: 2026-02-21