- frame caching

src/view.rs

@@ -1,9 +1,9 @@
 use crate::axes::{Ax, Axis, Operation, Shape, Slice, SliceInfoElemDef, slice_info};
 use crate::error::Error;
 use crate::metadata::Metadata;
-use crate::readers::{Dimensions, DynReader, Reader};
+use crate::readers::{Dimensions, DynReader, Frame, Reader};
 use crate::stats::MinMax;
-use indexmap::IndexMap;
+use indexmap::{Equivalent, IndexMap};
 use itertools::{Itertools, iproduct};
 use ndarray::{
     Array, Array0, Array1, Array2, ArrayD, Dimension, IntoDimension, Ix0, Ix1, Ix2, Ix5, IxDyn,
@@ -15,12 +15,13 @@ use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
 use std::any::type_name;
 use std::collections::{HashMap, HashSet};
-use std::fmt::{Display, Formatter};
+use std::fmt::{Debug, Display, Formatter};
 use std::hash::{Hash, Hasher};
 use std::iter::Sum;
 use std::marker::PhantomData;
 use std::ops::{AddAssign, Deref, Div};
-use std::path::Path;
+use std::path::{Path, PathBuf};
+use std::sync::{Arc, Mutex, OnceLock};
 
 fn idx_bnd(idx: isize, bnd: isize) -> Result<isize, Error> {
     if idx < -bnd {
@@ -62,9 +63,157 @@ impl<T> Number for T where
 {
 }
 
+/// maximum number of frames held in the cache
+const DEFAULT_FRAME_CACHE_SIZE: usize = 128;
+
+/// identity of the reader a frame was read from
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct ReaderKey {
+    name: String,
+    path: PathBuf,
+    series: usize,
+    position: usize,
+}
+
+type FrameKey = (ReaderKey, usize, usize, usize);
+
+/// borrowed view of [`FrameKey`] for cache lookups without allocation.
+/// hashes byte-identically to [`FrameKey`] (str/String and Path/PathBuf hash
+/// the same), so it only matches the same reader identity.
+#[derive(Debug, PartialEq, Eq)]
+struct FrameKeyRef<'a> {
+    name: &'a str,
+    path: &'a Path,
+    series: usize,
+    position: usize,
+    c: usize,
+    z: usize,
+    t: usize,
+}
+
+impl Hash for FrameKeyRef<'_> {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+        self.path.hash(state);
+        self.series.hash(state);
+        self.position.hash(state);
+        self.c.hash(state);
+        self.z.hash(state);
+        self.t.hash(state);
+    }
+}
+
+impl Equivalent<FrameKey> for FrameKeyRef<'_> {
+    fn equivalent(&self, key: &FrameKey) -> bool {
+        let (rk, c, z, t) = key;
+        self.name == rk.name
+            && self.path == rk.path.as_path()
+            && self.series == rk.series
+            && self.position == rk.position
+            && self.c == *c
+            && self.z == *z
+            && self.t == *t
+    }
+}
+
+impl FrameKeyRef<'_> {
+    fn to_owned(&self) -> FrameKey {
+        (
+            ReaderKey {
+                name: self.name.to_string(),
+                path: self.path.to_path_buf(),
+                series: self.series,
+                position: self.position,
+            },
+            self.c,
+            self.z,
+            self.t,
+        )
+    }
+}
+
+/// process-wide LRU cache of frames, shared between all views and threads
+static GLOBAL_FRAME_CACHE: OnceLock<FrameCache> = OnceLock::new();
+
+/// thread-safe LRU cache of frames read from the underlying reader
+struct FrameCache {
+    inner: Mutex<FrameCacheInner>,
+}
+
+struct FrameCacheInner {
+    map: IndexMap<FrameKey, Arc<Frame>>,
+    capacity: usize,
+}
+
+impl Default for FrameCache {
+    fn default() -> Self {
+        Self::new(DEFAULT_FRAME_CACHE_SIZE)
+    }
+}
+
+impl FrameCache {
+    fn new(capacity: usize) -> Self {
+        Self {
+            inner: Mutex::new(FrameCacheInner {
+                map: IndexMap::with_capacity(capacity),
+                capacity,
+            }),
+        }
+    }
+
+    fn global() -> &'static FrameCache {
+        GLOBAL_FRAME_CACHE.get_or_init(FrameCache::default)
+    }
+
+    fn capacity(&self) -> usize {
+        self.inner.lock().unwrap().capacity
+    }
+
+    fn set_capacity(&self, capacity: usize) {
+        let mut inner = self.inner.lock().unwrap();
+        inner.capacity = capacity;
+        while inner.map.len() > inner.capacity {
+            inner.map.shift_remove_index(0);
+        }
+    }
+
+    fn get<Q>(&self, key: &Q) -> Option<Arc<Frame>>
+    where
+        Q: ?Sized + Hash + Equivalent<FrameKey>,
+    {
+        let mut inner = self.inner.lock().unwrap();
+        if let Some(idx) = inner.map.get_index_of(key) {
+            let (key, frame) = inner.map.shift_remove_index(idx).unwrap();
+            inner.map.insert(key, frame.clone());
+            Some(frame)
+        } else {
+            None
+        }
+    }
+
+    fn insert(&self, key: FrameKey, frame: Arc<Frame>) {
+        let mut inner = self.inner.lock().unwrap();
+        inner.map.insert(key, frame);
+        while inner.map.len() > inner.capacity {
+            inner.map.shift_remove_index(0);
+        }
+    }
+
+    #[cfg(test)]
+    fn len(&self) -> usize {
+        self.inner.lock().unwrap().map.len()
+    }
+}
+
+impl Debug for FrameCache {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("FrameCache").finish_non_exhaustive()
+    }
+}
+
 /// sliceable view on an image file
 #[serde_as]
