base/
lru.rs

1// lru -- a map that forgets the least recently used entries: a limit on the number of entries
2// and one on their weight (bytes), ordered maps only (deterministic, no hashing); caches of
3// simulators and servers
4//
5//   Lru::new(max_len, max_weight)          0: no limit
6//   get(k) (counts as a use) / peek(k) / contains(k) / get_mut(k)
7//   insert(k, v, weight) -> evicted        the entries that had to go, oldest first (the new
8//                                          one itself when it alone is too heavy)
9//   remove(k) / pop_lru() / clear() / len() / weight() / keys() (oldest first)
10//
11// Every use moves an entry to the young end; eviction takes from the old end until both
12// limits hold.  Costs are O(log n) per call.
13
14use alloc::collections::BTreeMap;
15use alloc::vec::Vec;
16
17#[derive(Clone, Debug, PartialEq, Eq)]
18struct Slot<V> {
19    v: V,
20    tick: u64,
21    weight: u64,
22}
23
24/// a least-recently-used map
25#[derive(Clone, Debug, PartialEq, Eq)]
26pub struct Lru<K: Ord + Clone, V> {
27    map: BTreeMap<K, Slot<V>>,
28    /// use tick -> key, the oldest first
29    order: BTreeMap<u64, K>,
30    tick: u64,
31    weight: u64,
32    /// entries at most (0: no limit)
33    pub max_len: usize,
34    /// summed weight at most (0: no limit)
35    pub max_weight: u64,
36}
37
38impl<K: Ord + Clone, V> Lru<K, V> {
39    pub fn new(max_len: usize, max_weight: u64) -> Lru<K, V> {
40        Lru { map: BTreeMap::new(), order: BTreeMap::new(), tick: 0, weight: 0, max_len, max_weight }
41    }
42    pub fn len(&self) -> usize {
43        self.map.len()
44    }
45    pub fn is_empty(&self) -> bool {
46        self.map.is_empty()
47    }
48    /// the summed weight of the entries
49    pub fn weight(&self) -> u64 {
50        self.weight
51    }
52    pub fn contains(&self, k: &K) -> bool {
53        self.map.contains_key(k)
54    }
55    /// an entry without counting it as a use
56    pub fn peek(&self, k: &K) -> Option<&V> {
57        self.map.get(k).map(|s| &s.v)
58    }
59    fn touch(&mut self, k: &K) {
60        self.tick += 1;
61        let t = self.tick;
62        if let Some(s) = self.map.get_mut(k) {
63            self.order.remove(&s.tick);
64            s.tick = t;
65            self.order.insert(t, k.clone());
66        }
67    }
68    /// an entry; it becomes the most recently used
69    pub fn get(&mut self, k: &K) -> Option<&V> {
70        if !self.map.contains_key(k) {
71            return None;
72        }
73        self.touch(k);
74        self.map.get(k).map(|s| &s.v)
75    }
76    /// an entry to change in place; it becomes the most recently used
77    pub fn get_mut(&mut self, k: &K) -> Option<&mut V> {
78        if !self.map.contains_key(k) {
79            return None;
80        }
81        self.touch(k);
82        self.map.get_mut(k).map(|s| &mut s.v)
83    }
84    /// stores an entry (replacing one with the same key) as the most recently used; the
85    /// entries that had to go for the limits, oldest first
86    pub fn insert(&mut self, k: K, v: V, weight: u64) -> Vec<(K, V)> {
87        let _ = self.remove(&k);
88        self.tick += 1;
89        let t = self.tick;
90        self.order.insert(t, k.clone());
91        self.map.insert(k, Slot { v, tick: t, weight });
92        self.weight = self.weight.saturating_add(weight);
93        let mut out = Vec::new();
94        while (self.max_len > 0 && self.map.len() > self.max_len) || (self.max_weight > 0 && self.weight > self.max_weight) {
95            match self.pop_lru() {
96                Some(e) => out.push(e),
97                None => break,
98            }
99        }
100        out
101    }
102    pub fn remove(&mut self, k: &K) -> Option<V> {
103        let s = self.map.remove(k)?;
104        self.order.remove(&s.tick);
105        self.weight = self.weight.saturating_sub(s.weight);
106        Some(s.v)
107    }
108    /// takes the least recently used entry
109    pub fn pop_lru(&mut self) -> Option<(K, V)> {
110        let (&t, _) = self.order.iter().next()?;
111        let k = self.order.remove(&t)?;
112        let s = self.map.remove(&k)?;
113        self.weight = self.weight.saturating_sub(s.weight);
114        Some((k, s.v))
115    }
116    pub fn clear(&mut self) {
117        self.map.clear();
118        self.order.clear();
119        self.weight = 0;
120    }
121    /// the keys, least recently used first
122    pub fn keys(&self) -> impl Iterator<Item = &K> {
123        self.order.values()
124    }
125    /// the entries in key order (no use counted)
126    pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
127        self.map.iter().map(|(k, s)| (k, &s.v))
128    }
129}
base/lru.rs

base/
lru.rs
