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

---

📄 81% (0.81x) speedup for _infer_array_cast_type in codeflash/languages/java/instrumentation.py

⏱️ Runtime : 4.27 milliseconds → 2.36 milliseconds (best of 148 runs)

📝 Explanation and details

The optimized code achieves an 81% speedup (from 4.27ms to 2.36ms) by eliminating expensive generator overhead and reducing redundant operations.

Key Optimizations:

1. Replaced any() generator with direct string checks: The original code used any(method in line for method in assertion_methods) which creates a generator and performs up to 2 substring searches with early termination. The optimized version uses explicit "assertArrayEquals" not in line and "assertArrayNotEquals" not in line, which leverages Python's short-circuit evaluation to perform at most 2 substring searches but without generator overhead. Line profiler shows this reduced time from 4.53ms (47.5% of total) to 870μs (16.5% of total) - a 5.2x improvement on this single line.

2. Inverted match condition logic: Changed if match: to if not match: return None to enable early returns and reduce nesting, which slightly improves code path efficiency.

3. Removed module-level tuple creation: While the unused _ASSERTION_METHODS constant was added, it's not actually used in the optimized function, so the real benefit comes from eliminating the repeated tuple creation inside the function (503μs in original vs 0μs in optimized).

Performance Characteristics:

The optimization excels across all test scenarios:

• Fast-path cases (no assertion methods): 156-198% faster due to immediate rejection without generator overhead
• Regex matching cases: 40-60% faster from reduced overhead before regex execution
• Bulk processing tests: 78-83% faster when processing 1000+ lines, showing excellent scalability
• Edge cases (empty strings, long lines): Consistent 33-191% improvements

The regex search itself (_PRIMITIVE_ARRAY_PATTERN.search) remains the dominant cost at 42.3% of optimized runtime (vs 24.6% original), but only because we've eliminated the larger bottleneck of the any() generator expression.

Impact: This function appears to be called frequently in Java code instrumentation contexts (3,677 hits in profiling). The optimization significantly reduces overhead for every assertion line processed, making it particularly valuable in hot paths where Java test code is being analyzed or transformed.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 3677 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.instrumentation import _infer_array_cast_type

def test_no_assertion_method_returns_none():
    # If no assertion method name is present in the line, the function should return None.
    line = "some random code with new int[] and other tokens"
    codeflash_output = _infer_array_cast_type(line) # 1.46μs -> 571ns (156% faster)

def test_infers_int_array_for_assertarrayequals():
    # A typical assertion containing a primitive int array should return "int[]".
    line = 'assertArrayEquals("message", new int[] {1, 2, 3}, result);'
    codeflash_output = _infer_array_cast_type(line) # 3.80μs -> 2.39μs (58.6% faster)

@pytest.mark.parametrize(
    "primitive",
    ["int", "long", "double", "float", "short", "byte", "char", "boolean"],
)
def test_infers_all_primitives_for_both_assert_methods(primitive):
    # Ensure that every primitive type is recognized for both supported assertion methods.
    for method in ("assertArrayEquals", "assertArrayNotEquals"):
        # construct a line that contains the assertion method and a primitive array instantiation
        line = f'{method}("m", new {primitive}[] {{ }}, someOther);'
        # the function should return the primitive with [] appended
        codeflash_output = _infer_array_cast_type(line) # 44.7μs -> 29.0μs (54.1% faster)

def test_ignores_non_primitive_object_arrays():
    # When the array is an Object[] (non-primitive), the function should return None.
    line = 'assertArrayEquals("msg", new Object[] {"a"}, result);'
    codeflash_output = _infer_array_cast_type(line) # 2.94μs -> 1.65μs (77.7% faster)

def test_handles_extra_whitespace_and_tabs():
    # The regex allows extra spaces between "new", the type, and the brackets.
    line = "assertArrayEquals(  new   int   [   ]   , something );"
    codeflash_output = _infer_array_cast_type(line) # 3.69μs -> 2.36μs (56.0% faster)

