radiate_utils/

lru.rs

use foldhash::fast::RandomState;
use hashbrown::HashTable;
use hashbrown::hash_table::Entry;
use slotmap::{Key, SlotMap, new_key_type};
use std::borrow::Borrow;
use std::hash::{BuildHasher, Hash};

new_key_type! {
    struct LruKey;
}

#[derive(Copy, Clone, Default)]
struct LruListNode {
    more_recent: LruKey,
    less_recent: LruKey,
}

struct LruEntry<K, V> {
    key: K,
    value: V,
    list: LruListNode,
}

pub struct LruCache<K, V, S = RandomState> {
    table: HashTable<LruKey>,
    elements: SlotMap<LruKey, LruEntry<K, V>>,
    max_capacity: usize,
    most_recent: LruKey,
    least_recent: LruKey,
    build_hasher: S,
}

impl<K, V> LruCache<K, V> {
    pub fn with_capacity(capacity: usize) -> Self {
        Self::with_capacity_and_hasher(capacity, RandomState::default())
    }
}

impl<K, V, S> LruCache<K, V, S> {
    pub fn with_capacity_and_hasher(max_capacity: usize, build_hasher: S) -> Self {
        assert!(max_capacity > 0);
        Self {
            // Allocate one more capacity to prevent double-lookup or realloc
            // when doing get_or_insert when full.
            table: HashTable::with_capacity(max_capacity + 1),
            elements: SlotMap::with_capacity_and_key(max_capacity + 1),
            max_capacity,
            most_recent: LruKey::null(),
            least_recent: LruKey::null(),
            build_hasher,
        }
    }
}

impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
    fn lru_list_unlink(&mut self, lru_key: LruKey) {
        let list = self.elements[lru_key].list;
        if let Some(more_recent) = self.elements.get_mut(list.more_recent) {
            more_recent.list.less_recent = list.less_recent;
        } else {
            self.most_recent = list.less_recent;
        }

        if let Some(less_recent) = self.elements.get_mut(list.less_recent) {
            less_recent.list.more_recent = list.more_recent;
        } else {
            self.least_recent = list.more_recent;
        }
    }

    fn lru_list_insert_mru(&mut self, lru_key: LruKey) {
        let prev_most_recent_key = self.most_recent;
        self.most_recent = lru_key;
        if let Some(prev_most_recent) = self.elements.get_mut(prev_most_recent_key) {
            prev_most_recent.list.more_recent = lru_key;
        } else {
            self.least_recent = lru_key;
        }
        let list = &mut self.elements[lru_key].list;
        list.more_recent = LruKey::null();
        list.less_recent = prev_most_recent_key;
    }

    pub fn pop_lru(&mut self) -> Option<(K, V)> {
        if self.elements.is_empty() {
            return None;
        }
        let lru_key = self.least_recent;
        let hash = self.build_hasher.hash_one(&self.elements[lru_key].key);
        self.lru_list_unlink(lru_key);
        let lru_entry = self.elements.remove(lru_key).unwrap();
        self.table
            .find_entry(hash, |k| *k == lru_key)
            .unwrap()
            .remove();
        Some((lru_entry.key, lru_entry.value))
    }

    pub fn get<Q>(&mut self, key: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        let hash = self.build_hasher.hash_one(key);
        let lru_key = *self
            .table
            .find(hash, |lru_key| self.elements[*lru_key].key.borrow() == key)?;
        self.lru_list_unlink(lru_key);
        self.lru_list_insert_mru(lru_key);
        let lru_node = self.elements.get(lru_key).unwrap();
        Some(&lru_node.value)
    }

    /// Returns the old value, if any.
    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        let hash = self.build_hasher.hash_one(&key);
        match self.table.entry(
            hash,
            |lru_key| self.elements[*lru_key].key == key,
            |lru_key| self.build_hasher.hash_one(&self.elements[*lru_key].key),
        ) {
            Entry::Occupied(o) => {
                let lru_key = *o.get();
                self.lru_list_unlink(lru_key);
                self.lru_list_insert_mru(lru_key);
                Some(core::mem::replace(&mut self.elements[lru_key].value, value))
            }

            Entry::Vacant(v) => {
                let lru_entry = LruEntry {
                    key,
                    value,
                    list: LruListNode::default(),
                };
                let lru_key = self.elements.insert(lru_entry);
                v.insert(lru_key);
                self.lru_list_insert_mru(lru_key);
                if self.elements.len() > self.max_capacity {
                    self.pop_lru();
                }
                None
            }
        }
    }

    pub fn get_or_insert_with<Q, F>(&mut self, key: &Q, f: F) -> &mut V
    where
        F: FnOnce(&Q) -> V,
        K: Borrow<Q>,
        Q: Hash + Eq + ToOwned<Owned = K> + ?Sized,
    {
        enum Never {}
        let Ok(ret) = self.try_get_or_insert_with::<Q, Never, _>(key, |k| Ok(f(k)));
        ret
    }

    pub fn try_get_or_insert_with<Q, E, F: FnOnce(&Q) -> Result<V, E>>(
        &mut self,
        key: &Q,
        f: F,
    ) -> Result<&mut V, E>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ToOwned<Owned = K> + ?Sized,
    {
        let hash = self.build_hasher.hash_one(key);
        match self.table.entry(
            hash,
            |lru_key| self.elements[*lru_key].key.borrow() == key,
            |lru_key| self.build_hasher.hash_one(&self.elements[*lru_key].key),
        ) {
            Entry::Occupied(o) => {
                let lru_key = *o.get();
                if lru_key != self.most_recent {
                    self.lru_list_unlink(lru_key);
                    self.lru_list_insert_mru(lru_key);
                }
                Ok(&mut self.elements[lru_key].value)
            }

            Entry::Vacant(v) => {
                let lru_entry = LruEntry {
                    value: f(key)?,
                    key: key.to_owned(),
                    list: LruListNode::default(),
                };
                let lru_key = self.elements.insert(lru_entry);
                v.insert(lru_key);
                self.lru_list_insert_mru(lru_key);
                if self.elements.len() > self.max_capacity {
                    self.pop_lru();
                }
                Ok(&mut self.elements[lru_key].value)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::LruCache;
    #[test]
    fn test_lru_cache_basic() {
        let mut lru: LruCache<u32, &str> = LruCache::with_capacity(2);

        assert_eq!(lru.insert(1, "one"), None);
        assert_eq!(lru.insert(2, "two"), None);
        assert_eq!(lru.get(&1), Some(&"one"));
        assert_eq!(lru.insert(3, "three"), None); // Evicts key 2
        assert_eq!(lru.get(&2), None);
        assert_eq!(lru.get(&1), Some(&"one"));
        assert_eq!(lru.get(&3), Some(&"three"));
        assert_eq!(lru.insert(1, "uno"), Some("one")); // Update key 1
        assert_eq!(lru.get(&1), Some(&"uno"));
    }

    #[test]
    fn test_lru_cache_pop() {
        let mut lru: LruCache<u32, &str> = LruCache::with_capacity(2);

        assert_eq!(lru.insert(1, "one"), None);
        assert_eq!(lru.insert(2, "two"), None);
        assert_eq!(lru.pop_lru(), Some((1, "one"))); // Evicts key 1
        assert_eq!(lru.get(&1), None); // Key 1 should be gone
        assert_eq!(lru.get(&2), Some(&"two")); // Key 2 should still be present
    }
}