ds_ext/
map.rs

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