⚡️ Speed up method CommentMapper.get_comment by 30% in PR #537 (runtime-fixes-jul10)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #537

If you approve this dependent PR, these changes will be merged into the original PR branch runtime-fixes-jul10.

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

---

📄 30% (0.30x) speedup for CommentMapper.get_comment in codeflash/code_utils/edit_generated_tests.py

⏱️ Runtime : 7.33 milliseconds → 5.61 milliseconds (best of 203 runs)

📝 Explanation and details

Here’s a much faster, lower-overhead rewrite of your code, retaining all return values and logic, and keeping all output and function signatures the same. The key changes.

• The format_time and format_perf functions are inlined, branch logic is optimized, and string formatting is simplified.
• The performance_gain calculation is unchanged, but evaluated directly.
• The class CommentMapper keeps the same logic, but with attribute access and stack allocation slightly tightened, and avoids some redundant method lookups or assignments.
• No function is renamed or signature changed; all comments are preserved as requested.

Here's the optimized code.

Summary of speed-up.

• All for-loop and handler logic in format_time replaced with simple tight branch conditions (much faster for typical values).
• Redundant checks compressed, float division minimized.
• Speed-up for formatting applies both for time and percentage formatting (calls are now faster).
• Class methods leverage straight-in calculations to avoid extra function call layers as much as possible, and local variable assignment is reduced.

All output is identical to the original and all comments are fully preserved.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import ast

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.edit_generated_tests import CommentMapper


class GeneratedTests:
    """Minimal stub for testing purposes."""
    def __init__(self, behavior_file_path):
        self.behavior_file_path = behavior_file_path

class DummyPath:
    def __init__(self, path):
        self._path = path
    def with_suffix(self, suffix):
        return self
from codeflash.code_utils.edit_generated_tests import CommentMapper

# unit tests

# -------------------------------
# Basic Test Cases
# -------------------------------

def make_mapper(orig, opt):
    """Helper to create a CommentMapper with dummy path."""
    test = GeneratedTests(DummyPath("dummy.py"))
    return CommentMapper(test, orig, opt)

def test_basic_faster():
    # Optimized is faster than original
    orig = {"foo": 2_000_000}
    opt = {"foo": 1_000_000}
    mapper = make_mapper(orig, opt)
    codeflash_output = mapper.get_comment("foo"); comment = codeflash_output # 5.69μs -> 4.66μs (22.2% faster)

def test_basic_slower():
    # Optimized is slower than original
    orig = {"bar": 1_000_000}
    opt = {"bar": 2_000_000}
    mapper = make_mapper(orig, opt)
    codeflash_output = mapper.get_comment("bar"); comment = codeflash_output # 5.26μs -> 4.39μs (19.8% faster)

def test_basic_equal():
    # Optimized is same as original
    orig = {"baz": 1_000_000}
    opt = {"baz": 1_000_000}
    mapper = make_mapper(orig, opt)
    codeflash_output = mapper.get_comment("baz"); comment = codeflash_output # 4.96μs -> 4.15μs (19.6% faster)

def test_basic_microseconds():
    # Test microsecond formatting
    orig = {"test": 12_345}
    opt = {"test": 6_789}
    mapper = make_mapper(orig, opt)
    codeflash_output = mapper.get_comment("test"); comment = codeflash_output # 5.76μs -> 4.88μs (18.1% faster)

def test_basic_nanoseconds():
    # Test nanosecond formatting (no decimal)
    orig = {"nano": 999}
    opt = {"nano": 500}
    mapper = make_mapper(orig, opt)
    codeflash_output = mapper.get_comment("nano"); comment = codeflash_output # 3.81μs -> 3.68μs (3.54% faster)

# -------------------------------
# Edge Test Cases
# -------------------------------

def test_zero_optimized_runtime():
    # Optimized runtime is zero (should not divide by zero)
    orig = {"edge": 1_000_000}
    opt = {"edge": 0}
    mapper = make_mapper(orig, opt)
    codeflash_output = mapper.get_comment("edge"); comment = codeflash_output # 4.80μs -> 4.04μs (18.8% faster)

