polars_utils/
cache.rs

1use std::borrow::Borrow;
2use std::hash::{BuildHasher, Hash};
3
4use foldhash::fast::RandomState;
5use hashbrown::HashTable;
6use hashbrown::hash_table::Entry;
7use slotmap::{Key, SlotMap, new_key_type};
8
9/// A cached function that use `LruCache`.
10pub struct LruCachedFunc<T, R, F> {
11    func: F,
12    cache: LruCache<T, R>,
13}
14
15impl<T, R, F> LruCachedFunc<T, R, F>
16where
17    F: FnMut(T) -> R,
18    T: std::hash::Hash + Eq + Clone,
19    R: Copy,
20{
21    pub fn new(func: F, size: usize) -> Self {
22        Self {
23            func,
24            cache: LruCache::with_capacity(size),
25        }
26    }
27
28    pub fn eval(&mut self, x: T, use_cache: bool) -> R {
29        if use_cache {
30            *self
31                .cache
32                .get_or_insert_with(&x, |xr| (self.func)(xr.clone()))
33        } else {
34            (self.func)(x)
35        }
36    }
37}
38
39new_key_type! {
40    struct LruKey;
41}
42
43pub struct LruCache<K, V, S = RandomState> {
44    table: HashTable<LruKey>,
45    elements: SlotMap<LruKey, LruEntry<K, V>>,
46    max_capacity: usize,
47    most_recent: LruKey,
48    least_recent: LruKey,
49    build_hasher: S,
50}
51
52struct LruEntry<K, V> {
53    key: K,
54    value: V,
55    list: LruListNode,
56}
57
58#[derive(Copy, Clone, Default)]
59struct LruListNode {
60    more_recent: LruKey,
61    less_recent: LruKey,
62}
63
64impl<K, V> LruCache<K, V> {
65    pub fn with_capacity(capacity: usize) -> Self {
66        Self::with_capacity_and_hasher(capacity, RandomState::default())
67    }
68}
69
70impl<K, V, S> LruCache<K, V, S> {
71    pub fn with_capacity_and_hasher(max_capacity: usize, build_hasher: S) -> Self {
72        let max_capacity = max_capacity.max(1);
73        Self {
74            // Allocate one more capacity to prevent double-lookup or realloc
75            // when doing get_or_insert when full.
76            table: HashTable::with_capacity(max_capacity + 1),
77            elements: SlotMap::with_capacity_and_key(max_capacity + 1),
78            max_capacity,
79            most_recent: LruKey::null(),
80            least_recent: LruKey::null(),
81            build_hasher,
82        }
83    }
84
85    pub fn max_capacity(&self) -> usize {
86        self.max_capacity
87    }
88
89    pub fn len(&self) -> usize {
90        self.elements.len()
91    }
92
93    pub fn is_empty(&self) -> bool {
94        self.elements.is_empty()
95    }
96}
97
98impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
99    fn lru_list_unlink(&mut self, lru_key: LruKey) {
100        let list = self.elements[lru_key].list;
101        if let Some(more_recent) = self.elements.get_mut(list.more_recent) {
102            more_recent.list.less_recent = list.less_recent;
103        } else {
104            self.most_recent = list.less_recent;
105        }
106        if let Some(less_recent) = self.elements.get_mut(list.less_recent) {
107            less_recent.list.more_recent = list.more_recent;
108        } else {
109            self.least_recent = list.more_recent;
110        }
111    }
112
113    fn lru_list_insert_mru(&mut self, lru_key: LruKey) {
114        let prev_most_recent_key = self.most_recent;
115        self.most_recent = lru_key;
116        if let Some(prev_most_recent) = self.elements.get_mut(prev_most_recent_key) {
117            prev_most_recent.list.more_recent = lru_key;
118        } else {
119            self.least_recent = lru_key;
120        }
121        let list = &mut self.elements[lru_key].list;
122        list.more_recent = LruKey::null();
123        list.less_recent = prev_most_recent_key;
124    }
125
126    pub fn pop_lru(&mut self) -> Option<(K, V)> {
127        if self.elements.is_empty() {
128            return None;
129        }
130        let lru_key = self.least_recent;
131        let hash = self.build_hasher.hash_one(&self.elements[lru_key].key);
132        self.lru_list_unlink(lru_key);
133        let lru_entry = self.elements.remove(lru_key).unwrap();
134        self.table
135            .find_entry(hash, |k| *k == lru_key)
136            .unwrap()
137            .remove();
138        Some((lru_entry.key, lru_entry.value))
139    }
140
141    pub fn get<Q>(&mut self, key: &Q) -> Option<&V>
142    where
143        K: Borrow<Q>,
144        Q: Hash + Eq + ?Sized,
145    {
146        let hash = self.build_hasher.hash_one(key);
147        let lru_key = *self
148            .table
149            .find(hash, |lru_key| self.elements[*lru_key].key.borrow() == key)?;
150        self.lru_list_unlink(lru_key);
151        self.lru_list_insert_mru(lru_key);
152        let lru_node = self.elements.get(lru_key).unwrap();
153        Some(&lru_node.value)
154    }
155
156    /// Returns the old value, if any.
157    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
158        let hash = self.build_hasher.hash_one(&key);
159        match self.table.entry(
160            hash,
161            |lru_key| self.elements[*lru_key].key == key,
162            |lru_key| self.build_hasher.hash_one(&self.elements[*lru_key].key),
163        ) {
164            Entry::Occupied(o) => {
165                let lru_key = *o.get();
166                self.lru_list_unlink(lru_key);
167                self.lru_list_insert_mru(lru_key);
168                Some(core::mem::replace(&mut self.elements[lru_key].value, value))
169            },
170
171            Entry::Vacant(v) => {
172                let lru_entry = LruEntry {
173                    key,
174                    value,
175                    list: LruListNode::default(),
176                };
177                let lru_key = self.elements.insert(lru_entry);
178                v.insert(lru_key);
179                self.lru_list_insert_mru(lru_key);
180                if self.elements.len() > self.max_capacity {
181                    self.pop_lru();
182                }
183                None
184            },
185        }
186    }
187
188    pub fn get_or_insert_with<Q, F: FnOnce(&Q) -> V>(&mut self, key: &Q, f: F) -> &mut V
189    where
190        K: Borrow<Q>,
191        Q: Hash + Eq + ToOwned<Owned = K> + ?Sized,
192    {
193        enum Never {}
194        let Ok(ret) = self.try_get_or_insert_with::<Q, Never, _>(key, |k| Ok(f(k)));
195        ret
196    }
197
198    pub fn try_get_or_insert_with<Q, E, F: FnOnce(&Q) -> Result<V, E>>(
199        &mut self,
200        key: &Q,
201        f: F,
202    ) -> Result<&mut V, E>
203    where
204        K: Borrow<Q>,
205        Q: Hash + Eq + ToOwned<Owned = K> + ?Sized,
206    {
207        let hash = self.build_hasher.hash_one(key);
208        match self.table.entry(
209            hash,
210            |lru_key| self.elements[*lru_key].key.borrow() == key,
211            |lru_key| self.build_hasher.hash_one(&self.elements[*lru_key].key),
212        ) {
213            Entry::Occupied(o) => {
214                let lru_key = *o.get();
215                if lru_key != self.most_recent {
216                    self.lru_list_unlink(lru_key);
217                    self.lru_list_insert_mru(lru_key);
218                }
219                Ok(&mut self.elements[lru_key].value)
220            },
221
222            Entry::Vacant(v) => {
223                let lru_entry = LruEntry {
224                    value: f(key)?,
225                    key: key.to_owned(),
226                    list: LruListNode::default(),
227                };
228                let lru_key = self.elements.insert(lru_entry);
229                v.insert(lru_key);
230                self.lru_list_insert_mru(lru_key);
231                if self.elements.len() > self.max_capacity {
232                    self.pop_lru();
233                }
234                Ok(&mut self.elements[lru_key].value)
235            },
236        }
237    }
238}
polars_utils/cache.rs

polars_utils/
cache.rs
