1
use std::borrow::Borrow;
2
use std::collections::{btree_map, BTreeMap};
3
use std::fmt::Debug;
4
use std::hash::Hash;
5
use std::ops::RangeBounds;
6

            
7
use crate::lru::{EntryCache, EntryRef, IntoIter, LruCache, NodeId, Removed};
8
use crate::LruMap;
9

            
10
/// A Least Recently Used map with fixed capacity that stores keys using a
11
/// [`BTreeMap`] internally. Inserting and querying has similar performance to
12
/// using a [`BTreeMap`], but internally this data structure keeps track of the
13
/// order in which the keys were last touched.
14
///
15
/// When inserting a new key and the map is at-capacity, the least recently used
16
/// key will be evicted to make room for the new key.
17
///
18
/// To avoid `unsafe`, this crate must store each entry's key twice. This means
19
/// that `Key` must implement `Clone`. If you're using expensive-to-clone keys,
20
/// consider wrapping the key in an `Rc`/`Arc` or using an alternate LRU crate.
21
1
#[derive(Debug)]
22
#[must_use]
23
pub struct LruBTreeMap<Key, Value> {
24
    map: BTreeMap<Key, NodeId>,
25
    cache: LruCache<Key, Value>,
26
}
27

            
28
impl<Key, Value> LruBTreeMap<Key, Value>
29
where
30
    Key: Ord + Clone,
