ds_ext/
queue.rs

1//! A linked hash map ordered by insertion which can be reordered by swapping,
2//! useful as a simple priority queue (e.g. an LFU or LRU cache).
3
4use std::borrow::Borrow;
5use std::cell::{Ref, RefCell, RefMut};
6use std::collections::HashMap;
7use std::hash::Hash;
8use std::sync::Arc;
9use std::{fmt, mem};
10
11struct ItemState<K> {
12    prev: Option<Arc<K>>,
13    next: Option<Arc<K>>,
14}
15
16struct Item<K, V> {
17    key: Arc<K>,
18    value: V,
19    state: RefCell<ItemState<K>>,
20}
21
22impl<K, V> Item<K, V> {
23    #[inline]
24    fn state(&self) -> Ref<ItemState<K>> {
25        self.state.borrow()
26    }
27
28    #[inline]
29    fn state_mut(&self) -> RefMut<ItemState<K>> {
30        self.state.borrow_mut()
31    }
32}
33
34type Inner<K, V> = HashMap<Arc<K>, Item<K, V>>;
35
36/// An iterator over the contents of a [`LinkedHashMap`]
37pub struct IntoIter<K, V> {
38    queue: LinkedHashMap<K, V>,
39}
40
41impl<K: Eq + Hash + fmt::Debug, V> Iterator for IntoIter<K, V> {
42    type Item = (K, V);
43
44    fn next(&mut self) -> Option<Self::Item> {
45        self.queue.pop_first_entry()
46    }
47
48    fn size_hint(&self) -> (usize, Option<usize>) {
49        (self.queue.len(), Some(self.queue.len()))
50    }
51}
52
53impl<K: Eq + Hash + fmt::Debug, V> DoubleEndedIterator for IntoIter<K, V> {
54    fn next_back(&mut self) -> Option<Self::Item> {
55        self.queue.pop_last_entry()
56    }
57}
58
59/// An iterator over the entries in a [`LinkedHashMap`]
60pub struct Iter<'a, K, V> {
61    list: &'a Inner<K, V>,
62    next: Option<Arc<K>>,
63    last: Option<Arc<K>>,
64    size: usize,
65}
66
67impl<'a, K: Eq + Hash, V> Iterator for Iter<'a, K, V> {
68    type Item = (&'a K, &'a V);
69
70    fn next(&mut self) -> Option<Self::Item> {
71        let next = self.next.take()?;
72        let (key, item) = self.list.get_key_value(&*next).expect("next");
73
74        if self.last == Some(next) {
75            self.next = None;
76            self.last = None;
77        } else {
78            self.next = item.state().next.clone();
79        }
80
81        self.size -= 1;
82
83        Some((key, &item.value))
84    }
85
86    fn size_hint(&self) -> (usize, Option<usize>) {
87        (self.size, Some(self.size))
88    }
89}
90
91impl<'a, K: Eq + Hash, V> DoubleEndedIterator for Iter<'a, K, V> {
92    fn next_back(&mut self) -> Option<Self::Item> {
93        let last = self.last.take()?;
94        let (key, item) = self.list.get_key_value(&*last).expect("next");
95
96        if self.next == Some(last) {
97            self.next = None;
98            self.last = None;
99        } else {
100            self.last = item.state().prev.clone();
101        }
102
103        self.size -= 1;
104
105        Some((key, &item.value))
106    }
107}
108
109/// An iterator over the keys in a [`LinkedHashMap`]
110pub struct Keys<'a, K, V> {
111    inner: Iter<'a, K, V>,
112}
113
114impl<'a, K: Hash + Eq, V> Iterator for Keys<'a, K, V> {
115    type Item = &'a K;
116
117    fn next(&mut self) -> Option<Self::Item> {
118        self.inner.next().map(|(key, _value)| key)
119    }
120
121    fn size_hint(&self) -> (usize, Option<usize>) {
122        self.inner.size_hint()
123    }
124}
125
126impl<'a, K: Hash + Eq, V> DoubleEndedIterator for Keys<'a, K, V> {
127    fn next_back(&mut self) -> Option<Self::Item> {
128        self.inner.next_back().map(|(key, _value)| key)
129    }
130}
131
132/// An iterator over the values in a [`LinkedHashMap`]
133pub struct Values<'a, K, V> {
134    inner: Iter<'a, K, V>,
135}
136
137impl<'a, K: Eq + Hash, V> Iterator for Values<'a, K, V> {
138    type Item = &'a V;
139
140    fn next(&mut self) -> Option<Self::Item> {
141        self.inner.next().map(|(_key, value)| value)
142    }
143
144    fn size_hint(&self) -> (usize, Option<usize>) {
145        self.inner.size_hint()
146    }
147}
148
149impl<'a, K: Eq + Hash, V> DoubleEndedIterator for Values<'a, K, V> {
150    fn next_back(&mut self) -> Option<Self::Item> {
151        self.inner.next_back().map(|(_key, value)| value)
152    }
153}
154
155/// A hash map in insertion order which can be reordered using [`Self::bump`] and [`Self::swap`].
156pub struct LinkedHashMap<K, V> {
157    list: Inner<K, V>,
158    head: Option<Arc<K>>,
159    tail: Option<Arc<K>>,
160}
161
162impl<K: Clone + Eq + Hash, V: Clone> Clone for LinkedHashMap<K, V> {
163    fn clone(&self) -> Self {
164        let mut other = Self::with_capacity(self.list.capacity());
165
166        for (key, item) in &self.list {
167            let key = K::clone(&**key);
168            let value = V::clone(&item.value);
169            other.insert(key, value);
170        }
171
172        other
173    }
174}
175
176impl<K: Eq + Hash, V> LinkedHashMap<K, V> {
177    /// Construct a new [`LinkedHashMap`].
178    pub fn new() -> Self {
179        Self {
180            list: HashMap::new(),
181            head: None,
182            tail: None,
183        }
184    }
185
186    /// Construct a new [`LinkedHashMap`] with the given `capacity`.
187    pub fn with_capacity(capacity: usize) -> Self {
188        Self {
189            list: HashMap::with_capacity(capacity),
190            head: None,
191            tail: None,
192        }
193    }
194
195    /// If `key` is present, increase its priority by one and return `true`.
196    pub fn bump(&mut self, key: &K) -> bool {
197        let item = if let Some(item) = self.list.get(key) {
198            item
199        } else {
200            return false;
201        };
202
203        let mut item_state = item.state_mut();
204
205        if item_state.prev.is_none() {
206            // can't bump the first item
207            return true;
208        } else if item_state.next.is_none() && item_state.prev.is_some() {
209            // bump the last item
210
211            let prev_key = item_state.prev.as_ref().expect("prev key").clone();
212            let mut prev = self.list.get::<K>(&prev_key).expect("prev").state_mut();
213
214            mem::swap(&mut prev.next, &mut item_state.next); // set prev.next
215            mem::swap(&mut item_state.prev, &mut prev.prev); // set item.prev
216            mem::swap(&mut item_state.next, &mut prev.prev); // set item.next & prev.prev
217
218            self.tail = Some(prev_key)
219        } else {
220            // bump an item in the middle
221
222            let prev_key = item_state.prev.as_ref().expect("previous key").clone();
223            let mut prev = self.list.get::<K>(&prev_key).expect("prev").state_mut();
224
225            let next_key = item_state.next.as_ref().expect("next key").clone();
226            let mut next = self.list.get::<K>(&next_key).expect("next").state_mut();
227
228            mem::swap(&mut next.prev, &mut item_state.prev); // set next.prev
229            mem::swap(&mut item_state.prev, &mut prev.prev); // set item.prev
230            mem::swap(&mut prev.next, &mut item_state.next); // set prev.next
231
232            item_state.next = Some(prev_key);
233        }
234
235        if let Some(prev_key) = &item_state.prev {
236            let mut prev = self.list.get::<K>(prev_key).expect("prev").state_mut();
237            prev.next = Some(item.key.clone());
238        } else {
239            self.head = Some(item.key.clone());
240        }
241
242        std::mem::drop(item_state);
243
244        true
245    }
246
247    /// Remove all entries from this [`LinkedHashMap`].
248    pub fn clear(&mut self) {
249        self.list.clear();
250        self.head = None;
251        self.tail = None;
252    }
253
254    /// Return `true` if there is an entry for the given `key` in this [`LinkedHashMap`].
255    pub fn contains_key<Q>(&self, key: &Q) -> bool
256    where
257        Arc<K>: Borrow<Q>,
258        Q: Eq + Hash + ?Sized,
259    {
260        self.list.contains_key(key)
261    }
262
263    /// Consume the `iter` and insert all its elements into this [`LinkedHashMap`].
264    pub fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
265        for (key, value) in iter {
266            self.insert(key, value);
267        }
268    }
269
270    /// Borrow the value at the given `key`, if present.
271    pub fn get<Q>(&self, key: &Q) -> Option<&V>
272    where
273        Arc<K>: Borrow<Q>,
274        Q: Eq + Hash + ?Sized,
275    {
276        self.list.get(key).map(|item| &item.value)
277    }
278
279    /// Borrow the entry at the given `key`, if present.
280    pub fn get_key_value<Q>(&self, key: &Q) -> Option<(&K, &V)>
281    where
282        Arc<K>: Borrow<Q>,
283        Q: Eq + Hash + ?Sized,
284    {
285        self.list
286            .get_key_value(key)
287            .map(|(key, item)| (&**key, &item.value))
288    }
289
290    /// Borrow the value at the given `key` mutably, if present.
291    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
292    where
293        Arc<K>: Borrow<Q>,
294        Q: Eq + Hash + ?Sized,
295    {
296        self.list.get_mut(key).map(|item| &mut item.value)
297    }
298
299    /// Insert a new `value` at `key` and return the previous value, if any.
300    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
301        let old_value = self.remove(&key);
302
303        let key = Arc::new(key);
304        let mut next = Some(key.clone());
305        mem::swap(&mut self.head, &mut next);
306
307        if let Some(prev_key) = &next {
308            let mut prev = self.list.get::<K>(prev_key).expect("prev").state_mut();
309            prev.prev = Some(key.clone());
310        } else {
311            debug_assert!(self.tail.is_none());
312            self.tail = Some(key.clone());
313        }
314
315        let item = Item {
316            key: key.clone(),
317            value,
318            state: RefCell::new(ItemState { prev: None, next }),
319        };
320
321        assert!(self.list.insert(key, item).is_none());
322
323        old_value
324    }
325
326    /// Construct an iterator over the entries in this [`LinkedHashMap`].
327    pub fn iter(&self) -> Iter<K, V> {
328        Iter {
329            list: &self.list,
330            next: self.head.clone(),
331            last: self.tail.clone(),
332            size: self.len(),
333        }
334    }
335
336    /// Return `true` if this [`LinkedHashMap`] is empty.
337    pub fn is_empty(&self) -> bool {
338        self.list.is_empty()
339    }
340
341    /// Construct an iterator over keys of this [`LinkedHashMap`].
342    pub fn keys(&self) -> Keys<K, V> {
343        Keys { inner: self.iter() }
344    }
345
346    /// Return the size of this [`LinkedHashMap`].
347    pub fn len(&self) -> usize {
348        self.list.len()
349    }
350
351    /// Remove and return the first value in this [`LinkedHashMap`].
352    pub fn pop_first(&mut self) -> Option<V> {
353        if self.head.is_none() {
354            return None;
355        }
356
357        let item = self
358            .list
359            .remove(self.head.as_ref().expect("head"))
360            .expect("head");
361
362        Some(self.remove_inner(item))
363    }
364
365    /// Remove and return the first entry in this [`LinkedHashMap`].
366    pub fn pop_first_entry(&mut self) -> Option<(K, V)>
367    where
368        K: fmt::Debug,
369    {
370        if self.head.is_none() {
371            return None;
372        }
373
374        let (key, item) = self
375            .list
376            .remove_entry(self.head.as_ref().expect("head"))
377            .expect("head");
378
379        let value = self.remove_inner(item);
380        let key = Arc::try_unwrap(key).expect("key");
381        Some((key, value))
382    }
383
384    /// Remove and return the last value in this [`LinkedHashMap`].
385    pub fn pop_last(&mut self) -> Option<V> {
386        if self.tail.is_none() {
387            return None;
388        }
389
390        let item = self
391            .list
392            .remove(self.tail.as_ref().expect("tail"))
393            .expect("tail");
394
395        Some(self.remove_inner(item))
396    }
397
398    /// Remove and return the last entry in this [`LinkedHashMap`].
399    pub fn pop_last_entry(&mut self) -> Option<(K, V)>
400    where
401        K: fmt::Debug,
402    {
403        if self.tail.is_none() {
404            return None;
405        }
406
407        let (key, item) = self
408            .list
409            .remove_entry(self.tail.as_ref().expect("tail"))
410            .expect("tail");
411
412        let value = self.remove_inner(item);
413        let key = Arc::try_unwrap(key).expect("key");
414        Some((key, value))
415    }
416
417    fn remove_inner(&mut self, item: Item<K, V>) -> V {
418        let mut item_state = item.state_mut();
419
420        if item_state.prev.is_none() && item_state.next.is_none() {
421            // there was only one item and now the map is empty
422            self.head = None;
423            self.tail = None;
424        } else if item_state.prev.is_none() {
425            // the first item has been removed
426            self.head = item_state.next.clone();
427
428            let next_key = self.head.as_ref().expect("next key");
429            let mut next = self.list.get::<K>(&next_key).expect("next").state_mut();
430
431            mem::swap(&mut next.prev, &mut item_state.prev);
432        } else if item_state.next.is_none() {
433            // the last item has been removed
434            self.tail = item_state.prev.clone();
435
436            let prev_key = self.tail.as_ref().expect("previous key");
437            let mut prev = self.list.get::<K>(prev_key).expect("prev").state_mut();
438
439            mem::swap(&mut prev.next, &mut item_state.next);
440        } else {
441            // an item in the middle has been removed
442            let prev_key = item_state.prev.as_ref().expect("previous key");
443            let mut prev = self.list.get::<K>(prev_key).expect("prev").state_mut();
444
445            let next_key = item_state.next.as_ref().expect("next key");
446            let mut next = self.list.get::<K>(next_key).expect("next item").state_mut();
447
448            mem::swap(&mut next.prev, &mut item_state.prev);
449            mem::swap(&mut prev.next, &mut item_state.next);
450        }
451
452        std::mem::drop(item_state);
453
454        item.value
455    }
456
457    /// Remove an entry from this [`LinkedHashMap`] and return its value, if present.
458    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
459    where
460        Arc<K>: Borrow<Q>,
461        Q: Hash + Eq + ?Sized,
462    {
463        let item = self.list.remove(key)?;
464        Some(self.remove_inner(item))
465    }
466
467    /// Remove and return an entry from this [`LinkedHashMap`], if present.
468    pub fn remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V)>
469    where
470        K: fmt::Debug,
471        Arc<K>: Borrow<Q>,
472        Q: Hash + Eq + ?Sized,
473    {
474        let (key, item) = self.list.remove_entry(key)?;
475        let key = Arc::try_unwrap(key).expect("key");
476        Some((key, self.remove_inner(item)))
477    }
478
479    /// Swap the position of two keys in this [`LinkedHashMap`].
480    /// Returns `true` if both keys are present.
481    pub fn swap<Q>(&mut self, l: &Q, r: &Q) -> bool
482    where
483        Arc<K>: Borrow<Q>,
484        Q: Hash + Eq + ?Sized,
485    {
486        if l == r {
487            return self.contains_key(l) && self.contains_key(r);
488        }
489
490        let (l_key, l_item) = if let Some(entry) = self.list.get_key_value(l) {
491            entry
492        } else {
493            return false;
494        };
495
496        let (r_key, r_item) = if let Some(entry) = self.list.get_key_value(r) {
497            entry
498        } else {
499            return false;
500        };
501
502        if l_item.state().next.as_ref() == Some(r_key) {
503            let key = r_key.clone();
504            return self.bump(&key);
505        } else if r_item.state().next.as_ref() == Some(l_key) {
506            let key = l_key.clone();
507            return self.bump(&key);
508        } else {
509            let mut l_state = l_item.state_mut();
510            let mut r_state = r_item.state_mut();
511            mem::swap(&mut *l_state, &mut *r_state);
512        }
513
514        if self.head.as_ref() == Some(l_key) {
515            self.head = Some(r_key.clone());
516        } else if self.head.as_ref() == Some(r_key) {
517            self.head = Some(l_key.clone());
518        }
519
520        if self.tail.as_ref() == Some(l_key) {
521            self.tail = Some(r_key.clone());
522        } else if self.tail.as_ref() == Some(r_key) {
523            self.tail = Some(l_key.clone());
524        }
525
526        true
527    }
528
529    /// Construct an iterator over the values in this [`LinkedHashMap`].
530    pub fn values(&self) -> Values<K, V> {
531        Values { inner: self.iter() }
532    }
533}
534
535impl<K: Eq + Hash, V> FromIterator<(K, V)> for LinkedHashMap<K, V> {
536    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
537        let iter = iter.into_iter();
538        let mut map = match iter.size_hint() {
539            (_, Some(max)) => Self::with_capacity(max),
540            (min, None) if min > 0 => Self::with_capacity(min),
541            _ => Self::new(),
542        };
543
544        map.extend(iter);
545        map
546    }
547}
548
549impl<K: Eq + Hash + fmt::Debug, V> IntoIterator for LinkedHashMap<K, V> {
550    type Item = (K, V);
551    type IntoIter = IntoIter<K, V>;
552
553    fn into_iter(self) -> Self::IntoIter {
554        IntoIter { queue: self }
555    }
556}
557
558impl<'a, K: Eq + Hash, V> IntoIterator for &'a LinkedHashMap<K, V> {
559    type Item = (&'a K, &'a V);
560    type IntoIter = Iter<'a, K, V>;
561
562    fn into_iter(self) -> Self::IntoIter {
563        self.iter()
564    }
565}
566
567impl<K: Eq + Hash + fmt::Debug, V: fmt::Debug> fmt::Debug for LinkedHashMap<K, V> {
568    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
569        f.write_str("{")?;
570
571        for (key, value) in self.iter() {
572            write!(f, "{:?}: {:?}, ", key, value)?;
573        }
574
575        f.write_str("}")
576    }
577}
578
579#[allow(dead_code)]
580fn validate<K: Eq + Hash + fmt::Debug, V>(queue: &LinkedHashMap<K, V>) {
581    if queue.list.is_empty() {
582        assert!(queue.head.is_none(), "head is {:?}", queue.head);
583        assert!(queue.tail.is_none(), "tail is {:?}", queue.tail);
584    } else {
585        let first_key = queue.head.as_ref().expect("first key");
586        let first = queue.list.get::<K>(first_key).expect("first item");
587        assert_eq!(first.state().prev, None);
588
589        let last_key = queue.tail.as_ref().expect("last key");
590        let last = queue.list.get::<K>(last_key).expect("last item");
591        assert_eq!(last.state().next, None);
592    }
593
594    let mut size = 0;
595    let mut last = None;
596    let mut next = queue.head.clone();
597    while let Some(key) = next {
598        let item = queue.list.get::<K>(&key).expect("item");
599
600        let item_state = item.state.borrow();
601        assert_ne!(item_state.prev.as_ref(), Some(&key));
602        assert_ne!(item_state.next.as_ref(), Some(&key));
603
604        let prev_key = item_state.prev.as_ref();
605        assert_eq!(last.as_ref(), prev_key);
606
607        last = Some(key);
608        next = item.state.borrow().next.clone();
609        size += 1;
610    }
611
612    assert_eq!(size, queue.len());
613}
614
615#[allow(dead_code)]
616fn print_debug<K: fmt::Debug + Eq + Hash, V>(queue: &LinkedHashMap<K, V>) {
617    let mut next = queue.head.clone();
618
619    if next.is_some() {
620        println!();
621    }
622
623    while let Some(next_key) = next {
624        let item = queue.list.get::<K>(&next_key).expect("item").state();
625
626        if let Some(prev_key) = item.prev.as_ref() {
627            print!("{:?}-", prev_key);
628        }
629
630        print!("{:?}", next_key);
631
632        next = item.next.clone();
633        if let Some(next_key) = &next {
634            print!("-{:?}", next_key);
635        }
636
637        println!();
638    }
639}
640
641#[cfg(test)]
642mod tests {
643    use rand::{thread_rng, Rng};
644
645    use super::*;
646
647    #[test]
648    fn test_order() {
649        let mut queue = LinkedHashMap::new();
650        let expected: Vec<i32> = (0..10).collect();
651
652        for i in expected.iter() {
653            queue.insert(*i, i.to_string());
654            validate(&queue);
655        }
656
657        assert_eq!(queue.len(), expected.len());
658
659        let mut actual = Vec::with_capacity(expected.len());
660        for (i, s) in queue.iter() {
661            assert_eq!(&i.to_string(), s);
662            actual.push(i);
663        }
664
665        assert_eq!(actual.len(), expected.len());
666        assert!(actual
667            .iter()
668            .zip(expected.into_iter().rev())
669            .all(|(l, r)| **l == r))
670    }
671
672    #[test]
673    fn test_access() {
674        let mut queue = LinkedHashMap::new();
675        validate(&queue);
676
677        let mut rng = thread_rng();
678        for _ in 1..100_000 {
679            let i: i32 = rng.gen_range(0..1000);
680            queue.insert(i, i.to_string());
681            validate(&queue);
682
683            let mut size = 0;
684            for _ in queue.iter() {
685                size += 1;
686            }
687
688            assert_eq!(queue.len(), size);
689            assert!(!queue.is_empty());
690
691            while !queue.is_empty() {
692                let i: i32 = rng.gen_range(0..queue.len() as i32);
693                queue.bump(&i);
694                validate(&queue);
695
696                if queue.pop_first().is_some() {
697                    validate(&queue);
698                    size -= 1;
699                }
700
701                if !queue.is_empty() {
702                    let i: i32 = rng.gen_range(0..**queue.tail.as_ref().expect("tail"));
703                    queue.bump(&i);
704                    validate(&queue);
705                }
706
707                if queue.pop_last().is_some() {
708                    validate(&queue);
709                    size -= 1;
710                }
711
712                assert_eq!(queue.len(), size);
713            }
714
715            assert_eq!(queue.len(), 0);
716        }
717    }
718}
ds_ext/queue.rs

ds_ext/
queue.rs
