Rust LFU Leetcode

lfu.rs

use std::{
    borrow::Borrow,
    collections::{BTreeSet, HashMap},
    hash::Hash,
};

pub struct LFUCache(Lfu<i32, i32>);

// This is the Leetcode API we have to implement
impl LFUCache {
    pub fn new(capacity: i32) -> Self {
        Self(Lfu::with_capacity(capacity.try_into().unwrap()))
    }

    pub fn get(&mut self, key: i32) -> i32 {
        self.0.get(&key).copied().unwrap_or(-1)
    }

    pub fn put(&mut self, key: i32, value: i32) {
        self.0.insert(key, value);
    }
}

#[derive(Debug)]
pub struct Lfu<K, V> {
    values: HashMap<K, Entry<K, V>>,
    pqueue: BTreeSet<Entry<K, V>>,
    capacity: usize,
    calls: usize,
}

impl<K: Clone + Ord + Hash, V: Clone + Ord> Lfu<K, V> {
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            values: HashMap::with_capacity(capacity),
            pqueue: BTreeSet::new(),
            capacity,
            calls: 0,
        }
    }

    pub fn get<Q>(&mut self, k: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.calls += 1;
        let entry = self.values.get_mut(k)?;
        self.pqueue.remove(entry);
        entry.priority = self.calls;
        entry.uses += 1;
        self.pqueue.insert(entry.clone());
        Some(&entry.value)
    }

    pub fn insert(&mut self, k: K, v: V) {
        if self.capacity == 0 {
            return;
        }

        self.calls += 1;

        match self.values.get_mut(&k) {
            Some(entry) => {
                self.pqueue.remove(entry);
                entry.priority = self.calls;
                entry.uses += 1;
                entry.value = v;
                self.pqueue.insert(entry.clone());
            }

            None => {
                if self.values.len() >= self.capacity {
                    let least_frequent = self.pqueue.pop_first().unwrap();
                    self.values.remove(&least_frequent.key);
                }

                let entry = Entry {
                    key: k.clone(),
                    value: v,
                    uses: 1,
                    priority: self.calls,
                };
                self.pqueue.insert(entry.clone());
                self.values.insert(k, entry);
            }
        }
    }

    pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
        self.values.iter().map(|(key, val)| (key, &val.value))
    }

    pub fn keys(&self) -> impl Iterator<Item = &K> {
        self.values.keys()
    }
}

#[derive(Debug, Clone, Hash, Eq)]
struct Entry<K, V> {
    key: K,
    value: V,
    uses: usize,
    priority: usize,
}

impl<K: PartialEq, V> PartialEq for Entry<K, V> {
    fn eq(&self, other: &Entry<K, V>) -> bool {
        self.key == other.key
    }
}

impl<K: PartialEq, V: PartialEq> PartialOrd for Entry<K, V> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        (self.uses, self.priority).partial_cmp(&(other.uses, other.priority))
    }
}

impl<K: Eq, V: Eq> Ord for Entry<K, V> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        (self.uses, self.priority).cmp(&(other.uses, other.priority))
    }
}

#[cfg(test)]
mod test {
    use std::collections::HashSet;

    use crate::Lfu;

    #[test]
    fn test_empty_lfu() {
        let mut lfu = Lfu::<i32, i32>::with_capacity(10);
        assert_eq!(None, lfu.get(&123));
    }

    #[test]
    fn test_zero_capacity() {
        let mut lfu = Lfu::with_capacity(0);
        lfu.insert("lol", 1);
        assert_eq!(None, lfu.get("lol"));
    }

    #[test]
    fn test_removes_least_recent_inserted() {
        let mut lfu = Lfu::with_capacity(2);
        lfu.insert("lol", 1);
        lfu.insert("wut", 2);
        lfu.insert("wtf", 3);

        assert_eq!(
            lfu.keys().copied().collect::<HashSet<_>>(),
            HashSet::from(["wut", "wtf"])
        );
    }

    #[test]
    fn test_removes_least_recent_queried() {
        let mut lfu = Lfu::with_capacity(2);
        lfu.insert("lol", 1);
        lfu.insert("wut", 2);
        lfu.get("lol");
        lfu.insert("wtf", 3);

        assert_eq!(
            lfu.keys().copied().collect::<HashSet<_>>(),
            HashSet::from(["lol", "wtf"])
        );
    }

    #[test]
    fn leet_tests() {
        #[rustfmt::skip]
        leet_test(
            &["LFUCache", "put", "put", "get", "get", "get", "put", "put", "get", "get", "get", "get"],
            &[&[3], &[2, 2], &[1, 1], &[2], &[1], &[2], &[3, 3], &[4, 4], &[3], &[2], &[1], &[4]],
            &[None, None, None, Some(2), Some(1), Some(2), None, None, Some(-1), Some(2), Some(1), Some(4)],
        );

        #[rustfmt::skip]
        leet_test(
            &["LFUCache", "put", "put", "get", "put", "get", "get", "put", "get", "get", "get"],
            &[&[2], &[1,1], &[2,2], &[1], &[3,3], &[2], &[3], &[4,4], &[1], &[3], &[4]],
            &[None, None, None, Some(1), None, Some(-1), Some(3), None, Some(-1), Some(3), Some(4)],
        );
    }

    fn leet_test(commands: &[&'static str], args: &[&[i32]], expectations: &[Option<i32>]) {
        let mut lfu = Lfu::<i32, i32>::with_capacity(0);
        for ((command, arg), expectation) in commands
            .into_iter()
            .zip(args.into_iter())
            .zip(expectations.into_iter())
        {
            match command {
                &"LFUCache" => lfu = Lfu::with_capacity(arg[0] as usize),
                &"put" => lfu.insert(arg[0], arg[1]),
                &"get" => assert_eq!(expectation, &lfu.get(&arg[0]).copied().or_else(|| Some(-1))),
                _ => panic!("Invalid command {command}"),
            }
        }
    }
}