def test_zero_original_runtime():
    # Original runtime is zero, optimized is nonzero
    orig = {"edge": 0}
    opt = {"edge": 1_000_000}
    mapper = make_mapper(orig, opt)
    codeflash_output = mapper.get_comment("edge"); comment = codeflash_output # 4.74μs -> 4.03μs (17.7% faster)

def test_both_zero():
    # Both runtimes are zero
    orig = {"edge": 0}
    opt = {"edge": 0}
    mapper = make_mapper(orig, opt)
    codeflash_output = mapper.get_comment("edge"); comment = codeflash_output # 3.48μs -> 3.14μs (10.8% faster)

def test_large_values():
    # Very large nanosecond values (seconds formatting)
    orig = {"big": 12_345_678_901}
    opt = {"big": 1_234_567_890}
    mapper = make_mapper(orig, opt)
    codeflash_output = mapper.get_comment("big"); comment = codeflash_output # 5.86μs -> 4.75μs (23.4% faster)

def test_small_difference():
    # Very small difference between original and optimized
    orig = {"tiny": 1_000_001}
    opt = {"tiny": 1_000_000}
    mapper = make_mapper(orig, opt)
    codeflash_output = mapper.get_comment("tiny"); comment = codeflash_output # 5.32μs -> 4.43μs (20.2% faster)

def test_negative_runtime_raises():
    # Negative runtime should raise ValueError
    orig = {"neg": -1}
    opt = {"neg": 1}
    mapper = make_mapper(orig, opt)
    with pytest.raises(ValueError):
        mapper.get_comment("neg") # 5.13μs -> 5.11μs (0.391% faster)

def test_non_integer_runtime_raises():
    # Non-integer runtime should raise TypeError
    orig = {"bad": 1.5}
    opt = {"bad": 2}
    mapper = make_mapper(orig, opt)
    with pytest.raises(TypeError):
        mapper.get_comment("bad") # 5.01μs -> 5.12μs (2.17% slower)

# -------------------------------
# Large Scale Test Cases
# -------------------------------

def test_many_keys_all_faster():
    # Test with 1000 keys, all optimized are half original
    orig = {f"k{i}": 2_000_000 for i in range(1000)}
    opt = {f"k{i}": 1_000_000 for i in range(1000)}
    mapper = make_mapper(orig, opt)
    for i in range(1000):
        codeflash_output = mapper.get_comment(f"k{i}"); comment = codeflash_output # 1.87ms -> 1.45ms (29.0% faster)

def test_many_keys_mixed():
    # 1000 keys, half faster, half slower
    orig = {f"k{i}": 1_000_000 for i in range(1000)}
    opt = {f"k{i}": 500_000 if i % 2 == 0 else 2_000_000 for i in range(1000)}
    mapper = make_mapper(orig, opt)
    for i in range(1000):
        codeflash_output = mapper.get_comment(f"k{i}"); comment = codeflash_output # 1.91ms -> 1.48ms (29.2% faster)
        if i % 2 == 0:
            pass
        else:
            pass

def test_large_numbers_performance():
    # Test with large values to ensure no overflow or slowdowns
    orig = {f"k{i}": 999_999_999 for i in range(100)}
    opt = {f"k{i}": 1_000_000 for i in range(100)}
    mapper = make_mapper(orig, opt)
    for i in range(100):
        codeflash_output = mapper.get_comment(f"k{i}"); comment = codeflash_output # 192μs -> 144μs (32.5% faster)

def test_large_scale_edge_zeroes():
    # All original and optimized are zero
    orig = {f"k{i}": 0 for i in range(100)}
    opt = {f"k{i}": 0 for i in range(100)}
    mapper = make_mapper(orig, opt)
    for i in range(100):
        codeflash_output = mapper.get_comment(f"k{i}"); comment = codeflash_output # 120μs -> 106μs (13.7% faster)