def test_case_sensitivity_of_method_names():
    # Method name matching is case-sensitive; lowercased method should not be detected.
    line = "assertarrayequals(new int[] {1});"  # lowercased method name
    codeflash_output = _infer_array_cast_type(line) # 1.53μs -> 591ns (159% faster)

def test_empty_string_returns_none():
    # An empty input line should safely return None.
    codeflash_output = _infer_array_cast_type("") # 1.37μs -> 461ns (198% faster)

def test_none_input_raises_type_error():
    # Passing None (wrong type) should raise a TypeError when the function tries to iterate over it.
    with pytest.raises(TypeError):
        _infer_array_cast_type(None) # 4.02μs -> 3.02μs (33.2% faster)

def test_multiple_primitive_matches_returns_first():
    # If multiple primitive array constructions appear, the function should return the first matched primitive type.
    # Place 'long' first and 'int' later to ensure the first is returned.
    line = 'assertArrayEquals("m", new long[] {1L}, new int[] {1});'
    codeflash_output = _infer_array_cast_type(line) # 3.84μs -> 2.38μs (60.9% faster)

def test_no_space_after_new_prevents_match():
    # The regex requires whitespace after 'new'. 'newint[]' should not match.
    line = "assertArrayEquals(newint[] {1}, result);"
    codeflash_output = _infer_array_cast_type(line) # 2.83μs -> 1.49μs (89.2% faster)

def test_repeated_invocations_stability_large_loop():
    # Repeatedly call the function many times to ensure stability and consistent results.
    # Use 1000 iterations as a large-scale scenario to detect regressions that appear under load.
    for i in range(1000):
        # alternate between two forms to vary content
        if i % 2 == 0:
            line = f'assertArrayEquals("m{i}", new int[] {{ {i} }}, res{i});'
            expected = "int[]"
        else:
            line = f'assertArrayNotEquals({i}, new boolean[] {{ true }}, res{i});'
            expected = "boolean[]"
        # Each call must deterministically return the expected primitive array type.
        codeflash_output = _infer_array_cast_type(line) # 1.17ms -> 654μs (78.6% faster)

def test_many_different_lines_processed_quickly():
    # Construct 1000 different lines with various primitives mixed in and verify correct detection.
    primitives = ["int", "long", "double", "float", "short", "byte", "char", "boolean"]
    # Create 1000 lines cycling through the primitives and both assertion methods.
    lines = []
    expected_results = []
    for i in range(1000):
        prim = primitives[i % len(primitives)]
        method = "assertArrayEquals" if (i % 3) != 0 else "assertArrayNotEquals"
        # embed the primitive array near the start of the line
        line = f'{method}("m", new {prim}[] {{ }}, other{i});'
        lines.append(line)
        expected_results.append(f"{prim}[]")

    # Validate all lines produce expected results
    for line, expected in zip(lines, expected_results):
        codeflash_output = _infer_array_cast_type(line) # 1.16ms -> 637μs (82.8% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import re

import pytest
from codeflash.languages.java.instrumentation import _infer_array_cast_type

def test_basic_int_array_cast():
    """Test detection of int[] primitive array cast type."""
    line = "assertArrayEquals(new int[] {1, 2, 3}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 4.25μs -> 2.94μs (44.7% faster)

def test_basic_long_array_cast():
    """Test detection of long[] primitive array cast type."""
    line = "assertArrayEquals(new long[] {1L, 2L}, actual);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.87μs -> 2.56μs (51.4% faster)

def test_basic_double_array_cast():
    """Test detection of double[] primitive array cast type."""
    line = "assertArrayEquals(new double[] {1.0, 2.5}, value);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.82μs -> 2.42μs (57.4% faster)

def test_basic_float_array_cast():
    """Test detection of float[] primitive array cast type."""
    line = "assertArrayEquals(new float[] {1.0f}, arr);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.62μs -> 2.58μs (40.0% faster)

def test_basic_short_array_cast():
    """Test detection of short[] primitive array cast type."""
    line = "assertArrayEquals(new short[] {1, 2}, data);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.63μs -> 2.56μs (41.5% faster)

def test_basic_byte_array_cast():
    """Test detection of byte[] primitive array cast type."""
    line = "assertArrayEquals(new byte[] {1, 2}, bytes);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.59μs -> 2.52μs (42.6% faster)

