expiremap 0.1.0-beta.3

use indexmap::IndexMap;
use std::hash::Hash;

#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, Clone)]
pub struct ExpireMap<K, V, T>
where
    K: Eq + Hash + Clone,
{
    map: IndexMap<K, ExpiryEntry<V, T>>,
}

impl<K, V, T> Default for ExpireMap<K, V, T>
where
    K: Eq + Hash + Clone,
{
    fn default() -> Self {
        Self {
            map: IndexMap::default(),
        }
    }
}

impl<K, V, T> ExpireMap<K, V, T>
where
    K: Eq + Hash + Clone,
{
    pub fn new() -> Self {
        Self {
            map: IndexMap::new(),
        }
    }
}
impl<K, V, T> ExpireMap<K, V, T>
where
    K: Eq + Hash + Clone,
    T: Ord,
{
    /// Insert a key-value pair with expiry into the map.
    /// If a previous value for this key was present, the value and
    /// it's expiry time are returned.
    /// If no value was present for the key, `None` is returned.
    ///
    /// **NOTE:** Inserting an item with an expiry time > than the
    /// expiry of all other elements is O(1) complexity (amortized).
    /// Inserting an item with a smaller complexity is O(n) where n
    /// is the number of items with greater expiry time
    pub fn insert(&mut self, key: K, value: V, expiry: T) -> Option<ExpiryEntry<V, T>> {
        let insert_idx = self.map.partition_point(|_, entry| entry.expiry < expiry);
        let entry = ExpiryEntry::new(value, expiry);
        self.map.insert_before(insert_idx, key, entry).1
    }

    /// Return a reference to the value stored under this key
    /// if it exists. If the value does not exist, `None` is
    /// returned
    pub fn get(&self, key: &K) -> Option<&V> {
        self.map.get(key).map(|x| &x.value)
    }
    /// Remove and return all values from the map where the expiry
    /// time is less or equal the given expiry.
    pub fn expire<'a>(&'a mut self, expiry: &T) -> ExpiredIterator<'a, K, ExpiryEntry<V, T>> {
        let expiry_index = self.map.partition_point(|_, entry| entry.expiry <= *expiry);
        ExpiredIterator(self.map.drain(0..expiry_index))
    }

    /// Return the number of key-value pairs in the map.
    ///
    /// Computes in **O(1)** time.
    pub fn len(&self) -> usize {
        self.map.len()
    }

    /// Returns true if the map contains no elements.
    ///
    /// Computes in **O(1)** time.
    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    /// Remove a key from the map returning its associated value if it was in the map
    pub fn remove(&mut self, key: &K) -> Option<V> {
        self.map.shift_remove(key).map(|entry| entry.value)
    }

    pub fn into_values(self) -> ValuesIterator<K, V, T> {
        ValuesIterator(self.map.into_values())
    }
}

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct ExpiryEntry<V, T> {
    pub value: V,
    pub expiry: T,
}
impl<V, T> ExpiryEntry<V, T> {
    fn new(value: V, expiry: T) -> Self {
        Self { value, expiry }
    }
}

/// An iterator over all expired elements which where removed from the map
pub struct ExpiredIterator<'a, K, V>(indexmap::map::Drain<'a, K, V>);

impl<K, V> Iterator for ExpiredIterator<'_, K, V> {
    type Item = (K, V);

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

pub struct ValuesIterator<K, V, T>(indexmap::map::IntoValues<K, ExpiryEntry<V, T>>);

