evicting_cache_map 0.6.0

use ordered_hash_map::ordered_map::DefaultHashBuilder;
use ordered_hash_map::OrderedHashMap;

use core::borrow::Borrow;
use core::fmt;
use core::hash::Hash;
use core::ops::Fn;

mod iter;
pub use iter::{Iter, IterMut};

/// An evicting cache map.
///
/// Backed by a HashMap but maintains order and a maximum capacity. Upon inserting beyond capacity
/// the cache map will remove the eldest elements to make room. Each removal or prune will call the
/// provided prune hook with the values.
///
/// Being backed by a HashMap, this container has all the same requirements on its keys as
/// [`HashMap`]. Specifically, Keys must implement [`Eq`] and [`Hash`]. You can typically derive
/// these for types with `#[derive(PartialEq, Eq, hash)]`. If you implement these yourself you must
/// ensure the following holds:
///
/// ```text
/// k1 == k2 -> hash(k1) == hash(k2)
/// ```
/// You are also responsible to ensure Keys do not mutate such that their hash can change while in
/// the map. For more information, check [`HashMap`].
///
/// [`AHash`]: https://crates.io/crates/ahash
/// [`hashbrown`]: https://crates.io/crates/hashbrown
/// [`with_hasher`]: Self::with_hasher
/// [`with_capacity_and_hasher`]: Self::with_capacity_and_hasher
/// [`HashMap`]: std::collections::HashMap
/// [`RandomState`]: std::collections::hash_map::RandomState
pub struct EvictingCacheMap<K, V, const N: usize, F, S = DefaultHashBuilder>
where
    F: Fn(K, V),
{
    inner: OrderedHashMap<K, V, S>,
    func: F,
}

impl<K, V, const N: usize, F, S> fmt::Debug for EvictingCacheMap<K, V, N, F, S>
where
    K: fmt::Debug,
    V: fmt::Debug,
    F: Fn(K, V),
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.inner.fmt(f)
    }
}

impl<K, V, const N: usize> Default
    for EvictingCacheMap<K, V, N, fn(K, V) -> (), DefaultHashBuilder>
{
    fn default() -> Self {
        Self::new()
    }
}

impl<K, V, const N: usize> EvictingCacheMap<K, V, N, fn(K, V) -> (), DefaultHashBuilder> {
    /// Creates a new map
    ///
    /// Default prune hook simply drops K and V.
    pub fn new() -> Self {
        Self::with_prune_hook(|_, _| ())
    }
}

impl<K, V, const N: usize, F> EvictingCacheMap<K, V, N, F, DefaultHashBuilder>
where
    F: Fn(K, V),
{
    /// Create a new map with the provided prune hook
    pub fn with_prune_hook(func: F) -> Self {
        Self {
            inner: OrderedHashMap::with_capacity(N),
            func,
        }
    }
}

impl<K, V, const N: usize, S> EvictingCacheMap<K, V, N, fn(K, V) -> (), S> {
    /// Create a new map with provided hasher
    pub fn with_hasher(hash_builder: S) -> Self {
        Self::with_hasher_and_prune_hook(hash_builder, |_, _| ())
    }
}

impl<K, V, const N: usize, F, S> EvictingCacheMap<K, V, N, F, S>
where
    F: Fn(K, V),
{
    /// Create a new map with the provided hasher and prune hook
    pub fn with_hasher_and_prune_hook(hash_builder: S, func: F) -> Self {
        Self {
            inner: OrderedHashMap::with_capacity_and_hasher(N, hash_builder),
            func,
        }
    }
}