def test_basic_char_array_cast():
    """Test detection of char[] primitive array cast type."""
    line = "assertArrayEquals(new char[] {'a', 'b'}, chars);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.65μs -> 2.51μs (45.0% faster)

def test_basic_boolean_array_cast():
    """Test detection of boolean[] primitive array cast type."""
    line = "assertArrayEquals(new boolean[] {true, false}, flags);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.78μs -> 2.63μs (43.3% faster)

def test_no_array_assertion_returns_none():
    """Test that non-array assertion lines return None."""
    line = "assertEquals(5, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 1.62μs -> 652ns (149% faster)

def test_no_primitive_array_returns_none():
    """Test that assertion with object array returns None."""
    line = "assertArrayEquals(new String[] {\"a\", \"b\"}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.19μs -> 1.84μs (72.8% faster)

def test_extra_whitespace_before_bracket():
    """Test detection with extra whitespace before opening bracket."""
    line = "assertArrayEquals(new int  [  ] {1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.73μs -> 2.55μs (45.9% faster)

def test_no_whitespace_before_bracket():
    """Test detection with no whitespace before bracket."""
    line = "assertArrayEquals(new int[]{1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.69μs -> 2.51μs (46.7% faster)

def test_tabs_instead_of_spaces():
    """Test detection with tabs as whitespace."""
    line = "assertArrayEquals(new int\t[\t] {1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.60μs -> 2.40μs (49.6% faster)

def test_multiple_tabs_and_spaces():
    """Test detection with mixed tabs and spaces."""
    line = "assertArrayEquals(new long \t [ \t ] {}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.65μs -> 2.50μs (45.6% faster)

def test_newline_in_middle_of_declaration():
    """Test detection with newline in array type declaration."""
    line = "assertArrayEquals(new int\n[] {1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.66μs -> 2.50μs (46.6% faster)

def test_assertArrayEquals_with_message():
    """Test detection of primitive array with message argument."""
    line = 'assertArrayEquals("Arrays should be equal", new int[] {1, 2}, result);'
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.63μs -> 2.44μs (48.3% faster)

def test_assertArrayNotEquals_with_int_array():
    """Test detection with assertArrayNotEquals method."""
    line = "assertArrayNotEquals(new int[] {1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.81μs -> 2.61μs (45.6% faster)

def test_assertArrayNotEquals_with_long_array():
    """Test detection of long[] with assertArrayNotEquals."""
    line = "assertArrayNotEquals(new long[] {5L}, value);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.77μs -> 2.54μs (48.6% faster)

def test_empty_assertion_method_call():
    """Test line with assertion method but no array."""
    line = "assertArrayEquals(expected, actual);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 2.18μs -> 941ns (132% faster)

def test_assertion_with_partial_method_name():
    """Test that partial method name matches are detected."""
    line = "someMethodWithAssertArrayEqualsInside(new int[] {1}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 1.63μs -> 792ns (106% faster)

def test_multiple_new_keywords_in_line():
    """Test line with multiple 'new' keywords, one with primitive array."""
    line = "assertArrayEquals(someObject, new int[] {1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.82μs -> 2.48μs (53.6% faster)

def test_first_match_wins_with_multiple_arrays():
    """Test that first matching primitive array is returned when multiple exist."""
    line = "assertArrayEquals(new int[] {1}, new long[] {2L});"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.62μs -> 2.50μs (44.4% faster)

def test_case_sensitive_assertion_method():
    """Test that assertion method names are case-sensitive (lowercase)."""
    line = "AssertArrayEquals(new int[] {1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 1.55μs -> 581ns (167% faster)

def test_case_sensitive_primitive_type():
    """Test that primitive type names are case-sensitive (lowercase)."""
    line = "assertArrayEquals(new Int[] {1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 2.98μs -> 1.78μs (67.4% faster)

def test_empty_string():
    """Test with empty input string."""
    codeflash_output = _infer_array_cast_type(""); result = codeflash_output # 1.35μs -> 490ns (176% faster)