31
{
32
    /// Creates a new map with the maximum `capacity`.
33
    ///
34
    /// # Panics
35
    ///
36
    /// Panics if `capacity` is <= 1 or > `u32::MAX`.
37
8
    pub fn new(capacity: usize) -> Self {
38
8
        assert!(capacity > 1);
39
8
        assert!(capacity <= usize::try_from(u32::MAX).unwrap());
40
8
        Self {
41
8
            map: BTreeMap::new(),
42
8
            cache: LruCache::new(capacity),
43
8
        }
44
8
    }
45

            
46
    /// Returns the stored value for `key`, if present.
47
    ///
48
    /// This function touches the key, making it the most recently used key.
49
1010
    pub fn get<QueryKey>(&mut self, key: &QueryKey) -> Option<&Value>
50
1010
    where
51
1010
        QueryKey: Ord + ?Sized,
52
1010
        Key: Borrow<QueryKey>,
53
1010
    {
54
1010
        let node = self.map.get(key).copied();
55
1010
        node.map(|node| self.cache.get(node).value())
56
1010
    }
57

            
58
    /// Returns the stored value for `key`, if present.
59
    ///
60
    /// This function does not touch the key, preserving its current position in
61
    /// the lru cache.
62
1
    pub fn get_without_update<QueryKey>(&self, key: &QueryKey) -> Option<&Value>
63
1
    where
64
1
        QueryKey: Ord + ?Sized,
65
1
        Key: Borrow<QueryKey>,
66
1
    {
67
1
        self.map
68
1
            .get(key)
69
1
            .map(|node| self.cache.get_without_touch(*node).value())
70
1
    }
71

            
72
    /// Returns an [`EntryRef`] for `key`, if present.
73
    ///
74
    /// This function does not touch the key, preserving its current position in
75
    /// the lru cache. The [`EntryRef`] can touch the key, depending on which
76
    /// functions are used.
77
    ///
78
    /// ```rust
79
    /// use lrumap::{LruBTreeMap, LruMap, Removed};
80
    ///
81
    /// let mut lru = LruBTreeMap::new(3);
82
    /// lru.push(1, 1);
83
    /// lru.push(2, 2);
84
    /// lru.push(3, 3);
85
    ///
86
    /// // The cache has been updated once since entry 2 was touched.
87
    /// let mut entry = lru.entry(&2).unwrap();
88
    /// assert_eq!(entry.staleness(), 1);
89
    /// // Peeking the value will not update the entry's position.
90
    /// assert_eq!(entry.peek_value(), &2);
91
    /// assert_eq!(entry.staleness(), 1);
92
    /// // Querying the value or touching the entry will move it to the
93
    /// // front of the cache.
94
    /// assert_eq!(entry.value(), &2);
95
    /// assert_eq!(entry.staleness(), 0);
96
    ///
97
    /// assert_eq!(lru.head().unwrap().key(), &2);
98
    /// ```
99
9
    pub fn entry<QueryKey>(&mut self, key: &QueryKey) -> Option<EntryRef<'_, Self, Key, Value>>
100
9
    where
101
9
        QueryKey: Ord + ?Sized,
102
9
        Key: Borrow<QueryKey>,
103
9
    {
104
9
        self.map
105
9
            .get(key)
106
9
            .copied()
107
9
            .map(|node| EntryRef::new(self, node))
108
9
    }
109

            
110
    /// Inserts `value` for `key` into this map. If a value is already stored
111
    /// for this key, [`Removed::PreviousValue`] is returned with the previously
112
    /// stored value. If no value is currently stored and the map is full, the
113
    /// least recently used entry will be returned in [`Removed::Evicted`].
114
    /// Otherwise, `None` will be returned.
115
    ///
116
    /// This function touches the key, making it the most recently used key.
117
    ///
118
    /// ```rust
119
    /// use lrumap::{LruBTreeMap, LruMap, Removed};
120
    ///
121
    /// let mut lru = LruBTreeMap::new(3);
122
    /// lru.push(1, 1);
123
    /// lru.push(2, 2);
124
    /// lru.push(3, 3);
125
    ///
126
    /// // The cache is now full. The next push will evict an entry.
127
    /// let removed = lru.push(4, 4);
128
    /// assert_eq!(removed, Some(Removed::Evicted(1, 1)));
129
    ///
130
    /// // This leaves the cache with 4 as the most recent key, and 2 as the
131
    /// // least recent key.
132
    /// assert_eq!(lru.head().unwrap().key(), &4);
133
    /// assert_eq!(lru.tail().unwrap().key(), &2);
134
    /// ```
135
5018
    pub fn push(&mut self, key: Key, value: Value) -> Option<Removed<Key, Value>> {
136
5018
        // Create the new entry for this key/value pair, which also puts it at
137
5018
        // the front of the LRU
138
5018
        // let existing_entry = self.map.entry(key.clone());
139
5018
        let entry = self.map.entry(key.clone());
140

            
141
5018
        if let btree_map::Entry::Occupied(entry) = &entry {
142
1001
            let node_ref = *entry.get();
143
1001
            // Swap the value out.
144
1001
            let value = self.cache.get_mut(node_ref).replace_value(value);
145
1001

            
146
1001
            return Some(Removed::PreviousValue(value));
147
4017
        }
148
4017

            
149
4017
        // Key is not currently contained. Create a new node.
150
4017
        let (node, result) = self.cache.push(key, value);
151
4017

            
152
4017
        // Insert the node into the BTreeMap
153
4017
        entry.or_insert(node);
154

            
155
4017
        if let Some(Removed::Evicted(key, _)) = &result {
156
2006
            self.map.remove(key);
157
2014
        }
158

            
159
4017
        result
160
5018
    }
161

            
162
    /// Pushes all items from `iterator` into this map. If there are more
163
    /// entries in the iterator than capacity remaining, keys will be evicted as
164
    /// needed.
165
    ///
166
    /// This function is equivalent to a for loop calling [`Self::push()`].
167
    ///
168
    /// ```rust
169
    /// use lrumap::{LruBTreeMap, LruMap};
170
    ///
171
    /// let mut lru = LruBTreeMap::new(3);
172
    /// lru.extend([(1, 1), (2, 2), (3, 3), (4, 4)]);
173
    ///
174
    /// assert_eq!(lru.head().unwrap().key(), &4);
175
    /// assert_eq!(lru.tail().unwrap().key(), &2);
176
    /// ```
177
3
    pub fn extend<IntoIter: IntoIterator<Item = (Key, Value)>>(&mut self, iterator: IntoIter) {
178
18
        for (key, value) in iterator {
179
15
            let (node_id, removed) = self.cache.push(key.clone(), value);
180
15
            if let Some(Removed::Evicted(evicted_key, _)) = removed {
181
                self.map.remove(&evicted_key);
182
15
            }
183
15
            self.map.insert(key, node_id);
184
        }
185
3
    }
186

            
187
    /// Returns the most recently touched entry with a key within `range`.
188
    ///
189
    /// This function uses [`BTreeMap::range`] to identify all entries that
190
    /// match the given range. For each returned entry, the entry's
191
    /// [staleness](EntryRef::staleness) is compared, and the least stale entry
192
    /// is returned. If no keys match the range, `None` is returned.
193
    ///
194
    /// This function does not touch any keys, preserving the current order of
195
    /// the lru cache. The [`EntryRef`] returned can be used to peek, touch, or
196
    /// remove the entry.
197
    ///
198
    /// ```rust
199
    /// use lrumap::LruBTreeMap;
200
    ///
201
    /// let mut lru = LruBTreeMap::new(5);
202
    /// lru.extend([(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)]);
203
    ///
204
    /// assert_eq!(lru.most_recent_in_range(2..=4).unwrap().key(), &4);
205
    /// // Change the order by retrieving key 2.
206
    /// lru.get(&2);
207
    /// assert_eq!(lru.most_recent_in_range(2..=4).unwrap().key(), &2);
208
    /// ```
209
2
    pub fn most_recent_in_range<QueryKey, Range>(
210
2
        &mut self,
211
2
        range: Range,
212
2
    ) -> Option<EntryRef<'_, Self, Key, Value>>
213
2
    where
214
2
        QueryKey: Ord + ?Sized,
215
2
        Key: Borrow<QueryKey>,
216
2
        Range: RangeBounds<QueryKey>,
217
2
    {
218
6
        self.most_recent_in_range_where(range, |_, _| true)
219
2
    }
220

            
221
    /// Returns the most recently touched entry with a key within `range` that
222
    /// passes the `condition` check.
223
    ///
224
    /// This function uses [`BTreeMap::range`] to identify all entries that
225
    /// match the given range. Each key and value that matches is passed to
226
    /// `condition`. For each entry where `condition` returns true, the
227
    /// [staleness](EntryRef::staleness) is compared, and the least stale entry
228
    /// is returned. If no keys match the range, `None` is returned.
229
    ///
230
    /// This function does not touch any keys, preserving the current order of
231
    /// the lru cache. The [`EntryRef`] returned can be used to peek, touch, or
232
    /// remove the entry.
233
    ///
234
    /// ```rust
235
    /// use lrumap::LruBTreeMap;
236
    ///
237
    /// let mut lru = LruBTreeMap::<u32, u16>::new(5);
238
    /// lru.extend([(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)]);
239
    ///
240
    /// let condition = |key: &u32, value: &u16| key == &3 || value == &4;
241
    /// assert_eq!(
242
    ///     lru.most_recent_in_range_where(2..=4, condition)
243
    ///         .unwrap()
244
    ///         .key(),
245
    ///     &4
246
    /// );
247
    ///
248
    /// // Change the order by retrieving key 2. However, 2 doesn't meet the
249
    /// // condition, so the result is unchanged.
250
    /// lru.get(&2);
251
    /// assert_eq!(
252
    ///     lru.most_recent_in_range_where(2..=4, condition)
253
    ///         .unwrap()
254
    ///         .key(),
255
    ///     &4
256
    /// );
257
    ///
258
    /// // Request 3, moving it to the front. Since 3 matches the condition, the
259
    /// // result is now 3.
260
    /// lru.get(&3);
261
    /// assert_eq!(
262
    ///     lru.most_recent_in_range_where(2..=4, condition)
263
    ///         .unwrap()
264
    ///         .key(),
265
    ///     &3
266
    /// );
267
    /// ```
268
3
    pub fn most_recent_in_range_where<QueryKey, Range, Condition>(
269
3
        &mut self,
270
3
        range: Range,
271
3
        mut condition: Condition,
272
3
    ) -> Option<EntryRef<'_, Self, Key, Value>>
273
3
    where
274
3
        QueryKey: Ord + ?Sized,
275
3
        Key: Borrow<QueryKey>,
276
3
        Range: RangeBounds<QueryKey>,
277
3
        Condition: for<'key, 'value> FnMut(&'key Key, &'value Value) -> bool,
278
3
    {
279
3
        let mut closest_node = None;
280
3
        let mut closest_staleness = usize::MAX;
281
9
        for (_, &node_id) in self.map.range(range) {
282
9
            let node = self.cache.get_without_touch(node_id);
283
9
            if condition(node.key(), node.value()) {
284
8
                let staleness = self.cache.sequence().wrapping_sub(node.last_accessed());
285
8
                if staleness < closest_staleness {
286
6
                    closest_staleness = staleness;
287
6
                    closest_node = Some(node_id);
288
6
                }
289
1
            }
290
        }
291
3
        closest_node.map(|node| EntryRef::new(self, node))
292
3
    }
293
}
294

            
295
impl<Key, Value> LruMap<Key, Value> for LruBTreeMap<Key, Value>
296
where
297
    Key: Ord + Clone,
