defaultdict/
default_hashmap.rs

1#![deny(missing_docs)]
2
3use std::borrow::Borrow;
4use std::collections::hash_map::{
5    Drain, Entry, ExtractIf, HashMap, IntoIter, IntoKeys, IntoValues, Iter, IterMut, Keys,
6    RandomState, Values, ValuesMut,
7};
8use std::default::Default;
9use std::hash::{BuildHasher, Hash};
10use std::ops::Index;
11/// This struct mimicks the behaviour of a python defaultdict. This means alongside the traitbounds
12/// that apply on the key and value that are inherited from the [`HashMap`], it also requires the
13/// [`Default`] trait be implemented on the value type.
14#[derive(Clone, Debug)]
15pub struct DefaultHashMap<K, V, S = RandomState>
16where
17    K: Eq + Hash,
18    V: Default,
19{
20    _inner: HashMap<K, V, S>,
21    _default: V,
22}
23
24impl<K, V> DefaultHashMap<K, V, RandomState>
25where
26    K: Eq + Hash,
27    V: Default,
28{
29    /// Creates an empty [`DefaultHashMap`].
30    ///
31    /// # Example
32    /// ```
33    /// use defaultdict::DefaultHashMap;
34    ///
35    /// let map = DefaultHashMap::<i8, i8>::new();
36    ///
37    /// // This map is empty
38    /// let collected: Vec<(i8, i8)> = map.into_iter().collect();
39    /// let empty: Vec<(i8, i8)> = vec![];
40    ///
41    /// assert_eq!(empty, collected);
42    /// ```
43    #[must_use]
44    pub fn new() -> Self {
45        Self {
46            _inner: HashMap::new(),
47            _default: V::default(),
48        }
49    }
50}
51
52impl<K, V, S> DefaultHashMap<K, V, S>
53where
54    K: Eq + Hash,
55    V: Default,
56    S: BuildHasher,
57{
58    /// Returns the number of elements the map can hold without reallocating.
59    ///
60    /// This number is a lower bound; the `HashMap<K, V>` might be able to hold more, but is
61    /// guaranteed to be able to hold at least this many.
62    ///
63    /// # Example
64    /// ```
65    /// use defaultdict::DefaultHashMap;
66    ///
67    /// let mut map = DefaultHashMap::<i8, i8>::new();
68    /// map.insert(10, 20);
69    ///
70    /// println!("{}", map.capacity());
71    /// ```
72    #[inline]
73    pub fn capacity(&self) -> usize {
74        self._inner.capacity()
75    }
76
77    /// Clears the map, removing all key-value pairs. Keeps the allocated memory for reuse.
78    ///
79    /// # Example
80    /// ```
81    /// use defaultdict::DefaultHashMap;
82    ///
83    /// let mut map = DefaultHashMap::<i8, i8>::new();
84    /// map.insert(10, 20);
85    /// map.clear();
86    ///
87    /// // This map is empty
88    /// let collected: Vec<(i8, i8)> = map.into_iter().collect();
89    /// let empty: Vec<(i8, i8)> = vec![];
90    ///
91    /// assert_eq!(empty, collected);
92    /// ```
93    #[inline]
94    pub fn clear(&mut self) {
95        self._inner.clear()
96    }
97
98    /// Returns `true` if the key passed in exists in the HashMap.
99    ///
100    /// # Example
101    /// ```
102    /// use defaultdict::DefaultHashMap;
103    ///
104    /// let mut map = DefaultHashMap::<i8, i8>::new();
105    /// map.insert(10, 20);
106    ///
107    /// assert!(map.contains_key(&10));
108    /// ```
109    #[inline]
110    pub fn contains_key(&self, key: &K) -> bool
111    where
112        K: Eq + Hash,
113    {
114        self._inner.contains_key(key)
115    }
116
117    /// Clears the map, returning all key-value pairs as an iterator. Keeps the allocated memory for
118    /// reuse.
119    ///
120    /// If the returned iterator is dropped before being fully consumed, it drops the remaining
121    /// key-value pairs. The returned iterator keeps a mutable borrow on the map to optimize its
122    /// implementation.
123    ///
124    /// # Example
125    /// ```
126    /// use defaultdict::DefaultHashMap;
127    ///
128    /// let mut map = DefaultHashMap::<i8, i8>::new();
129    /// map.insert(10, 20);
130    ///
131    /// let contents: Vec<(i8, i8)> = map.drain().into_iter().collect();
132    /// assert_eq!(vec![(10, 20)], contents);
133    ///
134    ///
135    /// // The map is empty
136    /// let collected: Vec<(i8, i8)> = map.into_iter().collect();
137    /// let empty: Vec<(i8, i8)> = vec![];
138    ///
139    /// assert_eq!(empty, collected);
140    /// ```
141    #[inline]
142    pub fn drain(&mut self) -> Drain<K, V> {
143        self._inner.drain()
144    }
145
146    /// Gets the given key’s corresponding entry in the map for in-place manipulation.
147    ///
148    /// # Example
149    /// ```
150    /// use std::collections::hash_map::Entry;
151    /// use defaultdict::DefaultHashMap;
152    ///
153    /// let mut map = DefaultHashMap::<i8, i8>::new();
154    /// map.insert(10, 20);
155    ///
156    /// assert_eq!(&20, map.get(&10));
157    ///
158    /// if let Entry::Occupied(inner) = map.entry(10) {
159    ///     *inner.into_mut() += 10;
160    /// };
161    ///
162    /// assert_eq!(&30, map.get(&10));
163    /// ```
164    #[inline]
165    pub fn entry(&mut self, key: K) -> Entry<K, V> {
166        self._inner.entry(key)
167    }
168
169    /// Creates an iterator which uses a closure to determine if an element should be removed.
170    ///
171    /// If the closure returns true, the element is removed from the map and yielded. If the closure
172    /// returns false, or panics, the element remains in the map and will not be yielded.
173    ///
174    /// Note that extract_if lets you mutate every value in the filter closure, regardless of
175    /// whether you choose to keep or remove it.
176    pub fn extract_if<F>(&mut self, predicate: F) -> ExtractIf<'_, K, V, F>
177    where
178        F: FnMut(&K, &mut V) -> bool,
179    {
180        self._inner.extract_if(predicate)
181    }
182
183    /// Attempts to get mutable references to N values in the map at once.
184    ///
185    /// Returns an array of length N with the results of each query. For soundness, at most one
186    /// mutable reference will be returned to any value. None will be used if the key is missing.
187    pub fn get_disjoint_mut<Q, const N: usize>(&mut self, keys: [&Q; N]) -> [Option<&mut V>; N]
188    where
189        K: Borrow<Q>,
190        Q: Hash + Eq + ?Sized,
191    {
192        self._inner.get_disjoint_mut(keys)
193    }
194
195    /// Returns a reference to the value of the key passed in.
196    /// Because this hashmap mimicks the python defaultdict, it will also return a reference to a
197    /// value if the key is not present.
198    ///
199    /// The key type must implement [`Hash`] and [`Eq`].
200    ///
201    /// # Example
202    /// ```
203    /// use defaultdict::DefaultHashMap;
204    ///
205    /// let mut map = DefaultHashMap::<i8, i8>::new();
206    /// map.insert(10, 20);
207    ///
208    /// assert_eq!(&20, map.get(&10));
209    /// assert_eq!(&0, map.get(&20));
210    /// ```
211    #[must_use]
212    pub fn get<Q>(&self, key: &Q) -> &V
213    where
214        K: Borrow<Q>,
215        Q: Hash + Eq + ?Sized,
216    {
217        self._inner.get(key).unwrap_or(&self._default)
218    }
219
220    /// Returns the key-value pair corresponding to the supplied key.
221    /// The supplied key may be any borrowed form of the map’s key type, but [`Hash`] and [`Eq`] on
222    /// the borrowed form must match those for the key type.Returns a reference to the value of the
223    /// key passed in.
224    ///
225    /// # Example
226    /// ```
227    /// use defaultdict::DefaultHashMap;
228    ///
229    /// let mut map = DefaultHashMap::<i8, i8>::new();
230    /// map.insert(10, 20);
231    ///
232    /// let key_value: (&i8, &i8) = map.get_key_value(&10);
233    ///
234    /// assert_eq!((&10, &20), key_value);
235    /// ```
236    #[must_use]
237    pub fn get_key_value<'a>(&'a self, key: &'a K) -> (&'a K, &'a V)
238    where
239        K: Eq + Hash,
240    {
241        self._inner
242            .get_key_value(key)
243            .unwrap_or((key, &self._default))
244    }
245
246    /// Returns a mutable reference to the value corresponding to the key.
247    /// If the key is not present in the hashmap it will return the default value and insert it in
248    /// the map.
249    ///
250    /// The key may be any borrowed form of the map’s key type, but [`Hash`] and [`Eq`] on the
251    /// borrowed form must match those for the key type.
252    ///
253    /// # Example
254    /// ```
255    /// use defaultdict::DefaultHashMap;
256    ///
257    /// let mut map = DefaultHashMap::<i8, i8>::new();
258    /// let number = map.get_mut(&10);
259    ///
260    /// *number = 100;
261    ///
262    /// assert_eq!(&100, map.get(&10));
263    /// ```
264    #[must_use]
265    pub fn get_mut(&mut self, key: &K) -> &mut V
266    where
267        K: Hash + Eq + Clone,
268    {
269        let exists = self._inner.keys().any(|k| key == k);
270        if !exists {
271            self.insert(key.clone(), V::default());
272        }
273        self._inner.get_mut(key).unwrap()
274    }
275
276    /// Inserts a key value pair into the map. If the map did not have this key present, `None` is
277    /// returned.
278    ///
279    /// If the map had the key already present it will be overwritten.
280    ///
281    /// # Example
282    /// ```
283    /// use defaultdict::DefaultHashMap;
284    ///
285    /// let mut map = DefaultHashMap::<i8, i8>::new();
286    /// map.insert(10, 20);
287    ///
288    /// let old_value = map.insert(10, 30).unwrap();
289    ///
290    /// assert_eq!(&30, map.get(&10));
291    /// assert_eq!(20, old_value);
292    /// ```
293    #[inline]
294    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
295        self._inner.insert(key, value)
296    }
297
298    /// Creates a consuming iterator visiting all the keys in arbitrary order. The map cannot be
299    /// used after calling this. The iterator element type is `K`.
300    ///
301    /// # Example
302    /// ```
303    /// use defaultdict::DefaultHashMap;
304    ///
305    /// let mut map = DefaultHashMap::<i8, i8>::new();
306    /// map.insert(10, 20);
307    /// map.insert(30, 40);
308    ///
309    /// let mut collected: Vec<i8> = map.into_keys().collect();
310    /// collected.sort();
311    ///
312    /// assert_eq!(vec![10, 30], collected);
313    /// ```
314    #[inline]
315    pub fn into_keys(self) -> IntoKeys<K, V> {
316        self._inner.into_keys()
317    }
318
319    /// Creates a consuming iterator visiting all the values in arbitrary order. The map cannot be
320    /// used after calling this. The iterator element type is `V`.
321    ///
322    /// # Example
323    /// ```
324    /// use defaultdict::DefaultHashMap;
325    ///
326    /// let mut map = DefaultHashMap::<i8, i8>::new();
327    /// map.insert(10, 20);
328    /// map.insert(30, 40);
329    ///
330    /// let mut collected: Vec<i8> = map.into_values().collect();
331    /// collected.sort();
332    ///
333    /// assert_eq!(vec![20, 40], collected);
334    /// ```
335    #[inline]
336    pub fn into_values(self) -> IntoValues<K, V> {
337        self._inner.into_values()
338    }
339
340    /// Returns `true` if the map does not contain any keys.
341    ///
342    /// # Example
343    /// ```
344    /// use defaultdict::DefaultHashMap;
345    ///
346    /// let map = DefaultHashMap::<i8, i8>::new();
347    ///
348    /// assert!(map.is_empty());
349    /// ```
350    #[inline]
351    pub fn is_empty(&self) -> bool {
352        self._inner.is_empty()
353    }
354
355    /// Returns an iterator visiting all keys in arbitrary order. The iterator element type is
356    /// `&'a K`.
357    ///
358    /// # Example
359    /// ```
360    /// use defaultdict::DefaultHashMap;
361    ///
362    /// let mut map = DefaultHashMap::<i8, i8>::new();
363    /// map.insert(10, 20);
364    /// map.insert(11, 21);
365    /// map.insert(12, 22);
366    /// map.insert(13, 23);
367    ///
368    /// let mut collected: Vec<&i8> = map.keys().collect();
369    /// collected.sort();
370    ///
371    /// assert_eq!(vec![&10, &11, &12, &13], collected);
372    /// ```
373    #[inline]
374    pub fn keys(&self) -> Keys<'_, K, V> {
375        self._inner.keys()
376    }
377
378    /// Returns the length of the keys in the map.
379    ///
380    /// # Example
381    /// ```
382    /// use defaultdict::DefaultHashMap;
383    ///
384    /// let mut map = DefaultHashMap::<i8, i8>::new();
385    /// map.insert(10, 20);
386    /// map.insert(11, 21);
387    /// map.insert(12, 22);
388    /// map.insert(13, 23);
389    ///
390    /// assert_eq!(4, map.len());
391    /// ```
392    #[inline]
393    pub fn len(&self) -> usize {
394        self._inner.len()
395    }
396
397    /// Removes a key from the map, returning the value at the key if the key was previously in the
398    /// map. If the key is not present in the map it will return the default value.
399    ///
400    /// The key may be any borrowed form of the map’s key type, but [`Hash`] and [`Eq`] on the
401    /// borrowed form must match those for the key type.
402    ///
403    /// # Example
404    /// ```
405    /// use defaultdict::DefaultHashMap;
406    ///
407    /// let mut map = DefaultHashMap::<i8, i8>::new();
408    /// map.insert(10, 20);
409    ///
410    /// println!("{}", map.remove(&10));
411    ///
412    /// println!("{}", map.remove(&90));
413    ///
414    /// assert!(map.is_empty());
415    /// ```
416    #[must_use]
417    pub fn remove<Q>(&mut self, key: &Q) -> V
418    where
419        K: Borrow<Q>,
420        Q: Hash + Eq + ?Sized,
421    {
422        self._inner.remove(key).unwrap_or_default()
423    }
424
425    /// Removes a key from the map, returning the stored key and value if the key was previously in
426    /// the map. If the key is not present in the map, a default value will be returned.
427    ///
428    /// The key may be any borrowed form of the map’s key type, but [`Hash`] and [`Eq`] on the
429    /// borrowed form must match those for the key type.
430    ///
431    /// # Example
432    /// ```
433    /// use defaultdict::DefaultHashMap;
434    ///
435    /// let mut map = DefaultHashMap::<i8, i8>::new();
436    ///
437    /// for i in 0..10 {
438    ///     map.insert(i, 20);
439    /// }
440    ///
441    /// let entry = map.remove_entry(&0);
442    ///
443    /// let default_entry = map.remove_entry(&0);
444    ///
445    /// assert_eq!((1, 20), map.remove_entry(&1));
446    /// assert_eq!((1, 0), map.remove_entry(&1));
447    /// ```
448    #[must_use]
449    pub fn remove_entry(&mut self, key: &K) -> (K, V)
450    where
451        K: Clone,
452        V: Clone,
453    {
454        self._inner
455            .remove_entry(key)
456            .unwrap_or((key.clone(), self._default.to_owned()))
457    }
458
459    /// Retains only the elements specified by the predicate.
460    /// In other words, remove all pairs (k, v) for which f(&k, &mut v) returns false. The elements
461    /// are visited in unsorted (and unspecified) order.
462    ///
463    /// # Example
464    /// ```
465    /// use defaultdict::{DefaultHashMap, defaulthashmap};
466    ///
467    /// let mut map: DefaultHashMap<i8, i8> = defaulthashmap!();
468    ///
469    /// for i in 0..10 {
470    ///     map.insert(i, i);
471    /// }
472    ///
473    /// map.retain(|key, value| {
474    ///     key <= &2
475    /// });
476    ///
477    /// let mut collected: Vec<(i8, i8)> = map.into_iter().collect();
478    /// collected.sort();
479    /// let golden: Vec<(i8, i8)> = vec![(0, 0), (1, 1), (2, 2)];
480    ///
481    /// assert_eq!(golden, collected);
482    /// ```
483    pub fn retain<F>(&mut self, func: F)
484    where
485        F: FnMut(&K, &mut V) -> bool,
486    {
487        self._inner.retain(func);
488    }
489
490    /// Returns an iterator visiting all values in arbitrary order. The iterator element type is
491    /// &'a V.
492    ///
493    /// # Example
494    /// ```
495    /// use defaultdict::DefaultHashMap;
496    ///
497    /// let mut map = DefaultHashMap::<i8, i8>::new();
498    /// map.insert(10, 20);
499    /// map.insert(11, 21);
500    /// map.insert(12, 22);
501    /// map.insert(13, 23);
502    ///
503    /// let mut collected: Vec<&i8> = map.values().collect();
504    /// collected.sort();
505    /// let golden: Vec<&i8> = vec![&20, &21, &22, &23];
506    ///
507    /// assert_eq!(golden, collected);
508    /// ```
509    #[inline]
510    pub fn values(&self) -> Values<'_, K, V> {
511        self._inner.values()
512    }
513
514    /// Gets a mutable iterator over the values of the map, in order by key.
515    ///
516    /// # Example
517    /// ```
518    /// use defaultdict::DefaultHashMap;
519    ///
520    /// let mut map = DefaultHashMap::<i8, i8>::new();
521    ///
522    /// for i in 0..10 {
523    ///     map.insert(i, i);
524    /// }
525    ///
526    /// for value in map.values_mut() {
527    ///     *value += 1;
528    /// }
529    ///
530    /// let mut collected: Vec<&i8> = map.values().collect();
531    /// collected.sort();
532    /// let golden: Vec<&i8> = vec![&1, &2, &3, &4, &5, &6, &7, &8, &9, &10];
533    ///
534    /// assert_eq!(golden, collected);
535    /// ```
536    #[inline]
537    pub fn values_mut(&mut self) -> ValuesMut<K, V> {
538        self._inner.values_mut()
539    }
540
541    /// Creates an empty [`DefaultHashMap`] which will use the given hash builder to hash
542    /// keys.
543    ///
544    /// Warning: `hash_builder` is normally randomly generated, and
545    /// is designed to allow HashMaps to be resistant to attacks that
546    /// cause many collisions and very poor performance. Setting it
547    /// manually using this function can expose a DoS attack vector.
548    ///
549    /// The `hash_builder` passed should implement the [`BuildHasher`] trait for
550    /// the HashMap to be useful, see its documentation for details.
551    ///
552    /// # Examples
553    ///
554    /// ```
555    /// use defaultdict::DefaultHashMap;
556    /// use std::collections::hash_map::RandomState;
557    ///
558    /// let s = RandomState::new();
559    /// let mut map = DefaultHashMap::with_hasher(s);
560    /// map.insert(1, 2);
561    /// ```
562    #[inline]
563    pub fn with_hasher(hash_builder: S) -> Self {
564        DefaultHashMap {
565            _inner: HashMap::with_hasher(hash_builder),
566            _default: V::default(),
567        }
568    }
569}
570
571impl<K, V> Default for DefaultHashMap<K, V>
572where
573    K: Eq + Hash,
574    V: Default,
575{
576    fn default() -> Self {
577        Self::new()
578    }
579}
580
581impl<K, V, S> PartialEq for DefaultHashMap<K, V, S>
582where
583    K: Eq + Hash,
584    V: PartialEq + Default,
585    S: BuildHasher,
586{
587    fn eq(&self, other: &DefaultHashMap<K, V, S>) -> bool {
588        self._inner == other._inner && self._default == other._default
589    }
590}
591
592impl<K, V, S> Eq for DefaultHashMap<K, V, S>
593where
594    K: Eq + Hash,
595    V: Eq + Default,
596    S: BuildHasher,
597{
598}
599
600impl<K, V, S> IntoIterator for DefaultHashMap<K, V, S>
601where
602    K: Eq + Hash + Clone,
603    V: Default,
604    S: BuildHasher,
605{
606    type Item = (K, V);
607    type IntoIter = IntoIter<K, V>;
608
609    fn into_iter(self) -> Self::IntoIter {
610        self._inner.into_iter()
611    }
612}
613
614impl<'a, K, V, S> IntoIterator for &'a DefaultHashMap<K, V, S>
615where
616    K: Eq + Hash,
617    V: Default,
618    S: BuildHasher,
619{
620    type Item = (&'a K, &'a V);
621    type IntoIter = Iter<'a, K, V>;
622
623    fn into_iter(self) -> Self::IntoIter {
624        self._inner.iter()
625    }
626}
627
628impl<K, V, S> Index<&K> for DefaultHashMap<K, V, S>
629where
630    K: Eq + Hash,
631    V: Default,
632    S: BuildHasher,
633{
634    type Output = V;
635
636    fn index(&self, key: &K) -> &V {
637        self._inner.get(key).unwrap_or(&self._default)
638    }
639}
640
641impl<'a, K, V, S> IntoIterator for &'a mut DefaultHashMap<K, V, S>
642where
643    K: Eq + Hash,
644    V: Default,
645    S: BuildHasher,
646{
647    type Item = (&'a K, &'a mut V);
648    type IntoIter = IterMut<'a, K, V>;
649
650    fn into_iter(self) -> Self::IntoIter {
651        self._inner.iter_mut()
652    }
653}
654
655impl<K, V, S> From<HashMap<K, V, S>> for DefaultHashMap<K, V, S>
656where
657    K: Eq + Hash,
658    V: Default,
659    S: BuildHasher,
660{
661    fn from(hashmap: HashMap<K, V, S>) -> Self {
662        Self {
663            _inner: hashmap,
664            _default: V::default(),
665        }
666    }
667}
668
669impl<K, V, S> From<DefaultHashMap<K, V, S>> for HashMap<K, V, S>
670where
671    K: Eq + Hash,
672    V: Default,
673    S: BuildHasher,
674{
675    fn from(hashmap: DefaultHashMap<K, V, S>) -> Self {
676        hashmap._inner
677    }
678}
679
680impl<K, V, S> FromIterator<(K, V)> for DefaultHashMap<K, V, S>
681where
682    K: Eq + Hash,
683    V: Default,
684    S: BuildHasher + Default,
685{
686    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
687        let mut map = DefaultHashMap::with_hasher(Default::default());
688        for (k, v) in iter {
689            let _ = map.insert(k, v);
690        }
691        map
692    }
693}
694
695#[macro_export]
696/// A quick way to instantiate a HashMap.
697///
698/// A trailing comma is allowed but not required here
699///
700/// # Example
701/// ```
702/// use defaultdict::DefaultHashMap;
703///
704///
705/// let default_map: DefaultHashMap<i8, i8> = defaultdict::defaulthashmap!(1,2,3,);
706///
707/// let custom_map: DefaultHashMap<i8, i8> = defaultdict::defaulthashmap!(
708///     (1, 1),
709///     (2, 2),
710/// );
711/// ```
712macro_rules! defaulthashmap {
713
714    // match 1
715    ( ) => {
716        {
717            DefaultHashMap::new()
718        }
719    };
720
721    // match 2
722    ( $( ($key:expr, $val:expr) ),* $(,)? ) => {
723        {
724            let mut map = DefaultHashMap::new();
725            $(
726                let _ = map.insert($key, $val);
727            )*
728            map
729        }
730    };
731
732    // match 3
733    ( $( $key:expr ),* $(,)? ) => {
734        {
735            let mut map = DefaultHashMap::new();
736            $(
737                let _ = map.get_mut(&$key);
738            )*
739            map
740        }
741    };
742
743}