perf: O(1) chunk cache lookup with shared Arc buffers instead of O(n) scan+clone

The decompressed-chunk LRU cache was the hottest path in the read pipeline
(every chunked-dataset read goes through it) but did a linear scan through
up to 521 slots on every get/put, and a full buffer copy on every cache hit
(to_vec()/clone() of the whole decompressed chunk). chunked_read.rs then
cloned the buffer a second time just to insert it into the cache after
already having it in hand.

- Added a HashMap<ChunkCoord, usize> index alongside the LRU slots for O(1)
  lookup. Eviction uses swap_remove, so the swapped-in slot's index entry is
  fixed up on every eviction (covered by a dedicated test).
- CachedChunk.data is now Arc<CacheAlignedBuffer> — a cache hit is a
  refcount bump, not a copy. CacheAlignedBuffer gained a Sync impl (same
  soundness argument as its existing Send impl: access is only ever through
  borrow-checked &/&mut, like Vec<u8>) so Arc<CacheAlignedBuffer> is itself
  Send/Sync.
- put_decompressed/put_decompressed_aligned now return the Arc they just
  inserted (or the existing cached copy), so callers can reuse that
  allocation instead of holding a separate clone — eliminates the second
  copy in chunked_read.rs's three call sites, which now consume the
  Arc<CacheAlignedBuffer> (Deref's to &[u8], so downstream indexing/copy
  code is unchanged).
- prefetch_hint's doc comment now leads with "bookkeeping only, does not
  prefetch" instead of describing behavior it doesn't have.

Co-Authored-By: Claude Sonnet 5 <[email protected]>

crates/clawhdf5-format/src/chunk_cache.rs

@@ -16,6 +16,8 @@ use core::ops::{Deref, DerefMut};
 use alloc::collections::BTreeMap;
 #[cfg(feature = "std")]
 use std::collections::HashMap;
+#[cfg(feature = "std")]
+use std::sync::Arc;
 
 use crate::chunk_index::{ChunkIndex, ChunkLayout};
 use crate::chunked_read::ChunkInfo;
