⚡️ Speed up function _add_timing_instrumentation by 21% in PR #1390 (feat/kryo-serialization)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1390

If you approve this dependent PR, these changes will be merged into the original PR branch feat/kryo-serialization.

This PR will be automatically closed if the original PR is merged.

---

📄 21% (0.21x) speedup for _add_timing_instrumentation in codeflash/languages/java/instrumentation.py

⏱️ Runtime : 3.17 milliseconds → 2.61 milliseconds (best of 245 runs)

📝 Explanation and details

This optimization achieves a 21% runtime improvement by eliminating expensive character-by-character iteration when counting braces in Java method bodies.

Key Performance Changes:

1. Replaced per-character loop with str.count() calls: The original code iterated through every character in each line to count opening and closing braces, which consumed 35% of total runtime (18% + 17% from line profiler). The optimized version uses Python's built-in str.count("{") and str.count("}") methods, which are implemented in C and run significantly faster.

2. Direct appending eliminates intermediate list: The original code collected method body lines in a body_lines list, then later iterated through it again to append lines with indentation to the result. The optimized version directly appends each line to result as it's processed, removing both the intermediate storage and the second iteration.

3. Cached list length and lstrip() calls: Storing n = len(lines) upfront avoids repeated length calculations in tight loops (1,071 iterations), and storing lstripped = line.lstrip() avoids redundant string operations.

Why This Matters:

Based on the function references, _add_timing_instrumentation is called from test instrumentation workflows that process Java test classes. The test results show consistent 10-30% speedups across various scenarios:

• Simple single methods: 11-16% faster
• Methods with nested braces: 19-32% faster
• Large-scale scenarios (100+ methods): 10-15% faster
• Complex method bodies with deep nesting: 192% faster (most dramatic for brace-heavy code)

The optimization is particularly effective for:

• Test files with deeply nested control flow (loops, switches, conditionals) where brace counting dominates
• Large test suites where the function is called repeatedly on many methods
• Build/CI pipelines where test instrumentation happens frequently

Since this function processes every @Test method in Java source files during instrumentation, reducing its runtime by 21% provides compounding benefits in workflows that instrument multiple test classes or run in tight optimization loops.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import re
import textwrap

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

