torment/torment-core/src/utils.rs

use std::collections::{BTreeMap, BTreeSet};
use std::fmt::{Debug, Formatter};
use std::time::{Duration, Instant};

#[derive(Default)]
pub struct EphemeralMap<K: Ord + Clone, V> {
    entry_map: BTreeMap<K, EphemeralNode<V>>,
    expiry_map: BTreeMap<Instant, BTreeSet<K>>,
}

impl<K: Ord + Debug + Clone, V: Debug> Debug for EphemeralMap<K, V> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_map().entries(self.entry_map.iter()).finish()
    }
}

impl<K: Ord + Clone + Debug, V> EphemeralMap<K, V> {
    pub fn new() -> Self {
        EphemeralMap {
            entry_map: Default::default(),
            expiry_map: Default::default(),
        }
    }

    pub fn clean(&mut self) {
        let valid = self.expiry_map.split_off(&Instant::now());
        let expired = std::mem::replace(&mut self.expiry_map, valid);

        for (_, keys) in expired {
            for key in keys {
                self.entry_map.remove(&key);
            }
        }
    }

    pub fn insert(&mut self, key: K, value: V, expires_at: Instant) -> Option<EphemeralNode<V>> {
        let node = EphemeralNode { value, expires_at };
        let removed = self.remove(&key);

        self.add_expiry(key.clone(), expires_at);
        self.entry_map.insert(key, node);
        removed
    }

    pub fn remove(&mut self, key: &K) -> Option<EphemeralNode<V>> {
        if let Some(node) = self.entry_map.remove(key) {
            self.remove_expiry(key, node.expires_at);
            Some(node)
        } else {
            None
        }
    }

    pub fn remove_value(&mut self, key: &K) -> Option<V> {
        self.remove(key).map(|x| x.value)
    }

    fn add_expiry(&mut self, key: K, expires_at: Instant) {
        self.expiry_map
            .entry(expires_at)
            .or_insert_with(|| BTreeSet::new())
            .insert(key);
    }

    fn remove_expiry(&mut self, key: &K, expires_at: Instant) {
        let remove_node = if let Some(item) = self.expiry_map.get_mut(&expires_at) {
            item.remove(&key);
            item.len() == 0
        } else {
            false
        };

        if remove_node {
            self.expiry_map.remove(&expires_at);
        }
    }

    pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
        self.entry_map.iter().filter_map(|(key, entry)| {
            if entry.is_expired() {
                None
            } else {
                Some((key, &entry.value))
            }
        })
    }

    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut V)> {
        self.entry_map.iter_mut().filter_map(|(key, entry)| {
            if entry.is_expired() {
                None
            } else {
                Some((key, &mut entry.value))
            }
        })
    }

    pub fn len(&self) -> usize {
        self.iter().count()
    }

    pub fn contains_key(&self, key: &K) -> bool {
        self.get(key).is_some()
    }

    pub fn get_and_bump(&mut self, key: &K, from_now: Duration) -> Option<&V> {
        if self.get_node(key).is_none() {
            return None;
        }

        self.update_expiry(&key, Instant::now() + from_now);
        self.get(key)
    }

    pub fn get(&self, key: &K) -> Option<&V> {
        self.get_node(key).map(|x| &x.value)
    }

    pub fn get_and_bump_mut(&mut self, key: &K, from_now: Duration) -> Option<&mut V> {
        if self.get_node(key).is_none() {
            return None;
        }

        self.update_expiry(&key, Instant::now() + from_now);
        self.get_mut(key)
    }

    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
        self.get_node_mut(key).map(|x| &mut x.value)
    }

    pub fn get_node(&self, key: &K) -> Option<&EphemeralNode<V>> {
        let node = if let Some(node) = self.entry_map.get(&key) {
            node
        } else {
            return None;
        };

        if node.is_expired() {
            return None;
        }

        Some(node)
    }

    fn get_node_mut(&mut self, key: &K) -> Option<&mut EphemeralNode<V>> {
        let node = if let Some(node) = self.entry_map.get_mut(&key) {
            node
        } else {
            return None;
        };

        if node.is_expired() {
            return None;
        }

        Some(node)
    }

    pub fn update_expiry(&mut self, key: &K, expires_at: Instant) {
        let node = if let Some(node) = self.entry_map.get_mut(key) {
            node
        } else {
            return;
        };

        let old_expiry = node.expires_at;
        node.expires_at = expires_at;

        self.remove_expiry(key, old_expiry);
        self.add_expiry(key.clone(), expires_at);
    }

    pub fn entry(&mut self, key: K) -> EphemeralEntry<'_, K, V> {
        EphemeralEntry { map: self, key }
    }
}

pub struct EphemeralNode<V> {
    expires_at: Instant,
    value: V,
}

impl<V: Debug> Debug for EphemeralNode<V> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EphemeralNode")
            .field("expires_at", &self.expires_at)
            .field("value", &self.value)
            .finish()
    }
}

