⚡️ Speed up function _build_code_strings_for_language by 28% in PR #1473 (refactor/tree-sitter-instrumentation)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1473

If you approve this dependent PR, these changes will be merged into the original PR branch refactor/tree-sitter-instrumentation.

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

---

📄 28% (0.28x) speedup for _build_code_strings_for_language in codeflash/context/code_context_extractor.py

⏱️ Runtime : 51.7 milliseconds → 40.4 milliseconds (best of 74 runs)

📝 Explanation and details

This optimization achieves a 27% runtime improvement (51.7ms → 40.4ms) by eliminating redundant work in path resolution and regex compilation—two operations that were consuming over 70% of the original runtime.

Key optimizations:

1. Precompiled regex pattern (_RE_JAVADOC): The original code imported re and compiled the Javadoc pattern on every call to _strip_javadoc_comments. By moving the regex compilation to module scope, we eliminate repeated compilation overhead. Line profiler shows this function dropping from 444ns to 22ns per call—a ~20× improvement.

2. Cached path resolution: The original code called .resolve() repeatedly inside the helper loop (line taking 70.9% of total time). The optimization resolves project_root_path once upfront and reuses project_root_resolved throughout. For helpers, it now resolves each file_path once and reuses helper_file_resolved, avoiding 1,208+ redundant resolve operations per invocation.

3. Hoisted target file resolution: Similar to the project root, the target file path is resolved once and stored in target_file_resolved, eliminating duplicate work when computing target_relative_path.

4. List comprehensions in joins: Changed generator expressions to list comprehensions in "\n\n".join() calls. While this has minimal performance impact for small collections, it can improve performance for larger helper lists by avoiding iterator overhead.

Why this matters:

• Path resolution involves filesystem syscalls and is expensive—the profiler shows the original helper loop spending 125ms on path operations alone
• The optimization is particularly effective for the large-scale test case (1000 helpers), which sees a 39.4% speedup (32.3ms → 23.1ms)
• Tests with Javadoc stripping also benefit significantly (15.2% faster), as regex compilation overhead is eliminated
• The cross-file helpers test improves by 9.83% due to reduced path resolution overhead

Trade-offs:

Minor increases in some small test cases (1-4% slower) are within measurement noise and acceptable given the substantial gains in realistic workloads with multiple helpers and path operations.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from pathlib import Path  # used to construct file paths
from types import \
    SimpleNamespace  # simple attribute container for test objects

# imports
import pytest  # used for our unit tests
from codeflash.context.code_context_extractor import \
    _build_code_strings_for_language
from codeflash.models.models import \
    FunctionSource  # used to compare helper function sources

def test_combines_target_with_same_file_helper_and_preserves_language_and_relative_path():
    # Create a project root and a target file under it so relative_to succeeds
    project_root = Path("/project")
    target_path = Path("/project/src/target.java")

    # Build a simple same-file helper (same file path as target)
    helper = SimpleNamespace(
        file_path=target_path,
        qualified_name="pkg.Class.helper",
        name="helper",
        source_code="public void helper() {}",
    )

    # Build a minimal code_context with no imports and one helper
    code_context = SimpleNamespace(
        imports=[],  # no imports
        helper_functions=[helper],  # one helper in same file
        target_code="public void target() {}",
        read_only_context=None,  # no read-only context
    )

    # The function_to_optimize only needs file_path and language for this function
    function_to_optimize = SimpleNamespace(file_path=target_path, language="java")  # use non-python to avoid validation

    # Call the function under test
    code_strings, helper_function_sources, read_only_context = _build_code_strings_for_language(
        code_context, function_to_optimize, project_root
    ) # 68.1μs -> 68.7μs (0.812% slower)
    cs = code_strings[0]

    # The helper_function_sources should include one FunctionSource corresponding to the helper
    expected_fs = FunctionSource(
        file_path=helper.file_path,
        qualified_name=helper.qualified_name,
        fully_qualified_name=helper.qualified_name,
        only_function_name=helper.name,
        source_code=helper.source_code,
        jedi_definition=None,
    )

