⚡️ Speed up function get_optimized_code_for_module by 33% in PR #1772 (fix-markdown-code-path-lookup)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1772

If you approve this dependent PR, these changes will be merged into the original PR branch fix-markdown-code-path-lookup.

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

---

📄 33% (0.33x) speedup for get_optimized_code_for_module in codeflash/languages/code_replacer.py

⏱️ Runtime : 9.16 milliseconds → 6.86 milliseconds (best of 23 runs)

📝 Explanation and details

The hot path in Fallback 2 previously built a list comprehension with Path(path).name == target_name for every (path, code) pair, constructing many short-lived Path objects and allocating an intermediate list that profiler data showed consumed ~30% of total runtime. The optimized code replaces this with a single-pass loop using basename(path) (a C-level string operation) and tracks match_count inline, exiting early on duplicate basenames and avoiding all list/Path allocations. Logger guards (isEnabledFor) defer expensive f-string and list(...) construction until logging is actually enabled, eliminating overhead in production where debug/warning levels are often disabled. The combination yields a 33% speedup, with the largest gains in the test_large_scale_unique_basename_among_many case (683% faster) where basename scanning dominated.

✅ Correctness verification report:

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

⚙️ Click to see Existing Unit Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_languages/test_get_optimized_code_for_module.py::test_basename_fallback_ambiguous_returns_empty	629μs	627μs	0.297%✅
test_languages/test_get_optimized_code_for_module.py::test_basename_fallback_different_directory	27.5μs	19.0μs	44.9%✅
test_languages/test_get_optimized_code_for_module.py::test_basename_fallback_skips_non_matching_context_files	29.5μs	20.0μs	47.4%✅
test_languages/test_get_optimized_code_for_module.py::test_context_files_only_returns_empty	675μs	660μs	2.22%✅
test_languages/test_get_optimized_code_for_module.py::test_empty_markdown_returns_empty	574μs	576μs	-0.349%⚠️
test_languages/test_get_optimized_code_for_module.py::test_exact_match_preferred_over_basename	15.9μs	15.7μs	1.15%✅
test_languages/test_get_optimized_code_for_module.py::test_exact_path_match_picks_correct_file	16.5μs	15.8μs	4.32%✅
test_languages/test_get_optimized_code_for_module.py::test_exact_path_match_single_file	17.2μs	16.4μs	5.03%✅
test_languages/test_get_optimized_code_for_module.py::test_no_match_returns_empty	606μs	598μs	1.25%✅
test_languages/test_get_optimized_code_for_module.py::test_none_path_fallback_ignored_when_named_blocks_exist	766μs	727μs	5.40%✅
test_languages/test_get_optimized_code_for_module.py::test_none_path_fallback_single_block	21.8μs	14.6μs	49.5%✅

🌀 Click to see Generated Regression Tests

from pathlib import Path

# imports
import pytest  # used for our unit tests
from codeflash.languages.code_replacer import get_optimized_code_for_module
from codeflash.models.models import CodeStringsMarkdown

def test_exact_path_match_basic():
    # Create a real CodeStringsMarkdown instance (no special constructor args)
    optimized = CodeStringsMarkdown()
    # Pre-populate the cached file->code mapping so file_to_path() returns it
    optimized._cache["file_to_path"] = {"src/module.py": "print('hello')"}
    # Query with the exact same path string -> should return the corresponding code
    codeflash_output = get_optimized_code_for_module(Path("src/module.py"), optimized); result = codeflash_output # 8.38μs -> 7.96μs (5.16% faster)

def test_none_single_block_fallback():
    # Single code block with key "None" should be used for any requested module
    optimized = CodeStringsMarkdown()
    optimized._cache["file_to_path"] = {"None": "single-block-code"}
    codeflash_output = get_optimized_code_for_module(Path("any/thing.py"), optimized); result = codeflash_output # 12.7μs -> 8.04μs (58.0% faster)

def test_basename_match_single():
    # If there is a single entry whose basename matches the requested module name,
    # it should be returned even if the directory prefix differs.
    optimized = CodeStringsMarkdown()
    optimized._cache["file_to_path"] = {
        "other/path/target.py": "target-code",
        "another/file.py": "foo",
    }
    # Request a Path with same basename 'target.py' but different directory
    codeflash_output = get_optimized_code_for_module(Path("some/dir/target.py"), optimized); result = codeflash_output # 21.6μs -> 11.1μs (94.6% faster)

