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

---

📄 22% (0.22x) speedup for _add_behavior_instrumentation in codeflash/languages/java/instrumentation.py

⏱️ Runtime : 4.44 milliseconds → 3.63 milliseconds (best of 250 runs)

📝 Explanation and details

This optimization achieves a 22% runtime improvement (4.44ms → 3.63ms) by addressing three key performance bottlenecks:

Primary Optimization: Cached Regex Compilation (29.7% of optimized runtime)

The original code compiled the same regex pattern 202 times inside a loop (consuming 17.8% of runtime). The optimized version introduces:

@lru_cache(maxsize=128)
def _get_method_call_pattern(func_name: str):
    return re.compile(...)

This caches compiled patterns, eliminating redundant compilation. While the first call appears slower in the line profiler (9.3ms vs 8.3ms total), this is because it includes cache initialization overhead. Subsequent calls benefit from instant retrieval, making this optimization particularly valuable when:

• Instrumenting multiple test methods in sequence
• Processing classes with many @Test methods (e.g., the 50-method test shows 14.8% speedup)

Secondary Optimization: Efficient Brace Counting

The original code iterated character-by-character through method bodies (23.4% of runtime):

for ch in body_line:
    if ch == "{": brace_depth += 1
    elif ch == "}": brace_depth -= 1

The optimized version uses Python's built-in string methods:

open_count = body_line.count('{')
close_count = body_line.count('}')
brace_depth += open_count - close_count

This change shows dramatic improvements in tests with deeply nested structures:

• 10-level nested braces: 66.4% faster
• Large method bodies (100+ lines): 44.0% faster
• Methods with many variables (500+): 88.9% faster

Performance Characteristics

The optimization excels in scenarios common to Java test instrumentation:

• Multiple test methods: 11-15% speedup for classes with 30-100 test methods
• Complex method bodies: 29-44% speedup for methods with many nested structures or statements
• Sequential processing: Benefits accumulate when instrumenting multiple files due to regex caching

The minor slowdowns (3-9%) in trivial cases (empty methods, minimal source) are negligible compared to the substantial gains in realistic workloads, where Java test classes typically contain multiple complex test methods.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 50 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 re
import re as _re  # avoid shadowing re used in tests

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

def test_add_imports_before_class():
    # Minimal source with package and class, but no imports.
    src = "package com.example;\n\npublic class ExampleTest {\n}\n"
    codeflash_output = _add_behavior_instrumentation(src, "Example", "someFunc"); out = codeflash_output # 6.55μs -> 7.16μs (8.58% slower)
    # Ensure imports are placed before the class declaration (first occurrence of class)
    first_class_idx = out.find("public class")
    first_import_idx = out.find("import java.sql.Connection;")

def test_add_imports_after_existing_imports():
    # Source that already has an import line.
    src = "import java.util.List;\nimport java.util.Map;\npublic class ExampleTest {\n}\n"
    codeflash_output = _add_behavior_instrumentation(src, "Example", "someFunc"); out = codeflash_output # 6.23μs -> 6.53μs (4.60% slower)

def test_wrap_single_method_call_and_capture():
    # Create a @Test method with a simple target function call.
    src = """public class ExampleTest {
@Test
public void testSomething() {
    int x = obj.targetFunc(1);
}
}
"""
    codeflash_output = _add_behavior_instrumentation(src, "Example", "targetFunc"); out = codeflash_output # 33.0μs -> 28.3μs (16.5% faster)

def test_no_calls_serialize_null():
    src = """public class ExampleTest {
@Test
public void testNothing() {
    int a = 1 + 2; // no calls to target function
}
}
"""
    codeflash_output = _add_behavior_instrumentation(src, "Example", "targetFunc"); out = codeflash_output # 25.2μs -> 20.8μs (21.3% faster)

def test_multiple_calls_in_same_method():
    src = """public class ExampleTest {
@Test
public void testMultiple() {
    int a = calc.targetFunc(10);
    String s = new Helper().targetFunc("x");
}
}
"""
    codeflash_output = _add_behavior_instrumentation(src, "Example", "targetFunc"); out = codeflash_output # 37.5μs -> 32.5μs (15.2% faster)