@@ -64,6 +66,11 @@ pub struct CacheAlignedBuffer {
 
 // SAFETY: The raw pointer is exclusively owned — no aliasing.
 unsafe impl Send for CacheAlignedBuffer {}
+// SAFETY: `CacheAlignedBuffer` exposes its contents only via `&[u8]`/`&mut
+// [u8]` through the ordinary borrow-checked `Deref`/`DerefMut` impls below —
+// the same access pattern as `Vec<u8>`, which is `Sync`. Needed so
+// `Arc<CacheAlignedBuffer>` (used by the chunk cache) is itself `Send`.
+unsafe impl Sync for CacheAlignedBuffer {}
 
 impl CacheAlignedBuffer {
     /// Allocate a new cache-line-aligned buffer of exactly `len` bytes,
@@ -223,7 +230,9 @@ pub const DEFAULT_MAX_SLOTS: usize = 521;
 #[cfg(feature = "std")]
 struct CachedChunk {
     coord: ChunkCoord,
-    data: CacheAlignedBuffer,
+    /// Shared so a cache hit is a refcount bump, not a copy of the whole
+    /// (potentially large) decompressed chunk.
+    data: Arc<CacheAlignedBuffer>,
     /// Monotonically increasing access counter for LRU ordering.
     last_access: u64,
 }
@@ -267,6 +276,12 @@ struct CacheInner {
     /// LRU cache of decompressed chunk data.
     slots: Vec<CachedChunk>,
 
+    /// Coordinate -> index into `slots`, for O(1) lookup instead of a linear
+    /// scan. Kept in sync with `slots` on every insert/evict/clear — in
+    /// particular, `slots.swap_remove(i)` moves the last element into slot
+    /// `i`, so the moved element's index entry must be updated too.
+    slot_index: HashMap<ChunkCoord, usize>,
+
     /// Current total bytes of cached decompressed data.
     current_bytes: usize,
 
@@ -344,6 +359,7 @@ impl ChunkCache {
                 index: None,
                 index_addr: None,
                 slots: Vec::with_capacity(max_slots.min(64)),
+                slot_index: HashMap::with_capacity(max_slots.min(64)),
                 current_bytes: 0,
                 max_bytes,
                 max_slots,
@@ -375,6 +391,7 @@ impl ChunkCache {
         inner.chunk_index = None;
         inner.chunk_layout = None;
         inner.slots.clear();
+        inner.slot_index.clear();
         inner.current_bytes = 0;
         inner.last_coord = None;
         inner.index_addr = Some(addr);
@@ -477,8 +494,20 @@ impl ChunkCache {
 
     /// Try to get cached decompressed data for a chunk coordinate.
     ///
-    /// Returns a clone of the cache-line-aligned buffer.
+    /// O(1) lookup. Returns an owned copy for API compatibility with callers
+    /// that need a `Vec<u8>`; prefer [`Self::get_decompressed_aligned`] when
+    /// an `Arc`-shared buffer works for the caller, since that avoids the
+    /// copy entirely.
     pub fn get_decompressed(&self, coord: &[u64]) -> Option<Vec<u8>> {
+        self.get_decompressed_aligned(coord)
+            .map(|arc| arc.as_slice().to_vec())
+    }
+
+    /// Try to get a reference-counted clone of the aligned buffer for a chunk.
+    ///
+    /// O(1) index lookup; the clone is an `Arc` refcount bump, not a copy of
+    /// the underlying decompressed data.
+    pub fn get_decompressed_aligned(&self, coord: &[u64]) -> Option<Arc<CacheAlignedBuffer>> {
         let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
         inner.tick += 1;
         let tick = inner.tick;
@@ -500,36 +529,12 @@ impl ChunkCache {
         }
         inner.last_coord = Some(coord.to_vec());
 
-        let mut found = None;
+        let found = if let Some(&idx) = inner.slot_index.get(coord) {
-        for slot in inner.slots.iter_mut() {
+            inner.slots[idx].last_access = tick;
-            if slot.coord.as_slice() == coord {
+            Some(Arc::clone(&inner.slots[idx].data))
-                slot.last_access = tick;
-                found = Some(slot.data.to_vec());
-                break;
-            }
-        }
-        if let Some(ref data) = found {
-            inner.stats.hits += 1;
-            inner.stats.bytes_read += data.len() as u64;
         } else {
-            inner.stats.misses += 1;
+            None
-        }
+        };
-        found
-    }
-
-    /// Try to get a reference-counted clone of the aligned buffer for a chunk.
-    pub fn get_decompressed_aligned(&self, coord: &[u64]) -> Option<CacheAlignedBuffer> {
-        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
-        inner.tick += 1;
-        let tick = inner.tick;
-        let mut found = None;
-        for slot in inner.slots.iter_mut() {
-            if slot.coord.as_slice() == coord {
-                slot.last_access = tick;
-                found = Some(slot.data.clone());
-                break;
-            }
-        }
         if let Some(ref data) = found {
             inner.stats.hits += 1;
             inner.stats.bytes_read += data.len() as u64;
@@ -542,30 +547,39 @@ impl ChunkCache {
     /// Insert decompressed chunk data into the LRU cache.
     ///
     /// The data is stored in a [`CacheAlignedBuffer`] so subsequent reads
-    /// return cache-line-aligned memory.
+    /// return cache-line-aligned memory. Returns the `Arc`-shared buffer that
-    pub fn put_decompressed(&self, coord: ChunkCoord, data: Vec<u8>) {
+    /// is now cached (or already was), so the caller can reuse it directly
-        let aligned = CacheAlignedBuffer::from_slice(&data);
+    /// instead of holding a separate copy of the same data.
-        self.put_decompressed_aligned(coord, aligned);
+    pub fn put_decompressed(&self, coord: ChunkCoord, data: Vec<u8>) -> Arc<CacheAlignedBuffer> {
+        let aligned = CacheAlignedBuffer::from_vec(data);
+        self.put_decompressed_aligned(coord, aligned)
     }
 
     /// Insert an already-aligned buffer into the LRU cache.
-    pub fn put_decompressed_aligned(&self, coord: ChunkCoord, data: CacheAlignedBuffer) {
+    ///
+    /// Returns the `Arc`-shared buffer now held by the cache (the one just
+    /// inserted, or the existing cached copy if `coord` was already present).
+    pub fn put_decompressed_aligned(
+        &self,
+        coord: ChunkCoord,
+        data: CacheAlignedBuffer,
+    ) -> Arc<CacheAlignedBuffer> {
+        let data = Arc::new(data);
         let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
         let data_len = data.len();
 
-        // Don't cache if single chunk exceeds budget
+        // Don't cache if single chunk exceeds budget — still return the data
+        // to the caller, just don't retain it.
         if data_len > inner.max_bytes {
-            return;
+            return data;
         }
 
         // Check if already present
         inner.tick += 1;
         let tick = inner.tick;
-        for slot in inner.slots.iter_mut() {
+        if let Some(&idx) = inner.slot_index.get(&coord) {
-            if slot.coord == coord {
+            inner.slots[idx].last_access = tick;
-                slot.last_access = tick;
+            return Arc::clone(&inner.slots[idx].data); // already cached
-                return; // already cached
-            }
         }
 
         // Evict until we have room
@@ -581,16 +595,26 @@ impl ChunkCache {
                 .map(|(i, _)| i)
                 .unwrap();
             let removed = inner.slots.swap_remove(lru_idx);
+            inner.slot_index.remove(&removed.coord);
+            // swap_remove moved the former last element into `lru_idx` (unless
+            // it *was* the last element) — fix up that element's index entry.
+            if lru_idx < inner.slots.len() {
+                let moved_coord = inner.slots[lru_idx].coord.clone();
+                inner.slot_index.insert(moved_coord, lru_idx);
+            }
             inner.current_bytes -= removed.data.len();
             inner.stats.evictions += 1;
         }
 
         inner.current_bytes += data_len;
+        let new_idx = inner.slots.len();
+        inner.slot_index.insert(coord.clone(), new_idx);
         inner.slots.push(CachedChunk {
             coord,
-            data,
+            data: Arc::clone(&data),
             last_access: tick,
         });
+        data
     }
 
     /// Clear the entire cache (index + decompressed data).
@@ -599,6 +623,7 @@ impl ChunkCache {
         inner.index = None;
         inner.index_addr = None;
         inner.slots.clear();
+        inner.slot_index.clear();
         inner.current_bytes = 0;
         inner.tick = 0;
         inner.last_coord = None;
@@ -607,11 +632,13 @@ impl ChunkCache {
         inner.chunk_layout = None;
     }
 
-    /// Hint that the given chunk coordinates will be accessed soon.
+    /// Record that the given chunk coordinates are predicted to be accessed
+    /// soon (bookkeeping only).
     ///
-    /// Pre-populates the chunk index for these coordinates so that
+    /// This does **not** prefetch or pre-decompress anything — it only
-    /// subsequent lookups are O(1). This does NOT pre-decompress the
+    /// checks whether each coordinate is already in the chunk index and
-    /// chunks — it only ensures the index entries exist.
+    /// updates access-pattern stats accordingly. Real prefetching (e.g.
+    /// background pre-decompression) is not implemented.
     pub fn prefetch_hint(&self, next_coords: &[ChunkCoord]) {
         let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
         if inner.index.is_none() {
@@ -785,6 +812,50 @@ mod tests {
         assert_eq!(cache.cached_bytes(), 3);
     }
 
+    #[test]
+    fn slot_index_consistent_after_many_evictions() {
+        // Force repeated swap_remove evictions (small slot budget, many
+        // inserts) and confirm the coord -> slot index stays correct: every
+        // remaining coord must still resolve to its own data, not another
+        // slot's (which would happen if swap_remove's index fixup were wrong).
+        let cache = ChunkCache::with_capacity(1024 * 1024, 4); // max 4 slots
+
+        for i in 0..50u64 {
+            cache.put_decompressed(vec![i], vec![(i % 256) as u8; 8]);
+            // Interleave reads of a couple of earlier coords to churn LRU
+            // order (and thus which slot gets swap_remove'd) beyond simple
+            // FIFO eviction.
+            if i >= 2 {
+                let _ = cache.get_decompressed(&[i - 2]);
+            }
+        }
+
+        // Whatever remains in the cache (at most 4 slots) must return its
+        // own correct data.
+        for i in 0..50u64 {
+            if let Some(data) = cache.get_decompressed(&[i]) {
+                assert_eq!(
+                    data,
+                    vec![(i % 256) as u8; 8],
+                    "coord {i} returned wrong data after eviction churn"
+                );
+            }
+        }
+        assert!(cache.cached_chunk_count() <= 4);
+    }
+
+    #[test]
+    fn get_decompressed_aligned_shares_arc_on_hit() {
+        let cache = ChunkCache::new();
+        cache.put_decompressed(vec![0, 0], vec![9, 9, 9, 9]);
+        let a = cache.get_decompressed_aligned(&[0, 0]).unwrap();
+        let b = cache.get_decompressed_aligned(&[0, 0]).unwrap();
+        // A cache hit clones the Arc (refcount bump), not the underlying
+        // buffer — both handles point at the same allocation.
+        assert!(Arc::ptr_eq(&a, &b));
+        assert_eq!(a.as_slice(), &[9, 9, 9, 9]);
+    }
+
     // --- CacheAlignedBuffer tests ---
 
     #[test]

crates/clawhdf5-format/src/chunked_read.rs

@@ -9,6 +9,8 @@ use alloc::{format, vec, vec::Vec};
 use crate::chunk_cache::CacheAlignedBuffer;
 #[cfg(feature = "std")]
 use crate::chunk_cache::ChunkCache;
+#[cfg(feature = "std")]
+use std::sync::Arc;
 use crate::data_layout::DataLayout;
 use crate::dataspace::Dataspace;
 use crate::datatype::Datatype;
@@ -689,7 +691,7 @@ pub fn read_chunked_data_cached(
         let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
 
         // Try decompressed cache first
-        let decompressed = if let Some(cached) = cache.get_decompressed(&coord) {
+        let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
             cached
         } else {
             // Decompress from file
@@ -711,8 +713,7 @@ pub fn read_chunked_data_cached(
             } else {
                 raw_chunk.to_vec()
             };
-            cache.put_decompressed(coord, dec.clone());
+            cache.put_decompressed(coord, dec)
-            dec
         };
 
         let chunk_offsets: Vec<usize> = chunk_info
@@ -1055,7 +1056,7 @@ pub fn read_chunked_data_sweep(
         }
 
         // Try decompressed cache first
-        let decompressed = if let Some(cached) = cache.get_decompressed(&coord) {
+        let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
             cached
         } else {
             // Decompress from file
@@ -1077,8 +1078,7 @@ pub fn read_chunked_data_sweep(
             } else {
                 raw_chunk.to_vec()
             };
-            cache.put_decompressed(coord, dec.clone());
+            cache.put_decompressed(coord, dec)
-            dec
         };
 
         let chunk_offsets: Vec<usize> = chunk_info
@@ -1271,7 +1271,7 @@ pub fn read_chunked_data_indexed(
         .ok_or_else(|| FormatError::ChunkedReadError("chunk layout not available".into()))?;
 
     // Decompress chunks (using LRU cache where possible)
-    let mut chunk_buffers: Vec<CacheAlignedBuffer> = Vec::with_capacity(mappings_info.len());
+    let mut chunk_buffers: Vec<Arc<CacheAlignedBuffer>> = Vec::with_capacity(mappings_info.len());
     for (coord, file_offset, file_size, filter_mask) in &mappings_info {
         if let Some(cached) = cache.get_decompressed_aligned(coord) {
             chunk_buffers.push(cached);
@@ -1295,8 +1295,8 @@ pub fn read_chunked_data_indexed(
                 raw_chunk.to_vec()
             };
             let aligned = CacheAlignedBuffer::from_vec(decompressed);
-            cache.put_decompressed_aligned(coord.clone(), aligned.clone());
+            let arc = cache.put_decompressed_aligned(coord.clone(), aligned);
-            chunk_buffers.push(aligned);
+            chunk_buffers.push(arc);
         }
     }