def test_no_match_returns_empty():
    # Empty mapping should lead to empty-string fallback
    optimized = CodeStringsMarkdown()
    optimized._cache["file_to_path"] = {}
    codeflash_output = get_optimized_code_for_module(Path("nothing.py"), optimized); result = codeflash_output # 575μs -> 576μs (0.255% slower)

def test_multiple_basename_matches_returns_empty():
    # If multiple code blocks share the same basename, basename fallback must not choose one
    optimized = CodeStringsMarkdown()
    optimized._cache["file_to_path"] = {
        "dir1/dup.py": "code1",
        "dir2/dup.py": "code2",
    }
    codeflash_output = get_optimized_code_for_module(Path("dup.py"), optimized); result = codeflash_output # 610μs -> 598μs (1.98% faster)

def test_none_key_not_used_when_multiple_keys():
    # "None" key should only be used when it's the only entry in the mapping.
    optimized = CodeStringsMarkdown()
    optimized._cache["file_to_path"] = {
        "None": "should-not-be-used",
        "some/other.py": "other-code",
    }
    # Request a path that does not match any basename or exact key
    codeflash_output = get_optimized_code_for_module(Path("unmatched.py"), optimized); result = codeflash_output # 602μs -> 600μs (0.245% faster)

def test_exact_path_with_none_value_treated_as_missing_but_basename_match_returns_none():
    # If the exact path maps to None, module_optimized_code will be None and the
    # function should fall through to basename matching. If basename matches uniquely
    # and its associated value is None, the function will return None.
    optimized = CodeStringsMarkdown()
    optimized._cache["file_to_path"] = {
        "a/b.py": None,  # exact match exists but value is None
    }
    # Because basename matches uniquely (there's a single entry with same basename),
    # the function returns that entry's value, which is None.
    codeflash_output = get_optimized_code_for_module(Path("a/b.py"), optimized); result = codeflash_output # 18.7μs -> 10.9μs (71.5% faster)

def test_large_scale_exact_match_performance_and_correctness():
    # Build a large mapping (1000 entries) and ensure an exact match is returned
    optimized = CodeStringsMarkdown()
    mapping = {}
    for i in range(1000):
        mapping[f"pkg/sub{i}/file{i}.py"] = f"code{i}"
    # Insert a special exact key with a unique value
    mapping["special/path/target_module.py"] = "SPECIAL_CODE"
    optimized._cache["file_to_path"] = mapping
    codeflash_output = get_optimized_code_for_module(Path("special/path/target_module.py"), optimized); result = codeflash_output # 8.19μs -> 7.83μs (4.61% faster)

def test_large_scale_unique_basename_among_many():
    # Build many entries with unique basenames and one entry with the basename we will request
    optimized = CodeStringsMarkdown()
    mapping = {}
    for i in range(999):
        mapping[f"otherdir/sub{i}/unique_file_{i}.py"] = f"code_{i}"
    # Add one entry whose basename is 'unique_target.py' inside some directory
    mapping["deep/path/unique_target.py"] = "UNIQUE_TARGET_CODE"
    optimized._cache["file_to_path"] = mapping
    # Request with a different directory prefix but same basename — should match uniquely
    codeflash_output = get_optimized_code_for_module(Path("some/other/unique_target.py"), optimized); result = codeflash_output # 2.47ms -> 315μs (683% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from pathlib import Path

# imports
import pytest
from codeflash.languages.code_replacer import get_optimized_code_for_module
from codeflash.models.models import CodeString, CodeStringsMarkdown

def test_empty_code_strings_list():
    """Test with empty code_strings list."""
    markdown = CodeStringsMarkdown(code_strings=[])
    codeflash_output = get_optimized_code_for_module(Path("any/path.py"), markdown); result = codeflash_output # 603μs -> 598μs (0.822% faster)

def test_exact_match_with_multiple_fallback_candidates():
    """Test exact match is preferred even when basename matches exist."""
    code_strings = [
        CodeString(file_path=Path("exact/path.py"), code="exact"),
        CodeString(file_path=Path("other/path.py"), code="other"),
    ]
    markdown = CodeStringsMarkdown(code_strings=code_strings)
    codeflash_output = get_optimized_code_for_module(Path("exact/path.py"), markdown); result = codeflash_output # 18.1μs -> 17.6μs (2.96% faster)

def test_basename_match_with_multiple_files_no_match():
    """Test basename fallback is not used when multiple files have same basename."""
    code_strings = [
        CodeString(file_path=Path("dir1/module.py"), code="code1"),
        CodeString(file_path=Path("dir2/module.py"), code="code2"),
    ]
    markdown = CodeStringsMarkdown(code_strings=code_strings)
    codeflash_output = get_optimized_code_for_module(Path("dir3/module.py"), markdown); result = codeflash_output # 650μs -> 645μs (0.735% faster)

