⚡️ Speed up method SegmentLRUCache.set by 7%

[codeflash-ai]

📄 7% (0.07x) speedup for SegmentLRUCache.set in chromadb/segment/impl/manager/cache/cache.py

⏱️ Runtime : 1.81 milliseconds → 1.70 milliseconds (best of 153 runs)

📝 Explanation and details

The optimized code achieves a 6% speedup through two key data structure improvements that eliminate expensive O(n) operations:

1. Dict-based LRU History (O(1) operations)

• Replaces self.history = [] with self.history: Dict[uuid.UUID, None] = {}
• Leverages Python 3.7+ dict insertion order preservation for LRU tracking
• Before: self.history.remove(key) + self.history.append(key) = O(n) linear search + O(1) append
• After: self.history.pop(key) + self.history[key] = None = O(1) + O(1)

2. Cached Size Tracking (eliminates repeated computations)

• Adds self.key_sizes: Dict[uuid.UUID, int] = {} and self.total_size: int = 0
• Before: key_sizes = {key: self.size_func(key) for key in self.cache} + sum(key_sizes.values()) = O(n) dict comprehension + O(n) summation on every insertion
• After: Maintains running totals with O(1) updates: self.total_size += item_size

Performance Impact by Test Type:

• Single insertions: 15-30% faster (eliminates dict comprehension overhead)
• Large-scale operations: 6-8% faster (benefits compound with cache size)
• Variable-size scenarios: Most benefit from avoiding repeated size_func calls

The optimization is particularly effective for workloads with frequent cache operations on moderately-sized caches, where the O(n) operations in the original code become bottlenecks.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import threading
import uuid
from typing import Any, Callable, Dict, Optional

# imports
import pytest  # used for our unit tests
from chromadb.segment.impl.manager.cache.cache import SegmentLRUCache


# Minimal mock Segment type for testing
class MockSegment:
    def __init__(self, data):
        self.data = data

# ------------------------- UNIT TESTS -------------------------

# Helper functions for tests
def fixed_size_func(size):
    """Returns a size_func that always returns the given size."""
    return lambda key: size

def variable_size_func(sizes):
    """Returns a size_func that returns a size from a dict keyed by UUID."""
    return lambda key: sizes[key]

def make_segments(n):
    """Create n mock segments and UUIDs."""
    return [(uuid.uuid4(), MockSegment(f"segment_{i}")) for i in range(n)]

# ----------------- BASIC TEST CASES -----------------

def test_set_basic_insertion():
    """Test basic insertion of a single item."""
    cache = SegmentLRUCache(capacity=10, size_func=fixed_size_func(1))
    key, value = uuid.uuid4(), MockSegment("foo")
    cache.set(key, value) # 2.92μs -> 2.26μs (28.9% faster)


def test_set_duplicate_key_no_overwrite():
    """Test that setting a key already present does not overwrite value."""
    cache = SegmentLRUCache(capacity=5, size_func=fixed_size_func(1))
    k = uuid.uuid4()
    v1 = MockSegment("first")
    v2 = MockSegment("second")
    cache.set(k, v1) # 2.80μs -> 2.17μs (29.3% faster)
    cache.set(k, v2) # 540ns -> 539ns (0.186% faster)

def test_set_history_order():
    """Test that history is updated correctly on insert."""
    cache = SegmentLRUCache(capacity=10, size_func=fixed_size_func(1))
    k1, v1 = uuid.uuid4(), MockSegment("a")
    k2, v2 = uuid.uuid4(), MockSegment("b")
    cache.set(k1, v1) # 2.82μs -> 2.31μs (22.1% faster)
    cache.set(k2, v2) # 593ns -> 534ns (11.0% faster)

# ----------------- EDGE TEST CASES -----------------

def test_set_eviction_on_capacity_exceeded():
    """Test that the oldest item is evicted when capacity is exceeded."""
    cache = SegmentLRUCache(capacity=2, size_func=fixed_size_func(1))
    k1, v1 = uuid.uuid4(), MockSegment("first")
    k2, v2 = uuid.uuid4(), MockSegment("second")
    k3, v3 = uuid.uuid4(), MockSegment("third")
    cache.set(k1, v1) # 2.75μs -> 2.26μs (21.5% faster)
    cache.set(k2, v2) # 575ns -> 545ns (5.50% faster)
    cache.set(k3, v3) # 366ns -> 338ns (8.28% faster)

def test_set_eviction_with_variable_sizes():
    """Test eviction when items have variable sizes."""
    k1, k2, k3 = uuid.uuid4(), uuid.uuid4(), uuid.uuid4()
    sizes = {k1: 2, k2: 3, k3: 4}
    cache = SegmentLRUCache(capacity=5, size_func=variable_size_func(sizes))
    cache.set(k1, MockSegment("first")) # 2.78μs -> 2.40μs (15.7% faster)
    cache.set(k2, MockSegment("second")) # 584ns -> 534ns (9.36% faster)
    cache.set(k3, MockSegment("third")) # 356ns -> 367ns (3.00% slower)