-#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct View<D: Dimension, R: Reader = DynReader> {
     reader: R,
     /// same order as axes
@@ -176,6 +325,18 @@ impl<D: Dimension, R: Reader> View<D, R> {
         self.reader.series()
     }
 
+    /// set the maximum number of frames the process-wide cache may hold,
+    /// evicting the least recently used frames
+    pub fn with_cache_capacity(self, capacity: usize) -> Self {
+        FrameCache::global().set_capacity(capacity);
+        self
+    }
+
+    /// the maximum number of frames the process-wide cache may hold
+    pub fn cache_capacity(&self) -> usize {
+        FrameCache::global().capacity()
+    }
+
     fn with_operations(mut self, operations: IndexMap<Axis, Operation>) -> Self {
         self.operations = operations;
         self
@@ -744,13 +905,13 @@ impl<D: Dimension, R: Reader> View<D, R> {
             if let Some(i) = axes_out_idx[2] {
                 slice[i] = SliceInfoElem::Index(t)
             };
-            let frame = self.reader.get_frame(
+            let frame = self.get_cached_frame(
                 (c % shape.c as isize) as usize,
                 (z % shape.z as isize) as usize,
                 (t % shape.t as isize) as usize,
             )?;
 
-            let arr_frame: Array2<T> = frame.try_into()?;
+            let arr_frame: Array2<T> = frame.as_ref().try_into()?;
             let arr_frame = match xy_dim {
                 0 => {
                     if op_xy.contains_key(&Axis::X) && op_xy.contains_key(&Axis::Y) {
@@ -888,6 +1049,28 @@ impl<D: Dimension, R: Reader> View<D, R> {
             .collect())
     }
 
+    fn frame_key(&self, c: usize, z: usize, t: usize) -> FrameKeyRef<'_> {
+        FrameKeyRef {
+            name: self.reader.reader_name(),
+            path: self.reader.path(),
+            series: self.reader.series(),
+            position: self.reader.position(),
+            c,
+            z,
+            t,
+        }
+    }
+
+    fn get_cached_frame(&self, c: usize, z: usize, t: usize) -> Result<Arc<Frame>, Error> {
+        let key = self.frame_key(c, z, t);
+        if let Some(frame) = FrameCache::global().get(&key) {
+            return Ok(frame);
+        }
+        let frame = Arc::new(self.reader.get_frame(c, z, t)?);
+        FrameCache::global().insert(key.to_owned(), frame.clone());
+        Ok(frame)
+    }
+
     /// retrieve a single frame at czt, sliced accordingly
     pub fn get_frame<T, N>(&self, c: N, z: N, t: N) -> Result<Array2<T>, Error>
     where
@@ -1116,11 +1299,13 @@ to_bytes_vec_impl!(
 mod tests {
     use crate::axes::Axis;
     use crate::error::Error;
-    use crate::readers::{DynReader, Reader};
+    use crate::readers::{DynReader, Frame, Reader};
     use crate::stats::MinMax;
-    use crate::view::Item;
+    use crate::view::{FrameCache, Item, ReaderKey};
     use ndarray::{Array, Array4, Array5, NewAxis};
     use ndarray::{Array2, s};
+    use std::path::PathBuf;
+    use std::sync::Arc;
 
     fn open(file: &str) -> Result<DynReader, Error> {
         let path = std::env::current_dir()?
@@ -1388,4 +1573,60 @@ mod tests {
         assert_eq!(a.shape(), [1280, 1280]);
         Ok(())
     }
+
+    #[test]
+    fn frame_cache() -> Result<(), Error> {
+        let cache = FrameCache::new(2);
+        let key = |c: usize, z: usize, t: usize| {
+            (
+                ReaderKey {
+                    name: "test".to_string(),
+                    path: PathBuf::from("test.tif"),
+                    series: 0,
+                    position: 0,
+                },
+                c,
+                z,
+                t,
+            )
+        };
+        cache.insert(
+            key(0, 0, 0),
+            Arc::new(Frame::from(Array2::from_elem((2, 2), 1u16))),
+        );
+        cache.insert(
+            key(0, 0, 1),
+            Arc::new(Frame::from(Array2::from_elem((2, 2), 2u16))),
+        );
+        assert!(cache.get(&key(0, 0, 0)).is_some());
+        cache.insert(
+            key(0, 0, 2),
+            Arc::new(Frame::from(Array2::from_elem((2, 2), 3u16))),
+        );
+        assert!(cache.get(&key(0, 0, 0)).is_some());
+        assert!(cache.get(&key(0, 0, 1)).is_none());
+        assert!(cache.get(&key(0, 0, 2)).is_some());
+        assert_eq!(cache.len(), 2);
+        Ok(())
+    }
+
+    #[test]
+    fn as_array_caches_frames() -> Result<(), Error> {
+        let file = "tiffseq/YTL1841B2-2-1_1hr_DMSO_galinduction_1";
+        let reader = open(file)?;
+        let view = reader.view();
+        let shape = view.shape();
+        let (c, z, t) = (shape[0] - 1, shape[1] - 1, shape[2] - 1);
+        let cache = FrameCache::global();
+        let a: Array5<usize> = view.clone().try_into()?;
+        assert!(cache.get(&view.frame_key(c, z, t)).is_some());
+        let b: Array5<usize> = view.clone().try_into()?;
+        assert_eq!(a, b);
+        // a fresh reader on the same file shares the same cache key
+        let view2 = open(file)?.view();
+        assert_eq!(view.frame_key(c, z, t), view2.frame_key(c, z, t));
+        let c2: Array5<usize> = view2.clone().try_into()?;
+        assert_eq!(a, c2);
+        Ok(())
+    }
 }