Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 0 additions & 29 deletions codeflash/languages/java/instrumentation.py
Original file line number Diff line number Diff line change
Expand Up @@ -1325,32 +1325,3 @@ def instrument_generated_java_test(

logger.debug("Instrumented generated Java test for %s (mode=%s)", function_name, mode)
return modified_code


def _add_import(source: str, import_statement: str) -> str:
"""Add an import statement to the source.

Args:
source: The source code.
import_statement: The import to add.

Returns:
Source with import added.

"""
lines = source.splitlines(keepends=True)
insert_idx = 0

# Find the last import or package statement
for i, line in enumerate(lines):
stripped = line.strip()
if stripped.startswith(("import ", "package ")):
insert_idx = i + 1
elif stripped and not stripped.startswith("//") and not stripped.startswith("/*"):
# First non-import, non-comment line
if insert_idx == 0:
insert_idx = i
break

lines.insert(insert_idx, import_statement + "\n")
return "".join(lines)
234 changes: 223 additions & 11 deletions codeflash/languages/java/remove_asserts.py
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,15 @@ def __init__(
# Precompile regex to find next special character (quotes, parens, braces).
self._special_re = re.compile(r"[\"'{}()]")


# Precompile literal/cast regexes to avoid recompilation on each literal check.
self._LONG_LITERAL_RE = re.compile(r"^-?\d+[lL]$")
self._INT_LITERAL_RE = re.compile(r"^-?\d+$")
self._DOUBLE_LITERAL_RE = re.compile(r"^-?\d+\.\d*[dD]?$|^-?\d+[dD]$")
self._FLOAT_LITERAL_RE = re.compile(r"^-?\d+\.?\d*[fF]$")
self._CHAR_LITERAL_RE = re.compile(r"^'.'$|^'\\.'$")
self._cast_re = re.compile(r"^\((\w+)\)")

def transform(self, source: str) -> str:
"""Remove assertions from source code, preserving target function calls.

Expand Down Expand Up @@ -894,6 +903,138 @@ def _find_balanced_braces(self, code: str, open_brace_pos: int) -> tuple[str | N

return code[open_brace_pos + 1 : pos - 1], pos

def _infer_return_type(self, assertion: AssertionMatch) -> str:
"""Infer the Java return type from the assertion context.

For assertEquals(expected, actual) patterns, the expected literal determines the type.
For assertTrue/assertFalse, the result is boolean.
Falls back to Object when the type cannot be determined.
"""
method = assertion.assertion_method

# assertTrue/assertFalse always deal with boolean values
if method == "assertTrue" or method == "assertFalse":
return "boolean"

# assertNull/assertNotNull — keep Object (reference type)
if method == "assertNull" or method == "assertNotNull":
return "Object"

# For assertEquals/assertNotEquals/assertSame, try to infer from the expected literal
if method in JUNIT5_VALUE_ASSERTIONS:
return self._infer_type_from_assertion_args(assertion.original_text, method)

# For fluent assertions (assertThat), type inference is harder — keep Object
return "Object"

# Regex patterns for Java literal type inference
_LONG_LITERAL_RE = re.compile(r"^-?\d+[lL]$")
_INT_LITERAL_RE = re.compile(r"^-?\d+$")
_DOUBLE_LITERAL_RE = re.compile(r"^-?\d+\.\d*[dD]?$|^-?\d+[dD]$")
_FLOAT_LITERAL_RE = re.compile(r"^-?\d+\.?\d*[fF]$")
_CHAR_LITERAL_RE = re.compile(r"^'.'$|^'\\.'$")

def _infer_type_from_assertion_args(self, original_text: str, method: str) -> str:
"""Infer the return type from assertEquals/assertNotEquals expected value."""
# Extract the args portion from the assertion text
# Pattern: assertXxx( args... )
paren_idx = original_text.find("(")
if paren_idx < 0:
return "Object"

args_str = original_text[paren_idx + 1 :]
# Remove trailing ");", whitespace
args_str = args_str.rstrip()
if args_str.endswith(");"):
args_str = args_str[:-2]
elif args_str.endswith(")"):
args_str = args_str[:-1]

# Fast-path: only extract the first top-level argument instead of splitting all arguments.
first_arg = self._extract_first_arg(args_str)
if not first_arg:
return "Object"

# assertEquals has (expected, actual) or (expected, actual, message/delta)
# Some overloads have (message, expected, actual) in JUnit 4 but JUnit 5 uses (expected, actual[, message])
# Try the first argument as the expected value
expected = first_arg.strip()

return self._type_from_literal(expected)

def _type_from_literal(self, value: str) -> str:
"""Determine the Java type of a literal value."""
if value in ("true", "false"):
return "boolean"
if value == "null":
return "Object"
if self._FLOAT_LITERAL_RE.match(value):
return "float"
if self._DOUBLE_LITERAL_RE.match(value):
return "double"
if self._LONG_LITERAL_RE.match(value):
return "long"
if self._INT_LITERAL_RE.match(value):
return "int"
if self._CHAR_LITERAL_RE.match(value):
return "char"
Comment on lines +971 to +980
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚡️Codeflash found 64% (0.64x) speedup for JavaAssertTransformer._type_from_literal in codeflash/languages/java/remove_asserts.py

⏱️ Runtime : 9.46 milliseconds 5.77 milliseconds (best of 68 runs)

📝 Explanation and details

Runtime improvement (primary): The optimized _type_from_literal reduces the method runtime from 9.46 ms to 5.77 ms — a ~63% overall speedup. That throughput win is why this change was accepted.

What changed (specific optimizations)

  • Replaced the expensive compiled-regex.match calls for numeric/char detection with cheap, deterministic string operations:
    • Check suffixes by inspecting v[-1] (single-char index) instead of running a regex.
    • Use slicing (v[:-1]), simple sign-strip logic, str.isdigit(), and a single split(".", 1) for decimal validation.
    • Detect char literals with length checks and startswith/endswith rather than a regex.
  • Kept only the cast regex (_cast_re.match) for the cast-specific extraction; everything else uses C-level string methods.
  • Bound value to a local variable v early to reduce repeated attribute lookups (fewer attribute and name lookups).

Why this is faster (how it maps to Python performance)

  • Regex.match invokes the regex engine, produces match objects, and can do backtracking — it’s relatively expensive even when compiled. The original profiler shows large time spent on those regex checks (FLOAT/DOBULE/LONG/INT regex lines dominated the function's time).
  • String operations used here (indexing, slicing, isdigit, startswith, split) are implemented in C with minimal overhead and no match-object allocation. They therefore execute much faster and allocate less temporary memory.
  • Early-exit checks and ordering reduce average work: common simple cases (suffix checks, pure digits, decimals) are identified quickly with small, deterministic checks.

Profiler evidence

  • Original profile: the compiled-regex checks (FLOAT/DOUBLE/LONG/INT) were hot — together they accounted for a large fraction of wall time (e.g., FLOAT_LITERAL_RE.match showed ~24.5% alone).
  • Optimized profile: time shifts to cheap operations (v = value, v[-1] checks, isdigit, split) and the expensive regex is no longer invoked for normal numeric cases. Overall time per-call drops substantially (consistent with the measured runtime improvement).

Behavioral/compatibility notes and trade-offs

  • Functional behavior is preserved for the tested inputs (booleans, null, numeric suffixes, decimals, char and string detection, and cast extraction).
  • Some microbenchmarks show small regressions (certain f-suffix float cases or a few edge inputs), because the manual logic does more branching/splits in those particular paths. This is an acceptable trade-off given the substantial aggregate runtime benefit — the overall throughput and bulk tests improved dramatically.
  • The cast-detection is still done with the compiled regex, preserving the original extraction behavior for casts.

Workloads that benefit most

  • Hot paths that call _type_from_literal many times (parsing many literals, bulk transformations) gain the most — annotated large-scale tests show ~67–75% faster behavior.
  • Numeric-heavy workloads (ints, longs, doubles) see especially large wins, since they were the original regex-dominated cases.
  • Very small or extremely pathological single-case inputs may see no change or minor regressions; however the overall system throughput improves.

Summary

  • Primary win: 63% reduction in runtime (9.46 ms → 5.77 ms).
  • How: remove repeated regex.match calls in favor of C-level string ops and early exits, reducing engine overhead and allocations.
  • Trade-off: a few microbenchmarks are slightly slower due to extra branching, but the aggregate and hot-path performance improvements justify the change for workloads that exercise this method frequently.

Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests 12069 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_boolean_and_null_literals_basic():
    # Create a real instance of JavaAssertTransformer with a dummy function name.
    transformer = JavaAssertTransformer("dummy")
    # The literal "true" should be recognized as boolean.
    codeflash_output = transformer._type_from_literal("true") # 451ns -> 531ns (15.1% slower)
    # The literal "false" should be recognized as boolean.
    codeflash_output = transformer._type_from_literal("false") # 290ns -> 280ns (3.57% faster)
    # The literal "null" is mapped to the Java Object type.
    codeflash_output = transformer._type_from_literal("null") # 290ns -> 320ns (9.38% slower)

def test_numeric_and_string_and_char_literals_basic():
    # Instantiate the transformer to call the instance method under test.
    transformer = JavaAssertTransformer("fn")
    # Float literals with 'f' or 'F' suffix -> float
    codeflash_output = transformer._type_from_literal("1.0f") # 2.21μs -> 2.33μs (5.57% slower)
    codeflash_output = transformer._type_from_literal("2F") # 891ns -> 852ns (4.58% faster)
    # Double literals: decimal without 'f' but with '.' or with 'd' suffix -> double
    codeflash_output = transformer._type_from_literal("3.1415") # 1.45μs -> 982ns (48.0% faster)
    codeflash_output = transformer._type_from_literal("4.0d") # 1.09μs -> 862ns (26.7% faster)
    codeflash_output = transformer._type_from_literal("5D") # 982ns -> 511ns (92.2% faster)
    # Long literals end with L/l -> long
    codeflash_output = transformer._type_from_literal("123L") # 1.46μs -> 781ns (87.3% faster)
    codeflash_output = transformer._type_from_literal("-999l") # 1.48μs -> 932ns (59.0% faster)
    # Integer literals (no suffix, whole number) -> int
    codeflash_output = transformer._type_from_literal("42") # 1.39μs -> 591ns (136% faster)
    codeflash_output = transformer._type_from_literal("-7") # 1.29μs -> 571ns (126% faster)
    # Char literals: single char inside single quotes -> char
    codeflash_output = transformer._type_from_literal("'a'") # 1.20μs -> 1.27μs (5.58% slower)
    # Char escape like backslash + char should also be recognized as char.
    codeflash_output = transformer._type_from_literal("'\\n'") # 1.20μs -> 962ns (24.9% faster)
    # String values (start with double quote) -> String
    codeflash_output = transformer._type_from_literal('"hello"') # 1.35μs -> 721ns (87.7% faster)
    # Even a malformed string that merely starts with a quote is treated as String.
    codeflash_output = transformer._type_from_literal('"unterminated') # 1.01μs -> 982ns (3.05% faster)

def test_cast_expressions_and_fallback_behavior():
    transformer = JavaAssertTransformer("f")
    # Casts like (byte)0 should return the cast type (group 1 from the regex).
    codeflash_output = transformer._type_from_literal("(byte)0") # 3.98μs -> 3.81μs (4.49% faster)
    # Cast should preserve case (e.g., capitalized types).
    codeflash_output = transformer._type_from_literal("(Short)1") # 1.80μs -> 1.54μs (16.9% faster)
    # A user-defined type in a cast should also be extracted.
    codeflash_output = transformer._type_from_literal("(MyType)someValue") # 1.45μs -> 1.14μs (27.2% faster)
    # A value that does not match any pattern should fall back to Object.
    codeflash_output = transformer._type_from_literal("SomeRandomToken") # 1.21μs -> 952ns (27.3% faster)
    # Boolean is case-sensitive, so "False" (capitalized) is not recognized -> fallback Object.
    codeflash_output = transformer._type_from_literal("False") # 1.21μs -> 842ns (44.1% faster)

def test_empty_string_and_none_behavior():
    transformer = JavaAssertTransformer("x")
    # An empty string does not match any literal regex -> Object.
    codeflash_output = transformer._type_from_literal("")
    # Passing None is not a valid string; ensure it raises an AttributeError when startswith is used.
    # The implementation will attempt value.startswith('"') and thus should raise AttributeError for None.
    with pytest.raises(AttributeError):
        transformer._type_from_literal(None)  # type: ignore[arg-type]

def test_large_scale_many_literals_mixed_types():
    transformer = JavaAssertTransformer("big")
    literals = []
    expected = []
    # Build 1000 deterministic literals covering many categories in a round-robin fashion.
    for i in range(1000):
        mod = i % 8
        if mod == 0:
            # boolean true/false alternating
            val = "true" if (i % 2 == 0) else "false"
            typ = "boolean"
        elif mod == 1:
            # float with suffix f
            val = f"{i}.0f"
            typ = "float"
        elif mod == 2:
            # double with decimal point
            val = f"{i}.5"
            typ = "double"
        elif mod == 3:
            # long with L suffix
            val = f"{i}L"
            typ = "long"
        elif mod == 4:
            # int plain
            val = str(i)
            typ = "int"
        elif mod == 5:
            # char alternating between simple and escaped
            val = "'a'" if (i % 2 == 0) else "'\\t'"
            typ = "char"
        elif mod == 6:
            # string literal
            val = f"\"s{i}\""
            typ = "String"
        else:
            # cast to a named type
            val = f"(Custom{i})x"
            typ = f"Custom{i}"
        literals.append(val)
        expected.append(typ)
    # Check each literal's resolved type matches expected.
    for lit, exp in zip(literals, expected):
        codeflash_output = transformer._type_from_literal(lit); got = codeflash_output # 947μs -> 542μs (74.6% faster)

def test_large_scale_performance_consistency():
    transformer = JavaAssertTransformer("perf")
    # Run the same small set of literals 1000 times to exercise regex caching and ensure no state bleed.
    sample_literals = ['true', 'false', 'null', '1.0f', '2.5', '123L', '42', "'z'", '"hi"', "(byte)0"]
    # Expected results for one pass, determined statically.
    expected_single = ["boolean", "boolean", "Object", "float", "double", "long", "int", "char", "String", "byte"]
    # Execute 1000 iterations and ensure results remain identical each time.
    for _ in range(1000):
        for lit, exp in zip(sample_literals, expected_single):
            codeflash_output = transformer._type_from_literal(lit)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
import pytest  # used for our unit tests
from codeflash.languages.java.remove_asserts import JavaAssertTransformer

def test_boolean_and_null_literals():
    # Create a real instance of the transformer with a dummy function name.
    t = JavaAssertTransformer("dummy")
    # 'true' and 'false' must be detected as boolean.
    codeflash_output = t._type_from_literal("true") # 932ns -> 842ns (10.7% faster)
    codeflash_output = t._type_from_literal("false") # 221ns -> 230ns (3.91% slower)
    # 'null' must map to Object
    codeflash_output = t._type_from_literal("null") # 351ns -> 370ns (5.14% slower)

def test_string_and_char_literals_basic():
    t = JavaAssertTransformer("fn")
    # A value that starts with a double quote is considered a String.
    codeflash_output = t._type_from_literal('"hello"') # 3.55μs -> 2.83μs (25.1% faster)
    # Even if it only starts with a quote (unterminated), rule is to return String.
    codeflash_output = t._type_from_literal('"unterminated') # 1.32μs -> 1.68μs (21.4% slower)
    # Simple char literal: single character enclosed in single quotes -> char
    codeflash_output = t._type_from_literal("'a'") # 1.44μs -> 1.24μs (16.2% faster)
    # Escaped char (backslash + char) should also be recognized as char
    codeflash_output = t._type_from_literal("'\\n'") # 1.27μs -> 1.01μs (25.7% faster)

def test_numeric_literals_and_suffixes():
    t = JavaAssertTransformer("n")
    # Integers with no suffix
    codeflash_output = t._type_from_literal("0") # 2.42μs -> 1.59μs (51.6% faster)
    codeflash_output = t._type_from_literal("123") # 2.01μs -> 701ns (187% faster)
    codeflash_output = t._type_from_literal("-42") # 1.66μs -> 982ns (69.3% faster)
    # Long literals with 'L' or 'l' suffix
    codeflash_output = t._type_from_literal("123L") # 1.67μs -> 1.03μs (62.1% faster)
    codeflash_output = t._type_from_literal("-999l") # 1.52μs -> 1.04μs (46.2% faster)
    # Float literals require an f/F suffix (per the regex used)
    codeflash_output = t._type_from_literal("1.0f") # 722ns -> 1.26μs (42.8% slower)
    codeflash_output = t._type_from_literal("2F") # 561ns -> 641ns (12.5% slower)
    codeflash_output = t._type_from_literal("-3.14f") # 521ns -> 912ns (42.9% slower)
    # Double detection: decimals (with optional trailing d/D), and '1.' should be allowed
    codeflash_output = t._type_from_literal("1.0") # 1.09μs -> 832ns (31.2% faster)
    codeflash_output = t._type_from_literal("1.") # 832ns -> 591ns (40.8% faster)
    codeflash_output = t._type_from_literal("6D") # 851ns -> 531ns (60.3% faster)
    codeflash_output = t._type_from_literal("-2.71828") # 1.07μs -> 811ns (32.2% faster)

def test_char_and_nonchar_single_quote_inputs():
    t = JavaAssertTransformer("x")
    # If single-quoted content is more than one char (e.g. "'ab'") it's not matched by the char regex -> Object
    codeflash_output = t._type_from_literal("'ab'") # 3.32μs -> 4.18μs (20.6% slower)
    # Empty string should be treated as unknown -> Object
    codeflash_output = t._type_from_literal("") # 1.24μs -> 1.25μs (0.799% slower)
    # A capitalized boolean-like 'True' should NOT be recognized -> Object
    codeflash_output = t._type_from_literal("True") # 1.41μs -> 1.30μs (8.53% faster)

def test_cast_expressions_and_spacing():
    t = JavaAssertTransformer("c")
    # Direct cast without space
    codeflash_output = t._type_from_literal("(byte)0") # 4.02μs -> 4.19μs (4.08% slower)
    # Cast followed by a space before number still matches because regex only matches the prefix "(type)"
    codeflash_output = t._type_from_literal("(short) -1") # 1.83μs -> 1.42μs (29.0% faster)
    # Cast with a multi-character type name
    codeflash_output = t._type_from_literal("(MyType)someExpr") # 1.42μs -> 1.17μs (21.4% faster)
    # If the parentheses are not a cast (e.g. "(123)"), group will match '123' only if it is a word; '123' isn't a word -> Object
    # But because the cast regex only captures word characters for the type, "(123)5" should not match and fall back to Object
    codeflash_output = t._type_from_literal("(123)5") # 1.38μs -> 951ns (45.3% faster)

def test_numeric_precedence_and_ambiguities():
    t = JavaAssertTransformer("p")
    # The order in the implementation checks float (f/F), then double, then long, then int.
    # So "1f" should be float; "1d" or "1D" double; "1L" long
    codeflash_output = t._type_from_literal("1f") # 2.02μs -> 1.53μs (32.0% faster)
    codeflash_output = t._type_from_literal("1d") # 1.32μs -> 762ns (73.5% faster)
    codeflash_output = t._type_from_literal("1L") # 1.16μs -> 761ns (52.8% faster)
    # A decimal without suffix picks up as double even if it could be seen as int by digits portion.
    codeflash_output = t._type_from_literal("10.0") # 1.27μs -> 1.23μs (3.25% faster)
    # Negative long with capital L
    codeflash_output = t._type_from_literal("-100L") # 1.50μs -> 991ns (51.7% faster)

def test_large_scale_many_literals():
    t = JavaAssertTransformer("bulk")
    # Prepare a base set of representative literal strings and their expected types.
    base_pairs = [
        ("true", "boolean"),
        ("false", "boolean"),
        ("null", "Object"),
        ('"s"', "String"),
        ("100", "int"),
        ("-7", "int"),
        ("42L", "long"),
        ("99l", "long"),
        ("3.1415", "double"),
        ("2.", "double"),
        ("6D", "double"),
        ("7d", "double"),
        ("8.0f", "float"),
        ("9F", "float"),
        ("'z'", "char"),
        ("'\\t'", "char"),
        ("(byte)5", "byte"),
        ("(short) 6", "short"),
        ("'ab'", "Object"),  # not a char
        ("", "Object"),  # empty -> unknown
    ]
    # Build a large list (1000 elements) by repeating the base set until we reach 1000
    total = 1000
    values = []
    expected = []
    i = 0
    while len(values) < total:
        val, typ = base_pairs[i % len(base_pairs)]
        # Slightly vary numeric values to catch regex differences (append iteration index where appropriate)
        if val.isdigit() or (val.startswith("-") and val[1:].isdigit()):
            # change numeric magnitude to keep variety but type stays the same
            variant = str(int(val) + i)
            values.append(variant)
            # Determine expected type from the base mapping
            expected.append("int")
        elif val.endswith("L") or val.endswith("l"):
            variant = val[:-1] + str((i % 100) + 1) + val[-1]
            values.append(variant)
            expected.append("long")
        elif val.endswith("f") or val.endswith("F"):
            variant = val[:-1] + str((i % 10) + 1) + val[-1]
            values.append(variant)
            expected.append("float")
        elif val.endswith("d") or val.endswith("D") or "." in val:
            # Keep as double-like
            variant = str(float((i % 100) + 1))  # e.g. "1.0", "2.0"
            values.append(variant)
            expected.append("double")
        else:
            # non-numeric base values: keep them repeated
            values.append(val)
            expected.append(typ)
        i += 1

    # Now run the transformer over all values and assert types match expectations.
    # This checks consistent behavior over a large volume of inputs.
    for v, exp in zip(values, expected):
        codeflash_output = t._type_from_literal(v); got = codeflash_output # 970μs -> 578μs (67.8% faster)

def test_none_raises_type_error():
    t = JavaAssertTransformer("err")
    # Calling _type_from_literal with None should raise a TypeError when regex.match is attempted.
    with pytest.raises(TypeError):
        t._type_from_literal(None) # 3.85μs -> 4.54μs (15.2% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To test or edit this optimization locally git merge codeflash/optimize-pr1663-2026-02-25T22.18.13

Click to see suggested changes
Suggested change
if self._FLOAT_LITERAL_RE.match(value):
return "float"
if self._DOUBLE_LITERAL_RE.match(value):
return "double"
if self._LONG_LITERAL_RE.match(value):
return "long"
if self._INT_LITERAL_RE.match(value):
return "int"
if self._CHAR_LITERAL_RE.match(value):
return "char"
# Fast-path numeric suffix checks and simple digit/decimal validations
v = value
# Check float suffix (e.g., 1.23f or 10f)
if v and (v[-1] in "fF"):
core = v[:-1]
if core:
if core[0] == "-":
core2 = core[1:]
else:
core2 = core
if core2:
if "." in core2:
a, b = core2.split(".", 1)
if a.isdigit() and (b == "" or b.isdigit()):
return "float"
else:
if core2.isdigit():
return "float"
# Check double suffix (e.g., 1.23d or 10d)
if v and (v[-1] in "dD"):
core = v[:-1]
if core:
if core[0] == "-":
core2 = core[1:]
else:
core2 = core
if core2:
if "." in core2:
a, b = core2.split(".", 1)
if a.isdigit() and (b == "" or b.isdigit()):
return "double"
else:
if core2.isdigit():
return "double"
# Decimal without suffix implies double if it fits the pattern like 1. or 1.23
if "." in v:
core = v
if core and core[0] == "-":
core = core[1:]
a, b = core.split(".", 1)
if a.isdigit() and (b == "" or b.isdigit()):
return "double"
# Long suffix (e.g., 123L)
if v and (v[-1] in "lL"):
core = v[:-1]
if core:
if core[0] == "-":
core2 = core[1:]
else:
core2 = core
if core2.isdigit():
return "long"
# Integer literal (e.g., 123)
tmp = v
if tmp and tmp[0] == "-":
tmp2 = tmp[1:]
else:
tmp2 = tmp
if tmp2.isdigit():
return "int"
# Character literal ('a' or '\n')
# Matches patterns of length 3 (e.g. 'a') or length 4 (e.g. '\n')
if len(v) in (3, 4) and v.startswith("'") and v.endswith("'"):
if len(v) == 3:
return "char"
# len == 4: must be an escape like '\n'
if v[1] == "\\":
return "char"

Static Badge

if value.startswith('"'):
return "String"
# Cast expression like (byte)0, (short)1
cast_match = self._cast_re.match(value)
if cast_match:
return cast_match.group(1)
return "Object"

def _split_top_level_args(self, args_str: str) -> list[str]:
"""Split assertion arguments at top-level commas, respecting parens/strings/generics."""
# Fast-path: if there are no special delimiters that require parsing,
# we can use a simple split which is much faster for common simple cases.
if not self._special_re.search(args_str):
# Preserve original behavior of returning a list with the single unstripped string
# when there are no commas, otherwise split on commas.
if "," in args_str:
return args_str.split(",")
return [args_str]

args: list[str] = []
depth = 0
current: list[str] = []
i = 0
in_string = False
string_char = ""

while i < len(args_str):
ch = args_str[i]

if in_string:
current.append(ch)
if ch == "\\" and i + 1 < len(args_str):
i += 1
current.append(args_str[i])
elif ch == string_char:
in_string = False
elif ch in ('"', "'"):
in_string = True
string_char = ch
current.append(ch)
elif ch in ("(", "<", "[", "{"):
depth += 1
current.append(ch)
elif ch in (")", ">", "]", "}"):
depth -= 1
current.append(ch)
elif ch == "," and depth == 0:
args.append("".join(current))
current = []
else:
current.append(ch)
i += 1

if current:
args.append("".join(current))
Comment on lines +1007 to +1035
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚡️Codeflash found 96% (0.96x) speedup for JavaAssertTransformer._split_top_level_args in codeflash/languages/java/remove_asserts.py

⏱️ Runtime : 4.99 milliseconds 2.55 milliseconds (best of 250 runs)

📝 Explanation and details

Runtime improvement: The optimized _split_top_level_args cuts the call time from 4.99 ms to 2.55 ms (~95% faster by the reported metric), with especially large wins on long inputs (hundreds/thousands of args).

What changed (concrete optimizations)

  • Eliminated per-character list-building + join: the original appended each character to current and used "".join(current) on splits. The optimized code records slice boundaries (last..i) and appends substrings once, avoiding the frequent small-list allocations and the join cost.
  • Cached locals and precomputed values: binds args.append to args_append, assigns s = args_str and n = len(s) once. This avoids repeated attribute lookups and len() calls inside the hot loop.
  • Early fast-path for empty input: returns immediately for empty args_str.
  • Reduced branching and work in the hot loop:
    • Handles string escapes by advancing i by 2 and continuing, instead of building characters.
    • Uses a single-character membership string for open/close brackets (opens/closes) and checks ch in opens/closes — a cheap O(1) C-level operation.
    • Minimizes repeated current.append and other Python-level list ops by moving to index arithmetic.
  • Keeps loop indexing simple (while i < n, s[i]) and uses continues to reduce duplicate i += 1 logic.

Why this yields a speedup

  • The original implementation spent most time inside the while loop doing a lot of Python-level work per character: list appends for each char, join() calls on splits, and repeated attribute/len lookups. Those are high-overhead Python operations.
  • The optimized version reduces allocations (no per-char lists), reduces Python bytecode executed per character (fewer appends, fewer attribute accesses), and shifts work to C-level substring slicing and in-operator checks. That lowers CPU and memory churn and reduces the constant factor per character, which explains the big wins shown by the line profiler (significantly less time inside the main loop and fewer append/join costs).
  • Caching args.append and precomputing len(s) eliminates repeated attribute lookups inside the hot path, which matters at scale.

Behavioral and workload impact

  • Functionality preserved for the tested cases: strings, escapes, nested parentheses/generics/arrays/braces are still respected (see annotated tests). The optimizer excels on long inputs (e.g., 1000-element tests show 80–100%+ speedups) because savings are per-character/per-arg and accumulate.
  • Small-case trade-off: one microcase (single top-level comma) was marginally slower in a microbenchmark (annotated_tests shows ~12% slower). This is an acceptable trade-off because the optimization dramatically reduces runtime for typical and large inputs where this function is most likely to be on a hot path.
  • No new heavy dependencies were introduced; changes are local to the function and safe for callers.

When you get the most benefit

  • Hot paths that call this on long argument lists or in tight loops (e.g., splitting thousands of arguments or repeatedly parsing assertion arguments) will see the largest absolute and relative improvements.
  • Short inputs also generally benefit (small microbenchmarks mostly show faster times), but the biggest wins are for large/mixed/nested inputs as shown in the annotated tests.

Summary
The optimized implementation reduces per-character Python-level operations, avoids frequent small allocations and joins, and caches locals so the inner loop runs much lighter. Those combined micro-optimizations produce the reported runtime improvement (4.99 ms → 2.55 ms) and large speedups on real workloads, at the cost of a tiny micro-regression in one trivial input case.

Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests 47 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_split_basic_simple():
    # Create a real transformer instance using the real constructor
    t = JavaAssertTransformer("assert")
    # Simple comma-separated arguments should split into three items
    args = "a,b,c"
    codeflash_output = t._split_top_level_args(args); result = codeflash_output # 3.03μs -> 2.97μs (2.02% faster)

def test_split_respects_parentheses_and_whitespace():
    # Transformer instance
    t = JavaAssertTransformer("assert")
    # First arg contains a comma inside parentheses, which should not be split.
    # There is a top-level comma between the function call and the 'x' argument.
    args = "func(1,2), x"
    codeflash_output = t._split_top_level_args(args); result = codeflash_output # 4.88μs -> 3.83μs (27.5% faster)

def test_split_respects_generics_brackets_and_braces():
    t = JavaAssertTransformer("assert")
    # Generics use angle brackets which the splitter treats like parentheses (depth tracking).
    args = "List<String>, Map<K,V>"
    codeflash_output = t._split_top_level_args(args); res = codeflash_output # 7.08μs -> 4.86μs (45.7% faster)

def test_split_respects_strings_and_escapes():
    t = JavaAssertTransformer("assert")
    # Double-quoted string containing a comma should be preserved as one argument.
    args1 = '"a,b", c'
    codeflash_output = t._split_top_level_args(args1); r1 = codeflash_output # 3.44μs -> 2.87μs (19.9% faster)
    # Single-quoted string with an escaped single-quote and a comma inside should also be preserved.
    # Use a Python double-quoted literal so the inner backslash is represented correctly in the input string.
    args2 = "'don\\'t,split', y"
    codeflash_output = t._split_top_level_args(args2); r2 = codeflash_output # 3.58μs -> 2.33μs (53.2% faster)

def test_split_handles_brackets_and_braces():
    t = JavaAssertTransformer("assert")
    # Arrays and braces should be treated as nested structures; commas inside them are ignored.
    args = "[1,2], {a,b}, (x,y)"
    codeflash_output = t._split_top_level_args(args); res = codeflash_output # 6.62μs -> 4.66μs (42.2% faster)

def test_split_unmatched_parentheses_returns_single():
    t = JavaAssertTransformer("assert")
    # Unmatched parentheses means top-level commas may not exist; all characters remain in a single arg.
    args = "((a,b), c"  # note: one more '(' than ')'
    codeflash_output = t._split_top_level_args(args); res = codeflash_output # 3.78μs -> 3.02μs (25.2% faster)

def test_split_empty_string_returns_empty_list():
    t = JavaAssertTransformer("assert")
    # An empty input string should produce an empty list (no args).
    codeflash_output = t._split_top_level_args(""); res = codeflash_output # 742ns -> 581ns (27.7% faster)

def test_split_double_closing_angle_brackets_in_generics():
    t = JavaAssertTransformer("assert")
    # Handle cases with consecutive closing angle brackets (e.g., nested generics like Foo<Bar<Baz>>)
    args = "Map<A<B>, C>>, y"
    codeflash_output = t._split_top_level_args(args); res = codeflash_output # 5.39μs -> 3.92μs (37.6% faster)

def test_split_large_number_of_simple_args():
    t = JavaAssertTransformer("assert")
    # Build 1000 simple args: a0,a1,a2,...,a999
    n = 1000
    expected = [f"a{i}" for i in range(n)]
    big_args = ",".join(expected)
    # Split should return the original list
    codeflash_output = t._split_top_level_args(big_args); res = codeflash_output # 1.16ms -> 623μs (85.6% faster)

def test_split_large_mixed_nested_args():
    t = JavaAssertTransformer("assert")
    # Build 1000 args where even indices are nested function calls containing internal commas,
    # and odd indices are simple values. Joining with top-level commas must split into the 1000 parts.
    n = 1000
    parts = []
    for i in range(n):
        if i % 2 == 0:
            # nested call contains an internal comma which should NOT cause a top-level split
            parts.append(f"f({i},{i+1})")
        else:
            parts.append(f"v{i}")
    big = ",".join(parts)
    codeflash_output = t._split_top_level_args(big); res = codeflash_output # 1.86ms -> 926μs (100% 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  # used for our unit tests
from codeflash.languages.java.remove_asserts import JavaAssertTransformer

def test_simple_top_level_commas():
    # Create a real instance of the transformer using the actual constructor.
    t = JavaAssertTransformer("assertX")
    # Basic comma-separated tokens at top-level should split into separate fragments.
    args = "a,b,c"
    # Call the instance method; it's an instance method so we call on the real object.
    codeflash_output = t._split_top_level_args(args); parts = codeflash_output # 3.13μs -> 2.80μs (12.1% faster)

def test_preserve_whitespace_and_trailing_leading_commas():
    t = JavaAssertTransformer("assertX")
    # Whitespace around items should be preserved in the fragments.
    args = "  a , b  ,c "
    codeflash_output = t._split_top_level_args(args); parts = codeflash_output # 4.77μs -> 3.52μs (35.6% faster)

    # Trailing comma should result in omission of a final empty fragment (current is empty -> not appended).
    codeflash_output = t._split_top_level_args("a,b,"); parts_trailing = codeflash_output # 1.74μs -> 1.44μs (20.9% faster)

    # Leading comma results in the first fragment being an empty string, per implementation.
    codeflash_output = t._split_top_level_args(",a"); parts_leading = codeflash_output # 1.02μs -> 961ns (6.35% faster)

def test_parentheses_and_nested_commas_are_ignored_for_splitting():
    t = JavaAssertTransformer("assertX")
    # Commas inside parentheses should not split top-level arguments.
    args = "foo(bar,baz),qux"
    codeflash_output = t._split_top_level_args(args); parts = codeflash_output # 5.62μs -> 4.14μs (35.9% faster)

    # Multiple nested parentheses should also be respected.
    complex_args = "outer(inner(a,b), other(c,d)),last"
    codeflash_output = t._split_top_level_args(complex_args); complex_parts = codeflash_output # 8.34μs -> 4.61μs (80.9% faster)

def test_angle_brackets_generics_are_respected():
    t = JavaAssertTransformer("assertX")
    # Commas inside angle-bracket generics shouldn't split at top-level.
    args = "Map<String,Integer>, other"
    codeflash_output = t._split_top_level_args(args); parts = codeflash_output # 8.02μs -> 5.17μs (55.1% faster)

    # Deeper nested generics remain intact.
    args2 = "List<Map<String,List<Integer>>>, z"
    codeflash_output = t._split_top_level_args(args2); parts2 = codeflash_output # 8.43μs -> 4.62μs (82.5% faster)

def test_string_literals_with_commas_and_escaped_quotes():
    t = JavaAssertTransformer("assertX")
    # Build a quoted string containing commas and escaped internal quotes.
    part1 = '"He said \\"Hello, world\\""'
    # Compose the args string with that quoted string followed by a simple identifier.
    args = part1 + ",x"
    codeflash_output = t._split_top_level_args(args); parts = codeflash_output # 5.83μs -> 4.07μs (43.3% faster)

    # Single-quoted strings (char literals) with escaped characters should also be preserved.
    char_part = "'\\n'"
    args2 = char_part + ", other"
    codeflash_output = t._split_top_level_args(args2); parts2 = codeflash_output # 3.04μs -> 2.11μs (43.6% faster)

def test_braces_and_array_initializers():
    t = JavaAssertTransformer("assertX")
    # Commas inside braces should not produce top-level splits.
    args = "{1,2},3"
    codeflash_output = t._split_top_level_args(args); parts = codeflash_output # 3.58μs -> 3.04μs (17.8% faster)

    # Array initializer inside expression should be preserved as one fragment.
    args2 = "new int[]{1,2}, other"
    codeflash_output = t._split_top_level_args(args2); parts2 = codeflash_output # 5.70μs -> 3.34μs (70.9% faster)

def test_empty_input_and_single_comma_behavior():
    t = JavaAssertTransformer("assertX")
    # Empty input should return an empty list (no fragments).
    codeflash_output = t._split_top_level_args("") # 722ns -> 572ns (26.2% faster)

    # A single comma at top-level: implementation appends current (which is empty) once and
    # does not append another empty current at end -> produces a single empty string.
    codeflash_output = t._split_top_level_args(",") # 1.30μs -> 1.48μs (12.2% slower)

    # Two commas: the first comma will append empty, then the second comma will append empty again,
    # but because current will be empty at end there is no final append, resulting in two empties.
    codeflash_output = t._split_top_level_args(",,") # 1.08μs -> 1.06μs (1.88% faster)

def test_complex_mixed_constructs():
    t = JavaAssertTransformer("assertX")
    # Mix strings with commas, nested calls with commas and maps containing commas.
    part = 'assertEquals("expected, with,commas", List.of(1,2), Map.of("k","v"))'
    args = part + ", x"
    codeflash_output = t._split_top_level_args(args); parts = codeflash_output # 15.3μs -> 9.17μs (67.2% faster)

def test_large_scale_many_elements_with_strings_and_commas():
    t = JavaAssertTransformer("assertX")
    # Build 1000 elements; every 10th element is a quoted string containing a comma.
    parts = []
    for i in range(1000):
        if i % 10 == 0:
            # Quoted string containing a comma to test that commas within strings are not split.
            parts.append(f'"{i},{i+1}"')
        else:
            parts.append(f'item{i}')
    # Join into a single args string with top-level commas between parts.
    args = ",".join(parts)
    # Split using the method under test.
    codeflash_output = t._split_top_level_args(args); result = codeflash_output # 1.86ms -> 916μs (103% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To test or edit this optimization locally git merge codeflash/optimize-pr1663-2026-02-25T20.47.23

Click to see suggested changes
Suggested change
while i < len(args_str):
ch = args_str[i]
if in_string:
current.append(ch)
if ch == "\\" and i + 1 < len(args_str):
i += 1
current.append(args_str[i])
elif ch == string_char:
in_string = False
elif ch in ('"', "'"):
in_string = True
string_char = ch
current.append(ch)
elif ch in ("(", "<", "[", "{"):
depth += 1
current.append(ch)
elif ch in (")", ">", "]", "}"):
depth -= 1
current.append(ch)
elif ch == "," and depth == 0:
args.append("".join(current))
current = []
else:
current.append(ch)
i += 1
if current:
args.append("".join(current))
# Fast-path: empty input
if not args_str:
return args
s = args_str
n = len(s)
last = 0
args_append = args.append
# Use string membership checks for bracket characters
opens = "({[<"
closes = ")}]>"
while i < n:
ch = s[i]
if in_string:
# Respect escape sequences inside strings
if ch == "\\" and i + 1 < n:
i += 2
continue
if ch == string_char:
in_string = False
i += 1
continue
# Not in string
if ch == '"' or ch == "'":
in_string = True
string_char = ch
i += 1
continue
# Track nesting depth for parentheses/generics/arrays/blocks
if ch in opens:
depth += 1
i += 1
continue
if ch in closes:
depth -= 1
i += 1
continue
# Top-level comma splits arguments
if ch == "," and depth == 0:
args_append(s[last:i])
last = i + 1
i += 1
continue
i += 1
# Append the final segment if any
if last < n:
args_append(s[last:n])

Static Badge

return args

def _generate_replacement(self, assertion: AssertionMatch) -> str:
"""Generate replacement code for an assertion.

Expand All @@ -912,18 +1053,35 @@ def _generate_replacement(self, assertion: AssertionMatch) -> str:
if not assertion.target_calls:
return ""

# Infer the return type from assertion context to avoid Object→primitive cast errors
return_type = self._infer_return_type(assertion)

# Generate capture statements for each target call
replacements = []
replacements: list[str] = []
# For the first replacement, use the full leading whitespace
# For subsequent ones, strip leading newlines to avoid extra blank lines
base_indent = assertion.leading_whitespace.lstrip("\n\r")
for i, call in enumerate(assertion.target_calls):
self.invocation_counter += 1
var_name = f"_cf_result{self.invocation_counter}"
if i == 0:
replacements.append(f"{assertion.leading_whitespace}Object {var_name} = {call.full_call};")
else:
replacements.append(f"{base_indent}Object {var_name} = {call.full_call};")
leading_ws = assertion.leading_whitespace
base_indent = leading_ws.lstrip("\n\r")

# Use a local counter to minimize attribute write overhead in the loop.
inv = self.invocation_counter

calls = assertion.target_calls
# Handle first call explicitly to avoid a per-iteration branch
if calls:
inv += 1
var_name = "_cf_result" + str(inv)
replacements.append(f"{leading_ws}{return_type} {var_name} = {calls[0].full_call};")

# Handle remaining calls
for call in calls[1:]:
inv += 1
var_name = "_cf_result" + str(inv)
replacements.append(f"{base_indent}{return_type} {var_name} = {call.full_call};")


# Write back the counter
self.invocation_counter = inv

return "\n".join(replacements)

Expand All @@ -942,8 +1100,10 @@ def _generate_exception_replacement(self, assertion: AssertionMatch) -> str:
try { code(); } catch (IllegalArgumentException _cf_caught1) { ex = _cf_caught1; } catch (Exception _cf_ignored1) {}

"""
self.invocation_counter += 1
counter = self.invocation_counter
# Increment invocation counter once for this exception handling
inv = self.invocation_counter + 1
self.invocation_counter = inv
counter = inv
ws = assertion.leading_whitespace
base_indent = ws.lstrip("\n\r")

Expand Down Expand Up @@ -982,6 +1142,58 @@ def _generate_exception_replacement(self, assertion: AssertionMatch) -> str:
# Fallback: comment out the assertion
return f"{ws}// Removed assertThrows: could not extract callable"

def _extract_first_arg(self, args_str: str) -> str | None:
"""Extract the first top-level argument from args_str.

This is a lightweight alternative to splitting all top-level arguments;
it stops at the first top-level comma, respects nested delimiters and strings,
and avoids constructing the full argument list for better performance.
"""
n = len(args_str)
i = 0

# skip leading whitespace
while i < n and args_str[i].isspace():
i += 1
if i >= n:
return None

depth = 0
in_string = False
string_char = ""
cur: list[str] = []

while i < n:
ch = args_str[i]

if in_string:
cur.append(ch)
if ch == "\\" and i + 1 < n:
i += 1
cur.append(args_str[i])
elif ch == string_char:
in_string = False
elif ch in ('"', "'"):
in_string = True
string_char = ch
cur.append(ch)
elif ch in ("(", "<", "[", "{"):
depth += 1
cur.append(ch)
elif ch in (")", ">", "]", "}"):
depth -= 1
cur.append(ch)
elif ch == "," and depth == 0:
break
else:
cur.append(ch)
i += 1

# Trim trailing whitespace from the extracted argument
if not cur:
return None
return "".join(cur).rstrip()
Comment on lines +1162 to +1195
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚡️Codeflash found 37% (0.37x) speedup for JavaAssertTransformer._extract_first_arg in codeflash/languages/java/remove_asserts.py

⏱️ Runtime : 4.31 milliseconds 3.15 milliseconds (best of 240 runs)

📝 Explanation and details

Runtime improvement (primary): The optimized version reduces the function's median runtime from about 4.31 ms to 3.15 ms — a ~36% speedup overall — by removing per-character Python-level allocations and doing a single substring slice at the end.

What changed (concrete optimizations)

  • Removed per-character accumulation into a Python list (cur.append + "".join(cur)). Instead the optimized code records the start index and scans with integer indices, producing exactly one slice (args_str[start:i]) and a single rstrip() at the end.
  • Replaced the "in_string" state machine that appended every character with an inner loop that advances the index over quoted strings (skipping escapes) in larger leaps (i += 2 for escaped chars, and i += 1 for normal chars). This reduces Python bytecode executed per character inside strings.
  • Avoided many small Python-level operations inside the hot loop (fewer method calls and list appends). Condition checks for delimiters use explicit equality chains which avoid repeated list/tuple membership checks and repeated cur.append calls.
  • Kept nesting depth tracking but avoided storing all characters that are not needed until the final slice.

Why this yields a speedup (Python-level reasoning)

  • List append per character and the final join impose heavy Python overhead (many function calls, memory operations, and a large join cost). The original profiler shows substantial time attributed to cur.append and the final "".join(cur). The optimized approach replaces those many operations with a handful of integer increments and a single slice allocation — far less Python interpreter overhead.
  • Fewer Python-level objects and calls in the inner loop reduces garbage/allocator pressure and bytecode dispatch cost. Index arithmetic and slicing are lower-overhead than repeated append/join work at the same string size.
  • The optimized inner loop also avoids re-handling escaped characters via repeated append; jumping the index over escapes reduces loop iterations for quoted string segments.

Profiler evidence

  • Total profiled time for the function dropped (see provided line profiler results). The heavy per-character append/join lines in the original are gone or greatly reduced in the optimized trace; time moved into index scans and a single slice/rstrip operation. This aligns with the measured runtime improvement.

Behavioral changes & trade-offs

  • Behavior is preserved for the intended inputs (all regression tests pass and outputs match). There are a few tiny regressions on very small edge cases in the tests (e.g., whitespace-only or leading-comma inputs are marginally slower in a couple of microbenchmarks). These regressions are minor (single-digit percentage points) and are reasonable trade-offs given the consistent and significant overall runtime reduction, especially on realistic/hot workloads.
  • The change does not add new dependencies or alter public signatures.

Impact on workloads (where it matters)

  • Big wins for:
    • Long or deeply nested arguments (many characters scanned) because the per-character overhead is now much smaller.
    • High-throughput or repeated-invocation scenarios (see tests calling the function thousands of times) — the savings per call accumulate into large throughput gains.
  • Small inputs see smaller absolute gains; a couple of tiny cases in tests even slowed a bit, but those are microbenchmarks and don't offset overall throughput improvements for typical usage.

Tests & suitability

  • Annotated tests show consistent speedups across nested-parentheses, strings with escapes, generics, and long repeated loops. The optimization is therefore especially beneficial for realistic Java argument fragments (strings with nested delimiters, generics, escaped quotes) and when the transformer is used frequently.

Summary

  • Primary benefit: 36% faster runtime (4.31ms → 3.15ms).
  • How: replace many Python-level character appends and join with index-based scanning + one slice, plus an inner loop that skips quoted segments efficiently.
  • Trade-offs: negligible regressions on a few tiny edge microbenchmarks, acceptable given the large gains for real workloads and repeated calls.

Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests 2055 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_simple_and_whitespace_handling():
    # Create a real transformer instance with a sample function name.
    transformer = JavaAssertTransformer("assertEquals")

    # Simple first-arg extraction: basic comma-separated args.
    args = "a, b"
    # Expect the first top-level argument to be "a"
    codeflash_output = transformer._extract_first_arg(args) # 2.12μs -> 1.85μs (14.6% faster)

    # Leading whitespace should be skipped before extracting the first arg.
    args = "   foo.bar(), baz"
    # Expect "foo.bar()" with no leading spaces preserved.
    codeflash_output = transformer._extract_first_arg(args) # 3.53μs -> 2.56μs (37.5% faster)

    # If there's no comma (single argument), the whole trimmed string is returned.
    args = " singleArg "
    # Trailing and leading whitespace trimmed; internal preserved.
    codeflash_output = transformer._extract_first_arg(args) # 2.90μs -> 2.00μs (44.6% faster)

def test_basic_nested_delimiters_and_strings():
    transformer = JavaAssertTransformer("assertThat")

    # Nested parentheses: the inner comma should not end the top-level argument.
    args = "func(1, 2), other"
    # The extractor should include the entire func(1, 2) as the first arg.
    codeflash_output = transformer._extract_first_arg(args) # 4.33μs -> 3.32μs (30.5% faster)

    # Strings containing commas should not be split on those commas.
    args = '"hello, world", next'
    # The quoted string with the internal comma is preserved as a unit.
    codeflash_output = transformer._extract_first_arg(args) # 2.44μs -> 1.79μs (35.8% faster)

    # Char literal containing an escaped backslash (Java char '\\').
    # In Python literal we write it as "'\\\\'" to represent the two backslashes inside single quotes.
    args = "'\\\\', other"
    # Expect the char literal preserved exactly as in the input string.
    codeflash_output = transformer._extract_first_arg(args) # 1.23μs -> 962ns (28.2% faster)

    # Generics with commas inside angle brackets should be treated as nested depth.
    args = "Map<String, List<Integer>>, other"
    # The extractor should return the full generic type as first arg.
    codeflash_output = transformer._extract_first_arg(args) # 6.14μs -> 4.17μs (47.4% faster)

def test_edge_cases_empty_and_misplaced_commas_and_casts():
    transformer = JavaAssertTransformer("assertTrue")

    # Empty string should return None (nothing to extract).
    codeflash_output = transformer._extract_first_arg("") # 601ns -> 531ns (13.2% faster)

    # String with only whitespace should return None too.
    codeflash_output = transformer._extract_first_arg("   \t\n  ") # 932ns -> 991ns (5.95% slower)

    # Leading comma indicates an empty first argument -> return None.
    codeflash_output = transformer._extract_first_arg(", second") # 942ns -> 1.09μs (13.7% slower)

    # Unbalanced parentheses (no closing paren) - extractor should return whatever it saw.
    args = "(a, b"
    # It will not find a top-level comma (the comma is inside depth>0) so return the full string.
    codeflash_output = transformer._extract_first_arg(args) # 2.20μs -> 1.41μs (56.1% faster)

    # Leading cast should be preserved (cast uses parentheses but is part of the argument).
    args = "(MyType) obj.method(), rest"
    # The extractor should include the cast and the method call.
    codeflash_output = transformer._extract_first_arg(args) # 5.00μs -> 3.40μs (47.2% faster)

    # Braces and brackets count toward nesting and should not cause early split.
    args = "{x, y}, second"
    # The braces form a single argument even though there is a comma inside.
    codeflash_output = transformer._extract_first_arg(args) # 1.96μs -> 1.49μs (31.5% faster)

def test_string_with_escaped_quotes_inside():
    transformer = JavaAssertTransformer("assertThat")

    # Double-quoted string containing escaped quotes: "He said \"hi\""
    # Python literal must escape backslashes: "\"He said \\\"hi\\\"\", next"
    args = "\"He said \\\"hi\\\"\", other"
    # The extractor should treat the inner escaped quotes as part of the string and return the full quoted literal.
    codeflash_output = transformer._extract_first_arg(args) # 3.82μs -> 2.75μs (39.1% faster)

    # Single-quoted string containing an escaped single quote inside (char literal with escape).
    # Represent Java char literal '\'' in Python as "'\\''" - but to avoid quoting ambiguity we place overall string in double quotes:
    args = "'\\'' , following"
    # The extractor should include the escaped single-quote char literal.
    codeflash_output = transformer._extract_first_arg(args) # 1.72μs -> 1.39μs (23.7% faster)

def test_large_scale_nested_parentheses_and_repeated_invocations():
    transformer = JavaAssertTransformer("assertLarge")

    # Construct a very deep nested parentheses expression: (((...1...))) with depth 1000
    depth = 1000
    nested = "(" * depth + "1" + ")" * depth
    # Append other arguments after a top-level comma.
    args = nested + ", something_else"
    # Ensure the extractor returns the full nested expression as the first argument.
    codeflash_output = transformer._extract_first_arg(args) # 274μs -> 205μs (34.0% faster)

    # Now call the extractor repeatedly (1000 iterations) with slightly varying inputs
    # to ensure stability and performance under repeated use.
    for i in range(1000):
        # Alternate between a simple arg and a nested arg to exercise the code paths.
        if (i % 3) == 0:
            s = f"val{i}, rest"
            expected = f"val{i}"
        elif (i % 3) == 1:
            s = f"fn({i}, {i+1}), more"
            expected = f"fn({i}, {i+1})"
        else:
            # Use a quoted argument containing a comma to ensure string handling repeats correctly.
            s = f"\"x,{i}\", tail"
            expected = f"\"x,{i}\""
        # Each invocation must return the expected first argument deterministically.
        codeflash_output = transformer._extract_first_arg(s) # 1.86ms -> 1.36ms (36.4% 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  # used for our unit tests
from codeflash.languages.java.remove_asserts import JavaAssertTransformer

def test_basic_simple_first_arg_extraction():
    # Create a real transformer instance with a dummy function name.
    t = JavaAssertTransformer("assertX")
    # A simple argument list: no nesting, should return the first token trimmed.
    codeflash_output = t._extract_first_arg("foo, bar, baz") # 2.79μs -> 2.34μs (19.2% faster)
    # Leading whitespace should be ignored and trailing whitespace trimmed.
    codeflash_output = t._extract_first_arg("   singleArg   ") # 4.01μs -> 2.87μs (39.8% faster)
    # No comma means the whole (trimmed) string is returned.
    codeflash_output = t._extract_first_arg("onlyOneArg") # 2.77μs -> 1.96μs (41.3% faster)

def test_edge_empty_and_whitespace_inputs_return_none():
    t = JavaAssertTransformer("assertX")
    # Empty string -> None
    codeflash_output = t._extract_first_arg("") # 621ns -> 601ns (3.33% faster)
    # String with only whitespace -> None
    codeflash_output = t._extract_first_arg("   \t\n  ") # 1.03μs -> 1.03μs (0.097% faster)
    # Leading comma means there's no first top-level argument -> None
    codeflash_output = t._extract_first_arg(",b") # 1.00μs -> 1.05μs (4.75% slower)
    # Comma immediately after whitespace still yields no first argument
    codeflash_output = t._extract_first_arg("   , after") # 871ns -> 922ns (5.53% slower)

def test_nested_parentheses_brackets_braces_and_generics_handled():
    t = JavaAssertTransformer("assertX")
    # Parentheses: the comma inside nested parens should not split top-level args.
    s = "a(b, c), d"
    codeflash_output = t._extract_first_arg(s) # 3.76μs -> 3.00μs (25.4% faster)
    # Brackets and braces (array initializers / anonymous blocks) should be respected.
    s2 = "new int[]{1, 2, 3}, other"
    codeflash_output = t._extract_first_arg(s2) # 4.80μs -> 3.43μs (40.0% faster)
    # Angle brackets for generics: commas inside generics shouldn't split.
    s3 = "Map<String, List<Integer>>, other"
    codeflash_output = t._extract_first_arg(s3) # 6.07μs -> 4.03μs (50.7% faster)

def test_strings_with_commas_and_escaped_quotes_are_respected():
    t = JavaAssertTransformer("assertX")
    # Double-quoted string containing a comma: should be treated as a single arg.
    s = '"He said, \\"hello\\"", trailing'
    codeflash_output = t._extract_first_arg(s) # 4.32μs -> 3.03μs (42.7% faster)
    # Single-quoted string containing a comma: should also be preserved.
    s2 = "'a,b', rest"
    codeflash_output = t._extract_first_arg(s2) # 1.50μs -> 1.23μs (22.0% faster)
    # Escaped backslash sequences inside string should not break parsing.
    s3 = '"path\\\\to\\\\file,withcomma", next'
    codeflash_output = t._extract_first_arg(s3) # 3.14μs -> 1.99μs (57.3% faster)

def test_char_literals_and_edge_literals():
    t = JavaAssertTransformer("assertX")
    # Simple char literal followed by other args.
    codeflash_output = t._extract_first_arg("'a', b") # 2.44μs -> 1.93μs (26.4% faster)
    # Escaped char literal like newline should be kept intact.
    codeflash_output = t._extract_first_arg("'\\n', other") # 1.39μs -> 1.06μs (31.1% faster)
    # A lone char (no comma) should be returned trimmed.
    codeflash_output = t._extract_first_arg("'z'") # 901ns -> 721ns (25.0% faster)

def test_casts_and_parenthetical_prefixes():
    t = JavaAssertTransformer("assertX")
    # A cast at the front uses parentheses which must be balanced before splitting.
    s = "(Type) obj.method(1, 2), somethingElse"
    # The first top-level argument should include the cast and the method invocation.
    codeflash_output = t._extract_first_arg(s) # 6.85μs -> 5.02μs (36.5% faster)
    # Multiple nested casts and parentheses should still work.
    s2 = "((A) b).c(d, e), tail"
    codeflash_output = t._extract_first_arg(s2) # 3.73μs -> 2.67μs (39.3% faster)

def test_no_top_level_comma_returns_full_trimmed_arg():
    t = JavaAssertTransformer("assertX")
    # Complex expression without any top-level comma should be returned entirely.
    complex_expr = "someFunc(1, new int[]{1,2,3}, Map.<K,V>of(k,v))   "
    codeflash_output = t._extract_first_arg(complex_expr) # 12.5μs -> 8.82μs (41.9% faster)

def test_leading_and_trailing_whitespace_trimming():
    t = JavaAssertTransformer("assertX")
    # Leading whitespace is skipped at start; trailing whitespace trimmed from extracted arg.
    codeflash_output = t._extract_first_arg("   alpha  , beta") # 3.97μs -> 3.12μs (27.3% faster)
    codeflash_output = t._extract_first_arg("   alpha   ") # 2.75μs -> 1.96μs (39.8% faster)

def test_large_scale_many_items_first_arg_from_long_list():
    t = JavaAssertTransformer("assertX")
    # Build a large comma-separated list of function calls (1000 items).
    # The first item itself contains nested parentheses and commas that shouldn't be split.
    first_item = "complexFn(" + ",".join(str(i) for i in range(10)) + ")"  # nested commas inside
    rest = ",".join(f"fn{i}()" for i in range(1, 1000))
    long_args = first_item + "," + rest
    # The extractor should return exactly the first complex item (with nested commas intact).
    codeflash_output = t._extract_first_arg(long_args) # 8.50μs -> 6.31μs (34.6% faster)

def test_large_scale_loop_many_iterations_stability():
    t = JavaAssertTransformer("assertX")
    # Call the extractor 1000 times with predictable inputs to ensure stability.
    for i in range(1000):
        s = f"value{i},other{ i }"
        # Each call must deterministically return the correct first argument.
        codeflash_output = t._extract_first_arg(s) # 2.03ms -> 1.48ms (37.2% faster)

def test_complex_mixture_of_all_cases():
    t = JavaAssertTransformer("assertX")
    # A deliberately complex mixture of generics, arrays, strings, casts, and nested calls.
    complex_arg = (
        ' (A<B>) map.get("key,stillInString", arr[0], new HashMap<String, List<Integer>>() {{ put(1, List.of(2)); }} ) '
        + ", tail"
    )
    # Ensure extraction yields the whole complex first argument trimmed.
    codeflash_output = t._extract_first_arg(complex_arg); extracted = codeflash_output # 23.1μs -> 15.8μs (45.9% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To test or edit this optimization locally git merge codeflash/optimize-pr1663-2026-02-25T22.31.35

Click to see suggested changes
Suggested change
in_string = False
string_char = ""
cur: list[str] = []
while i < n:
ch = args_str[i]
if in_string:
cur.append(ch)
if ch == "\\" and i + 1 < n:
i += 1
cur.append(args_str[i])
elif ch == string_char:
in_string = False
elif ch in ('"', "'"):
in_string = True
string_char = ch
cur.append(ch)
elif ch in ("(", "<", "[", "{"):
depth += 1
cur.append(ch)
elif ch in (")", ">", "]", "}"):
depth -= 1
cur.append(ch)
elif ch == "," and depth == 0:
break
else:
cur.append(ch)
i += 1
# Trim trailing whitespace from the extracted argument
if not cur:
return None
return "".join(cur).rstrip()
# Record start index of the extracted argument instead of building a list
start = i
while i < n:
ch = args_str[i]
if ch == '"' or ch == "'":
string_char = ch
# include opening quote; advance inside string while handling escapes
i += 1
while i < n:
ch2 = args_str[i]
if ch2 == "\\" and i + 1 < n:
i += 2
elif ch2 == string_char:
i += 1
break
else:
i += 1
elif ch == "(" or ch == "<" or ch == "[" or ch == "{":
depth += 1
i += 1
elif ch == ")" or ch == ">" or ch == "]" or ch == "}":
depth -= 1
i += 1
elif ch == "," and depth == 0:
# slice up to but not including the comma
res = args_str[start:i].rstrip()
if not res:
return None
return res
else:
i += 1
# If we reached the end without encountering a top-level comma
res = args_str[start:i].rstrip()
if not res:
return None
return res

Static Badge



def transform_java_assertions(source: str, function_name: str, qualified_name: str | None = None) -> str:
"""Transform Java test code by removing assertions and capturing function calls.
Expand Down
20 changes: 17 additions & 3 deletions codeflash/languages/java/test_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,17 @@
# Allows: letters, digits, underscores, dots, and dollar signs (inner classes)
_VALID_JAVA_CLASS_NAME = re.compile(r"^[a-zA-Z_$][a-zA-Z0-9_$.]*$")

# Skip validation/analysis plugins that reject generated instrumented files
# (e.g. Apache Rat rejects missing license headers, Checkstyle rejects naming, etc.)
_MAVEN_VALIDATION_SKIP_FLAGS = [
"-Drat.skip=true",
"-Dcheckstyle.skip=true",
"-Dspotbugs.skip=true",
"-Dpmd.skip=true",
"-Denforcer.skip=true",
"-Djapicmp.skip=true",
]


def _validate_java_class_name(class_name: str) -> bool:
"""Validate that a string is a valid Java class name.
Expand Down Expand Up @@ -416,8 +427,9 @@ def run_behavioral_tests(
add_jacoco_plugin_to_pom(pom_path)
coverage_xml_path = get_jacoco_xml_path(project_root)

# Use a minimum timeout of 60s for Java builds (120s when coverage is enabled due to verify phase)
min_timeout = 120 if enable_coverage else 60
# Use a minimum timeout of 60s for Java builds (300s when coverage is enabled due to verify phase
# which runs full compilation + instrumentation + test execution in multi-module projects)
min_timeout = 300 if enable_coverage else 60
effective_timeout = max(timeout or 300, min_timeout)

if enable_coverage:
Expand Down Expand Up @@ -493,6 +505,7 @@ def _compile_tests(
return subprocess.CompletedProcess(args=["mvn"], returncode=-1, stdout="", stderr="Maven not found")

cmd = [mvn, "test-compile", "-e", "-B"] # Show errors but not verbose output; -B for batch mode (no ANSI colors)
cmd.extend(_MAVEN_VALIDATION_SKIP_FLAGS)

if test_module:
cmd.extend(["-pl", test_module, "-am"])
Expand Down Expand Up @@ -1447,6 +1460,7 @@ def _run_maven_tests(
# JaCoCo's report goal is bound to the verify phase to get post-test execution data
maven_goal = "verify" if enable_coverage else "test"
cmd = [mvn, maven_goal, "-fae", "-B"] # Fail at end to run all tests; -B for batch mode (no ANSI colors)
cmd.extend(_MAVEN_VALIDATION_SKIP_FLAGS)

# Add --add-opens flags for Java 16+ module system compatibility.
# The codeflash-runtime Serializer uses Kryo which needs reflective access to
Expand Down Expand Up @@ -1483,7 +1497,7 @@ def _run_maven_tests(
# -am = also make dependencies
# -DfailIfNoTests=false allows dependency modules without tests to pass
# -DskipTests=false overrides any skipTests=true in pom.xml
cmd.extend(["-pl", test_module, "-am", "-DfailIfNoTests=false", "-DskipTests=false"])
cmd.extend(["-pl", test_module, "-am", "-DfailIfNoTests=false", "-Dsurefire.failIfNoSpecifiedTests=false", "-DskipTests=false"])

if test_filter:
# Validate test filter to prevent command injection
Expand Down
8 changes: 7 additions & 1 deletion codeflash/verification/parse_test_output.py
Original file line number Diff line number Diff line change
Expand Up @@ -805,7 +805,13 @@ def parse_test_xml(
if class_name is not None and class_name.startswith(test_module_path):
test_class = class_name[len(test_module_path) + 1 :] # +1 for the dot, gets Unittest class name

loop_index = int(testcase.name.split("[ ")[-1][:-2]) if testcase.name and "[" in testcase.name else 1
loop_index = 1
if testcase.name and "[" in testcase.name:
bracket_content = testcase.name.rsplit("[", 1)[-1].rstrip("]").strip()
try:
loop_index = int(bracket_content)
except ValueError:
loop_index = 1

timed_out = False
if len(testcase.result) > 1:
Expand Down
2 changes: 1 addition & 1 deletion tests/test_languages/test_java/test_instrumentation.py
Original file line number Diff line number Diff line change
Expand Up @@ -1324,7 +1324,7 @@ def test_instrument_generated_test_behavior_mode(self):
}
}
}
Object _cf_result1 = _cf_result1_1;
int _cf_result1 = (int)_cf_result1_1;
}
}
"""
Expand Down
Loading
Loading