def test_set_eviction_callback_called():
    """Test that callback is called on eviction."""
    evicted = []
    def callback(key, value):
        evicted.append((key, value.data))
    cache = SegmentLRUCache(capacity=2, size_func=fixed_size_func(1), callback=callback)
    k1, v1 = uuid.uuid4(), MockSegment("first")
    k2, v2 = uuid.uuid4(), MockSegment("second")
    k3, v3 = uuid.uuid4(), MockSegment("third")
    cache.set(k1, v1) # 2.77μs -> 2.20μs (25.5% faster)
    cache.set(k2, v2) # 633ns -> 549ns (15.3% faster)
    cache.set(k3, v3) # 357ns -> 341ns (4.69% faster)

def test_set_zero_capacity():
    """Test that no items can be inserted if capacity is zero."""
    cache = SegmentLRUCache(capacity=0, size_func=fixed_size_func(1))
    k1, v1 = uuid.uuid4(), MockSegment("first")
    cache.set(k1, v1) # 2.86μs -> 2.30μs (24.4% faster)

def test_set_item_larger_than_capacity():
    """Test that an item larger than capacity cannot be inserted."""
    cache = SegmentLRUCache(capacity=1, size_func=fixed_size_func(2))
    k1, v1 = uuid.uuid4(), MockSegment("big")
    cache.set(k1, v1) # 2.87μs -> 2.25μs (27.9% faster)






def test_set_large_eviction_callback():
    """Test callback is called correct number of times during large evictions."""
    N = 1000
    capacity = 500
    evicted = []
    def callback(key, value):
        evicted.append(key)
    cache = SegmentLRUCache(capacity=capacity, size_func=fixed_size_func(1), callback=callback)
    keys_vals = make_segments(N)
    for k, v in keys_vals:
        cache.set(k, v) # 345μs -> 324μs (6.63% faster)
    # Evicted keys should be the first N-capacity inserted keys
    expected_evicted = [kv[0] for kv in keys_vals[:N-capacity]]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import threading
import uuid
from typing import Any, Callable, Dict, Optional

# imports
import pytest
from chromadb.segment.impl.manager.cache.cache import SegmentLRUCache


class Segment:
    """Dummy Segment class for testing."""
    def __init__(self, data: Any):
        self.data = data

# unit tests

# Helper functions for tests
def fixed_size_func(size):
    """Returns a function that always returns a fixed size."""
    return lambda key: size

def varying_size_func(sizes_dict):
    """Returns a function that returns sizes based on a dict."""
    return lambda key: sizes_dict[key]

def make_segments(num, data_prefix="data"):
    """Create a list of Segment objects with unique data."""
    return [Segment(f"{data_prefix}_{i}") for i in range(num)]

def make_keys(num):
    """Create a list of unique UUID keys."""
    return [uuid.uuid4() for _ in range(num)]

# --- 1. Basic Test Cases ---

def test_set_adds_new_item():
    """Test that set adds a new item to the cache."""
    cache = SegmentLRUCache(capacity=10, size_func=fixed_size_func(1))
    key = uuid.uuid4()
    seg = Segment("A")
    cache.set(key, seg) # 3.76μs -> 2.98μs (26.2% faster)

def test_set_does_not_overwrite_existing_key():
    """Test that set does not overwrite existing key."""
    cache = SegmentLRUCache(capacity=10, size_func=fixed_size_func(1))
    key = uuid.uuid4()
    seg1 = Segment("A")
    seg2 = Segment("B")
    cache.set(key, seg1) # 3.02μs -> 2.62μs (15.2% faster)
    cache.set(key, seg2) # 574ns -> 563ns (1.95% faster)

def test_set_eviction_on_capacity():
    """Test that items are evicted when capacity is exceeded."""
    cache = SegmentLRUCache(capacity=2, size_func=fixed_size_func(1))
    keys = make_keys(3)
    segs = make_segments(3)
    cache.set(keys[0], segs[0]) # 3.00μs -> 2.42μs (23.5% faster)
    cache.set(keys[1], segs[1]) # 576ns -> 523ns (10.1% faster)
    cache.set(keys[2], segs[2]) # 350ns -> 344ns (1.74% faster)

def test_set_history_order():
    """Test that history is updated in LRU order."""
    cache = SegmentLRUCache(capacity=3, size_func=fixed_size_func(1))
    keys = make_keys(3)
    segs = make_segments(3)
    for k, s in zip(keys, segs):
        cache.set(k, s) # 3.77μs -> 3.15μs (20.0% faster)
    # Add a new key to cause eviction
    new_key = uuid.uuid4()
    new_seg = Segment("new")
    cache.set(new_key, new_seg) # 363ns -> 359ns (1.11% faster)