298
{
299
5
    fn new(capacity: usize) -> Self {
300
5
        Self::new(capacity)
301
5
    }
302

            
303
7
    fn len(&self) -> usize {
304
7
        self.cache.len()
305
7
    }
306

            
307
9
    fn head(&mut self) -> Option<EntryRef<'_, Self, Key, Value>> {
308
9
        self.cache.head().map(|node| EntryRef::new(self, node))
309
9
    }
310

            
311
5
    fn tail(&mut self) -> Option<EntryRef<'_, Self, Key, Value>> {
312
5
        self.cache.tail().map(|node| EntryRef::new(self, node))
313
5
    }
314

            
315
4
    fn iter(&self) -> crate::lru::Iter<'_, Key, Value> {
316
4
        self.cache.iter()
317
4
    }
318

            
319
8
    fn get<QueryKey>(&mut self, key: &QueryKey) -> Option<&Value>
320
8
    where
321
8
        QueryKey: Ord + Hash + Eq + ?Sized,
322
8
        Key: Borrow<QueryKey> + Ord + Eq + Hash,
323
8
    {
324
8
        self.get(key)
325
8
    }
326

            
327
1
    fn get_without_update<QueryKey>(&self, key: &QueryKey) -> Option<&Value>
328
1
    where
329
1
        QueryKey: Ord + Hash + Eq + ?Sized,