def test_annotations_and_signature_collection():
    src = """public class ExampleTest {
@Test
@Another
public
void
testWeird()
{
    this.targetFunc();
}
}
"""
    codeflash_output = _add_behavior_instrumentation(src, "Example", "targetFunc"); out = codeflash_output # 30.8μs -> 27.2μs (13.0% faster)

def test_nested_braces_and_brace_counting():
    src = """public class ExampleTest {
@Test
public void testNested() {
    if (true) {
        while (false) {
            new Foo().target();
            bar.targetFunc(5); // targetFunc is the one we're instrumenting
        }
    }
}
}
"""
    codeflash_output = _add_behavior_instrumentation(src, "Example", "targetFunc"); out = codeflash_output # 40.9μs -> 34.4μs (18.9% faster)

def test_call_split_across_lines_not_matched():
    src = """public class ExampleTest {
@Test
public void testSplit() {
    obj.
    targetFunc(1);
}
}
"""
    codeflash_output = _add_behavior_instrumentation(src, "Example", "targetFunc"); out = codeflash_output # 27.4μs -> 24.2μs (13.2% faster)

def test_multiple_test_methods_iteration_ids():
    src = """public class ExampleTest {
@Test
public void first() {
    a.targetFunc();
}
@Test
public void second() {
    b.targetFunc();
}
}
"""
    codeflash_output = _add_behavior_instrumentation(src, "Example", "targetFunc"); out = codeflash_output # 48.1μs -> 42.7μs (12.8% faster)

def test_large_scale_many_test_methods():
    # Generate 100 test methods to keep under the 1000 step and element constraints.
    n = 100
    methods = []
    for i in range(1, n + 1):
        methods.append(f"@Test\npublic void t{i}() {{\n    obj.targetFunc({i});\n}}")
    class_body = "public class LargeTest {\n" + "\n".join(methods) + "\n}\n"
    codeflash_output = _add_behavior_instrumentation(class_body, "Large", "targetFunc"); out = codeflash_output # 1.37ms -> 1.20ms (14.1% faster)

def test_existing_imports_not_duplicated():
    src = "import java.sql.Connection;\npublic class ExampleTest {\n}\n"
    codeflash_output = _add_behavior_instrumentation(src, "Example", "someFunc"); out = codeflash_output # 5.97μs -> 6.39μs (6.56% slower)

def test_function_name_in_comments_and_strings_ignored():
    src = '''public class ExampleTest {
@Test
public void testComments() {
    // targetFunc should not be captured when in comments: targetFunc()
    String s = "targetFunc(1)";
}
}
'''
    codeflash_output = _add_behavior_instrumentation(src, "Example", "targetFunc"); out = codeflash_output # 39.8μs -> 33.9μs (17.3% faster)

def test_empty_source_handled_gracefully():
    src = ""
    codeflash_output = _add_behavior_instrumentation(src, "Empty", "fn"); out = codeflash_output # 2.76μs -> 3.20μs (13.9% slower)
# 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 \
    _add_behavior_instrumentation