impl<V> EphemeralNode<V> {
    pub fn is_expired(&self) -> bool {
        self.expires_at <= Instant::now()
    }

    pub fn expires_at(&self) -> Instant {
        self.expires_at
    }

    pub fn value(&self) -> &V {
        &self.value
    }

    pub fn value_mut(&mut self) -> &mut V {
        &mut self.value
    }

    pub fn into_value(self) -> V {
        self.value
    }
}

pub struct EphemeralEntry<'map, K: Ord + Clone, V> {
    map: &'map mut EphemeralMap<K, V>,
    key: K,
}

impl<'map, K: Ord + Clone + Debug, V> EphemeralEntry<'map, K, V> {
    pub fn or_insert_with<F: FnOnce() -> V>(self, expires_at: Instant, f: F) -> &'map mut V {
        self.or_insert(expires_at, f())
    }

    pub fn or_insert(self, expires_at: Instant, default: V) -> &'map mut V {
        self.map.insert(self.key.clone(), default, expires_at);
        self.map.get_mut(&self.key).unwrap()
    }

    pub fn and_modify<F: FnOnce(&mut V)>(self, f: F) -> Self {
        if let Some(item) = self.map.get_mut(&self.key) {
            f(item)
        }

        self
    }
}

impl<'map, K: Ord + Clone + Debug, V: Default> EphemeralEntry<'map, K, V> {
    pub fn or_insert_default(self, expires_at: Instant) -> &'map mut V {
        self.or_insert(expires_at, Default::default())
    }
}

#[derive(Default)]
pub struct EphemeralSet<K: Clone + Ord>(EphemeralMap<K, ()>);

impl<K: Debug + Clone + Ord> Debug for EphemeralSet<K> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_map()
            .entries(
                self.0
                    .entry_map
                    .iter()
                    .map(|(key, node)| (key, node.expires_at)),
            )
            .finish()
    }
}

impl<K: Clone + Ord + Debug> EphemeralSet<K> {
    pub fn len(&self) -> usize {
        self.0.len()
    }

    pub fn new() -> Self {
        EphemeralSet(EphemeralMap::new())
    }

    pub fn clean(&mut self) {
        self.0.clean()
    }

    pub fn insert(&mut self, key: K, expires_at: Instant) -> bool {
        self.0.insert(key, (), expires_at).is_none()
    }

    pub fn remove(&mut self, key: &K) -> bool {
        self.0.remove(key).is_some()
    }

    pub fn contains(&self, key: &K) -> bool {
        self.0.get_node(key).is_some()
    }

    pub fn remove_first(&mut self) -> Option<K> {
        loop {
            let first_key = self.0.entry_map.keys().next().cloned();
            let item = if let Some(key) = first_key {
                key
            } else {
                break;
            };

            if self.contains(&item) {
                self.remove(&item);
                return Some(item);
            }
        }

        None
    }
}

#[derive(Default)]
pub struct EphemeralQueueMap<K: Clone + Ord, V>(EphemeralMap<K, Vec<EphemeralNode<V>>>);

impl<K: Debug + Clone + Ord, V: Debug> Debug for EphemeralQueueMap<K, V> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_map().entries(self.0.entry_map.iter()).finish()
    }
}

impl<K: Clone + Ord + Debug, V> EphemeralQueueMap<K, V> {
    pub fn clean(&mut self) {
        self.0.clean();
    }

    pub fn insert(&mut self, key: K, value: V, expires_at: Instant) {
        let update_expiry = if let Some(node) = self.0.get_node_mut(&key) {
            let mut old = std::mem::replace(&mut node.value, vec![]);
            old.push(EphemeralNode { value, expires_at });

            while let Some(old_node) = old.pop() {
                if old_node.is_expired() {
                    continue;
                }

                node.value.push(old_node)
            }

            node.expires_at < expires_at
        } else {
            self.0.insert(
                key.clone(),
                vec![EphemeralNode { value, expires_at }],
                expires_at,
            );
            false
        };

        if update_expiry {
            self.0.update_expiry(&key, expires_at);
        }
    }

    pub fn remove(&mut self, key: &K) {
        self.0.remove(key);
    }

    pub fn pop(&mut self, key: &K) -> Option<V> {
        if let Some(node) = self.0.get_node_mut(key) {
            let old_values = std::mem::replace(&mut node.value, vec![]);
            let mut selected = None;
            let mut expiry: Option<Instant> = None;
            for value in old_values {
                if value.is_expired() {
                    continue;
                }

                if selected.is_none() {
                    selected = Some(value);
                    continue;
                }

                if *expiry.get_or_insert(value.expires_at) > value.expires_at {
                    expiry = Some(value.expires_at)
                }

                node.value.push(value)
            }

            if !node.value.is_empty() {
                self.0.update_expiry(key, expiry.unwrap());
            } else {
                self.0.remove(key);
            }

            selected.map(|x| x.value)
        } else {
            None
        }
    }
}