def test_no_test_annotations_returns_identical_source():
    # Basic: When there are no @Test annotations, source should be unchanged.
    src = textwrap.dedent(
        """\
        public class Example {
            public void helper() {
                // nothing to do
            }
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "Example", "helper"); out = codeflash_output # 3.14μs -> 3.06μs (2.65% faster)

def test_single_test_method_inserts_instrumentation_and_preserves_body_order():
    # Basic: One @Test method with a simple single-line body.
    src = textwrap.dedent(
        """\
        public class MyClass {
            @Test
            public void testSomething() {
                System.out.println("hello");
            }
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "MyClass", "testSomething"); out = codeflash_output # 11.9μs -> 10.3μs (15.1% faster)

    # The original body line should still be present (once) and should appear after the start println
    start_pos = out.find("!$######")
    body_pos = out.find('System.out.println("hello");')

def test_multiple_annotations_and_multiline_signature_preserved_and_instrumented():
    # Edge: Additional annotations should be preserved and multi-line method signature handled.
    src = textwrap.dedent(
        """\
        public class Multi {
            @Test
            @CustomAnnotation(value = "x")
            public void
            testMultiLine()
            throws Exception
            {
                // inner action
            }
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "Multi", "testMultiLine"); out = codeflash_output # 12.3μs -> 11.0μs (12.2% faster)

def test_nested_braces_in_method_body_are_handled_correctly():
    # Edge: Method body contains nested braces, ensure correct brace counting and that nested content appears
    src = textwrap.dedent(
        """\
        public class Nested {
            @Test
            public void testNested() {
                for (int i = 0; i < 2; i++) {
                    if (i % 2 == 0) {
                        System.out.println(i);
                    }
                }
            }
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "Nested", "testNested"); out = codeflash_output # 15.3μs -> 11.6μs (32.2% faster)

    # Confirm that the finally block for iter_id 1 exists and follows the try
    idx_try = out.find("try {")
    idx_finally = out.find("} finally {")

def test_missing_opening_brace_results_in_timing_start_inserted_but_no_finally():
    # Edge: @Test present but no opening '{' for the method in the source.
    src = textwrap.dedent(
        """\
        public class Incomplete {
            @Test
            public void testNoBrace()
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "Incomplete", "testNoBrace"); out = codeflash_output # 7.46μs -> 7.05μs (5.83% faster)

def test_large_scale_many_tests_unique_iter_ids_and_counts():
    # Large Scale: Generate a source with many (@Test) methods to validate scalability and uniqueness of iter ids.
    # Use 100 methods (keeps under 1000).
    num_methods = 100
    methods = []
    for i in range(1, num_methods + 1):
        methods.append(textwrap.dedent(
            f"""\
            @Test
            public void test{i}() {{
                // body {i}
            }}"""
        ))
    class_src = "public class Bulk {\n" + "\n".join(methods) + "\n}"
    codeflash_output = _add_timing_instrumentation(class_src, "Bulk", "testUnknown"); out = codeflash_output # 436μs -> 396μs (10.1% faster)

    # Verify that we have exactly num_methods occurrences of loop variable declarations with increasing indices
    # E.g., int _cf_loop1, int _cf_loop2, ..., int _cf_loop100
    for idx in range(1, num_methods + 1):
        pattern = f"int _cf_loop{idx} = Integer.parseInt(System.getenv(\"CODEFLASH_LOOP_INDEX\"));"

    # There should be exactly num_methods start marker println occurrences (one per method) and exactly num_methods end markers.
    # Start marker uses "######$!" at end; end marker uses "######!".
    count_start_markers = out.count("######$!")
    count_end_markers = out.count("######!")

    # Ensure ordering: for the first method, the loop declaration appears before the first start println
    first_loop_idx = out.find("int _cf_loop1")
    first_start_idx = out.find("!$######")

def test_instrumentation_preserves_additional_annotation_blocks_and_spacing():
    # Basic/Edge: Ensure that blocks of annotations are preserved and spacing is not lost between annotation and signature.
    src = textwrap.dedent(
        """\
        public class Spacing {
            @Test
            @Repeat(3)
            @Another
            public void spacedTest() {
                doSomething();
            }
        }"""
    )
    codeflash_output = _add_timing_instrumentation(src, "Spacing", "spacedTest"); out = codeflash_output # 12.5μs -> 11.1μs (13.2% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import os

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

def test_basic_single_test_method_instrumentation():
    """Test that a single @Test method is properly instrumented with timing code."""
    source = """public class TestExample {
    @Test
    public void testSimple() {
        int x = 5;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testSimple"); result = codeflash_output # 12.2μs -> 11.0μs (11.1% faster)

def test_basic_method_signature_preserved():
    """Test that the method signature is preserved after instrumentation."""
    source = """public class TestExample {
    @Test
    public void testSimple() {
        int x = 5;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testSimple"); result = codeflash_output # 11.9μs -> 10.5μs (13.9% faster)

def test_basic_class_and_func_names_in_output():
    """Test that class and function names are properly embedded in timing markers."""
    source = """public class MyTest {
    @Test
    public void myFunc() {
        x = 1;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "MyTest", "myFunc"); result = codeflash_output # 11.5μs -> 10.4μs (11.1% faster)

def test_basic_multiple_test_methods():
    """Test that multiple @Test methods are each instrumented independently."""
    source = """public class TestExample {
    @Test
    public void testFirst() {
        int x = 1;
    }
    
    @Test
    public void testSecond() {
        int y = 2;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testMethod"); result = codeflash_output # 19.1μs -> 16.8μs (13.5% faster)

def test_basic_single_line_method_body():
    """Test instrumentation of a method with a single-line body."""
    source = """public class TestExample {
    @Test
    public void test() {
        assert true;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "test"); result = codeflash_output # 11.2μs -> 9.95μs (12.7% faster)

def test_basic_multiline_method_body():
    """Test instrumentation of a method with multiple lines in the body."""
    source = """public class TestExample {
    @Test
    public void test() {
        int a = 1;
        int b = 2;
        int c = a + b;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "test"); result = codeflash_output # 13.1μs -> 10.7μs (22.4% faster)

def test_edge_empty_method_body():
    """Test instrumentation of a method with an empty body."""
    source = """public class TestExample {
    @Test
    public void testEmpty() {
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testEmpty"); result = codeflash_output # 10.2μs -> 8.97μs (13.3% faster)

def test_edge_method_with_multiple_annotations():
    """Test instrumentation when method has multiple annotations."""
    source = """public class TestExample {
    @Test
    @Ignore
    @DisplayName("My Test")
    public void testMultiAnnot() {
        x = 1;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testMultiAnnot"); result = codeflash_output # 12.0μs -> 10.6μs (13.4% faster)

def test_edge_test_annotation_with_parameters():
    """Test instrumentation when @Test has parameters."""
    source = """public class TestExample {
    @Test(timeout = 1000)
    public void testWithParams() {
        int x = 5;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testWithParams"); result = codeflash_output # 11.3μs -> 9.80μs (15.4% faster)

def test_edge_nested_braces_in_method_body():
    """Test instrumentation when method body contains nested braces."""
    source = """public class TestExample {
    @Test
    public void testNested() {
        if (true) {
            int x = 1;
        }
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testNested"); result = codeflash_output # 13.0μs -> 10.9μs (19.1% faster)

def test_edge_special_characters_in_names():
    """Test instrumentation when class/function names contain underscores."""
    source = """public class Test_Example_Class {
    @Test
    public void test_Function_Name() {
        int x = 1;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "Test_Example_Class", "test_Function_Name"); result = codeflash_output # 11.3μs -> 10.0μs (12.6% faster)

def test_edge_method_with_parameters():
    """Test instrumentation of a test method that has parameters."""
    source = """public class TestExample {
    @Test
    public void testWithParams(int param1, String param2) {
        assert param1 > 0;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testWithParams"); result = codeflash_output # 11.4μs -> 9.84μs (16.1% faster)

def test_edge_method_with_throws_clause():
    """Test instrumentation of a test method with throws clause."""
    source = """public class TestExample {
    @Test
    public void testWithThrows() throws Exception {
        int x = 1;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testWithThrows"); result = codeflash_output # 11.3μs -> 9.80μs (15.3% faster)

def test_edge_method_with_braces_in_string():
    """Test instrumentation when method body contains braces inside strings."""
    source = """public class TestExample {
    @Test
    public void testStringBraces() {
        String s = "test { braces }";
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testStringBraces"); result = codeflash_output # 11.9μs -> 9.86μs (20.9% faster)

def test_edge_indented_class():
    """Test instrumentation when the test method is indented within a class."""
    source = """    public class TestExample {
        @Test
        public void testIndented() {
            int x = 1;
        }
    }"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testIndented"); result = codeflash_output # 11.9μs -> 10.3μs (15.1% faster)

def test_edge_no_test_annotation():
    """Test that methods without @Test annotation are not instrumented."""
    source = """public class TestExample {
    public void normalMethod() {
        int x = 1;
    }
    
    @Test
    public void testMethod() {
        int y = 2;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "test"); result = codeflash_output # 13.0μs -> 11.2μs (16.1% faster)

def test_edge_non_ascii_characters_in_body():
    """Test that non-ASCII characters in method body are preserved."""
    source = """public class TestExample {
    @Test
    public void testUnicode() {
        String s = "café";
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testUnicode"); result = codeflash_output # 12.7μs -> 10.8μs (17.9% faster)

def test_edge_very_long_method_name():
    """Test instrumentation with very long class and method names."""
    long_class = "VeryLongTestClassNameThatIsActuallyQuiteLong"
    long_func = "testWithAVeryLongFunctionNameThatExplainsWhatItDoes"
    
    source = f"""public class {long_class} {{
    @Test
    public void {long_func}() {{
        int x = 1;
    }}
}}"""
    
    codeflash_output = _add_timing_instrumentation(source, long_class, long_func); result = codeflash_output # 11.5μs -> 10.2μs (13.3% faster)

def test_edge_method_with_only_comment_in_body():
    """Test instrumentation when method body contains only comments."""
    source = """public class TestExample {
    @Test
    public void testOnlyComment() {
        // This is a comment
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testOnlyComment"); result = codeflash_output # 11.5μs -> 10.1μs (14.0% faster)

def test_edge_class_name_with_package():
    """Test instrumentation when class name contains package notation."""
    source = """public class TestExample {
    @Test
    public void testMethod() {
        int x = 1;
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "com.example.TestExample", "testMethod"); result = codeflash_output # 11.2μs -> 9.92μs (13.2% faster)

def test_large_scale_many_test_methods():
    """Test instrumentation of a class with many @Test methods."""
    # Create a source with 100 test methods
    lines = ["public class TestExample {"]
    for i in range(100):
        lines.append(f"    @Test")
        lines.append(f"    public void test{i}() {{")
        lines.append(f"        int x = {i};")
        lines.append(f"    }}")
        lines.append("")
    lines.append("}")
    
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "test"); result = codeflash_output # 494μs -> 430μs (14.7% faster)

def test_large_scale_large_method_body():
    """Test instrumentation of a method with a large body."""
    # Create method with 200 lines in body
    lines = ["public class TestExample {", "    @Test", "    public void testLargeBody() {"]
    for i in range(200):
        lines.append(f"        int var{i} = {i};")
    lines.append("    }")
    lines.append("}")
    
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testLargeBody"); result = codeflash_output # 173μs -> 76.6μs (127% faster)

def test_large_scale_deeply_nested_braces():
    """Test instrumentation with deeply nested brace structures."""
    # Create method with deeply nested braces
    lines = ["public class TestExample {", "    @Test", "    public void testDeepNesting() {"]
    indent_level = 2
    for i in range(20):
        lines.append("    " * indent_level + "if (true) {")
        indent_level += 1
    lines.append("    " * indent_level + "int x = 1;")
    for i in range(20):
        indent_level -= 1
        lines.append("    " * indent_level + "}")
    lines.append("    }")
    lines.append("}")
    
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testDeepNesting"); result = codeflash_output # 76.6μs -> 26.2μs (192% faster)

def test_large_scale_mixed_methods():
    """Test instrumentation with a mix of test and non-test methods."""
    # Create class with 50 @Test methods and 50 non-test methods
    lines = ["public class TestExample {"]
    for i in range(50):
        lines.append(f"    @Test")
        lines.append(f"    public void test{i}() {{")
        lines.append(f"        int x = {i};")
        lines.append(f"    }}")
        lines.append(f"    public void helper{i}() {{")
        lines.append(f"        int y = {i};")
        lines.append(f"    }}")
        lines.append("")
    lines.append("}")
    
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "test"); result = codeflash_output # 284μs -> 249μs (14.0% faster)
    # Non-test helper methods should remain unchanged
    for i in range(50):
        pass

def test_large_scale_very_long_source_file():
    """Test instrumentation on a very large source file."""
    # Create source with 500 total lines (mix of test and other code)
    lines = [
        "package com.example;",
        "import java.util.*;",
        "public class LargeTest {"
    ]
    
    # Add various class members and 50 test methods
    for i in range(50):
        lines.append(f"    private static final int CONSTANT_{i} = {i};")
    
    lines.append("")
    
    for i in range(50):
        lines.append(f"    @Test")
        lines.append(f"    public void testMethod{i}() {{")
        lines.append(f"        assertEquals({i}, CONSTANT_{i});")
        lines.append(f"    }}")
        lines.append("")
    
    lines.append("}")
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "LargeTest", "test"); result = codeflash_output # 285μs -> 228μs (24.5% faster)

def test_large_scale_complex_method_bodies():
    """Test instrumentation of methods with complex control flow."""
    source = """public class TestExample {
    @Test
    public void testComplex() {
        for (int i = 0; i < 100; i++) {
            if (i % 2 == 0) {
                switch (i) {
                    case 0:
                        int x = 0;
                        break;
                    case 2:
                        int y = 2;
                        break;
                    default:
                        int z = i;
                }
            }
        }
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testComplex"); result = codeflash_output # 24.7μs -> 15.9μs (55.6% faster)

def test_large_scale_variable_naming_no_conflicts():
    """Test that generated variable names don't conflict with user code."""
    # Create method that uses variables with similar names to generated ones
    source = """public class TestExample {
    @Test
    public void testNoConflict() {
        int _cf_loop = 5;
        int _cf_i = 10;
        int _cf_start = 20;
        for (int i = 0; i < 10; i++) {
            int x = _cf_loop + _cf_i + _cf_start;
        }
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testNoConflict"); result = codeflash_output # 17.2μs -> 12.5μs (37.7% faster)

def test_large_scale_stress_output_length():
    """Test that output size remains reasonable with many instrumented methods."""
    # Create 200 test methods to stress output size
    lines = ["public class StressTest {"]
    for i in range(200):
        lines.append(f"    @Test")
        lines.append(f"    public void test{i}() {{")
        lines.append(f"        assertTrue(true);")
        lines.append(f"    }}")
    lines.append("}")
    
    source = "\n".join(lines)
    codeflash_output = _add_timing_instrumentation(source, "StressTest", "test"); result = codeflash_output # 1.03ms -> 873μs (17.9% faster)

def test_large_scale_whitespace_preservation():
    """Test that significant whitespace variations are handled correctly."""
    # Create source with inconsistent indentation and spacing
    source = """public class TestExample {
	@Test
	public void testWhitespace() {
  int x=1;
    int     y  =  2;
			int z = 3;
	}
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "TestExample", "testWhitespace"); result = codeflash_output # 13.2μs -> 11.6μs (13.8% faster)

def test_integration_complete_class_with_multiple_features():
    """Test a complete realistic test class with various features."""
    source = """public class CompleteTest {
    private static final String TEST_DATA = "data";
    
    @Before
    public void setUp() {
        // Setup
    }
    
    @Test
    public void testFirst() {
        String result = processData(TEST_DATA);
        assertEquals("expected", result);
    }
    
    @Test
    public void testSecond() {
        try {
            doSomething();
        } catch (Exception e) {
            fail("Should not throw");
        }
    }
    
    @After
    public void tearDown() {
        // Cleanup
    }
    
    private String processData(String input) {
        return input.toUpperCase();
    }
}"""
    
    codeflash_output = _add_timing_instrumentation(source, "CompleteTest", "test"); result = codeflash_output # 29.8μs -> 23.8μs (25.1% 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-pr1390-2026-02-06T19.01.20 and push.