class TestBasicFunctionality:
    """Basic test cases verifying fundamental functionality."""

    def test_adds_required_imports(self):
        """Test that required SQL imports are added to source code."""
        source = "public class TestClass {\n}"
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "testFunc"); result = codeflash_output # 5.12μs -> 5.40μs (5.26% slower)

    def test_preserves_existing_imports(self):
        """Test that existing imports are preserved in output."""
        source = "import org.junit.Test;\npublic class TestClass {\n}"
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "testFunc"); result = codeflash_output # 5.71μs -> 5.90μs (3.14% slower)

    def test_identifies_test_annotation(self):
        """Test that @Test annotated methods are properly instrumented."""
        source = """public class TestClass {
    @Test
    public void testAdd() {
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 24.9μs -> 20.9μs (19.2% faster)

    def test_simple_method_with_function_call(self):
        """Test instrumentation of a simple method with one function call."""
        source = """public class TestClass {
    @Test
    public void testMethod() {
        Calculator calc = new Calculator();
        int result = calc.add(1, 2);
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 35.8μs -> 30.9μs (15.7% faster)

    def test_empty_method_body(self):
        """Test instrumentation of a test method with empty body."""
        source = """public class TestClass {
    @Test
    public void testEmpty() {
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "someFunc"); result = codeflash_output # 22.7μs -> 19.2μs (18.1% faster)

    def test_method_without_target_function_call(self):
        """Test that methods without the target function call are still instrumented."""
        source = """public class TestClass {
    @Test
    public void testNoTarget() {
        int x = 5;
        int y = 10;
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "targetFunc"); result = codeflash_output # 25.0μs -> 21.1μs (18.8% faster)

    def test_class_name_parameter_used_correctly(self):
        """Test that the class_name parameter is correctly used in instrumentation."""
        source = """public class MyTestClass {
    @Test
    public void testSomething() {
        MyClass obj = new MyClass();
        obj.func();
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "MyTestClass", "func"); result = codeflash_output # 32.3μs -> 27.7μs (16.6% faster)

    def test_function_name_parameter_used_correctly(self):
        """Test that the func_name parameter is correctly used in instrumentation."""
        source = """public class TestClass {
    @Test
    public void testReverse() {
        String obj = new String("hello");
        obj.toLowerCase();
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "toLowerCase"); result = codeflash_output # 32.0μs -> 27.0μs (18.7% faster)

    def test_multiple_function_calls_in_same_line(self):
        """Test handling of multiple function calls on the same line."""
        source = """public class TestClass {
    @Test
    public void testMultiple() {
        calc.add(1, 2); calc.add(3, 4);
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 31.6μs -> 27.3μs (15.7% faster)

class TestEdgeCases:
    """Edge case tests for unusual or extreme conditions."""

    def test_no_test_annotation(self):
        """Test that methods without @Test annotation are not instrumented."""
        source = """public class TestClass {
    public void regularMethod() {
        calc.add(1, 2);
    }
    
    @Test
    public void testMethod() {
        calc.add(3, 4);
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 31.2μs -> 27.4μs (13.6% faster)
        
        # Only one instrumentation should be added (for the @Test method)
        count_iter_id = result.count("_cf_iter1")

    def test_nested_braces_in_method(self):
        """Test correct handling of nested braces (if statements, loops, etc.)."""
        source = """public class TestClass {
    @Test
    public void testNested() {
        if (true) {
            calc.add(1, 2);
        }
        for (int i = 0; i < 5; i++) {
            calc.add(i, i+1);
        }
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 38.5μs -> 32.4μs (18.8% faster)

    def test_function_name_as_substring_of_other_calls(self):
        """Test that function name matching doesn't capture partial matches."""
        source = """public class TestClass {
    @Test
    public void testSubstring() {
        obj.add(1, 2);
        obj.addAll(3, 4);
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 33.0μs -> 28.7μs (14.8% faster)

    def test_method_with_new_keyword_receiver(self):
        """Test instrumentation where method is called on a newly created object."""
        source = """public class TestClass {
    @Test
    public void testNewKeyword() {
        new StringUtils().reverse("hello");
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "reverse"); result = codeflash_output # 29.6μs -> 25.6μs (15.7% faster)

    def test_method_with_this_keyword_receiver(self):
        """Test instrumentation where method is called using 'this'."""
        source = """public class TestClass {
    @Test
    public void testThis() {
        this.execute();
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "execute"); result = codeflash_output # 29.1μs -> 25.2μs (15.5% faster)

    def test_method_with_no_imports_in_source(self):
        """Test adding imports when source has no existing imports."""
        source = """public class TestClass {
    @Test
    public void test() {
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "func"); result = codeflash_output # 21.4μs -> 18.6μs (15.1% faster)
        lines = result.split("\n")
        import_index = next(i for i, line in enumerate(lines) if "import java.sql" in line)
        class_index = next(i for i, line in enumerate(lines) if "public class" in line)

    def test_method_indentation_preservation(self):
        """Test that method body indentation is preserved."""
        source = """public class TestClass {
    @Test
    public void testIndent() {
        int x = 5;
            int y = 10;
                int z = 15;
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "func"); result = codeflash_output # 25.6μs -> 21.1μs (21.5% faster)

    def test_function_call_with_complex_arguments(self):
        """Test function calls with complex argument expressions."""
        source = """public class TestClass {
    @Test
    public void testComplex() {
        calc.add(getValue() + 5, (100 - 50) * 2);
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 29.8μs -> 26.3μs (13.6% faster)

    def test_multiple_test_methods_in_class(self):
        """Test instrumentation of multiple @Test methods in one class."""
        source = """public class TestClass {
    @Test
    public void testOne() {
        calc.add(1, 2);
    }
    
    @Test
    public void testTwo() {
        calc.add(3, 4);
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 48.3μs -> 42.4μs (14.0% faster)

    def test_method_with_annotations_and_modifiers(self):
        """Test method with multiple annotations and modifiers."""
        source = """public class TestClass {
    @Test
    @Ignore
    @Timeout(5000)
    public void testAnnotated() {
        calc.add(1, 2);
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 28.4μs -> 26.2μs (8.21% faster)

    def test_empty_source_with_only_class_declaration(self):
        """Test behavior with minimal source code."""
        source = "public class TestClass { }"
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "func"); result = codeflash_output # 4.19μs -> 4.57μs (8.28% slower)

    def test_function_name_not_in_source(self):
        """Test when target function name never appears in source."""
        source = """public class TestClass {
    @Test
    public void testMethod() {
        int x = 5;
        System.out.println(x);
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "nonexistentFunc"); result = codeflash_output # 25.7μs -> 21.2μs (21.7% faster)

    def test_method_with_try_catch_block(self):
        """Test instrumentation of method containing try-catch block."""
        source = """public class TestClass {
    @Test
    public void testTryCatch() {
        try {
            calc.add(1, 2);
        } catch (Exception e) {
            System.out.println("Error");
        }
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 34.5μs -> 29.1μs (18.5% faster)

    def test_special_characters_in_class_name(self):
        """Test with class names containing numbers."""
        source = """public class Test2Class {
    @Test
    public void test() {
        obj.func();
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "Test2Class", "func"); result = codeflash_output # 28.9μs -> 25.9μs (11.7% faster)

    def test_function_name_with_special_regex_chars(self):
        """Test function names that are special regex characters."""
        source = """public class TestClass {
    @Test
    public void testDot() {
        obj.add(1, 2);
    }
}"""
        # Using a function name that contains a dot (though unlikely in Java)
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 28.3μs -> 25.0μs (13.3% faster)

class TestLargeScale:
    """Large scale test cases for performance and scalability."""

    def test_large_method_body(self):
        """Test instrumentation of method with large body."""
        # Create a method with 100 lines of code
        lines = ["public class TestClass {", "    @Test", "    public void testLarge() {"]
        for i in range(100):
            if i % 5 == 0:
                lines.append(f"        calc.add({i}, {i+1});")
            else:
                lines.append(f"        int var{i} = {i};")
        lines.append("    }")
        lines.append("}")
        source = "\n".join(lines)
        
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 168μs -> 117μs (44.0% faster)
        
        # Count how many add() calls were captured
        add_count = result.count("_cf_result1_")

    def test_many_function_calls_in_line(self):
        """Test method with many function calls on single line."""
        # Create a line with 50 function calls
        call_string = "; ".join([f"calc.add({i}, {i+1})" for i in range(50)])
        source = f"""public class TestClass {{
    @Test
    public void testMany() {{
        {call_string};
    }}
}}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 114μs -> 88.8μs (29.3% faster)
        
        # All 50 calls should be captured
        for i in range(1, 51):
            pass

    def test_many_test_methods(self):
        """Test class with many @Test methods."""
        # Create 50 test methods
        methods = []
        for i in range(50):
            methods.append(f"""    @Test
    public void test{i}() {{
        calc.add({i}, {i+1});
    }}
""")
        source = "public class TestClass {\n" + "\n".join(methods) + "\n}"
        
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 696μs -> 606μs (14.8% faster)
        
        # Each method should be instrumented with its own iteration ID
        for i in range(1, 51):
            pass

    def test_deeply_nested_braces(self):
        """Test method with deeply nested brace structures."""
        # Create nested if statements
        nested = "public class TestClass {\n    @Test\n    public void testNested() {\n"
        for i in range(10):
            nested += "        " * (i + 1) + "if (true) {\n"
        nested += "        " * 11 + "calc.add(1, 2);\n"
        for i in range(10):
            nested += "        " * (10 - i) + "}\n"
        nested += "    }\n}"
        
        codeflash_output = _add_behavior_instrumentation(source=nested, class_name="TestClass", func_name="add"); result = codeflash_output # 72.3μs -> 43.5μs (66.4% faster)

    def test_large_class_with_mixed_methods(self):
        """Test large class with mix of test and regular methods."""
        # Create class with 30 methods total (20 regular, 10 test)
        methods = []
        for i in range(20):
            methods.append(f"""    public void regular{i}() {{
        int x = {i};
    }}
""")
        for i in range(10):
            methods.append(f"""    @Test
    public void test{i}() {{
        calc.add({i}, {i+1});
    }}
""")
        source = "public class TestClass {\n" + "\n".join(methods) + "\n}"
        
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 181μs -> 163μs (11.0% faster)

    def test_method_with_very_long_line(self):
        """Test handling of very long source lines."""
        # Create a very long line with function calls
        long_call = "calc.add(a, b); " * 60
        source = f"""public class TestClass {{
    @Test
    public void testLongLine() {{
        {long_call}
    }}
}}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 126μs -> 97.4μs (29.8% faster)
        
        # All 60 calls should be captured
        for i in range(1, 61):
            pass

    def test_large_number_of_imports(self):
        """Test handling many existing imports."""
        # Create source with many imports
        imports = "\n".join([f"import com.package{i}.Class{i};" for i in range(100)])
        source = f"""{imports}
public class TestClass {{
    @Test
    public void test() {{
        calc.add(1, 2);
    }}
}}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 67.0μs -> 64.8μs (3.45% faster)
        
        # All original imports should be preserved
        for i in range(100):
            pass

    def test_performance_with_many_variable_declarations(self):
        """Test performance with many variable declarations in method."""
        # Create method with 500 variable declarations
        lines = ["public class TestClass {", "    @Test", "    public void testManyVars() {"]
        for i in range(500):
            lines.append(f"        int var{i} = {i};")
        lines.append("        calc.add(1, 2);")
        lines.append("    }")
        lines.append("}")
        source = "\n".join(lines)
        
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 549μs -> 290μs (88.9% faster)
        # All original declarations should be preserved
        for i in range(500):
            pass

    def test_method_with_string_containing_target_function(self):
        """Test that function names in string literals are not captured."""
        source = """public class TestClass {
    @Test
    public void testString() {
        String msg = "Please call add(1, 2) method";
        calc.add(3, 4);
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 36.6μs -> 31.7μs (15.3% faster)

    def test_multiple_files_instrumentation_sequence(self):
        """Test that multiple calls to function maintain correct iteration IDs."""
        # Simulate instrumenting multiple test methods in sequence
        source1 = """public class TestClass1 {
    @Test
    public void test1() {
        calc.add(1, 2);
    }
}"""
        source2 = """public class TestClass2 {
    @Test
    public void test2() {
        calc.add(3, 4);
    }
}"""
        
        codeflash_output = _add_behavior_instrumentation(source1, "TestClass1", "add"); result1 = codeflash_output # 27.7μs -> 24.4μs (13.7% faster)
        codeflash_output = _add_behavior_instrumentation(source2, "TestClass2", "add"); result2 = codeflash_output # 19.0μs -> 17.0μs (12.2% faster)

    def test_return_value_serialization_correctness(self):
        """Test that return values are properly serialized."""
        source = """public class TestClass {
    @Test
    public void testReturn() {
        int result = calc.add(1, 2);
        String text = utils.reverse("hello");
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 31.2μs -> 26.3μs (18.8% faster)
        # Both should use String.valueOf for serialization
        serialize_count = result.count("String.valueOf")

    def test_instrumentation_code_integrity(self):
        """Test that all instrumentation code elements are present."""
        source = """public class TestClass {
    @Test
    public void test() {
        calc.add(1, 2);
    }
}"""
        codeflash_output = _add_behavior_instrumentation(source, "TestClass", "add"); result = codeflash_output # 26.7μs -> 23.4μs (14.3% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

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