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

---

📄 12% (0.12x) speedup for instrument_generated_java_test in codeflash/languages/java/instrumentation.py

⏱️ Runtime : 4.07 milliseconds → 3.63 milliseconds (best of 201 runs)

📝 Explanation and details

This optimization achieves a 12% runtime improvement by targeting the most expensive operations in Java test instrumentation code. The key performance gains come from:

1. Efficient Brace Counting (40% of original time eliminated)
The original code iterates character-by-character through each line checking if ch == "{" or ch == "}", accounting for ~39% of the function's runtime. The optimized version replaces this with body_line.count("{") - body_line.count("}"), leveraging Python's C-implemented string methods that are 10-20x faster than Python loops for character counting.

2. Local Variable Caching
Critical loop variables like lines_local, res_append, and append_body are cached before hot loops. This eliminates repeated attribute lookups (e.g., result.append → res_append) which Python must perform on every call. The profiler shows the inner while loop runs 1,158 times per invocation - these small savings compound significantly.

3. Precompiled Regex Pattern
The instrument_generated_java_test function now precompiles the regex pattern once rather than calling re.sub with a raw pattern string. This eliminates the per-call compilation overhead, reducing the regex operation from 36.4% to 42.6% of that function's total time (though with better absolute performance).

4. Optimized String Concatenation
Using a constant body_prefix = " " instead of repeatedly creating " " + bl reduces string allocation overhead in the body line loop.

Test Results Show Consistent Improvements:

• Simple instrumentation cases: 17-47% faster (e.g., test_behavior_mode_basic_rename_only: 36.6% faster)
• Complex nested brace scenarios: 24-32% faster (e.g., test_large_method_with_nested_complexity: 31.8% faster)
• Large-scale tests with 200+ methods: 5-10% faster, still meaningful at scale

The optimization particularly excels when processing methods with significant body content or nested structures (where brace counting dominates), making it valuable for real-world Java test generation workflows that involve complex test methods with loops, conditionals, and exception handling.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 46 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

from __future__ import annotations

import logging
import re

# imports
import pytest  # used for our unit tests
from codeflash.languages.java.instrumentation import \
    instrument_generated_java_test

def test_behavior_mode_basic_rename_only():
    # Minimal Java test source with a public class and a single @Test method.
    original = (
        "package com.example;\n"
        "public class SampleTest {\n"
        "    @Test\n"
        "    public void testSomething() {\n"
        "        // original body\n"
        "    }\n"
        "}\n"
    )

    # Call the instrumentation in 'behavior' mode.
    codeflash_output = instrument_generated_java_test(original, "testSomething", "com.example.SampleTest", mode="behavior"); out = codeflash_output # 17.4μs -> 12.7μs (36.6% faster)

def test_performance_mode_single_test_instrumentation():
    # Create a test with one @Test method with the opening brace on the same line.
    original = (
        "class PerfClass {\n"
        "    @Test\n"
        "    public void doWork() {\n"
        "        int x = 1;\n"
        "        x++;\n"
        "    }\n"
        "}\n"
    )

    codeflash_output = instrument_generated_java_test(original, "doWork", "Pkg.PerfClass.doWork", mode="performance"); out = codeflash_output # 28.4μs -> 23.6μs (20.4% faster)

def test_no_class_found_returns_input_unchanged(caplog):
    caplog.set_level("WARNING")

    original = (
        "// This file intentionally has no class declaration\n"
        "/* @Test */\n"
        "void helper() {}\n"
    )
    codeflash_output = instrument_generated_java_test(original, "helper", "helper", mode="performance"); out = codeflash_output # 513μs -> 517μs (0.744% slower)

def test_multiple_annotations_and_multiline_signature():
    original = (
        "class MultiAnn {\n"
        "    @Test\n"
        "    @SomeOther\n"
        "    public void complexTest()\n"
        "        throws Exception\n"
        "    {\n"
        "        doSomething();\n"
        "    }\n"
        "}\n"
    )

    codeflash_output = instrument_generated_java_test(original, "complexTest", "pkg.MultiAnn.complexTest", mode="performance"); out = codeflash_output # 34.4μs -> 28.6μs (20.1% faster)