def test_empty_code_string_value():
    """Test code string with empty code value."""
    code_string = CodeString(file_path=Path("empty.py"), code="")
    markdown = CodeStringsMarkdown(code_strings=[code_string])
    codeflash_output = get_optimized_code_for_module(Path("empty.py"), markdown); result = codeflash_output # 16.5μs -> 16.2μs (1.55% faster)

def test_whitespace_only_code():
    """Test code string with whitespace only."""
    code_string = CodeString(file_path=Path("whitespace.py"), code="   \n\t  ")
    markdown = CodeStringsMarkdown(code_strings=[code_string])
    codeflash_output = get_optimized_code_for_module(Path("whitespace.py"), markdown); result = codeflash_output # 14.1μs -> 14.2μs (1.19% slower)

def test_path_with_dots():
    """Test path containing multiple dots."""
    code_string = CodeString(file_path=Path("src/module.test.py"), code="dots in name")
    markdown = CodeStringsMarkdown(code_strings=[code_string])
    codeflash_output = get_optimized_code_for_module(Path("src/module.test.py"), markdown); result = codeflash_output # 16.6μs -> 16.7μs (0.420% slower)

def test_very_long_code_string():
    """Test very long code string."""
    long_code = "x = 1\n" * 1000
    code_string = CodeString(file_path=Path("long.py"), code=long_code)
    markdown = CodeStringsMarkdown(code_strings=[code_string])
    codeflash_output = get_optimized_code_for_module(Path("long.py"), markdown); result = codeflash_output # 17.1μs -> 16.5μs (3.58% faster)

def test_large_code_blocks_performance():
    """Test performance with large code blocks."""
    large_code = "def function():\n    pass\n" * 100
    code_strings = [
        CodeString(file_path=Path(f"src/module_{i}.py"), code=large_code)
        for i in range(100)
    ]
    markdown = CodeStringsMarkdown(code_strings=code_strings)
    codeflash_output = get_optimized_code_for_module(Path("src/module_50.py"), markdown); result = codeflash_output # 113μs -> 108μs (4.81% faster)

To edit these changes git checkout codeflash/optimize-pr1772-2026-03-05T13.23.24 and push.

[claude]

Claude finished @codeflash-ai[bot]'s task in 3m 40s —— View job

---

PR Review Summary

• Triage PR scope
• Lint and typecheck changed files
• Resolve stale review threads
• Code review
• Duplicate detection
• Test coverage

---

Prek Checks

Two lint issues found and fixed (committed as style: fix ruff TC003 and suppress PTH119 in code_replacer):

1. TC003 — pathlib.Path was imported at module level but only used in type annotations. Moved to TYPE_CHECKING block (safe because from __future__ import annotations is present). This was a pre-existing issue surfaced because the file was modified.

2. PTH119 — ruff wanted os.path.basename() replaced with Path.name. Suppressed with # noqa: PTH119 since using Path(path).name would allocate Path objects in the hot loop, defeating the purpose of the optimization.

Mypy: no issues.

---

Code Review

The optimization is correct and well-implemented. Key observations:

Single-pass loop with early exit (code_replacer.py:45–54): When a second basename match is found, the loop breaks immediately. match_count will be 2 after the break, but correctness is preserved — if match_count != 1 we don't return the candidate anyway.

# type: ignore[return-value] at line 60: Necessary and correct — mypy cannot prove match_code is not None when match_count == 1, but the logic guarantees it.

Minor nit: The isEnabledFor(logging.WARNING) guard at code_replacer.py:62 is technically correct but unusual. logger.warning is almost always enabled; the guard exists only to prevent list(file_to_code_context.keys()) construction in the rare case where warnings are suppressed. This doesn't affect correctness and the performance gain from the guard in the warning path is negligible compared to the rest of the optimization — but it's not harmful either.

No bugs or correctness issues found.

---

Duplicate Detection

No duplicates detected. function_optimizer.py:1507 does a similar file_to_path() lookup but does not include the basename-fallback logic, so it is distinct.

---

Test Coverage

All 11 existing unit tests pass. The get_optimized_code_for_module function has 100% functional coverage per the PR's own test report. Overall file coverage is 42% because replace_function_definitions_for_language and _extract_function_from_code are not exercised by these tests — this is pre-existing, not a regression.

---

| Branch