def test_large_scale_varied():
    # 100 keys, with increasing original and decreasing optimized
    orig = {f"k{i}": (i+1)*1_000_000 for i in range(100)}
    opt = {f"k{i}": (100-i)*1_000_000 for i in range(100)}
    mapper = make_mapper(orig, opt)
    for i in range(100):
        codeflash_output = mapper.get_comment(f"k{i}"); comment = codeflash_output # 204μs -> 161μs (26.2% faster)
        # Check that the comment contains the correct units and "faster" or "slower"
        o_val = (i+1)*1_000_000
        n_val = (100-i)*1_000_000
        if n_val > o_val:
            pass
        else:
            pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

import ast

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.edit_generated_tests import CommentMapper


class DummyPath:
    """Dummy path object for testing."""
    def with_suffix(self, suffix):
        return self

class GeneratedTests:
    """Minimal stub for GeneratedTests."""
    def __init__(self, behavior_file_path):
        self.behavior_file_path = behavior_file_path
from codeflash.code_utils.edit_generated_tests import CommentMapper

# unit tests

@pytest.fixture
def mapper_factory():
    """Fixture to create a CommentMapper with provided runtimes."""
    def make_mapper(orig, opt):
        test = GeneratedTests(DummyPath())
        return CommentMapper(test, orig, opt)
    return make_mapper

# ----------------------
# Basic Test Cases
# ----------------------

def test_basic_faster_case(mapper_factory):
    """Test a basic case where optimized is faster than original."""
    orig = {"A": 2_000_000}
    opt = {"A": 1_000_000}
    mapper = mapper_factory(orig, opt)
    codeflash_output = mapper.get_comment("A"); comment = codeflash_output # 5.06μs -> 4.15μs (22.0% faster)

def test_basic_slower_case(mapper_factory):
    """Test a basic case where optimized is slower than original."""
    orig = {"B": 1_000_000}
    opt = {"B": 2_000_000}
    mapper = mapper_factory(orig, opt)
    codeflash_output = mapper.get_comment("B"); comment = codeflash_output # 5.06μs -> 4.34μs (16.6% faster)

def test_basic_equal_case(mapper_factory):
    """Test a basic case where runtimes are equal."""
    orig = {"C": 1_000_000}
    opt = {"C": 1_000_000}
    mapper = mapper_factory(orig, opt)
    codeflash_output = mapper.get_comment("C"); comment = codeflash_output # 4.96μs -> 3.92μs (26.6% faster)

def test_basic_ns_boundary(mapper_factory):
    """Test formatting at nanosecond boundary (< 1_000)."""
    orig = {"D": 999}
    opt = {"D": 1}
    mapper = mapper_factory(orig, opt)
    codeflash_output = mapper.get_comment("D"); comment = codeflash_output # 4.00μs -> 3.72μs (7.53% faster)

def test_basic_us_boundary(mapper_factory):
    """Test formatting at microsecond boundary (just below 1_000_000)."""
    orig = {"E": 999_999}
    opt = {"E": 1_000}
    mapper = mapper_factory(orig, opt)
    codeflash_output = mapper.get_comment("E"); comment = codeflash_output # 6.13μs -> 5.03μs (21.9% faster)
    # 999_999ns = 1_000.00μs, 1_000ns = 1.00μs
    perf_gain = (999_999 - 1_000) / 1_000 * 100

# ----------------------
# Edge Test Cases
# ----------------------

def test_zero_optimized_runtime(mapper_factory):
    """Test with optimized runtime == 0 (should not divide by zero)."""
    orig = {"F": 10_000}
    opt = {"F": 0}
    mapper = mapper_factory(orig, opt)
    codeflash_output = mapper.get_comment("F"); comment = codeflash_output # 5.03μs -> 4.34μs (15.9% faster)

def test_zero_original_runtime(mapper_factory):
    """Test with original runtime == 0, optimized > 0."""
    orig = {"G": 0}
    opt = {"G": 1_000}
    mapper = mapper_factory(orig, opt)
    codeflash_output = mapper.get_comment("G"); comment = codeflash_output # 5.01μs -> 4.44μs (12.9% faster)

def test_large_runtime_values(mapper_factory):
    """Test with very large runtime values (seconds)."""
    orig = {"H": 10_000_000_000}
    opt = {"H": 5_000_000_000}
    mapper = mapper_factory(orig, opt)
    codeflash_output = mapper.get_comment("H"); comment = codeflash_output # 5.87μs -> 4.76μs (23.4% faster)

