LRU O1

lru_o1_safe.rs

use std::{
    cell::{Ref, RefCell},
    collections::HashMap,
    fmt::Debug,
    rc::Rc,
};

/// LRU cache
struct LRU {
    cache: HashMap<String, Rc<RefCell<ListNode>>>,
    queue: DoubleLinkedList,
    max_len: usize,
}

#[derive(Debug)]
struct ListNode {
    key: String,
    value: String,
    prev: Option<Rc<RefCell<ListNode>>>,
    next: Option<Rc<RefCell<ListNode>>>,
}

impl ListNode {
    pub fn new(key: String, value: String) -> Self {
        Self {
            key,
            value,
            prev: None,
            next: None,
        }
    }
}

#[derive(Debug)]
struct DoubleLinkedList {
    head: Option<Rc<RefCell<ListNode>>>,
    tail: Option<Rc<RefCell<ListNode>>>,
}

impl DoubleLinkedList {
    pub fn new() -> Self {
        Self {
            head: None,
            tail: None,
        }
    }

    fn get_head(&self) -> Option<Rc<RefCell<ListNode>>> {
        self.head.clone()
    }

    fn get_tail(&self) -> Option<Rc<RefCell<ListNode>>> {
        self.tail.clone()
    }

    pub fn add_front_node(&mut self, node: Rc<RefCell<ListNode>>) {
        let head = self.get_head();
        if let Some(ref h) = head {
            h.borrow_mut().prev = Some(node.clone());
        }

        node.borrow_mut().prev = None;
        node.borrow_mut().next = head;

        self.head = Some(node);
    }

    pub fn add_back_node(&mut self, node: Rc<RefCell<ListNode>>) {
        let tail = self.get_tail();
        if let Some(ref t) = tail {
            t.borrow_mut().next = Some(node.clone());
        }

        node.borrow_mut().prev = tail;
        node.borrow_mut().next = None;

        self.tail = Some(node);
    }

    pub fn remove(&mut self, target: Rc<RefCell<ListNode>>) {
        let prev = target.borrow().prev.clone();
        let next = target.borrow().next.clone();

        match (prev, next) {
            (Some(prev), Some(next)) => {
                prev.borrow_mut().next = Some(next.clone());
                next.borrow_mut().prev = Some(prev);
            }
            (Some(prev), None) => {
                prev.borrow_mut().next.take();
                self.tail.replace(prev);
            }
            (None, Some(next)) => {
                next.borrow_mut().prev.take();
                self.head.replace(next);
            }
            (None, None) => {
                self.head.take();
                self.tail.take();
            }
        }
    }

    pub fn move_head_to_front(&mut self, target: Rc<RefCell<ListNode>>) {
        if let Some(ref head) = self.get_head() {
            if !Rc::ptr_eq(head, &target) {
                self.remove(target.clone());
                self.add_front_node(target);
            }
        }
    }
}

impl LRU {
    /// `max_size` - maximum number of items that this cache can hold.
    pub fn new(max_len: usize) -> Self {
        Self {
            cache: HashMap::new(),
            queue: DoubleLinkedList::new(),
            max_len,
        }
    }

    /// If `key` is in LRU then this will return associated value.
    ///
    /// Should be O(1)
    /// Required Option<&String> return type is hard, because String value live inside dynamically borrowed RefCell
    pub fn get(&mut self, key: &str) -> Option<Ref<String>> {
        match self.cache.get(key) {
            Some(node) => {
                self.queue.move_head_to_front(node.clone());
                let n = node.borrow();
                let s = Ref::map(n, |v| &v.value);
                Some(s)
            }
            None => None,
        }
    }

    /// Insert a new `value` with given `key`. If LRU already has a value
    /// associated with given key, the value is replaced (and old value is returned).
    ///
    /// If LRU cache already has `max_size` elements, then least recently used (get/put) key is
    /// removed from cache to keep it at `max_size` size after insertion.
    ///
    /// Should be O(1)
    pub fn insert(&mut self, key: String, value: String) -> Option<String> {
        let mut old_val: Option<String> = None;
        let node = match self.cache.get(&key) {
            Some(node) => {
                node.borrow_mut().value = value;
                self.queue.remove(node.clone());
                self.queue.add_front_node(node.clone());
                node.clone()
            }
            None => {
                let node = Rc::new(RefCell::new(ListNode::new(key.clone(), value)));

                if self.cache.len() == self.max_len {
                    let tail = self.queue.get_tail();
                    if let Some(tail) = tail {
                        self.cache.remove(&tail.borrow().key);
                        self.queue.remove(tail);
                    }
                }

                let old = self.cache.insert(key, node.clone()).map(Rc::try_unwrap);

                if let Some(Ok(v)) = old {
                    old_val = Some(v.into_inner().value);
                }

                self.queue.add_front_node(node.clone());

                node
            }
        };

        if self.queue.tail.is_none() {
            self.queue.add_back_node(node.clone());
        }

        old_val
    }
}

fn main() {
    let mut lru = LRU::new(2);
    lru.insert("a".to_string(), "a".to_string());
    lru.insert("b".to_string(), "b".to_string());
    assert_eq!(*lru.get("a").unwrap(), "a".to_string());
    lru.insert("c".to_string(), "c".to_string());
    assert!(lru.get("b").is_none());
    assert_eq!(*lru.get("c").unwrap(), "c".to_string());
    lru.insert("c".to_string(), "C".to_string());
    assert_eq!(*lru.get("c").unwrap(), "C".to_string());
}