use crate::cache::dataset::{CachedDataset, WarmingStrategy};
use crate::Dataset;
use tenflowers_core::Result;
pub trait CacheExt<T>: Dataset<T> + Sized {
fn cached(self, capacity: usize) -> CachedDataset<T, Self>
where
T: Clone + Send + Sync + 'static,
{
CachedDataset::new(self, capacity)
}
fn cached_with_warming(
self,
capacity: usize,
strategy: WarmingStrategy,
) -> Result<CachedDataset<T, Self>>
where
T: Clone + Send + Sync + 'static,
{
let cached = CachedDataset::new(self, capacity);
let indices = strategy.generate_indices(cached.len());
cached.warm_cache(&indices)?;
Ok(cached)
}
}
impl<T, D: Dataset<T>> CacheExt<T> for D {}