1
use std::borrow::Borrow;
2
#[cfg(not(feature = "hashbrown"))]
3
use std::collections::{hash_map, hash_map::RandomState as DefaultState, HashMap};
4
use std::fmt::Debug;
5
use std::hash::{BuildHasher, Hash};
6

            
7
#[cfg(feature = "hashbrown")]
8
use hashbrown::{
9
    hash_map::{self, DefaultHashBuilder as DefaultState},
10
    HashMap,
11
};
12

            
13
use crate::lru::{EntryCache, EntryRef, IntoIter, LruCache, NodeId, Removed};
14
use crate::LruMap;
15

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

            
34
impl<Key, Value> LruHashMap<Key, Value, DefaultState>
35
where
36
    Key: Hash + Eq + Clone,
37
{
38
    /// Creates a new map with the maximum `capacity`.
39
    ///
40
    /// # Panics
41
    ///
42
    /// Panics if `capacity` is <= 1.
43
7
    pub fn new(capacity: usize) -> Self {
44
7
        assert!(capacity > 1);
45
7
        Self {
46
7
            map: HashMap::with_capacity(capacity),
47
7
            cache: LruCache::new(capacity),
48
7
        }
49
7
    }
50
}
51

            
52
impl<Key, Value, State> LruHashMap<Key, Value, State>
53
where
54
    Key: Hash + Eq + Clone,
55
    State: BuildHasher,