def test_prepends_imports_to_target_code_when_present():
    # Create project root and target
    project_root = Path("/proj")
    target_path = Path("/proj/main/File.java")

    # No helpers
    code_context = SimpleNamespace(
        imports=["import a.b.C;", "import x.y.Z;"],
        helper_functions=[],
        target_code="class X {}",
        read_only_context="readonly",
    )

    function_to_optimize = SimpleNamespace(file_path=target_path, language="java")

    # Call function
    code_strings, helper_function_sources, read_only_context = _build_code_strings_for_language(
        code_context, function_to_optimize, project_root
    ) # 54.2μs -> 55.0μs (1.45% slower)
    cs = code_strings[0]

def test_includes_cross_file_helpers_as_separate_code_strings_and_relative_paths():
    # Project root is parent of both target and helper files
    project_root = Path("/repo")
    target_path = Path("/repo/src/Target.java")
    helper_path = Path("/repo/lib/Helper.java")

    # Cross-file helper (different file)
    helper = SimpleNamespace(
        file_path=helper_path,
        qualified_name="pkg.Helper.doIt",
        name="doIt",
        source_code="public void doIt() {}",
    )

    # Code context with the cross-file helper
    code_context = SimpleNamespace(
        imports=[],
        helper_functions=[helper],
        target_code="public void target() {}",
        read_only_context=None,
    )

    function_to_optimize = SimpleNamespace(file_path=target_path, language="java")

    code_strings, helper_function_sources, _ = _build_code_strings_for_language(
        code_context, function_to_optimize, project_root
    ) # 100μs -> 91.2μs (9.83% faster)

    # Find which CodeString corresponds to helper by file_path (relative)
    helper_cs = next(cs for cs in code_strings if cs.file_path == Path("lib/Helper.java"))

    # helper_function_sources should include the helper's FunctionSource
    expected_fs = FunctionSource(
        file_path=helper.file_path,
        qualified_name=helper.qualified_name,
        fully_qualified_name=helper.qualified_name,
        only_function_name=helper.name,
        source_code=helper.source_code,
        jedi_definition=None,
    )

def test_excludes_cross_file_helpers_when_flag_false_and_excludes_same_file_helpers_when_flag_false():
    # Project root and paths
    project_root = Path("/repo2")
    target_path = Path("/repo2/src/Target2.java")
    helper_same = SimpleNamespace(
        file_path=target_path,
        qualified_name="pkg.Target2.helper",
        name="helper",
        source_code="/* helper same */",
    )
    helper_cross = SimpleNamespace(
        file_path=Path("/repo2/lib/Cross.java"),
        qualified_name="pkg.Cross.fn",
        name="fn",
        source_code="// cross helper",
    )

    # Both helpers present
    code_context = SimpleNamespace(
        imports=[],
        helper_functions=[helper_same, helper_cross],
        target_code="target content",
        read_only_context=None,
    )

    function_to_optimize = SimpleNamespace(file_path=target_path, language="java")

    # Exclude same-file helpers and cross-file helpers (both False)
    code_strings, helper_function_sources, _ = _build_code_strings_for_language(
        code_context,
        function_to_optimize,
        project_root,
        include_cross_file_helpers=False,
        include_same_file_helpers=False,
    ) # 55.6μs -> 57.9μs (4.09% slower)
    cs = code_strings[0]