def test_small_runtime_values(mapper_factory):
    """Test with very small runtime values (1ns, 2ns)."""
    orig = {"I": 2}
    opt = {"I": 1}
    mapper = mapper_factory(orig, opt)
    codeflash_output = mapper.get_comment("I"); comment = codeflash_output # 3.80μs -> 3.62μs (4.98% faster)

def test_formatting_precision(mapper_factory):
    """Test formatting at various boundaries for significant digits."""
    orig = {"J": 123_456_789}
    opt = {"J": 12_345_678}
    mapper = mapper_factory(orig, opt)
    codeflash_output = mapper.get_comment("J"); comment = codeflash_output # 5.65μs -> 4.48μs (26.1% faster)
    # 123_456_789ns = 123.46ms, 12_345_678ns = 12.35ms
    perf_gain = (123_456_789 - 12_345_678) / 12_345_678 * 100

def test_non_integer_runtime_raises(mapper_factory):
    """Test that non-integer runtimes raise TypeError."""
    orig = {"K": "not_an_int"}
    opt = {"K": 1_000}
    mapper = mapper_factory(orig, opt)
    with pytest.raises(TypeError):
        mapper.get_comment("K") # 3.55μs -> 3.12μs (13.9% faster)

def test_negative_runtime_raises(mapper_factory):
    """Test that negative runtimes raise ValueError."""
    orig = {"L": -100}
    opt = {"L": 1_000}
    mapper = mapper_factory(orig, opt)
    with pytest.raises(ValueError):
        mapper.get_comment("L") # 5.10μs -> 5.05μs (0.990% faster)

def test_missing_key_raises(mapper_factory):
    """Test that missing match_key raises KeyError."""
    orig = {"M": 1_000}
    opt = {"M": 2_000}
    mapper = mapper_factory(orig, opt)
    with pytest.raises(KeyError):
        mapper.get_comment("not_in_dict") # 1.66μs -> 1.69μs (1.77% slower)

# ----------------------
# Large Scale Test Cases
# ----------------------

def test_many_keys_scalability(mapper_factory):
    """Test with a large number of keys to check scalability and correctness."""
    # 1000 keys, original = 1_000_000, optimized = 500_000
    n = 1000
    orig = {f"K{i}": 1_000_000 for i in range(n)}
    opt = {f"K{i}": 500_000 for i in range(n)}
    mapper = mapper_factory(orig, opt)
    for i in range(n):
        codeflash_output = mapper.get_comment(f"K{i}"); comment = codeflash_output # 1.89ms -> 1.43ms (32.8% faster)

def test_many_keys_varied_runtimes(mapper_factory):
    """Test with a large number of keys and varied runtimes."""
    n = 500
    orig = {f"K{i}": (i+1)*1_000 for i in range(n)}
    opt = {f"K{i}": (n-i)*1_000 for i in range(n)}
    mapper = mapper_factory(orig, opt)
    for i in range(n):
        key = f"K{i}"
        codeflash_output = mapper.get_comment(key); comment = codeflash_output # 999μs -> 730μs (36.9% faster)
        orig_val = orig[key]
        opt_val = opt[key]
        perf_gain = (orig_val - opt_val) / opt_val * 100 if opt_val != 0 else 0.0
        status = "slower" if opt_val > orig_val else "faster"

def test_large_values_and_small_values(mapper_factory):
    """Test with a mix of very large and very small runtime values."""
    orig = {
        "big": 999_999_999,
        "small": 1,
        "mid": 100_000
    }
    opt = {
        "big": 1_000_000_000,
        "small": 2,
        "mid": 50_000
    }
    mapper = mapper_factory(orig, opt)
    # big: slower, small: faster, mid: faster
    codeflash_output = mapper.get_comment("big") # 6.39μs -> 5.17μs (23.6% faster)
    codeflash_output = mapper.get_comment("small") # 2.12μs -> 2.05μs (3.41% faster)
    codeflash_output = mapper.get_comment("mid") # 3.19μs -> 2.56μs (24.7% 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-pr537-2025-07-11T04.11.46 and push.