impl<K, V, const N: usize, F> EvictingCacheMap<K, V, N, F>
where
    K: Eq + Hash,
    F: Fn(K, V),
{
    #[inline(always)]
    fn prune_impl(&mut self) {
        if let Some((k, v)) = self.inner.pop_front_entry() {
            (self.func)(k, v);
        }
    }

    #[inline(always)]
    fn prune_with_impl(&mut self, f: &dyn Fn(K, V)) {
        if let Some((k, v)) = self.inner.pop_front_entry() {
            (f)(k, v);
        }
    }

    /// Prune eldest entry
    #[inline]
    fn prune_one(&mut self) {
        if self.inner.len() == N {
            self.prune_impl();
        }
    }

    /// Prune a key
    ///
    /// If a prune needs to happen it will prune the given key if it exists, otherwise prunes
    /// eldest entry.
    #[inline]
    fn prune_one_key(&mut self, key: &K) {
        if self.inner.len() == N {
            if !self.inner.contains_key(key) {
                self.prune_impl();
            } else if let Some((k, v)) = self.inner.remove_entry(key) {
                (self.func)(k, v);
            }
        }
    }

    /// Remove elements and call prune hook
    ///
    /// Removes `len` elements from the cache and calls the prune hook for each. If `len` is
    /// greater than the len() of the cache it will leave it empty.
    #[inline]
    pub fn prune_by(&mut self, len: usize) {
        for _ in 0..if len < self.len() { len } else { self.len() } {
            self.prune_impl();
        }
    }

    /// Remove elements and calls provided prune hook
    ///
    /// Removes `len` elements from the cache and calls the prune hook for each. If `len` is
    /// greater than the len() of th cache it will leave it empty.
    #[inline]
    pub fn prune_by_with(&mut self, len: usize, f: impl Fn(K, V)) {
        for _ in 0..if len < self.len() { len } else { self.len() } {
            self.prune_with_impl(&f);
        }
    }

    /// Remove elements and call the prune hook
    ///
    /// Removes elements until the cache is the provided size.
    #[inline]
    pub fn prune_to(&mut self, len: usize) {
        self.prune_by(self.len().saturating_sub(len));
    }

    /// Removeelements and call the provided prune hook
    ///
    /// Removes elements until the cache is the provided size.
    #[inline]
    pub fn prune_to_with(&mut self, len: usize, f: impl Fn(K, V)) {
        self.prune_by_with(self.len().saturating_sub(len), f);
    }

    /// Promote the element if it exists
    ///
    /// This moves an element to the top of the cache as-if removed and re-inserted but does not
    /// call the prune hook.
    #[inline]
    pub fn promote<Q>(&mut self, key: &Q)
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.inner.move_to_back(key);
    }

    /// Checks if cache contains this key
    ///
    /// Does not promote the key.
    pub fn contains_key<Q>(&self, key: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.inner.contains_key(key)
    }

    /// Insert element to the top of the cache
    ///
    /// If the container is full this will prune a single element, calling the prune hook.
    #[inline]
    pub fn insert(&mut self, key: K, value: V) {
        self.prune_one_key(&key);
        self.inner.insert(key, value);
    }

    /// Get a reference to the corresponding element and promotes it to the top of the cache.
    #[inline]
    pub fn get<Q>(&mut self, key: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.promote(key);
        self.get_no_promote(key)
    }

    /// Get a reference to the corresponding element without promoting it.
    #[inline]
    pub fn get_no_promote<Q>(&self, key: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.inner.get(key)
    }

    /// Get a mutable reference to the corresponding element and promote it to the top of the
    /// cache.
    #[inline]
    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.promote(key);
        self.get_mut_no_promote(key)
    }

    /// Get a mutable reference to the corresponding element without promoting it.
    #[inline]
    pub fn get_mut_no_promote<Q>(&mut self, key: &Q) -> Option<&mut V>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.inner.get_mut(key)
    }

    /// Get a reference to the key-value pair and promote it to the top of the cache.
    #[inline]
    pub fn get_key_value<Q>(&mut self, key: &Q) -> Option<(&K, &V)>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.promote(key);
        self.get_key_value_no_promote(key)
    }

    /// Get a reference to the key-value pair without promoting it.
    #[inline]
    pub fn get_key_value_no_promote<Q>(&self, key: &Q) -> Option<(&K, &V)>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.inner.get_key_value(key)
    }

    /// Get a mutable reference to the key-value pair and promote it to the top of the cache.
    #[inline]
    pub fn get_key_value_mut<Q>(&mut self, key: &Q) -> Option<(&K, &mut V)>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.promote(key);
        self.get_key_value_mut_no_promote(key)
    }

    /// Get a mutable reference to the key-value pair without promoting it.
    #[inline]
    pub fn get_key_value_mut_no_promote<Q>(&mut self, key: &Q) -> Option<(&K, &mut V)>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.inner.get_key_value_mut(key)
    }

    /// Remove and return a value from the map.
    ///
    /// Does not call the prune hook.
    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.inner.remove(key)
    }

    /// Remove and return a key-value pair from the map.
    ///
    /// Does not call the prune hook.
    pub fn remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V)>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.inner.remove_entry(key)
    }

    /// Retrieve the current number of elements in the cache.
    pub fn len(&self) -> usize {
        self.inner.len()
    }

    /// Check whether the cache is empty.
    pub fn is_empty(&self) -> bool {
        self.inner.is_empty()
    }

    /// Clear all elements from the map and call the prune hook for each element.
    pub fn clear(&mut self) {
        self.inner.drain().for_each(|(k, v)| (self.func)(k, v));
    }

    /// Iterator over the values in the map
    pub fn iter(&self) -> Iter<'_, K, V> {
        self.into_iter()
    }

    /// Mutable Iterator for the values in the map
    ///
    /// Note: mutating a value through this iterator does not promote keys.
    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
        self.into_iter()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::mpsc;

    #[test]
    fn basic_operation() {
        let mut map: EvictingCacheMap<String, u8, 4, _> = EvictingCacheMap::default();
        map.insert("One".to_string(), 1);
        map.insert("Two".to_string(), 2);
        map.insert("Three".to_string(), 3);
        map.insert("Four".to_string(), 4);
        map.insert("Five".to_string(), 5);

        assert!(!map.contains_key("One"));
        assert!(map.contains_key("Two"));
        assert!(map.contains_key("Three"));
        assert!(map.contains_key("Four"));
        assert!(map.contains_key("Five"));
    }

    #[test]
    fn promote() {
        let (tx, rx) = mpsc::channel();
        let mut map: EvictingCacheMap<String, u8, 2, _> =
            EvictingCacheMap::with_prune_hook(|k, v| tx.send((k, v)).unwrap());
        map.insert("One".into(), 1);
        map.insert("Two".into(), 2);
        map.promote("One");
        map.insert("Three".into(), 3);

        assert!(map.contains_key("One"));
        assert!(!map.contains_key("Two"));
        assert!(map.contains_key("Three"));

        drop(map);
        drop(tx); // Drop senders so rx closes
        assert_eq!(rx.recv(), Ok(("Two".into(), 2)));
        assert!(rx.recv().is_err());
    }

    #[test]
    fn clear() {
        let (tx, rx) = mpsc::channel();
        let mut map: EvictingCacheMap<String, u8, 2, _> =
            EvictingCacheMap::with_prune_hook(|k, v| tx.send((k, v)).unwrap());
        map.insert("A".into(), 1);
        map.clear();
        assert_eq!(rx.recv(), Ok(("A".into(), 1)));
    }

    #[test]
    fn prune() {
        let (tx, rx) = mpsc::channel();
        let mut map: EvictingCacheMap<String, u8, 4, _> =
            EvictingCacheMap::with_prune_hook(|k, v| tx.send((k, v)).unwrap());
        map.insert("A".into(), 1);
        map.insert("B".into(), 2);
        map.insert("C".into(), 3);
        map.insert("D".into(), 4);
        map.prune_to(3);
        assert_eq!(rx.recv(), Ok(("A".into(), 1)));
        assert_eq!(rx.try_recv(), Err(mpsc::TryRecvError::Empty));
        map.prune_by(1);
        assert_eq!(rx.recv(), Ok(("B".into(), 2)));
        assert_eq!(rx.try_recv(), Err(mpsc::TryRecvError::Empty));
    }

    #[test]
    fn prune_with() {
        let (tx, rx) = mpsc::channel();
        let mut map: EvictingCacheMap<String, u8, 4, _> =
            EvictingCacheMap::with_prune_hook(|k, v| tx.send((k, v)).unwrap());

        map.insert("A".into(), 1);
        map.insert("B".into(), 2);
        map.insert("C".into(), 3);
        map.insert("D".into(), 4);
        map.prune_to_with(3, |k, v| assert_eq!((k, v), ("A".into(), 1)));
        assert!(rx.try_recv().is_err()); // We should have bypassed the registered prune_hook
        map.prune_by_with(1, |k, v| assert_eq!((k, v), ("B".into(), 2)));
        assert!(rx.try_recv().is_err());
    }

    #[test]
    fn get_ops() {
        let (tx, rx) = mpsc::channel();
        let mut map: EvictingCacheMap<i32, i32, 4, _> =
            EvictingCacheMap::with_prune_hook(|k, v| tx.send((k, v)).unwrap());
        map.insert(1, 1);
        map.insert(2, 2);
        map.insert(3, 3);
        map.insert(4, 4);

        assert!(!map.is_empty());

        map.get(&1);
        map.prune_by(1);
        assert_eq!(rx.recv(), Ok((2, 2)));

        map.get_no_promote(&3);
        map.prune_by(1);
        assert_eq!(rx.recv(), Ok((3, 3)));

        if let Some(v) = map.get_mut(&4) {
            *v += 1i32;
        };
        map.prune_by(2);
        assert_eq!(rx.recv(), Ok((1, 1)));
        assert_eq!(rx.recv(), Ok((4, 5)));

        assert!(map.is_empty());
    }

    #[test]
    fn iterators() {
        let mut map: EvictingCacheMap<i32, i32, 4, _> = EvictingCacheMap::new();

        map.insert(1, 1);
        map.insert(2, 2);
        map.insert(3, 3);
        map.insert(4, 4);

        for (_, v) in map.iter_mut() {
            *v += 1;
        }
        map.iter().for_each(|(k, v)| assert_eq!(k + 1, *v));
    }
}