- as_array_dyn lru caching

.gitignore

@@ -75,6 +75,7 @@ docs/_build/
 /tests/files/*
 AGENTS.md
 .agentbridge
+.agent-work
 
 py/ndbioimage/jassets
 py/ndbioimage/deps

src/view.rs

@@ -13,7 +13,7 @@ use num::traits::ToBytes;
 use num::{Bounded, FromPrimitive, ToPrimitive, Zero};
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
-use std::any::type_name;
+use std::any::{Any, type_name};
 use std::collections::{HashMap, HashSet};
 use std::fmt::{Debug, Display, Formatter};
 use std::hash::{Hash, Hasher};
@@ -48,11 +48,22 @@ fn slc_bnd(idx: isize, bnd: isize) -> Result<isize, Error> {
 }
 
 pub trait Number:
-    'static + AddAssign + Bounded + Clone + Div<Self, Output = Self> + FromPrimitive + PartialOrd + Zero
+    'static
+    + Send
+    + Sync
+    + AddAssign
+    + Bounded
+    + Clone
+    + Div<Self, Output = Self>
+    + FromPrimitive
+    + PartialOrd
+    + Zero
 {
 }
 impl<T> Number for T where
     T: 'static
+        + Send
+        + Sync
         + AddAssign
         + Bounded
         + Clone
@@ -66,6 +77,9 @@ impl<T> Number for T where
 /// maximum number of frames held in the cache
 const DEFAULT_FRAME_CACHE_SIZE: usize = 128;
 
+/// maximum number of materialized arrays held in the cache
+const DEFAULT_ARRAY_CACHE_SIZE: usize = 2;
+
 /// identity of the reader a frame was read from
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 struct ReaderKey {
@@ -75,17 +89,52 @@ struct ReaderKey {
     position: usize,
 }
 
-type FrameKey = (ReaderKey, usize, usize, usize);
+/// borrowed view of [`ReaderKey`] for cache lookups without allocation.
-
+/// hashes byte-identically to [`ReaderKey`] (str/String and Path/PathBuf hash
-/// 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> {
+struct ReaderKeyRef<'a> {
     name: &'a str,
     path: &'a Path,
     series: usize,
     position: usize,
+}
+
+impl Hash for ReaderKeyRef<'_> {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+        self.path.hash(state);
+        self.series.hash(state);
+        self.position.hash(state);
+    }
+}
+
+impl Equivalent<ReaderKey> for ReaderKeyRef<'_> {
+    fn equivalent(&self, key: &ReaderKey) -> bool {
+        self.name == key.name
+            && self.path == key.path.as_path()
+            && self.series == key.series
+            && self.position == key.position
+    }
+}
+
+impl ReaderKeyRef<'_> {
+    fn to_owned(&self) -> ReaderKey {
+        ReaderKey {
+            name: self.name.to_string(),
+            path: self.path.to_path_buf(),
+            series: self.series,
+            position: self.position,
+        }
+    }
+}
+
+type FrameKey = (ReaderKey, usize, usize, usize);
+
+/// borrowed view of [`FrameKey`] for cache lookups without allocation.
+#[derive(Debug, PartialEq, Eq)]
+struct FrameKeyRef<'a> {
+    reader: ReaderKeyRef<'a>,
     c: usize,
     z: usize,
     t: usize,
@@ -93,10 +142,7 @@ struct FrameKeyRef<'a> {
 
 impl Hash for FrameKeyRef<'_> {
     fn hash<H: Hasher>(&self, state: &mut H) {
-        self.name.hash(state);
+        self.reader.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);
@@ -106,29 +152,78 @@ impl Hash for FrameKeyRef<'_> {
 impl Equivalent<FrameKey> for FrameKeyRef<'_> {
     fn equivalent(&self, key: &FrameKey) -> bool {
         let (rk, c, z, t) = key;
-        self.name == rk.name
+        self.reader.equivalent(rk) && self.c == *c && self.z == *z && self.t == *t
-            && 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 {
-        (
+        (self.reader.to_owned(), self.c, self.z, self.t)
-            ReaderKey {
+    }
-                name: self.name.to_string(),
+}
-                path: self.path.to_path_buf(),
+
-                series: self.series,
+/// identity of a materialized array in the process-wide cache
-                position: self.position,
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-            },
+struct ArrayKey {
-            self.c,
+    reader: ReaderKey,
-            self.z,
+    dtype: &'static str,
-            self.t,
+    slice: Vec<SliceInfoElem>,
-        )
+    axes: Vec<Axis>,
+    operations: Vec<(Axis, Operation)>,
+}
+
+/// borrowed view of [`ArrayKey`] for cache lookups without allocation.
+#[derive(Debug, PartialEq, Eq)]
+struct ArrayKeyRef<'a> {
+    reader: ReaderKeyRef<'a>,
+    dtype: &'static str,
+    slice: &'a [SliceInfoElem],
+    axes: &'a [Axis],
+    operations: &'a IndexMap<Axis, Operation>,
+}
+
+impl Hash for ArrayKeyRef<'_> {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.reader.hash(state);
+        self.dtype.hash(state);
+        self.slice.hash(state);
+        self.axes.hash(state);
+        self.operations.len().hash(state);
+        for (ax, op) in self.operations.iter() {
+            ax.hash(state);
+            op.hash(state);
+        }
+    }
+}
+
+impl Equivalent<ArrayKey> for ArrayKeyRef<'_> {
+    fn equivalent(&self, key: &ArrayKey) -> bool {
+        self.reader.equivalent(&key.reader)
+            && self.dtype == key.dtype
+            && self.slice == key.slice
+            && self.axes == key.axes
+            && self.operations.len() == key.operations.len()
+            && self
+                .operations
+                .iter()
+                .zip(&key.operations)
+                .all(|((ax, op), (key_ax, key_op))| ax == key_ax && op == key_op)
+    }
+}
+
+impl ArrayKeyRef<'_> {
+    fn to_owned(&self) -> ArrayKey {
+        ArrayKey {
+            reader: self.reader.to_owned(),
+            dtype: self.dtype,
+            slice: self.slice.to_vec(),
+            axes: self.axes.to_vec(),
+            operations: self
+                .operations
+                .iter()
+                .map(|(ax, op)| (*ax, op.clone()))
+                .collect(),
+        }
     }
 }
 
@@ -211,6 +306,86 @@ impl Debug for FrameCache {
     }
 }
 
+/// process-wide LRU cache of materialized arrays, shared between all views and threads
+static GLOBAL_ARRAY_CACHE: OnceLock<ArrayCache> = OnceLock::new();
+
+/// thread-safe LRU cache of materialized arrays produced by `as_array_dyn`
+struct ArrayCache {
+    inner: Mutex<ArrayCacheInner>,
+}
+
+struct ArrayCacheInner {
+    map: IndexMap<ArrayKey, Arc<dyn Any + Send + Sync>>,
+    capacity: usize,
+}
+
+impl Default for ArrayCache {
+    fn default() -> Self {
+        Self::new(DEFAULT_ARRAY_CACHE_SIZE)
+    }
+}
+
+impl ArrayCache {
+    fn new(capacity: usize) -> Self {
+        Self {
+            inner: Mutex::new(ArrayCacheInner {
+                map: IndexMap::with_capacity(capacity),
+                capacity,
+            }),
+        }
+    }
+
+    fn global() -> &'static ArrayCache {
+        GLOBAL_ARRAY_CACHE.get_or_init(ArrayCache::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<T, Q>(&self, key: &Q) -> Option<Arc<ArrayD<T>>>
+    where
+        Q: ?Sized + Hash + Equivalent<ArrayKey>,
+        T: Any + Send + Sync,
+    {
+        let mut inner = self.inner.lock().unwrap();
+        if let Some(idx) = inner.map.get_index_of(key) {
+            let (key, array) = inner.map.shift_remove_index(idx).unwrap();
+            inner.map.insert(key, array.clone());
+            array.downcast::<ArrayD<T>>().ok()
+        } else {
+            None
+        }
+    }
+
+    fn insert<T: Any + Send + Sync>(&self, key: ArrayKey, array: ArrayD<T>) {
+        let mut inner = self.inner.lock().unwrap();
+        inner.map.insert(key, Arc::new(array));
+        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 ArrayCache {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ArrayCache").finish_non_exhaustive()
+    }
+}
+
 /// sliceable view on an image file
 #[serde_as]
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
@@ -337,6 +512,18 @@ impl<D: Dimension, R: Reader> View<D, R> {
         FrameCache::global().capacity()
     }
 
+    /// set the maximum number of materialized arrays the process-wide cache
+    /// may hold, evicting the least recently used arrays
+    pub fn with_array_cache_capacity(self, capacity: usize) -> Self {
+        ArrayCache::global().set_capacity(capacity);
+        self
+    }
+
+    /// the maximum number of materialized arrays the process-wide cache may hold
+    pub fn array_cache_capacity(&self) -> usize {
+        ArrayCache::global().capacity()
+    }
+
     fn with_operations(mut self, operations: IndexMap<Axis, Operation>) -> Self {
         self.operations = operations;
         self
@@ -799,6 +986,10 @@ impl<D: Dimension, R: Reader> View<D, R> {
         Array1<T>: MinMax<Output = Array0<T>>,
         Array2<T>: MinMax<Output = Array1<T>>,
     {
+        let key = self.array_key::<T>();
+        if let Some(arr) = ArrayCache::global().get(&key) {
+            return Ok(arr.as_ref().clone());
+        }
         let mut op_xy = IndexMap::new();
         if let Some((&ax, op)) = self.operations.first()
             && ((ax == Axis::X) || (ax == Axis::Y))
@@ -954,7 +1145,7 @@ impl<D: Dimension, R: Reader> View<D, R> {
                     }
                 }
                 _ => {
-                    panic!("xy cannot be 3d or more");
+                    unreachable!("xy cannot be 3d or more");
                 }
             };
             if let Some((_, op)) = op_czt.first() {
@@ -1020,6 +1211,7 @@ impl<D: Dimension, R: Reader> View<D, R> {
             let m = T::from_usize(n).unwrap_or_else(|| T::zero());
             out.take().unwrap().mapv(|x| x / m.clone())
         };
+        ArrayCache::global().insert(key.to_owned(), array.clone());
         Ok(array)
     }
 
@@ -1051,16 +1243,33 @@ impl<D: Dimension, R: Reader> View<D, R> {
 
     fn frame_key(&self, c: usize, z: usize, t: usize) -> FrameKeyRef<'_> {
         FrameKeyRef {
-            name: self.reader.reader_name(),
+            reader: ReaderKeyRef {
-            path: self.reader.path(),
+                name: self.reader.reader_name(),
-            series: self.reader.series(),
+                path: self.reader.path(),
-            position: self.reader.position(),
+                series: self.reader.series(),
+                position: self.reader.position(),
+            },
             c,
             z,
             t,
         }
     }
 
+    fn array_key<T: 'static>(&self) -> ArrayKeyRef<'_> {
+        ArrayKeyRef {
+            reader: ReaderKeyRef {
+                name: self.reader.reader_name(),
+                path: self.reader.path(),
+                series: self.reader.series(),
+                position: self.reader.position(),
+            },
+            dtype: type_name::<T>(),
+            slice: &self.slice,
+            axes: &self.axes,
+            operations: &self.operations,
+        }
+    }
+
     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) {
@@ -1297,14 +1506,17 @@ to_bytes_vec_impl!(
 
 #[cfg(test)]
 mod tests {
-    use crate::axes::Axis;
+    use crate::axes::{Axis, Operation};
     use crate::error::Error;
     use crate::readers::{DynReader, Frame, Reader};
     use crate::stats::MinMax;
-    use crate::view::{FrameCache, Item, ReaderKey};
+    use crate::view::{
+        ArrayCache, ArrayKey, ArrayKeyRef, FrameCache, Item, ReaderKey, ReaderKeyRef,
+    };
+    use indexmap::IndexMap;
     use ndarray::{Array, Array4, Array5, NewAxis};
-    use ndarray::{Array2, s};
+    use ndarray::{Array2, ArrayD, IxDyn, SliceInfoElem, s};
-    use std::path::PathBuf;
+    use std::path::{Path, PathBuf};
     use std::sync::Arc;
 
     fn open(file: &str) -> Result<DynReader, Error> {
@@ -1610,6 +1822,51 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn array_cache() -> Result<(), Error> {
+        let cache = ArrayCache::new(2);
+        let key = |dtype: &'static str, index: isize, axes: &[Axis]| ArrayKey {
+            reader: ReaderKey {
+                name: "test".to_string(),
+                path: PathBuf::from("test.tif"),
+                series: 0,
+                position: 0,
+            },
+            dtype,
+            slice: vec![SliceInfoElem::Index(index)],
+            axes: axes.to_vec(),
+            operations: vec![],
+        };
+        let k0 = key("u16", 0, &[Axis::T]);
+        let k1 = key("u16", 1, &[Axis::T]);
+        cache.insert(k0.clone(), ArrayD::<u16>::zeros(IxDyn(&[2, 2])));
+        cache.insert(k1.clone(), ArrayD::<u16>::zeros(IxDyn(&[2, 2])));
+        assert!(cache.get::<u16, _>(&k0).is_some());
+        cache.insert(
+            key("u16", 2, &[Axis::T]),
+            ArrayD::<u16>::zeros(IxDyn(&[2, 2])),
+        );
+        assert!(cache.get::<u16, _>(&k0).is_some());
+        assert!(cache.get::<u16, _>(&k1).is_none());
+        assert_eq!(cache.len(), 2);
+        let ops: IndexMap<Axis, Operation> = IndexMap::new();
+        let borrowed = ArrayKeyRef {
+            reader: ReaderKeyRef {
+                name: "test",
+                path: Path::new("test.tif"),
+                series: 0,
+                position: 0,
+            },
+            dtype: "u16",
+            slice: &[SliceInfoElem::Index(0)],
+            axes: &[Axis::T],
+            operations: &ops,
+        };
+        assert!(cache.get::<f64, _>(&borrowed).is_none());
+        assert!(cache.get::<u16, _>(&borrowed).is_some());
+        Ok(())
+    }
+
     #[test]
     fn as_array_caches_frames() -> Result<(), Error> {
         let file = "tiffseq/YTL1841B2-2-1_1hr_DMSO_galinduction_1";