use crate::{Key, KeyValue, Value};
#[cfg(feature = "serialize")]
use serde::{Deserialize, Serialize};
use std::collections::hash_map::Entry;
use std::collections::{HashMap, LinkedList};
#[cfg_attr(feature = "serialize", derive(Deserialize, Serialize))]
#[derive(Clone, Debug, PartialEq)]
pub struct EvictedHashMap {
map: HashMap<Key, Value>,
evict_list: LinkedList<Key>,
max_len: u32,
dropped_count: u32,
}
impl EvictedHashMap {
pub fn new(max_len: u32, capacity: usize) -> Self {
EvictedHashMap {
map: HashMap::with_capacity(capacity),
evict_list: LinkedList::new(),
max_len,
dropped_count: 0,
}
}
pub fn insert(&mut self, item: KeyValue) {
let KeyValue { key, value } = item;
let mut already_exists = false;
match self.map.entry(key.clone()) {
Entry::Occupied(mut occupied) => {
occupied.insert(value);
already_exists = true;
}
Entry::Vacant(entry) => {
entry.insert(value);
}
}
if already_exists {
self.move_key_to_front(key);
} else {
self.evict_list.push_front(key);
}
if self.evict_list.len() as u32 > self.max_len {
self.remove_oldest();
self.dropped_count += 1;
}
}
pub fn len(&self) -> usize {
self.map.len()
}
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
pub fn dropped_count(&self) -> u32 {
self.dropped_count
}
pub fn iter(&self) -> Iter<'_> {
Iter(self.map.iter())
}
pub fn get(&self, key: &Key) -> Option<&Value> {
self.map.get(key)
}
fn move_key_to_front(&mut self, key: Key) {
if self.evict_list.is_empty() {
self.evict_list.push_front(key);
} else if self.evict_list.front() == Some(&key) {
} else {
let key_idx = self
.evict_list
.iter()
.position(|k| k == &key)
.expect("key must exist in evicted hash map, this is a bug");
let mut tail = self.evict_list.split_off(key_idx);
let item = tail.pop_front().unwrap();
self.evict_list.push_front(item);
self.evict_list.append(&mut tail);
}
}
fn remove_oldest(&mut self) {
if let Some(oldest_item) = self.evict_list.pop_back() {
self.map.remove(&oldest_item);
}
}
}
#[derive(Debug)]
pub struct IntoIter(std::collections::hash_map::IntoIter<Key, Value>);
impl Iterator for IntoIter {
type Item = (Key, Value);
fn next(&mut self) -> Option<Self::Item> {
self.0.next()
}
}
impl IntoIterator for EvictedHashMap {
type Item = (Key, Value);
type IntoIter = IntoIter;
fn into_iter(self) -> Self::IntoIter {
IntoIter(self.map.into_iter())
}
}
impl<'a> IntoIterator for &'a EvictedHashMap {
type Item = (&'a Key, &'a Value);
type IntoIter = Iter<'a>;
fn into_iter(self) -> Self::IntoIter {
Iter(self.map.iter())
}
}
#[derive(Debug)]
pub struct Iter<'a>(std::collections::hash_map::Iter<'a, Key, Value>);
impl<'a> Iterator for Iter<'a> {
type Item = (&'a Key, &'a Value);
fn next(&mut self) -> Option<Self::Item> {
self.0.next()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
#[test]
fn insert_over_capacity_test() {
let max_len = 10;
let mut map = EvictedHashMap::new(max_len, max_len as usize);
for i in 0..=max_len {
map.insert(Key::new(i.to_string()).bool(true))
}
assert_eq!(map.dropped_count, 1);
assert_eq!(map.len(), max_len as usize);
assert_eq!(
map.map.keys().cloned().collect::<HashSet<_>>(),
(1..=max_len)
.map(|i| Key::new(i.to_string()))
.collect::<HashSet<_>>()
);
}
}