330
1
        Key: Borrow<QueryKey> + Ord + Eq + Hash,
331
1
    {
332
1
        self.get_without_update(key)
333
1
    }
334

            
335
9
    fn entry<QueryKey>(&mut self, key: &QueryKey) -> Option<EntryRef<'_, Self, Key, Value>>
336
9
    where
337
9
        QueryKey: Ord + Hash + Eq + ?Sized,
338
9
        Key: Borrow<QueryKey> + Ord + Eq + Hash,
339
9
    {
340
9
        self.entry(key)
341
9
    }
342

            
343
15
    fn push(&mut self, key: Key, value: Value) -> Option<Removed<Key, Value>> {
344
15
        self.push(key, value)
345
15
    }
346

            
347
2
    fn extend<IntoIter: IntoIterator<Item = (Key, Value)>>(&mut self, iterator: IntoIter) {
348
2
        self.extend(iterator);
349
2
    }
350
}
351

            
352
impl<Key, Value> EntryCache<Key, Value> for LruBTreeMap<Key, Value>
353
where
354
    Key: Ord + Clone,
355
{
356
46
    fn cache(&self) -> &LruCache<Key, Value> {
357
46
        &self.cache
358
46
    }
359

            
360
2
    fn cache_mut(&mut self) -> &mut LruCache<Key, Value> {
361
2
        &mut self.cache
362
2
    }
363

            
364
6
    fn remove(&mut self, node: NodeId) -> ((Key, Value), Option<NodeId>, Option<NodeId>) {
365
6
        let ((key, value), next, previous) = self.cache.remove(node);
366
6
        self.map.remove(&key);
367
6
        ((key, value), next, previous)
368
6
    }
369
}
370

            
371
impl<Key, Value> IntoIterator for LruBTreeMap<Key, Value>
372
where
373
    Key: Ord + Clone,
374
{
375
    type IntoIter = IntoIter<Key, Value>;
376
    type Item = (Key, Value);
377

            
378
1
    fn into_iter(self) -> Self::IntoIter {
379
1
        IntoIter::from(self.cache)
380
1
    }
381
}
382

            
383
1
#[test]
384
1
fn most_recent_in_range_test() {
385
1
    let mut lru = LruBTreeMap::new(5);
386
1
    lru.extend([(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)]);
387
1

            
388
1
    assert_eq!(lru.most_recent_in_range(2..=4).unwrap().key(), &4);
389
1
    lru.get(&2);
390
1
    assert_eq!(lru.most_recent_in_range(2..=4).unwrap().key(), &2);
391
1
    assert_eq!(
392
3
        lru.most_recent_in_range_where(2..=4, |key: &u32, _value: &u16| key != &2)
393
1
            .unwrap()
394
1
            .key(),
395
1
        &4
396
1
    );
397
1
}