56
{
57
    /// Creates a new map with the maximum `capacity` and `hasher`.
58
    ///
59
    /// # Panics
60
    ///
61
    /// Panics if `capacity` is <= 1
62
    pub fn with_hasher(capacity: usize, hasher: State) -> Self {
63
        assert!(capacity > 1);
64
        Self {
65
            map: HashMap::with_capacity_and_hasher(capacity, hasher),
66
            cache: LruCache::new(capacity),
67
        }
68
    }
69

            
70
    /// Returns the stored value for `key`, if present.
71
    ///
72
    /// This function touches the key, making it the most recently used key.
73
1009
    pub fn get<QueryKey>(&mut self, key: &QueryKey) -> Option<&Value>
74
1009
    where
75
1009
        QueryKey: Hash + Eq + ?Sized,
76
1009
        Key: Borrow<QueryKey>,
77
1009
    {
78
1009
        let node = self.map.get(key).copied();
79
1009
        node.map(|node| self.cache.get(node).value())
80
1009
    }
81

            
82
    /// Returns the stored value for `key`, if present.
83
    ///
84
    /// This function touches the key, making it the most recently used key.
85
    pub fn get_mut<QueryKey>(&mut self, key: &QueryKey) -> Option<&mut Value>
86
    where
87
        QueryKey: Hash + Eq + ?Sized,
88
        Key: Borrow<QueryKey>,
89
    {
90
        let node = self.map.get(key).copied();
91
        node.map(|node| self.cache.get_mut(node).value_mut())
92
    }
93

            
94
    /// Returns the stored value for `key`, if present.
95
    ///
96
    /// This function does not touch the key, preserving its current position in
97
    /// the lru cache.
98
1
    pub fn get_without_update<QueryKey>(&self, key: &QueryKey) -> Option<&Value>
99
1
    where
100
1
        QueryKey: Hash + Eq + ?Sized,
101
1
        Key: Borrow<QueryKey>,
102
1
    {
103
1
        self.map
104
1
            .get(key)
105
1
            .map(|node| self.cache.get_without_touch(*node).value())
106
1
    }
107

            
108
    /// Returns an [`EntryRef`] for `key`, if present.
109
    ///
110
    /// This function does not touch the key, preserving its current position in
111
    /// the lru cache. The [`EntryRef`] can touch the key, depending on which
112
    /// functions are used.
113
    ///
114
    /// ```rust
115
    /// use lrumap::{LruHashMap, LruMap, Removed};
116
    ///
117
    /// let mut lru = LruHashMap::new(3);
118
    /// lru.push(1, 1);
119
    /// lru.push(2, 2);
120
    /// lru.push(3, 3);
121
    ///
122
    /// // The cache has been updated once since entry 2 was touched.
123
    /// let mut entry = lru.entry(&2).unwrap();
124
    /// assert_eq!(entry.staleness(), 1);
125
    /// // Peeking the value will not update the entry's position.
126
    /// assert_eq!(entry.peek_value(), &2);
127
    /// assert_eq!(entry.staleness(), 1);
128
    /// // Querying the value or touching the entry will move it to the
129
    /// // front of the cache.
130
    /// assert_eq!(entry.value(), &2);
131
    /// assert_eq!(entry.staleness(), 0);
132
    ///
133
    /// assert_eq!(lru.head().unwrap().key(), &2);
134
    /// ```
135
9
    pub fn entry<QueryKey>(&mut self, key: &QueryKey) -> Option<EntryRef<'_, Self, Key, Value>>
136
9
    where
137
9
        QueryKey: Hash + Eq + ?Sized,
138
9
        Key: Borrow<QueryKey>,
139
9
    {
140
9
        self.map
141
9
            .get(key)
142
9
            .copied()
143
9
            .map(|node| EntryRef::new(self, node))
144
9
    }
145

            
146
    /// Inserts `value` for `key` into this map. If a value is already stored
147
    /// for this key, [`Removed::PreviousValue`] is returned with the previously
148
    /// stored value. If no value is currently stored and the map is full, the
149
    /// least recently used entry will be returned in [`Removed::Evicted`].
150
    /// Otherwise, `None` will be returned.
151
    ///
152
    /// This function touches the key, making it the most recently used key.
153
    ///
154
    /// ```rust
155
    /// use lrumap::{LruHashMap, LruMap, Removed};
156
    ///
157
    /// let mut lru = LruHashMap::new(3);
158
    /// lru.push(1, 1);
159
    /// lru.push(2, 2);
160
    /// lru.push(3, 3);
161
    ///
162
    /// // The cache is now full. The next push will evict an entry.
163
    /// let removed = lru.push(4, 4);
164
    /// assert_eq!(removed, Some(Removed::Evicted(1, 1)));
165
    ///
166
    /// // This leaves the cache with 4 as the most recent key, and 2 as the
167
    /// // least recent key.
168
    /// assert_eq!(lru.head().unwrap().key(), &4);
169
    /// assert_eq!(lru.tail().unwrap().key(), &2);
170
    /// ```
171
5028
    pub fn push(&mut self, key: Key, value: Value) -> Option<Removed<Key, Value>> {
172
5028
        // Create the new entry for this key/value pair, which also puts it at
173
5028
        // the front of the LRU
174
5028
        // let existing_entry = self.map.entry(key.clone());
175
5028
        let entry = self.map.entry(key.clone());
176

            
177
5028
        if let hash_map::Entry::Occupied(entry) = &entry {
178
1001
            let node_ref = *entry.get();
179
1001
            // Swap the value out.
180
1001
            let value = self.cache.get_mut(node_ref).replace_value(value);
181
1001

            
182
1001
            return Some(Removed::PreviousValue(value));
183
4027
        }
184
4027

            
185
4027
        // Key is not currently contained. Create a new node.
186
4027
        let (node, result) = self.cache.push(key, value);
187
4027

            
188
4027
        // Insert the node
189
4027
        entry.or_insert(node);
190

            
191
4027
        if let Some(Removed::Evicted(key, _)) = &result {
192
2006
            self.map.remove(key);
193
2024
        }
194

            
195
4027
        result
196
5028
    }
197

            
198
    /// Pushes all items from `iterator` into this map. If there are more
199
    /// entries in the iterator than capacity remaining, keys will be evicted as
200
    /// needed.
201
    ///
202
    /// This function is equivalent to a for loop calling [`Self::push()`].
203
    ///
204
    /// ```rust
205
    /// use lrumap::{LruHashMap, LruMap};
206
    ///
207
    /// let mut lru = LruHashMap::new(3);
208
    /// lru.extend([(1, 1), (2, 2), (3, 3), (4, 4)]);
209
    ///
210
    /// assert_eq!(lru.head().unwrap().key(), &4);
211
    /// assert_eq!(lru.tail().unwrap().key(), &2);
212
    /// ```
213
2
    pub fn extend<IntoIter: IntoIterator<Item = (Key, Value)>>(&mut self, iterator: IntoIter) {
214
12
        for (key, value) in iterator {
215
10
            self.push(key, value);
216
10
        }
217
2
    }
218
}
219

            
220
impl<Key, Value> LruMap<Key, Value> for LruHashMap<Key, Value, DefaultState>
221
where
222
    Key: Hash + Eq + Clone,