def test_nested_braces_inside_method_preserved_and_indented():
    original = (
        "class NestedBraces {\n"
        "    @Test\n"
        "    public void nested() {\n"
        "        if (true) {\n"
        "            while (x < 5) {\n"
        "                call();\n"
        "            }\n"
        "        }\n"
        "    }\n"
        "}\n"
    )

    codeflash_output = instrument_generated_java_test(original, "nested", "pkg.NestedBraces.nested", mode="performance"); out = codeflash_output # 33.3μs -> 26.7μs (24.7% faster)

def test_behavior_mode_no_performance_markers():
    original = (
        "class BehaviorOnly {\n"
        "    @Test\n"
        "    public void bTest() {\n"
        "        // behavior logic\n"
        "    }\n"
        "}\n"
    )

    codeflash_output = instrument_generated_java_test(original, "bTest", "pkg.BehaviorOnly.bTest", mode="behavior"); out = codeflash_output # 16.1μs -> 11.5μs (40.0% faster)

def test_large_scale_many_tests_performance():
    # Generate a class with 200 simple @Test methods (keeps under 1000 as requested)
    methods = []
    for i in range(1, 201):
        methods.append("    @Test")
        methods.append(f"    public void test{i}() {{")
        methods.append(f"        // body {i}")
        methods.append("    }")
    body = "\n".join(methods)
    original = "class BigClass {\n" + body + "\n}\n"

    codeflash_output = instrument_generated_java_test(original, "testX", "pkg.BigClass.testX", mode="performance"); out = codeflash_output # 1.27ms -> 1.21ms (4.94% faster)

def test_class_detection_without_public_keyword():
    original = (
        "package p;\n"
        "class Simple {\n"
        "    @Test\n"
        "    void t() {}\n"
        "}\n"
    )

    codeflash_output = instrument_generated_java_test(original, "t", "p.Simple.t", mode="behavior"); out = codeflash_output # 16.4μs -> 11.5μs (43.0% faster)

def test_method_opening_brace_on_separate_line():
    original = (
        "class BraceLater {\n"
        "    @Test\n"
        "    public void later()\n"
        "    {\n"
        "        doIt();\n"
        "    }\n"
        "}\n"
    )

    codeflash_output = instrument_generated_java_test(original, "later", "pkg.BraceLater.later", mode="performance"); out = codeflash_output # 28.3μs -> 24.0μs (17.9% faster)

def test_annotations_before_test_are_preserved_and_handled():
    original = (
        "class AnnotationsBefore {\n"
        "    @Before\n"
        "    @Test\n"
        "    public void check() {\n"
        "        work();\n"
        "    }\n"
        "}\n"
    )

    codeflash_output = instrument_generated_java_test(original, "check", "pkg.AnnotationsBefore.check", mode="performance"); out = codeflash_output # 27.8μs -> 23.1μs (20.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 re

import pytest
from codeflash.languages.java.instrumentation import \
    instrument_generated_java_test