def test_set_with_callback_on_eviction():
    """Test that callback is called on eviction."""
    evicted = []
    def cb(key, segment):
        evicted.append((key, segment.data))
    cache = SegmentLRUCache(capacity=2, size_func=fixed_size_func(1), callback=cb)
    keys = make_keys(3)
    segs = make_segments(3)
    cache.set(keys[0], segs[0]) # 2.86μs -> 2.29μs (24.5% faster)
    cache.set(keys[1], segs[1]) # 572ns -> 518ns (10.4% faster)
    cache.set(keys[2], segs[2]) # 364ns -> 352ns (3.41% faster)

# --- 2. Edge Test Cases ---

def test_set_zero_capacity():
    """Test that nothing is stored if capacity is zero."""
    cache = SegmentLRUCache(capacity=0, size_func=fixed_size_func(1))
    key = uuid.uuid4()
    seg = Segment("A")
    cache.set(key, seg) # 2.98μs -> 2.42μs (23.1% faster)

def test_set_item_larger_than_capacity():
    """Test that item larger than capacity is not stored."""
    cache = SegmentLRUCache(capacity=2, size_func=fixed_size_func(3))
    key = uuid.uuid4()
    seg = Segment("big")
    cache.set(key, seg) # 2.95μs -> 2.43μs (21.6% faster)


def test_set_varying_sizes():
    """Test eviction with varying item sizes."""
    sizes = {}
    keys = make_keys(4)
    sizes[keys[0]] = 1
    sizes[keys[1]] = 2
    sizes[keys[2]] = 3
    sizes[keys[3]] = 2
    cache = SegmentLRUCache(capacity=5, size_func=varying_size_func(sizes))
    segs = make_segments(4)
    cache.set(keys[0], segs[0]) # 2.89μs -> 2.22μs (30.1% faster)
    cache.set(keys[1], segs[1]) # 612ns -> 555ns (10.3% faster)
    cache.set(keys[2], segs[2]) # 368ns -> 337ns (9.20% faster)
    cache.set(keys[3], segs[3]) # 334ns -> 323ns (3.41% faster)


def test_set_history_removal_and_append():
    """Test that history is updated correctly when re-adding existing key."""
    cache = SegmentLRUCache(capacity=3, size_func=fixed_size_func(1))
    keys = make_keys(3)
    segs = make_segments(3)
    for k, s in zip(keys, segs):
        cache.set(k, s) # 5.00μs -> 4.21μs (18.7% faster)
    # Re-add key[0] (should not update cache, but history should move key[0] to end)
    cache.set(keys[0], segs[0]) # 347ns -> 328ns (5.79% faster)

# --- 3. Large Scale Test Cases ---



def test_set_large_items_varying_sizes():
    """Test inserting many items with varying sizes, checking eviction."""
    N = 100
    keys = make_keys(N)
    segs = make_segments(N)
    sizes = {k: (i % 5) + 1 for i, k in enumerate(keys)}  # Sizes 1-5
    total_size = sum(sizes[k] for k in keys)
    capacity = 250
    cache = SegmentLRUCache(capacity=capacity, size_func=varying_size_func(sizes))
    for k, s in zip(keys, segs):
        cache.set(k, s) # 37.3μs -> 34.7μs (7.64% faster)

def test_set_performance_under_load():
    """Test that set does not take excessive time for large number of items."""
    import time
    N = 900
    cache = SegmentLRUCache(capacity=900, size_func=fixed_size_func(1))
    keys = make_keys(N)
    segs = make_segments(N)
    start = time.time()
    for k, s in zip(keys, segs):
        cache.set(k, s) # 309μs -> 290μs (6.57% faster)
    duration = time.time() - start

def test_set_eviction_with_large_items():
    """Test eviction when inserting large items into a nearly full cache."""
    keys = make_keys(10)
    segs = make_segments(10)
    sizes = {k: 1 for k in keys}
    capacity = 10
    cache = SegmentLRUCache(capacity=capacity, size_func=varying_size_func(sizes))
    for k, s in zip(keys, segs):
        cache.set(k, s) # 6.65μs -> 5.75μs (15.6% faster)
    # Now insert a large item
    big_key = uuid.uuid4()
    big_seg = Segment("BIG")
    sizes[big_key] = 8
    cache.size_func = varying_size_func(sizes)
    cache.set(big_key, big_seg) # 362ns -> 341ns (6.16% faster)
    # Should evict enough items to fit big_seg
    total_size = sum(sizes[k] for k in cache.cache)
# 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-SegmentLRUCache.set-mh7jc82p and push.