223
{
224
5
    fn new(capacity: usize) -> Self {
225
5
        Self::new(capacity)
226
5
    }
227

            
228
7
    fn len(&self) -> usize {
229
7
        self.cache.len()
230
7
    }
231

            
232
9
    fn head(&mut self) -> Option<EntryRef<'_, Self, Key, Value>> {
233
9
        self.cache.head().map(|node| EntryRef::new(self, node))
234
9
    }
235

            
236
5
    fn tail(&mut self) -> Option<EntryRef<'_, Self, Key, Value>> {
237
5
        self.cache.tail().map(|node| EntryRef::new(self, node))
238
5
    }
239

            
240
8
    fn get<QueryKey>(&mut self, key: &QueryKey) -> Option<&Value>
241
8
    where
242
8
        QueryKey: Ord + Hash + Eq + ?Sized,
243
8
        Key: Borrow<QueryKey> + Ord + Hash + Eq,
244
8
    {
245
8
        self.get(key)
246
8
    }
247

            
248
1
    fn get_without_update<QueryKey>(&self, key: &QueryKey) -> Option<&Value>
249
1
    where
250
1
        QueryKey: Ord + Hash + Eq + ?Sized,
251
1
        Key: Borrow<QueryKey> + Ord + Hash + Eq,
252
1
    {
253
1
        self.get_without_update(key)
254
1
    }
255

            
256
9
    fn entry<QueryKey>(&mut self, key: &QueryKey) -> Option<EntryRef<'_, Self, Key, Value>>
257
9
    where
258
9
        QueryKey: Ord + Hash + Eq + ?Sized,
259
9
        Key: Borrow<QueryKey> + Ord + Hash + Eq,
260
9
    {
261
9
        self.entry(key)
262
9
    }
263

            
264
15
    fn push(&mut self, key: Key, value: Value) -> Option<Removed<Key, Value>> {
265
15
        self.push(key, value)
266
15
    }
267

            
268
4
    fn iter(&self) -> crate::lru::Iter<'_, Key, Value> {
269
4
        self.cache.iter()
270
4
    }
271

            
272
2
    fn extend<IntoIter: IntoIterator<Item = (Key, Value)>>(&mut self, iterator: IntoIter) {
273
2
        self.extend(iterator);
274
2
    }
275
}
276

            
277
impl<Key, Value, State> EntryCache<Key, Value> for LruHashMap<Key, Value, State>
278
where
279
    Key: Hash + Eq + Clone,
280
    State: BuildHasher,
281
{
282
43
    fn cache(&self) -> &LruCache<Key, Value> {
283
43
        &self.cache
284
43
    }
285

            
286
2
    fn cache_mut(&mut self) -> &mut LruCache<Key, Value> {
287
2
        &mut self.cache
288
2
    }
289

            
290
6
    fn remove(&mut self, node: NodeId) -> ((Key, Value), Option<NodeId>, Option<NodeId>) {
291
6
        let ((key, value), next, previous) = self.cache.remove(node);
292
6
        self.map.remove(&key);
293
6
        ((key, value), next, previous)
294
6
    }
295
}
296

            
297
impl<Key, Value, State> IntoIterator for LruHashMap<Key, Value, State>
298
where
299
    Key: Hash + Eq + Clone,
300
    State: BuildHasher,
301
{
302
    type IntoIter = IntoIter<Key, Value>;
303
    type Item = (Key, Value);
304

            
305
1
    fn into_iter(self) -> Self::IntoIter {
306
1
        IntoIter::from(self.cache)
307
1
    }
308
}