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}