use core::marker::Unpin;
use std::hash::Hash;
use futures::stream::{TryStream, TryStreamExt};
use linked_hash_map::LinkedHashMap;
pub struct Cache<K> {
map: LinkedHashMap<K, u64>,
size: u64,
}
impl<K> Cache<K>
where
K: Eq + Hash,
{
pub fn new() -> Self {
Cache {
map: LinkedHashMap::new(),
size: 0,
}
}
pub fn size(&self) -> u64 {
self.size
}
pub fn len(&self) -> usize {
self.map.len()
}
pub async fn from_stream<S>(mut stream: S) -> Result<Self, S::Error>
where
S: TryStream<Ok = (K, u64)> + Unpin,
{
let mut cache = Cache::new();
while let Some((oid, len)) = stream.try_next().await? {
cache.push(oid, len);
}
Ok(cache)
}
pub fn pop(&mut self) -> Option<(K, u64)> {
if let Some((k, v)) = self.map.pop_front() {
self.size -= v;
Some((k, v))
} else {
None
}
}
pub fn remove(&mut self, key: &K) -> Option<u64> {
self.map.remove(key).inspect(|size| {
self.size -= size;
})
}
pub fn get(&self, key: &K) -> Option<u64> {
self.map.get(key).cloned()
}
pub fn get_refresh(&mut self, key: &K) -> Option<u64> {
self.map.get_refresh(key).cloned()
}
pub fn push(&mut self, key: K, value: u64) -> Option<u64> {
self.size += value;
if let Some(old_value) = self.map.insert(key, value) {
self.size -= old_value;
Some(old_value)
} else {
None
}
}
}