def test_whitespace_only_string():
    """Test with whitespace-only input string."""
    codeflash_output = _infer_array_cast_type("   \t\t   "); result = codeflash_output # 1.31μs -> 461ns (185% faster)

def test_assertion_method_at_end_of_line():
    """Test when assertion method appears at the very end."""
    line = "assertArrayEquals"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 2.22μs -> 982ns (126% faster)

def test_primitive_array_only():
    """Test when only primitive array pattern exists without assertion."""
    line = "new int[] {1, 2}"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 1.31μs -> 451ns (191% faster)

def test_very_long_line_with_int_array():
    """Test with a very long line containing int array assertion."""
    line = "assertArrayEquals(" + "x" * 1000 + ", new int[] {1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 4.27μs -> 3.10μs (37.9% faster)

def test_special_characters_before_assertion():
    """Test with special characters before assertion method name."""
    line = "// Comment: assertArrayEquals(new int[] {1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.69μs -> 2.62μs (41.0% faster)

def test_string_literal_containing_assertion():
    """Test that assertion method name in string literal still matches."""
    line = 'String msg = "assertArrayEquals"; assertArrayEquals(new int[] {1}, r);'
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.62μs -> 2.48μs (45.6% faster)

def test_assertion_in_comment():
    """Test that assertion method name in comment is still matched."""
    line = "// assertArrayEquals(new int[] {1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.71μs -> 2.41μs (53.5% faster)

def test_unicode_characters_in_line():
    """Test line with unicode characters."""
    line = "assertArrayEquals(new int[] {1, 2}, résultat); // ñoño"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.64μs -> 2.57μs (41.3% faster)

def test_single_element_array():
    """Test primitive array with single element."""
    line = "assertArrayEquals(new int[] {42}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.44μs -> 2.54μs (35.0% faster)

def test_empty_primitive_array():
    """Test empty primitive array declaration."""
    line = "assertArrayEquals(new int[] {}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.50μs -> 2.38μs (46.6% faster)

def test_multiline_string_escaped_newlines():
    """Test with escaped newlines in string context (still single line)."""
    line = "assertArrayEquals(new int[] {1, 2}, result); // escape \\n test"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.59μs -> 2.43μs (47.4% faster)

def test_large_array_with_many_elements():
    """Test detection in line with very large array initialization."""
    large_array_content = ", ".join(str(i) for i in range(500))
    line = f"assertArrayEquals(new int[] {{{large_array_content}}}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.61μs -> 2.22μs (62.2% faster)

def test_long_method_call_chain_with_assertion():
    """Test with long method call chain containing the assertion."""
    line = "obj.method1().method2().method3().assertArrayEquals(new long[] {1L, 2L}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.82μs -> 2.62μs (46.0% faster)

def test_multiple_assertions_on_line():
    """Test line with multiple assertion method names (both match)."""
    line = "assertArrayEquals(new int[] {1}, r); assertArrayNotEquals(new long[] {2L}, s);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 3.53μs -> 2.39μs (47.3% faster)

def test_repeated_pattern_searching():
    """Test function called repeatedly with similar inputs (1000 iterations)."""
    inputs = [
        ("assertArrayEquals(new int[] {1, 2}, result);", "int[]"),
        ("assertArrayEquals(new long[] {1L}, result);", "long[]"),
        ("assertArrayEquals(new double[] {1.0}, result);", "double[]"),
        ("assertArrayEquals(new float[] {1.0f}, result);", "float[]"),
        ("assertArrayEquals(new short[] {1}, result);", "short[]"),
        ("assertArrayEquals(new byte[] {1}, result);", "byte[]"),
        ("assertArrayEquals(new char[] {'a'}, result);", "char[]"),
        ("assertArrayEquals(new boolean[] {true}, result);", "boolean[]"),
        ("assertEquals(5, result);", None),
        ("assertArrayEquals(new String[] {\"a\"}, result);", None),
    ]
    
    # Run 1000 iterations through the test patterns
    for _ in range(100):
        for line, expected in inputs:
            codeflash_output = _infer_array_cast_type(line); result = codeflash_output

