use lru::LruCache;
use std::hash::Hash;
use std::num::NonZeroUsize;
use std::sync::Mutex;
use super::Chunker;
use crate::UPOS;
pub struct CachedChunker<C: Chunker> {
inner: C,
cache: Mutex<LruCache<CacheKey, Vec<bool>>>,
}
impl<C: Chunker> CachedChunker<C> {
pub fn new(inner: C, capacity: NonZeroUsize) -> Self {
Self {
inner,
cache: Mutex::new(LruCache::new(capacity)),
}
}
}
impl<C: Chunker> Chunker for CachedChunker<C> {
fn chunk_sentence(&self, sentence: &[String], tags: &[Option<UPOS>]) -> Vec<bool> {
let key = CacheKey::new(sentence, tags);
if let Ok(mut read) = self.cache.try_lock()
&& let Some(result) = read.get(&key)
{
return result.clone();
};
let result = self.inner.chunk_sentence(sentence, tags);
if let Ok(mut cache) = self.cache.try_lock() {
cache.put(key, result.clone());
}
result
}
}
#[derive(Hash, PartialEq, Eq)]
struct CacheKey {
sentence: Vec<String>,
tags: Vec<Option<UPOS>>,
}
impl CacheKey {
fn new(sentence: &[String], tags: &[Option<UPOS>]) -> Self {
Self {
sentence: sentence.to_vec(),
tags: tags.to_vec(),
}
}
}