impl<K, V, T> Iterator for ValuesIterator<K, V, T> {
    type Item = V;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|x| x.value)
    }
}

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

    /// Check we can retrieve an inserted value
    #[test]
    fn insert_and_get() {
        let mut map = ExpireMap::new();
        map.insert("key", "HelloWorld", 42);
        let value = map.get(&"key");
        assert!(matches!(value, Some(&"HelloWorld")))
    }

    /// Check insert returns the previous value and expiry
    #[test]
    fn insert_retrieves() {
        let mut map = ExpireMap::new();
        map.insert("key", "HelloWorld", 42);
        let prev = map.insert("key", "HelloMom", 11);
        assert!(matches!(
            prev,
            Some(ExpiryEntry {
                value: "HelloWorld",
                expiry: 42
            })
        ))
    }

    // check we can not get a value after it was expired
    #[test]
    fn get_after_expired() {
        let mut map = ExpireMap::new();
        map.insert("key", "HelloWorld", 42);
        map.expire(&42);
        assert!(map.get(&"key").is_none())
    }

    // check we can get a value after it did not expire
    #[test]
    fn get_not_expired() {
        let mut map = ExpireMap::new();
        map.insert("key", "HelloWorld", 42);
        map.expire(&41);
        assert!(map.get(&"key").is_some())
    }

    /// Check the insertion order does not matter for expiry
    #[test]
    fn expiry_insertion_order() {
        let mut map = ExpireMap::new();
        map.insert("key0", "HelloWorld", 5);
        map.insert("key1", "HelloWorld", 2);
        map.insert("key2", "HelloWorld", 7);
        map.insert("key3", "HelloWorld", -22);
        map.expire(&4);
        for k in ["key1", "key3"] {
            assert!(map.get(&k).is_none())
        }
        for k in ["key0", "key2"] {
            assert!(map.get(&k).is_some())
        }
    }

    /// Check expired keys are returned as iterator
    #[test]
    fn expired_as_iterator() {
        let mut map = ExpireMap::new();
        let vals = ["foo", "bar", "baz"];
        for (i, x) in vals.into_iter().enumerate() {
            map.insert(x, x, i);
        }
        let expired: Vec<(&str, ExpiryEntry<&str, usize>)> = map.expire(&3).collect();
        let expected: Vec<(&str, ExpiryEntry<&str, usize>)> = vals
            .into_iter()
            .enumerate()
            .map(|(i, x)| (x, ExpiryEntry::new(x, i)))
            .collect();
        assert_eq!(expired, expected)
    }

    /// Check that `is_empty` works with key expiry
    #[test]
    fn is_empty() {
        let mut map = ExpireMap::new();

        assert!(map.is_empty());

        let vals = ["foo", "bar", "baz"];
        for (i, x) in vals.into_iter().enumerate() {
            map.insert(x, x, i);
        }

        assert!(!map.is_empty());

        let _ = map.expire(&1);
        assert!(!map.is_empty());

        let _ = map.expire(&2);
        assert!(map.is_empty());
    }

    /// Check that `len` works with key expiry
    #[test]
    fn len() {
        let mut map = ExpireMap::new();

        assert_eq!(0, map.len());

        let vals = ["foo", "bar", "baz"];
        for (i, x) in vals.into_iter().enumerate() {
            map.insert(x, x, i);
        }

        assert_eq!(3, map.len());

        let _ = map.expire(&1);
        assert_eq!(1, map.len());

        let _ = map.expire(&2);
        assert_eq!(0, map.len());
    }

    /// Check remove
    #[test]
    fn test_remove() {
        let mut map = ExpireMap::new();
        let vals = ["foo", "bar", "baz"];
        for (i, x) in vals.into_iter().enumerate() {
            map.insert(x, x, i);
        }

        let entry = map.remove(&"foo");
        assert_eq!(entry.unwrap(), "foo");

        assert!(map.remove(&"doesnotexist").is_none());
    }

    /// Check expiry works correctly after remove
    #[test]
    fn test_remove_expiry() {
        let mut map = ExpireMap::new();

        assert_eq!(0, map.len());
        let vals = ["foo", "bar", "baz"];
        for (i, x) in vals.into_iter().enumerate() {
            map.insert(x, x, i);
        }
        map.remove(&"bar");
        let _ = map.expire(&1);
        assert_eq!(1, map.len());
    }

    /// Check into_values returns all values
    #[test]
    fn test_into_values() {
        let mut map = ExpireMap::new();

        assert_eq!(0, map.len());
        let vals = [("foo", "FOO"), ("bar", "BAR"), ("baz", "BAZ")];
        for (i, x) in vals.into_iter().enumerate() {
            map.insert(x.0, x.1, i);
        }
        let values: Vec<_> = map.into_values().collect();
        assert_eq!(values, vec!["FOO", "BAR", "BAZ"]);
    }

    /// Check dropping the expire_iterator also removes entries i.e. the iterator does not have
    /// to be consumed
    #[test]
    fn test_expired_iterator_drop_removes() {
        let mut map = ExpireMap::new();
        let vals = [("foo", "FOO"), ("bar", "BAR"), ("baz", "BAZ")];
        for (i, x) in vals.into_iter().enumerate() {
            map.insert(x.0, x.1, i);
        }
        let it = map.expire(&usize::MAX);
        drop(it);
        assert!(map.is_empty())
    }
}