class TestInstrumentGeneratedJavaTestBasic:
    """Basic test cases for instrument_generated_java_test function."""

    def test_behavior_mode_class_rename(self):
        """Test that behavior mode renames class with __perfinstrumented suffix."""
        test_code = "public class MyTestClass { }"
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod", "com.example.TestClass", "behavior"
        ); result = codeflash_output # 14.7μs -> 10.5μs (40.0% faster)

    def test_performance_mode_class_rename(self):
        """Test that performance mode renames class with __perfonlyinstrumented suffix."""
        test_code = "public class MyTestClass { }"
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod", "com.example.TestClass", "performance"
        ); result = codeflash_output # 15.4μs -> 11.9μs (29.8% faster)

    def test_simple_test_method_behavior_mode(self):
        """Test behavior mode with simple @Test method."""
        test_code = """public class SimpleTest {
    @Test
    public void testSimple() {
        System.out.println("Hello");
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testSimple", "com.example.SimpleTest", "behavior"
        ); result = codeflash_output # 16.1μs -> 11.6μs (38.4% faster)

    def test_simple_test_method_performance_mode(self):
        """Test performance mode adds timing instrumentation."""
        test_code = """public class SimpleTest {
    @Test
    public void testMethod() {
        int x = 5;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod", "com.example.SimpleTest", "performance"
        ); result = codeflash_output # 27.7μs -> 23.1μs (19.9% faster)

    def test_preserves_method_body_behavior_mode(self):
        """Test that behavior mode preserves method body without modification."""
        test_code = """public class PreserveTest {
    @Test
    public void testPreserve() {
        System.out.println("test");
        int result = 10 + 20;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testPreserve", "com.example.PreserveTest", "behavior"
        ); result = codeflash_output # 15.9μs -> 11.9μs (33.4% faster)

    def test_class_name_extraction_with_no_public_keyword(self):
        """Test class name extraction works with class keyword even without public."""
        test_code = "class MyTest { }"
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod", "com.example.MyTest", "behavior"
        ); result = codeflash_output # 13.2μs -> 8.92μs (47.9% faster)

    def test_preserves_annotations_in_performance_mode(self):
        """Test that @Test annotation is preserved in performance mode."""
        test_code = """public class AnnotationTest {
    @Test
    public void testAnnotation() {
        int x = 1;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testAnnotation", "com.example.AnnotationTest", "performance"
        ); result = codeflash_output # 27.4μs -> 23.0μs (19.0% faster)

    def test_preserves_multiple_annotations(self):
        """Test that multiple annotations are preserved in instrumented code."""
        test_code = """public class MultiAnnotationTest {
    @Test
    @DisplayName("My Test")
    public void testMulti() {
        int x = 1;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testMulti", "com.example.MultiAnnotationTest", "performance"
        ); result = codeflash_output # 27.8μs -> 23.1μs (20.4% faster)

    def test_missing_class_definition_returns_unchanged(self):
        """Test that code without class definition is returned unchanged."""
        test_code = "// Just a comment without class"
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod", "com.example.Test", "behavior"
        ); result = codeflash_output # 519μs -> 516μs (0.583% faster)

class TestInstrumentGeneratedJavaTestEdgeCases:
    """Edge case tests for instrument_generated_java_test function."""

    def test_class_name_with_underscores(self):
        """Test handling of class names containing underscores."""
        test_code = "public class My_Test_Class { }"
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod", "com.example.My_Test_Class", "behavior"
        ); result = codeflash_output # 15.1μs -> 10.3μs (46.4% faster)

    def test_class_name_with_numbers(self):
        """Test handling of class names containing numbers."""
        test_code = "public class Test123Class { }"
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod", "com.example.Test123Class", "behavior"
        ); result = codeflash_output # 13.9μs -> 9.57μs (45.5% faster)

    def test_method_with_nested_braces(self):
        """Test handling of method with nested code blocks (if/for statements)."""
        test_code = """public class NestedTest {
    @Test
    public void testNested() {
        for (int i = 0; i < 10; i++) {
            if (i > 5) {
                System.out.println(i);
            }
        }
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testNested", "com.example.NestedTest", "performance"
        ); result = codeflash_output # 35.1μs -> 28.4μs (23.6% faster)

    def test_multiple_test_methods(self):
        """Test handling of multiple @Test methods in same class."""
        test_code = """public class MultiMethodTest {
    @Test
    public void testFirst() {
        int x = 1;
    }
    
    @Test
    public void testSecond() {
        int y = 2;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testFirst", "com.example.MultiMethodTest", "performance"
        ); result = codeflash_output # 36.7μs -> 32.1μs (14.2% faster)

    def test_empty_method_body(self):
        """Test handling of method with empty body."""
        test_code = """public class EmptyTest {
    @Test
    public void testEmpty() {
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testEmpty", "com.example.EmptyTest", "performance"
        ); result = codeflash_output # 25.3μs -> 20.8μs (21.5% faster)

    def test_method_with_throws_clause(self):
        """Test handling of method with throws declaration."""
        test_code = """public class ThrowsTest {
    @Test
    public void testThrows() throws Exception {
        int x = 1;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testThrows", "com.example.ThrowsTest", "performance"
        ); result = codeflash_output # 26.7μs -> 22.2μs (20.5% faster)

    def test_method_with_parameters(self):
        """Test handling of method with parameters."""
        test_code = """public class ParamTest {
    @Test
    public void testParams(String name, int value) {
        System.out.println(name);
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testParams", "com.example.ParamTest", "performance"
        ); result = codeflash_output # 27.9μs -> 22.8μs (22.1% faster)

    def test_class_with_inner_classes(self):
        """Test handling of code with inner classes."""
        test_code = """public class OuterTest {
    @Test
    public void testOuter() {
        int x = 1;
    }
    
    class InnerTest {
        void innerMethod() {
            int y = 2;
        }
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testOuter", "com.example.OuterTest", "performance"
        ); result = codeflash_output # 29.7μs -> 25.7μs (15.4% faster)

    def test_performance_mode_includes_correct_variable_names(self):
        """Test that performance mode uses correct variable names in timing code."""
        test_code = """public class VarTest {
    @Test
    public void testVar() {
        int x = 1;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testVar", "com.example.VarTest", "performance"
        ); result = codeflash_output # 25.9μs -> 21.7μs (19.5% faster)

    def test_behavior_mode_no_timing_instrumentation(self):
        """Test that behavior mode does NOT include timing instrumentation."""
        test_code = """public class NoTimingTest {
    @Test
    public void testNoTiming() {
        int x = 1;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testNoTiming", "com.example.NoTimingTest", "behavior"
        ); result = codeflash_output # 15.4μs -> 11.0μs (40.3% faster)

    def test_class_name_already_has_suffix(self):
        """Test handling when class name already contains the suffix pattern."""
        test_code = "public class TestClass__perfinstrumented { }"
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod", "com.example.TestClass__perfinstrumented", "behavior"
        ); result = codeflash_output # 14.2μs -> 10.0μs (41.2% faster)

    def test_special_characters_in_string_literals(self):
        """Test that string literals with special characters are preserved."""
        test_code = """public class SpecialTest {
    @Test
    public void testSpecial() {
        String s = "Hello: World!@#$%";
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testSpecial", "com.example.SpecialTest", "performance"
        ); result = codeflash_output # 28.1μs -> 22.8μs (23.0% faster)

    def test_method_with_comments(self):
        """Test that comments within methods are preserved."""
        test_code = """public class CommentTest {
    @Test
    public void testComment() {
        // This is a comment
        int x = 1; // inline comment
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testComment", "com.example.CommentTest", "performance"
        ); result = codeflash_output # 29.3μs -> 23.9μs (22.8% faster)

    def test_whitespace_preservation_in_method_body(self):
        """Test that indentation and whitespace in original method body is preserved."""
        test_code = """public class WhitespaceTest {
    @Test
    public void testWhitespace() {
        int x = 1;
        int y = 2;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testWhitespace", "com.example.WhitespaceTest", "performance"
        ); result = codeflash_output # 27.6μs -> 22.7μs (21.9% faster)

class TestInstrumentGeneratedJavaTestLargeScale:
    """Large scale test cases for instrument_generated_java_test function."""

    def test_large_method_body(self):
        """Test handling of method with large body (many statements)."""
        # Create a method with 100 statements
        statements = "\n        ".join([f"int var{i} = {i};" for i in range(100)])
        test_code = f"""public class LargeTest {{
    @Test
    public void testLarge() {{
        {statements}
    }}
}}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testLarge", "com.example.LargeTest", "performance"
        ); result = codeflash_output # 151μs -> 102μs (47.1% faster)
        # Verify all statements are preserved
        for i in range(100):
            pass

    def test_large_method_with_nested_complexity(self):
        """Test handling of method with deeply nested control structures."""
        # Create nested loops and conditions
        nested_code = """int result = 0;
        for (int i = 0; i < 10; i++) {
            for (int j = 0; j < 10; j++) {
                if (i > j) {
                    try {
                        result += i * j;
                    } catch (Exception e) {
                        result -= 1;
                    }
                }
            }
        }"""
        test_code = f"""public class ComplexTest {{
    @Test
    public void testComplex() {{
        {nested_code}
    }}
}}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testComplex", "com.example.ComplexTest", "performance"
        ); result = codeflash_output # 44.6μs -> 33.9μs (31.8% faster)

    def test_multiple_test_methods_large_count(self):
        """Test handling of class with many @Test methods."""
        # Create 50 test methods
        methods = "\n    ".join(
            [
                f"""@Test
    public void testMethod{i}() {{
        int x = {i};
    }}"""
                for i in range(50)
            ]
        )
        test_code = f"""public class ManyMethodsTest {{
    {methods}
}}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod0", "com.example.ManyMethodsTest", "performance"
        ); result = codeflash_output # 327μs -> 297μs (10.2% faster)

    def test_very_long_class_name(self):
        """Test handling of very long qualified class names."""
        long_class_name = "A" * 100 + "Test"
        test_code = f'public class {long_class_name} {{ }}'
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod", f"com.example.{long_class_name}", "behavior"
        ); result = codeflash_output # 14.2μs -> 10.00μs (41.8% faster)

    def test_method_with_long_body_and_many_braces(self):
        """Test method with hundreds of braces to stress brace counting."""
        # Create deeply nested structure with many braces
        body = "{"
        for i in range(50):
            body += f"{{ int x{i} = {i}; "
        for i in range(50):
            body += "}"
        body += "}"
        test_code = f"""public class BraceTest {{
    @Test
    public void testBraces() {body}
}}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testBraces", "com.example.BraceTest", "performance"
        ); result = codeflash_output # 44.5μs -> 40.5μs (9.98% faster)

    def test_performance_mode_timing_markers_format(self):
        """Test that performance mode creates correct timing marker format."""
        test_code = """public class MarkerTest {
    @Test
    public void testMarker() {
        int x = 1;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testMarker", "com.example.MarkerTest", "performance"
        ); result = codeflash_output # 26.6μs -> 21.8μs (22.3% faster)

    def test_class_with_mixed_method_types(self):
        """Test class with both @Test and non-@Test methods."""
        test_code = """public class MixedTest {
    @Test
    public void testMethod() {
        int x = 1;
    }
    
    public void helperMethod() {
        int y = 2;
    }
    
    @Test
    public void anotherTest() {
        int z = 3;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod", "com.example.MixedTest", "performance"
        ); result = codeflash_output # 37.6μs -> 33.1μs (13.6% faster)

    def test_performance_mode_with_original_class_name_in_markers(self):
        """Test that performance mode uses original class name in timing markers."""
        test_code = """public class OriginalName {
    @Test
    public void testMethod() {
        int x = 1;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testMethod", "com.example.OriginalName", "performance"
        ); result = codeflash_output # 25.9μs -> 21.4μs (21.0% faster)

    def test_function_name_with_special_case(self):
        """Test function name handling with mixed case."""
        test_code = """public class FunctionNameTest {
    @Test
    public void testMySpecialFunction() {
        int x = 1;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testMySpecialFunction", "com.example.Test", "performance"
        ); result = codeflash_output # 26.3μs -> 22.0μs (19.6% faster)

    def test_large_string_literals_in_method(self):
        """Test handling of large string literals within method body."""
        large_string = "a" * 500
        test_code = f"""public class LargeStringTest {{
    @Test
    public void testString() {{
        String s = "{large_string}";
    }}
}}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testString", "com.example.LargeStringTest", "performance"
        ); result = codeflash_output # 43.9μs -> 27.6μs (59.3% faster)

    def test_method_with_many_local_variables(self):
        """Test method declaring many local variables."""
        variables = "\n        ".join([f"int var{i} = {i};" for i in range(200)])
        test_code = f"""public class ManyVarsTest {{
    @Test
    public void testManyVars() {{
        {variables}
    }}
}}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testManyVars", "com.example.ManyVarsTest", "performance"
        ); result = codeflash_output # 281μs -> 185μs (52.4% faster)
        # All variables should be preserved
        for i in range(200):
            pass

    def test_mode_parameter_case_sensitivity(self):
        """Test that mode parameter respects case sensitivity."""
        test_code = """public class CaseTest {
    @Test
    public void testCase() {
        int x = 1;
    }
}"""
        codeflash_output = instrument_generated_java_test(
            test_code, "testCase", "com.example.CaseTest", "behavior"
        ); result_behavior = codeflash_output # 15.5μs -> 11.1μs (40.0% faster)
        codeflash_output = instrument_generated_java_test(
            test_code, "testCase", "com.example.CaseTest", "performance"
        ); result_performance = codeflash_output # 20.2μs -> 17.2μs (17.2% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr1199-2026-02-03T09.37.58 and push.

[KRRT7]

Closing stale bot PR.