tensor4all_core/cached_function/mod.rs
1//! Cached function wrapper for expensive function evaluations.
2//!
3//! [`CachedFunction`] wraps a user-supplied function `Fn(&[I]) -> V` with
4//! thread-safe memoization. On first call for a given multi-index, the
5//! function is evaluated and the result is cached; subsequent calls return
6//! the cached value.
7//!
8//! The internal cache key type is automatically selected based on the total
9//! index space size (up to 1024 bits by default). For larger index spaces,
10//! use [`CachedFunction::with_key_type`] with a custom [`CacheKey`]
11//! implementation.
12//!
13//! # Examples
14//!
15//! ```
16//! use tensor4all_core::CachedFunction;
17//!
18//! let cf = CachedFunction::new(
19//! |idx: &[usize]| idx[0] + idx[1],
20//! &[3, 4],
21//! ).unwrap();
22//!
23//! assert_eq!(cf.eval(&[1, 2]).unwrap(), 3);
24//! assert_eq!(cf.num_evals(), 1);
25//! assert_eq!(cf.eval(&[1, 2]).unwrap(), 3);
26//! assert_eq!(cf.num_cache_hits(), 1);
27//! ```
28
29pub mod cache_key;
30pub mod error;
31pub mod index_int;
32
33use std::collections::HashMap;
34use std::sync::atomic::{AtomicUsize, Ordering};
35use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
36
37use bnum::types::{U1024, U256, U512};
38
39use cache_key::CacheKey;
40use index_int::IndexInt;
41
42/// Compute total bits needed to represent the index space.
43pub(crate) fn total_bits(local_dims: &[usize]) -> u32 {
44 local_dims
45 .iter()
46 .map(|&d| {
47 if d <= 1 {
48 0
49 } else {
50 ((d - 1) as u64).ilog2() + 1
51 }
52 })
53 .sum()
54}
55
56/// Compute mixed-radix coefficients for flat index computation.
57///
58/// Returns `Err(CacheKeyError::Overflow)` if the index space overflows key type `K`.
59pub(crate) fn compute_coeffs<K: CacheKey>(
60 local_dims: &[usize],
61) -> Result<Vec<K>, error::CacheKeyError> {
62 let bits = total_bits(local_dims);
63 if bits > K::BITS_COUNT {
64 return Err(error::CacheKeyError::Overflow {
65 total_bits: bits,
66 max_bits: K::BITS_COUNT,
67 key_type: std::any::type_name::<K>(),
68 });
69 }
70
71 let mut coeffs = Vec::with_capacity(local_dims.len());
72 let mut prod = K::ONE;
73 for &d in local_dims {
74 coeffs.push(prod.clone());
75 let dim = K::from_usize(d);
76 prod = prod
77 .checked_mul(dim)
78 .ok_or_else(|| error::CacheKeyError::Overflow {
79 total_bits: bits,
80 max_bits: K::BITS_COUNT,
81 key_type: std::any::type_name::<K>(),
82 })?;
83 }
84
85 Ok(coeffs)
86}
87
88fn read_lock<T>(lock: &RwLock<T>) -> RwLockReadGuard<'_, T> {
89 lock.read().unwrap_or_else(|poisoned| poisoned.into_inner())
90}
91
92fn write_lock<T>(lock: &RwLock<T>) -> RwLockWriteGuard<'_, T> {
93 lock.write()
94 .unwrap_or_else(|poisoned| poisoned.into_inner())
95}
96
97/// Compute flat index from multi-index and coefficients.
98fn flat_index<K: CacheKey, I: IndexInt>(idx: &[I], coeffs: &[K]) -> K {
99 idx.iter().zip(coeffs).fold(K::ZERO, |acc, (&i, c)| {
100 acc.wrapping_add(
101 c.clone()
102 .checked_mul(K::from_usize(i.to_usize()))
103 .unwrap_or(K::ZERO),
104 )
105 })
106}
107
108/// Internal cache with automatically selected key type.
109enum InnerCache<V> {
110 U64 {
111 cache: RwLock<HashMap<u64, V>>,
112 coeffs: Vec<u64>,
113 },
114 U128 {
115 cache: RwLock<HashMap<u128, V>>,
116 coeffs: Vec<u128>,
117 },
118 U256 {
119 cache: RwLock<HashMap<U256, V>>,
120 coeffs: Vec<U256>,
121 },
122 U512 {
123 cache: RwLock<HashMap<U512, V>>,
124 coeffs: Vec<U512>,
125 },
126 U1024 {
127 cache: RwLock<HashMap<U1024, V>>,
128 coeffs: Vec<U1024>,
129 },
130}
131
132impl<V: Clone + Send + Sync> InnerCache<V> {
133 /// Create a new cache, automatically selecting the key type.
134 fn new(local_dims: &[usize]) -> Result<Self, error::CacheKeyError> {
135 let bits = total_bits(local_dims);
136 if bits <= 64 {
137 Ok(Self::U64 {
138 cache: RwLock::new(HashMap::new()),
139 coeffs: compute_coeffs::<u64>(local_dims)?,
140 })
141 } else if bits <= 128 {
142 Ok(Self::U128 {
143 cache: RwLock::new(HashMap::new()),
144 coeffs: compute_coeffs::<u128>(local_dims)?,
145 })
146 } else if bits <= 256 {
147 Ok(Self::U256 {
148 cache: RwLock::new(HashMap::new()),
149 coeffs: compute_coeffs::<U256>(local_dims)?,
150 })
151 } else if bits <= 512 {
152 Ok(Self::U512 {
153 cache: RwLock::new(HashMap::new()),
154 coeffs: compute_coeffs::<U512>(local_dims)?,
155 })
156 } else if bits <= 1024 {
157 Ok(Self::U1024 {
158 cache: RwLock::new(HashMap::new()),
159 coeffs: compute_coeffs::<U1024>(local_dims)?,
160 })
161 } else {
162 Err(error::CacheKeyError::Overflow {
163 total_bits: bits,
164 max_bits: 1024,
165 key_type: "auto",
166 })
167 }
168 }
169
170 fn get<I: IndexInt>(&self, idx: &[I]) -> Option<V> {
171 match self {
172 Self::U64 { cache, coeffs } => read_lock(cache).get(&flat_index(idx, coeffs)).cloned(),
173 Self::U128 { cache, coeffs } => read_lock(cache).get(&flat_index(idx, coeffs)).cloned(),
174 Self::U256 { cache, coeffs } => read_lock(cache).get(&flat_index(idx, coeffs)).cloned(),
175 Self::U512 { cache, coeffs } => read_lock(cache).get(&flat_index(idx, coeffs)).cloned(),
176 Self::U1024 { cache, coeffs } => {
177 read_lock(cache).get(&flat_index(idx, coeffs)).cloned()
178 }
179 }
180 }
181
182 fn insert<I: IndexInt>(&self, idx: &[I], value: V) {
183 match self {
184 Self::U64 { cache, coeffs } => {
185 write_lock(cache).insert(flat_index(idx, coeffs), value);
186 }
187 Self::U128 { cache, coeffs } => {
188 write_lock(cache).insert(flat_index(idx, coeffs), value);
189 }
190 Self::U256 { cache, coeffs } => {
191 write_lock(cache).insert(flat_index(idx, coeffs), value);
192 }
193 Self::U512 { cache, coeffs } => {
194 write_lock(cache).insert(flat_index(idx, coeffs), value);
195 }
196 Self::U1024 { cache, coeffs } => {
197 write_lock(cache).insert(flat_index(idx, coeffs), value);
198 }
199 }
200 }
201
202 fn contains<I: IndexInt>(&self, idx: &[I]) -> bool {
203 match self {
204 Self::U64 { cache, coeffs } => read_lock(cache).contains_key(&flat_index(idx, coeffs)),
205 Self::U128 { cache, coeffs } => read_lock(cache).contains_key(&flat_index(idx, coeffs)),
206 Self::U256 { cache, coeffs } => read_lock(cache).contains_key(&flat_index(idx, coeffs)),
207 Self::U512 { cache, coeffs } => read_lock(cache).contains_key(&flat_index(idx, coeffs)),
208 Self::U1024 { cache, coeffs } => {
209 read_lock(cache).contains_key(&flat_index(idx, coeffs))
210 }
211 }
212 }
213
214 fn len(&self) -> usize {
215 match self {
216 Self::U64 { cache, .. } => read_lock(cache).len(),
217 Self::U128 { cache, .. } => read_lock(cache).len(),
218 Self::U256 { cache, .. } => read_lock(cache).len(),
219 Self::U512 { cache, .. } => read_lock(cache).len(),
220 Self::U1024 { cache, .. } => read_lock(cache).len(),
221 }
222 }
223
224 fn clear(&self) {
225 match self {
226 Self::U64 { cache, .. } => write_lock(cache).clear(),
227 Self::U128 { cache, .. } => write_lock(cache).clear(),
228 Self::U256 { cache, .. } => write_lock(cache).clear(),
229 Self::U512 { cache, .. } => write_lock(cache).clear(),
230 Self::U1024 { cache, .. } => write_lock(cache).clear(),
231 }
232 }
233
234 fn key_type_name(&self) -> &'static str {
235 match self {
236 Self::U64 { .. } => "u64",
237 Self::U128 { .. } => "u128",
238 Self::U256 { .. } => "U256",
239 Self::U512 { .. } => "U512",
240 Self::U1024 { .. } => "U1024",
241 }
242 }
243}
244
245/// Type-erased cache interface for custom key types.
246trait DynCache<V>: Send + Sync {
247 fn get(&self, idx: &[usize]) -> Option<V>;
248 fn insert(&self, idx: &[usize], value: V);
249 fn contains(&self, idx: &[usize]) -> bool;
250 fn len(&self) -> usize;
251 fn clear(&self);
252}
253
254/// Generic cache for user-specified key types.
255struct GenericCache<K: CacheKey, V> {
256 cache: RwLock<HashMap<K, V>>,
257 coeffs: Vec<K>,
258}
259
260impl<K: CacheKey, V: Clone + Send + Sync> GenericCache<K, V> {
261 fn new(local_dims: &[usize]) -> Result<Self, error::CacheKeyError> {
262 Ok(Self {
263 cache: RwLock::new(HashMap::new()),
264 coeffs: compute_coeffs::<K>(local_dims)?,
265 })
266 }
267}
268
269impl<K: CacheKey, V: Clone + Send + Sync> DynCache<V> for GenericCache<K, V> {
270 fn get(&self, idx: &[usize]) -> Option<V> {
271 let key = flat_index::<K, usize>(idx, &self.coeffs);
272 read_lock(&self.cache).get(&key).cloned()
273 }
274
275 fn insert(&self, idx: &[usize], value: V) {
276 let key = flat_index::<K, usize>(idx, &self.coeffs);
277 write_lock(&self.cache).insert(key, value);
278 }
279
280 fn contains(&self, idx: &[usize]) -> bool {
281 let key = flat_index::<K, usize>(idx, &self.coeffs);
282 read_lock(&self.cache).contains_key(&key)
283 }
284
285 fn len(&self) -> usize {
286 read_lock(&self.cache).len()
287 }
288
289 fn clear(&self) {
290 write_lock(&self.cache).clear();
291 }
292}
293
294/// Internal backend: auto-selected enum or custom type-erased cache.
295enum CacheBackend<V: Clone + Send + Sync + 'static> {
296 Auto(InnerCache<V>),
297 Custom(Box<dyn DynCache<V>>),
298}
299
300impl<V: Clone + Send + Sync + 'static> CacheBackend<V> {
301 fn get<I: IndexInt>(&self, idx: &[I]) -> Option<V> {
302 match self {
303 Self::Auto(inner) => inner.get(idx),
304 Self::Custom(cache) => {
305 let usize_idx: Vec<usize> = idx.iter().map(|&i| i.to_usize()).collect();
306 cache.get(&usize_idx)
307 }
308 }
309 }
310
311 fn insert<I: IndexInt>(&self, idx: &[I], value: V) {
312 match self {
313 Self::Auto(inner) => inner.insert(idx, value),
314 Self::Custom(cache) => {
315 let usize_idx: Vec<usize> = idx.iter().map(|&i| i.to_usize()).collect();
316 cache.insert(&usize_idx, value);
317 }
318 }
319 }
320
321 fn contains<I: IndexInt>(&self, idx: &[I]) -> bool {
322 match self {
323 Self::Auto(inner) => inner.contains(idx),
324 Self::Custom(cache) => {
325 let usize_idx: Vec<usize> = idx.iter().map(|&i| i.to_usize()).collect();
326 cache.contains(&usize_idx)
327 }
328 }
329 }
330
331 fn len(&self) -> usize {
332 match self {
333 Self::Auto(inner) => inner.len(),
334 Self::Custom(cache) => cache.len(),
335 }
336 }
337
338 fn clear(&self) {
339 match self {
340 Self::Auto(inner) => inner.clear(),
341 Self::Custom(cache) => cache.clear(),
342 }
343 }
344
345 fn key_type_name(&self) -> &'static str {
346 match self {
347 Self::Auto(inner) => inner.key_type_name(),
348 Self::Custom(_) => "custom",
349 }
350 }
351}
352
353type BatchFunc<I, V> = dyn Fn(&[Vec<I>]) -> Vec<V> + Send + Sync;
354
355/// A wrapper that caches function evaluations for multi-index inputs.
356///
357/// Thread-safe: all methods take `&self`. Multiple threads can call `eval`
358/// concurrently.
359///
360/// # Type parameters
361///
362/// - `V` - cached value type
363/// - `F` - single-evaluation function `Fn(&[I]) -> V`
364/// - `I` - index element type (default `usize`); use `u8` for quantics
365///
366/// # Examples
367///
368/// ```
369/// use tensor4all_core::CachedFunction;
370///
371/// // Cache a 2-site function with local dimensions [3, 4]
372/// let cf = CachedFunction::new(
373/// |idx: &[usize]| (idx[0] * 4 + idx[1]) as f64,
374/// &[3, 4],
375/// ).unwrap();
376///
377/// // First call evaluates and caches
378/// let v00 = cf.eval(&[0, 0]).unwrap();
379/// assert_eq!(v00, 0.0);
380/// assert_eq!(cf.num_evals(), 1);
381/// assert_eq!(cf.num_cache_hits(), 0);
382///
383/// // Second call uses cache
384/// let v00_again = cf.eval(&[0, 0]).unwrap();
385/// assert_eq!(v00_again, 0.0);
386/// assert_eq!(cf.num_cache_hits(), 1);
387///
388/// let v12 = cf.eval(&[1, 2]).unwrap();
389/// assert_eq!(v12, 6.0); // 1*4 + 2
390/// ```
391pub struct CachedFunction<V, F, I = usize>
392where
393 I: IndexInt,
394 V: Clone + Send + Sync + 'static,
395 F: Fn(&[I]) -> V + Send + Sync,
396{
397 func: F,
398 batch_func: Option<Box<BatchFunc<I, V>>>,
399 cache: CacheBackend<V>,
400 local_dims: Vec<usize>,
401 num_evals: AtomicUsize,
402 num_cache_hits: AtomicUsize,
403 _phantom: std::marker::PhantomData<I>,
404}
405
406impl<V, F, I> CachedFunction<V, F, I>
407where
408 I: IndexInt,
409 V: Clone + Send + Sync + 'static,
410 F: Fn(&[I]) -> V + Send + Sync,
411{
412 /// Create a new cached function with automatic key selection (up to 1024 bits).
413 ///
414 /// # Errors
415 ///
416 /// Returns an error when the cache dimensions are invalid (a shape mismatch)
417 /// /// or the construction fails.
418 ///
419 /// # Examples
420 ///
421 /// ```
422 /// use tensor4all_core::CachedFunction;
423 ///
424 /// let cf = CachedFunction::new(|idx: &[usize]| idx[0] + idx[1], &[2, 3]).unwrap();
425 /// assert_eq!(cf.eval(&[1, 2]).unwrap(), 3);
426 /// assert_eq!(cf.num_sites(), 2);
427 /// assert_eq!(cf.local_dims(), &[2, 3]);
428 /// ```
429 pub fn new(func: F, local_dims: &[usize]) -> Result<Self, error::CacheKeyError> {
430 Ok(Self {
431 func,
432 batch_func: None,
433 cache: CacheBackend::Auto(InnerCache::new(local_dims)?),
434 local_dims: local_dims.to_vec(),
435 num_evals: AtomicUsize::new(0),
436 num_cache_hits: AtomicUsize::new(0),
437 _phantom: std::marker::PhantomData,
438 })
439 }
440
441 /// Create with a batch function for efficient multi-point evaluation.
442 ///
443 /// The batch function is used for cache misses during [`eval_batch`](Self::eval_batch)
444 /// calls, enabling amortized cost when evaluating many indices at once
445 /// (e.g., batch FFI calls or vectorized computations).
446 ///
447 /// # Errors
448 ///
449 /// Returns an error when the batch configuration is invalid (an
450 /// /// invalid-configuration failure).
451 ///
452 /// # Examples
453 ///
454 /// ```
455 /// use tensor4all_core::CachedFunction;
456 ///
457 /// let cf = CachedFunction::with_batch(
458 /// |idx: &[usize]| idx[0] * 10 + idx[1],
459 /// |indices: &[Vec<usize>]| indices.iter().map(|idx| idx[0] * 10 + idx[1]).collect(),
460 /// &[3, 4],
461 /// ).unwrap();
462 ///
463 /// let results = cf.eval_batch(&[vec![0, 1], vec![2, 3]]).unwrap();
464 /// assert_eq!(results, vec![1, 23]);
465 /// assert_eq!(cf.num_evals(), 2);
466 /// ```
467 pub fn with_batch<B>(
468 func: F,
469 batch_func: B,
470 local_dims: &[usize],
471 ) -> Result<Self, error::CacheKeyError>
472 where
473 B: Fn(&[Vec<I>]) -> Vec<V> + Send + Sync + 'static,
474 {
475 Ok(Self {
476 func,
477 batch_func: Some(Box::new(batch_func)),
478 cache: CacheBackend::Auto(InnerCache::new(local_dims)?),
479 local_dims: local_dims.to_vec(),
480 num_evals: AtomicUsize::new(0),
481 num_cache_hits: AtomicUsize::new(0),
482 _phantom: std::marker::PhantomData,
483 })
484 }
485
486 /// Create with an explicit key type for index spaces larger than 1024 bits.
487 ///
488 /// # Errors
489 ///
490 /// Returns an error when the key type configuration is invalid (an
491 /// /// invalid-configuration failure).
492 ///
493 /// # Example
494 ///
495 /// ```
496 /// use bnum::types::U2048;
497 /// use tensor4all_core::{CacheKey, CachedFunction};
498 ///
499 /// #[derive(Clone, Hash, PartialEq, Eq)]
500 /// struct U2048Key(U2048);
501 ///
502 /// impl CacheKey for U2048Key {
503 /// const BITS_COUNT: u32 = 2048;
504 /// const ZERO: Self = Self(U2048::ZERO);
505 /// const ONE: Self = Self(U2048::ONE);
506 ///
507 /// fn from_usize(v: usize) -> Self {
508 /// Self(U2048::from(v as u64))
509 /// }
510 ///
511 /// fn checked_mul(self, rhs: Self) -> Option<Self> {
512 /// self.0.checked_mul(rhs.0).map(Self)
513 /// }
514 ///
515 /// fn wrapping_add(self, rhs: Self) -> Self {
516 /// Self(self.0.wrapping_add(rhs.0))
517 /// }
518 /// }
519 ///
520 /// let local_dims = vec![2usize; 1025];
521 /// let cf = CachedFunction::with_key_type::<U2048Key>(
522 /// |idx: &[usize]| idx.iter().sum::<usize>(),
523 /// &local_dims,
524 /// ).unwrap();
525 /// let zeros = vec![0usize; 1025];
526 ///
527 /// assert_eq!(cf.eval(&zeros).unwrap(), 0);
528 /// assert_eq!(cf.key_type(), "custom");
529 /// ```
530 pub fn with_key_type<K: CacheKey>(
531 func: F,
532 local_dims: &[usize],
533 ) -> Result<Self, error::CacheKeyError> {
534 Ok(Self {
535 func,
536 batch_func: None,
537 cache: CacheBackend::Custom(Box::new(GenericCache::<K, V>::new(local_dims)?)),
538 local_dims: local_dims.to_vec(),
539 num_evals: AtomicUsize::new(0),
540 num_cache_hits: AtomicUsize::new(0),
541 _phantom: std::marker::PhantomData,
542 })
543 }
544
545 /// Create with explicit key type and batch function.
546 ///
547 /// Combines [`with_key_type`](Self::with_key_type) and
548 /// [`with_batch`](Self::with_batch) for index spaces larger than 1024
549 /// bits that also benefit from batch evaluation.
550 ///
551 /// # Errors
552 ///
553 /// Returns an error when the configuration is invalid (an
554 /// /// invalid-configuration failure).
555 ///
556 /// # Examples
557 ///
558 /// ```
559 /// use tensor4all_core::CachedFunction;
560 ///
561 /// // Use u128 key type with batch support
562 /// let cf = CachedFunction::with_key_type_and_batch::<u128, _>(
563 /// |idx: &[usize]| idx.iter().sum::<usize>(),
564 /// |indices: &[Vec<usize>]| indices.iter().map(|idx| idx.iter().sum()).collect(),
565 /// &[2, 3, 4],
566 /// ).unwrap();
567 ///
568 /// let results = cf.eval_batch(&[vec![0, 0, 0], vec![1, 2, 3]]).unwrap();
569 /// assert_eq!(results, vec![0, 6]);
570 /// ```
571 pub fn with_key_type_and_batch<K: CacheKey, B>(
572 func: F,
573 batch_func: B,
574 local_dims: &[usize],
575 ) -> Result<Self, error::CacheKeyError>
576 where
577 B: Fn(&[Vec<I>]) -> Vec<V> + Send + Sync + 'static,
578 {
579 Ok(Self {
580 func,
581 batch_func: Some(Box::new(batch_func)),
582 cache: CacheBackend::Custom(Box::new(GenericCache::<K, V>::new(local_dims)?)),
583 local_dims: local_dims.to_vec(),
584 num_evals: AtomicUsize::new(0),
585 num_cache_hits: AtomicUsize::new(0),
586 _phantom: std::marker::PhantomData,
587 })
588 }
589
590 /// Evaluate at a given index, using cache if available.
591 ///
592 /// On the first call for a given index, the wrapped function is invoked and
593 /// the result is cached. Subsequent calls with the same index return the
594 /// cached value. This method is thread-safe.
595 ///
596 /// # Errors
597 /// Returns [`error::CacheKeyError::InvalidIndexLength`] for a wrong-rank index or
598 /// [`error::CacheKeyError::IndexOutOfBounds`] for an invalid coordinate.
599 ///
600 /// # Examples
601 ///
602 /// ```
603 /// use tensor4all_core::CachedFunction;
604 ///
605 /// let cf = CachedFunction::new(|idx: &[usize]| idx[0] * idx[1], &[5, 5]).unwrap();
606 /// assert_eq!(cf.eval(&[3, 4]).unwrap(), 12);
607 /// assert_eq!(cf.num_evals(), 1);
608 ///
609 /// // Cache hit
610 /// assert_eq!(cf.eval(&[3, 4]).unwrap(), 12);
611 /// assert_eq!(cf.num_evals(), 1);
612 /// assert_eq!(cf.num_cache_hits(), 1);
613 /// ```
614 pub fn eval(&self, idx: &[I]) -> Result<V, error::CacheKeyError> {
615 self.validate_index(idx)?;
616 if let Some(value) = self.cache.get(idx) {
617 self.num_cache_hits.fetch_add(1, Ordering::Relaxed);
618 return Ok(value);
619 }
620
621 self.num_evals.fetch_add(1, Ordering::Relaxed);
622 let value = (self.func)(idx);
623 self.cache.insert(idx, value.clone());
624 Ok(value)
625 }
626
627 /// Evaluate bypassing the cache.
628 ///
629 /// The result is neither read from nor stored in the cache, and
630 /// evaluation counters are not updated. Useful for verification or
631 /// when the caller intentionally wants a fresh evaluation.
632 ///
633 /// # Errors
634 /// Returns [`error::CacheKeyError::InvalidIndexLength`] for a wrong-rank index or
635 /// [`error::CacheKeyError::IndexOutOfBounds`] for an invalid coordinate.
636 ///
637 /// # Examples
638 ///
639 /// ```
640 /// use tensor4all_core::CachedFunction;
641 ///
642 /// let cf = CachedFunction::new(|idx: &[usize]| idx[0] + 1, &[4]).unwrap();
643 /// assert_eq!(cf.eval_no_cache(&[2]).unwrap(), 3);
644 /// assert_eq!(cf.cache_size(), 0);
645 /// assert_eq!(cf.num_evals(), 0);
646 /// ```
647 pub fn eval_no_cache(&self, idx: &[I]) -> Result<V, error::CacheKeyError> {
648 self.validate_index(idx)?;
649 Ok((self.func)(idx))
650 }
651
652 /// Evaluate at multiple indices. Uses batch function for cache misses if available.
653 ///
654 /// Returns results in the same order as the input indices.
655 ///
656 /// # Errors
657 /// Returns an index validation error for malformed coordinates or
658 /// [`error::CacheKeyError::BatchResultLength`] when a batch callback returns the
659 /// wrong number of values.
660 ///
661 /// # Examples
662 ///
663 /// ```
664 /// use tensor4all_core::CachedFunction;
665 ///
666 /// let cf = CachedFunction::new(|idx: &[usize]| idx[0] * 2 + idx[1], &[2, 2]).unwrap();
667 /// let results = cf.eval_batch(&[vec![0, 0], vec![0, 1], vec![1, 0]]).unwrap();
668 /// assert_eq!(results, vec![0, 1, 2]);
669 /// ```
670 pub fn eval_batch(&self, indices: &[Vec<I>]) -> Result<Vec<V>, error::CacheKeyError> {
671 if indices.is_empty() {
672 return Ok(Vec::new());
673 }
674 for index in indices {
675 self.validate_index(index)?;
676 }
677
678 let mut results: Vec<Option<V>> = Vec::with_capacity(indices.len());
679 let mut miss_positions: Vec<usize> = Vec::new();
680 let mut miss_indices: Vec<Vec<I>> = Vec::new();
681
682 for (pos, idx) in indices.iter().enumerate() {
683 if let Some(value) = self.cache.get(idx) {
684 self.num_cache_hits.fetch_add(1, Ordering::Relaxed);
685 results.push(Some(value));
686 } else {
687 results.push(None);
688 miss_positions.push(pos);
689 miss_indices.push(idx.clone());
690 }
691 }
692
693 if miss_indices.is_empty() {
694 return Ok(results.into_iter().flatten().collect());
695 }
696
697 self.num_evals
698 .fetch_add(miss_indices.len(), Ordering::Relaxed);
699 let miss_values = if let Some(batch_func) = self.batch_func.as_ref() {
700 batch_func(&miss_indices)
701 } else {
702 miss_indices.iter().map(|idx| (self.func)(idx)).collect()
703 };
704 if miss_values.len() != miss_indices.len() {
705 return Err(error::CacheKeyError::BatchResultLength {
706 expected: miss_indices.len(),
707 got: miss_values.len(),
708 });
709 }
710
711 for (i, pos) in miss_positions.iter().enumerate() {
712 self.cache.insert(&miss_indices[i], miss_values[i].clone());
713 results[*pos] = Some(miss_values[i].clone());
714 }
715
716 Ok(results.into_iter().flatten().collect())
717 }
718
719 fn validate_index(&self, idx: &[I]) -> std::result::Result<(), error::CacheKeyError> {
720 if idx.len() != self.local_dims.len() {
721 return Err(error::CacheKeyError::InvalidIndexLength {
722 expected: self.local_dims.len(),
723 got: idx.len(),
724 });
725 }
726 for (axis, (&value, &dim)) in idx.iter().zip(&self.local_dims).enumerate() {
727 let value = value.to_usize();
728 if value >= dim {
729 return Err(error::CacheKeyError::IndexOutOfBounds {
730 axis,
731 index: value,
732 dim,
733 });
734 }
735 }
736 Ok(())
737 }
738
739 /// Get the local dimensions.
740 ///
741 /// # Examples
742 ///
743 /// ```
744 /// use tensor4all_core::CachedFunction;
745 ///
746 /// let cf = CachedFunction::new(|idx: &[usize]| 0, &[3, 4, 5]).unwrap();
747 /// assert_eq!(cf.local_dims(), &[3, 4, 5]);
748 /// ```
749 pub fn local_dims(&self) -> &[usize] {
750 &self.local_dims
751 }
752
753 /// Get the number of sites (length of the multi-index).
754 ///
755 /// # Examples
756 ///
757 /// ```
758 /// use tensor4all_core::CachedFunction;
759 ///
760 /// let cf = CachedFunction::new(|idx: &[usize]| 0, &[2, 3]).unwrap();
761 /// assert_eq!(cf.num_sites(), 2);
762 /// ```
763 pub fn num_sites(&self) -> usize {
764 self.local_dims.len()
765 }
766
767 /// Get the number of function evaluations (cache misses).
768 ///
769 /// # Examples
770 ///
771 /// ```
772 /// use tensor4all_core::CachedFunction;
773 ///
774 /// let cf = CachedFunction::new(|idx: &[usize]| idx[0], &[4]).unwrap();
775 /// cf.eval(&[0]).unwrap();
776 /// cf.eval(&[1]).unwrap();
777 /// cf.eval(&[0]).unwrap(); // cache hit, not a new eval
778 /// assert_eq!(cf.num_evals(), 2);
779 /// ```
780 pub fn num_evals(&self) -> usize {
781 self.num_evals.load(Ordering::Relaxed)
782 }
783
784 /// Get the number of cache hits.
785 ///
786 /// # Examples
787 ///
788 /// ```
789 /// use tensor4all_core::CachedFunction;
790 ///
791 /// let cf = CachedFunction::new(|idx: &[usize]| idx[0], &[4]).unwrap();
792 /// cf.eval(&[0]).unwrap();
793 /// assert_eq!(cf.num_cache_hits(), 0);
794 /// cf.eval(&[0]).unwrap();
795 /// assert_eq!(cf.num_cache_hits(), 1);
796 /// ```
797 pub fn num_cache_hits(&self) -> usize {
798 self.num_cache_hits.load(Ordering::Relaxed)
799 }
800
801 /// Get total calls (evaluations + cache hits).
802 ///
803 /// # Examples
804 ///
805 /// ```
806 /// use tensor4all_core::CachedFunction;
807 ///
808 /// let cf = CachedFunction::new(|idx: &[usize]| idx[0], &[4]).unwrap();
809 /// cf.eval(&[0]).unwrap();
810 /// cf.eval(&[1]).unwrap();
811 /// cf.eval(&[0]).unwrap(); // cache hit
812 /// assert_eq!(cf.total_calls(), 3);
813 /// assert_eq!(cf.total_calls(), cf.num_evals() + cf.num_cache_hits());
814 /// ```
815 pub fn total_calls(&self) -> usize {
816 self.num_evals() + self.num_cache_hits()
817 }
818
819 /// Get cache hit ratio (0.0 when no calls have been made).
820 ///
821 /// Returns `num_cache_hits() / total_calls()` as a value in `[0.0, 1.0]`.
822 ///
823 /// # Examples
824 ///
825 /// ```
826 /// use tensor4all_core::CachedFunction;
827 ///
828 /// let cf = CachedFunction::new(|idx: &[usize]| idx[0], &[4]).unwrap();
829 /// assert_eq!(cf.cache_hit_ratio(), 0.0); // no calls yet
830 ///
831 /// cf.eval(&[0]).unwrap();
832 /// cf.eval(&[0]).unwrap(); // cache hit
833 /// assert!((cf.cache_hit_ratio() - 0.5).abs() < 1e-10);
834 /// ```
835 pub fn cache_hit_ratio(&self) -> f64 {
836 let total = self.total_calls();
837 if total == 0 {
838 0.0
839 } else {
840 self.num_cache_hits() as f64 / total as f64
841 }
842 }
843
844 /// Clear the cache.
845 ///
846 /// # Examples
847 ///
848 /// ```
849 /// use tensor4all_core::CachedFunction;
850 ///
851 /// let cf = CachedFunction::new(|idx: &[usize]| idx[0], &[4]).unwrap();
852 /// cf.eval(&[2]).unwrap();
853 /// assert_eq!(cf.cache_size(), 1);
854 /// cf.clear_cache();
855 /// assert_eq!(cf.cache_size(), 0);
856 /// ```
857 pub fn clear_cache(&self) {
858 self.cache.clear();
859 }
860
861 /// Number of cached entries.
862 ///
863 /// # Examples
864 ///
865 /// ```
866 /// use tensor4all_core::CachedFunction;
867 ///
868 /// let cf = CachedFunction::new(|idx: &[usize]| idx[0], &[4]).unwrap();
869 /// assert_eq!(cf.cache_size(), 0);
870 /// cf.eval(&[0]).unwrap();
871 /// cf.eval(&[1]).unwrap();
872 /// assert_eq!(cf.cache_size(), 2);
873 /// cf.eval(&[0]).unwrap(); // cache hit, no new entry
874 /// assert_eq!(cf.cache_size(), 2);
875 /// ```
876 pub fn cache_size(&self) -> usize {
877 self.cache.len()
878 }
879
880 /// Check if an index is cached.
881 ///
882 /// # Examples
883 ///
884 /// ```
885 /// use tensor4all_core::CachedFunction;
886 ///
887 /// let cf = CachedFunction::new(|idx: &[usize]| idx[0], &[4]).unwrap();
888 /// assert!(!cf.is_cached(&[1]));
889 /// cf.eval(&[1]).unwrap();
890 /// assert!(cf.is_cached(&[1]));
891 /// ```
892 pub fn is_cached(&self, idx: &[I]) -> bool {
893 self.cache.contains(idx)
894 }
895
896 /// Internal key type name (for debugging).
897 ///
898 /// Returns `"u64"`, `"u128"`, `"U256"`, `"U512"`, `"U1024"` for
899 /// automatically selected types, or `"custom"` when constructed with
900 /// [`with_key_type`](Self::with_key_type).
901 ///
902 /// # Examples
903 ///
904 /// ```
905 /// use tensor4all_core::CachedFunction;
906 ///
907 /// // Small index space uses u64
908 /// let cf = CachedFunction::new(|idx: &[usize]| 0, &[2, 3]).unwrap();
909 /// assert_eq!(cf.key_type(), "u64");
910 /// ```
911 pub fn key_type(&self) -> &'static str {
912 self.cache.key_type_name()
913 }
914}
915
916#[cfg(test)]
917mod tests;