def test_strip_javadoc_removes_javadoc_from_target_helpers_and_read_only_context():
    # Project root and paths
    project_root = Path("/pj")
    target_path = Path("/pj/src/Thing.java")
    helper_path = Path("/pj/lib/HelperThing.java")

    # Insert Javadoc comments (/** ... */) in target, helper, and read-only context
    target_code = "public void target() {}/** javadoc target */"
    helper = SimpleNamespace(
        file_path=helper_path,
        qualified_name="pkg.HelperThing.h",
        name="h",
        source_code="/** helper javadoc */\npublic void h() {}",
    )
    read_only = "Some read only /** javadoc ro */ content"

    code_context = SimpleNamespace(
        imports=["import z;"],
        helper_functions=[helper],
        target_code=target_code,
        read_only_context=read_only,
    )

    function_to_optimize = SimpleNamespace(file_path=target_path, language="java")

    # Request stripping of Javadoc comments
    code_strings, helper_function_sources, ro_context = _build_code_strings_for_language(
        code_context,
        function_to_optimize,
        project_root,
        strip_javadoc=True,
        include_cross_file_helpers=True,
    ) # 110μs -> 95.6μs (15.2% faster)
    # Ensure helper javadoc removed from its CodeString
    helper_cs = next(cs for cs in code_strings if cs.file_path == Path("lib/HelperThing.java"))

    # Helper function sources should still be present (strip_javadoc does not affect FunctionSource.source_code)
    expected_fs = FunctionSource(
        file_path=helper.file_path,
        qualified_name=helper.qualified_name,
        fully_qualified_name=helper.qualified_name,
        only_function_name=helper.name,
        source_code=helper.source_code,
        jedi_definition=None,
    )

def test_when_file_not_under_project_root_target_relative_path_falls_back_to_original_path():
    # Use a project root that does NOT contain the target file, so relative_to will raise ValueError
    project_root = Path("/not_the_root")
    target_path = Path("/some/other/place/Target.java")

    code_context = SimpleNamespace(
        imports=[],
        helper_functions=[],
        target_code="x",
        read_only_context=None,
    )

    function_to_optimize = SimpleNamespace(file_path=target_path, language="java")

    code_strings, _, _ = _build_code_strings_for_language(code_context, function_to_optimize, project_root) # 59.6μs -> 60.8μs (1.95% slower)

def test_large_number_of_helpers_scalability_and_grouping():
    # This test creates 1000 helpers spanning 1000 unique files (plus a target)
    project_root = Path("/large_project")
    target_path = Path("/large_project/app/TargetLarge.java")

    num_helpers = 1000  # as required by the specification for large-scale tests
    helpers = []
    for i in range(num_helpers):
        # Each helper in its own file under project_root to force a separate CodeString per helper
        helper_file = Path(f"/large_project/helpers/helper_{i}.java")
        h = SimpleNamespace(
            file_path=helper_file,
            qualified_name=f"helpers.Helper{i}.fn",
            name=f"fn_{i}",
            source_code=f"// helper {i}\nvoid fn_{i}() {{}}",
        )
        helpers.append(h)

    # Compose the code context with many helpers
    code_context = SimpleNamespace(
        imports=[],
        helper_functions=helpers,
        target_code="public class TargetLarge {}",
        read_only_context="readonly_context",
    )

    function_to_optimize = SimpleNamespace(file_path=target_path, language="java")

    # Invoke the function; this should complete in reasonable time and return many CodeStrings
    code_strings, helper_function_sources, ro = _build_code_strings_for_language(
        code_context, function_to_optimize, project_root, include_cross_file_helpers=True
    ) # 32.3ms -> 23.1ms (39.4% faster)

    # Each helper file must be represented as a CodeString with the correct relative path and content
    helper_paths_expected = {Path(f"helpers/helper_{i}.java") for i in range(num_helpers)}
    helper_paths_found = {cs.file_path for cs in code_strings[1:]}  # skip the first which is the target
    # Spot-check a few helper FunctionSource entries are present
    sample_indices = [0, num_helpers // 2, num_helpers - 1]
    for idx in sample_indices:
        expected_fs = FunctionSource(
            file_path=helpers[idx].file_path,
            qualified_name=helpers[idx].qualified_name,
            fully_qualified_name=helpers[idx].qualified_name,
            only_function_name=helpers[idx].name,
            source_code=helpers[idx].source_code,
            jedi_definition=None,
        )
# 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-pr1473-2026-02-13T00.58.09 and push.

[codeflash-ai]

This PR has been automatically closed because the original PR #1473 by HeshamHM28 was closed.