def test_stress_test_random_order():
    """Test with 500 variations of assertion patterns."""
    primitive_types = ["int", "long", "double", "float", "short", "byte", "char", "boolean"]
    assertion_methods = ["assertArrayEquals", "assertArrayNotEquals"]
    
    # Test 500 combinations
    test_count = 0
    for i, prim_type in enumerate(primitive_types * 63):  # 504 repetitions
        method = assertion_methods[i % 2]
        line = f"{method}(new {prim_type}[] {{1, 2}}, result);"
        codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 597μs -> 329μs (81.3% faster)
        test_count += 1

def test_performance_with_long_nonsensical_prefix():
    """Test that function handles long content before the assertion."""
    prefix = "x = " + " + ".join(f"value{i}" for i in range(200)) + "; "
    line = prefix + "assertArrayEquals(new int[] {1, 2}, result);"
    codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 5.43μs -> 4.06μs (33.8% faster)

def test_all_primitive_types_in_sequence():
    """Test all 8 primitive array types in a deterministic sequence."""
    test_cases = [
        ("assertArrayEquals(new int[] {}, r);", "int[]"),
        ("assertArrayEquals(new long[] {}, r);", "long[]"),
        ("assertArrayEquals(new double[] {}, r);", "double[]"),
        ("assertArrayEquals(new float[] {}, r);", "float[]"),
        ("assertArrayEquals(new short[] {}, r);", "short[]"),
        ("assertArrayEquals(new byte[] {}, r);", "byte[]"),
        ("assertArrayEquals(new char[] {}, r);", "char[]"),
        ("assertArrayEquals(new boolean[] {}, r);", "boolean[]"),
    ]
    
    for line, expected in test_cases:
        codeflash_output = _infer_array_cast_type(line); result = codeflash_output # 12.9μs -> 7.75μs (67.1% faster)

def test_negative_cases_in_bulk():
    """Test 100 lines that should return None."""
    negative_cases = [
        "assertEquals(5, result);",
        "assertThat(result).isNotNull();",
        "assertTrue(result);",
        "assertFalse(flag);",
        "assertNull(value);",
        "assertNotNull(value);",
        "assertArrayEquals(new String[] {\"a\"}, result);",
        "assertArrayEquals(new Object[] {obj}, result);",
        "assertArrayEquals(expected, actual);",
        "SomeOtherClass.someMethod(new int[] {1}, result);",
    ]
    
    # Repeat to get 100 test cases
    for _ in range(10):
        for line in negative_cases:
            codeflash_output = _infer_array_cast_type(line); result = codeflash_output
# 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-20T07.48.46 and push.

[claude]

PR Review Summary

Prek Checks

• Fixed: 1 linting issue (W293 blank-line-with-whitespace in instrumentation.py)
• Status: All prek checks passing after fix (committed and pushed as style: auto-fix linting issues)

Mypy

• 3 pre-existing mypy errors in instrumentation.py (all from the omni-java base branch, not introduced by this PR):
  • Line 865: no-redef - result_parts redefined after early return
  • Line 921: var-annotated - test_methods needs type annotation
  • Line 1069: assignment - instrument_existing_test returns str | None but assigned to str

Code Review

No issues found. The optimization is clean and correct:

• Replaces any(method in line for method in assertion_methods) with direct in checks, eliminating generator overhead
• Inverts match condition for early return pattern
• Moves assertion methods tuple to module-level _ASSERTION_METHODS constant (unused by optimized code but harmless)
• Logic is equivalent to the original implementation

Test Coverage

File	Stmts	Miss	Coverage
codeflash/languages/java/instrumentation.py	483	85	82%

• The optimized function _infer_array_cast_type (lines 241-265): The early-return path (no assertion method found) is covered. Lines 260-265 (regex match + return) are not covered by tests — this is pre-existing and not introduced by this PR.
• This file is new relative to main (part of the omni-java feature branch), so no main-branch comparison is applicable.
• Overall: 3221 passed, 30 failed (all failures are pre-existing/environment-related — Maven dependency resolution, tracer attribute errors).

---

Last updated: 2026-02-20