countmap/lib.rs
1#![deny(missing_docs)]
2//! A map that helps counting elements.
3
4extern crate num_traits;
5
6use num_traits::identities::{One, Zero};
7
8use std::borrow::Borrow;
9use std::collections::HashMap;
10use std::collections::hash_map::{Drain, IntoIter, Iter, IterMut, Keys, RandomState, Values};
11use std::hash::{BuildHasher, Hash};
12use std::iter::FromIterator;
13use std::ops::{Add, Index};
14
15/// A count map is a hash map where the value field is a constantly incremented counter. If a key
16/// is inserted for the first time, the counter is set to 1. Every subsequent insert will increment
17/// the counter by 1. This implementation just acts as a decorator around a `HashMap` and exposes
18/// almost the complete API of `HashMap` except things like `iter_mut()` or `get_mut()` since it
19/// doesn't make sense in this use case.
20#[derive(Clone, Debug)]
21pub struct CountMap<K, C = u64, S = RandomState>
22where
23 K: Eq + Hash,
24 // C: Unsigned,
25 S: BuildHasher,
26{
27 map: HashMap<K, C, S>,
28}
29
30impl<K, C> CountMap<K, C, RandomState>
31where
32 K: Eq + Hash,
33 // C: Unsigned,
34{
35 /// Creates an empty `CountMap`.
36 ///
37 /// # Examples
38 /// ```
39 /// use countmap::CountMap;
40 ///
41 /// let mut count_map: CountMap<&str> = CountMap::new();
42 /// ```
43 pub fn new() -> Self {
44 Self::default()
45 }
46
47 /// Creates an empty `CountMap` with the specified capacity.
48 ///
49 /// The created map can hold at least `cap` elements before reallocating. If `cap` is `0` the
50 /// map will not allocate.
51 ///
52 /// # Examples
53 /// ```
54 /// use countmap::CountMap;
55 ///
56 /// let mut count_map: CountMap<&str> = CountMap::with_capacity(10);
57 /// ```
58 pub fn with_capacity(cap: usize) -> Self {
59 Self { map: HashMap::with_capacity(cap) }
60 }
61}
62
63impl<K, C, S> CountMap<K, C, S>
64where
65 K: Eq + Hash,
66 C: One + Zero + Copy + Clone + Add<Output = C>,
67 S: BuildHasher,
68{
69 /// Creates an empty `CountMap` which will use the given hash builder to hash keys.
70 ///
71 /// The created map has the default initial capacity.
72 ///
73 /// Warning: `hash_builder` is normally randomly generated, and is designed to allow HashMaps
74 /// to be resistant to attacks that cause many collisions and very poor performance. Setting it
75 /// manually using this function can expose a DoS attack vector.
76 ///
77 /// # Examples
78 /// ```
79 /// use std::collections::hash_map::RandomState;
80 /// use countmap::CountMap;
81 ///
82 /// let s = RandomState::new();
83 /// let mut map: CountMap<_, u16> = CountMap::with_hasher(s);
84 /// map.insert_or_increment("foo");
85 /// ```
86 pub fn with_hasher(hash_builder: S) -> Self {
87 Self { map: HashMap::with_hasher(hash_builder) }
88 }
89
90 /// Creates an empty `CountMap` with the specified capacity, using hash_builder to hash the
91 /// keys.
92 ///
93 /// The count map will be able to hold at least `capacity` elements without reallocating. If
94 /// `capacity` is 0, the hash map will not allocate.
95 ///
96 /// Warning: `hash_builder` is normally randomly generated, and is designed to allow HashMaps
97 /// to be resistant to attacks that cause many collisions and very poor performance. Setting it
98 /// manually using this function can expose a DoS attack vector.
99 ///
100 /// # Examples
101 /// ```
102 /// use std::collections::hash_map::RandomState;
103 /// use countmap::CountMap;
104 ///
105 /// let s = RandomState::new();
106 /// let mut map: CountMap<_, u16> = CountMap::with_capacity_and_hasher(10, s);
107 /// map.insert_or_increment("foo");
108 /// ```
109 pub fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> Self {
110 Self { map: HashMap::with_capacity_and_hasher(capacity, hash_builder) }
111 }
112
113 /// Returns a reference to the map's `BuildHasher`.
114 pub fn hasher(&self) -> &S {
115 self.map.hasher()
116 }
117
118 /// Returns the number of elements the map can hold without reallocating.
119 ///
120 /// This number is a lower bound; the `CountMap<K>` might be able to hold more, but is
121 /// guaranteed to be able to hold at least this many.
122 ///
123 /// # Examples
124 /// ```
125 /// use countmap::CountMap;
126 ///
127 /// let map: CountMap<&str> = CountMap::with_capacity(100);
128 /// assert!(map.capacity() >= 100);
129 /// ```
130 pub fn capacity(&self) -> usize {
131 self.map.capacity()
132 }
133
134 /// Reserves capacity for at least `additional` more elements to be inserted in the `CountMap`.
135 /// The collection ma reserve more space to avoid frequent reallocations.
136 ///
137 /// # Panics
138 /// Panics if the new allocation size overflows usize.
139 ///
140 /// # Examples
141 /// ```
142 /// use countmap::CountMap;
143 ///
144 /// let mut map: CountMap<&str> = CountMap::with_capacity(5);
145 /// map.reserve(10);
146 /// assert!(map.capacity() >= 15);
147 /// ```
148 pub fn reserve(&mut self, additional: usize) {
149 self.map.reserve(additional)
150 }
151
152 /// Shrinks the capacity of the map as much as possible. It will drop down as much as possible
153 /// while maintaining the internal rules and possibly leaving some space in accordance with the
154 /// resize policy.
155 ///
156 /// # Examples
157 /// ```
158 /// use countmap::CountMap;
159 ///
160 /// let mut map: CountMap<_, u16> = CountMap::with_capacity(100);
161 /// map.insert_or_increment("foo");
162 /// map.insert_or_increment("bar");
163 /// assert!(map.capacity() >= 100);
164 /// map.shrink_to_fit();
165 /// assert!(map.capacity() >= 2);
166 /// ```
167 pub fn shrink_to_fit(&mut self) {
168 self.map.shrink_to_fit()
169 }
170
171 /// An iterator visiting all keys in arbitrary order. The iterator element type is `&'a K`.
172 ///
173 /// # Examples
174 /// ```
175 /// use countmap::CountMap;
176 ///
177 /// let mut map: CountMap<_, u16> = CountMap::new();
178 /// map.insert_or_increment("foo");
179 /// map.insert_or_increment("bar");
180 /// map.insert_or_increment("foo");
181 ///
182 /// for key in map.keys() {
183 /// println!("{}", key);
184 /// }
185 /// ```
186 pub fn keys(&self) -> Keys<K, C> {
187 self.map.keys()
188 }
189
190 /// An iterator visiting all values in arbitrary order. The iterator element type is `&'a V`.
191 ///
192 /// # Examples
193 /// ```
194 /// use countmap::CountMap;
195 ///
196 /// let mut map: CountMap<_, u16> = CountMap::new();
197 /// map.insert_or_increment("foo");
198 /// map.insert_or_increment("bar");
199 /// map.insert_or_increment("foo");
200 ///
201 /// for val in map.values() {
202 /// println!("{}", val);
203 /// }
204 /// ```
205 pub fn values(&self) -> Values<K, C> {
206 self.map.values()
207 }
208
209 /// Inserts or increments an element by 1 in the `CountMap`. The new value of the counter is
210 /// returned.
211 ///
212 /// # Examples
213 /// ```
214 /// use countmap::CountMap;
215 ///
216 /// let mut count_map: CountMap<_, u16> = CountMap::new();
217 ///
218 /// assert_eq!(count_map.insert_or_increment("foo"), 1);
219 /// assert_eq!(count_map.insert_or_increment("foo"), 2);
220 /// assert_eq!(count_map.insert_or_increment("bar"), 1);
221 /// ```
222 pub fn insert_or_increment(&mut self, element: K) -> C {
223 self.insert_or_increment_by(element, C::one())
224 }
225
226 /// Inserts or increments an element by the specified difference in the `CountMap`. The new
227 /// value of the counter is returned.
228 ///
229 /// # Examples
230 /// ```
231 /// use countmap::CountMap;
232 ///
233 /// let mut count_map: CountMap<&str> = CountMap::new();
234 ///
235 /// assert_eq!(count_map.insert_or_increment_by("foo", 5), 5);
236 /// assert_eq!(count_map.insert_or_increment_by("foo", 2), 7);
237 /// assert_eq!(count_map.insert_or_increment_by("bar", 1), 1);
238 /// ```
239 pub fn insert_or_increment_by(&mut self, element: K, diff: C) -> C {
240 let count = self.map.entry(element).or_insert(C::zero());
241 // *count += diff;
242 *count = *count + diff;
243 // *count = count.add(diff);
244 *count
245 }
246
247 /// Increments an existing element in the `CountMap` by 1. Returns an `Option` with the new
248 /// value of the counter or `None` if the element doesn't exist.
249 ///
250 /// # Examples
251 /// ```
252 /// use countmap::CountMap;
253 ///
254 /// let mut count_map: CountMap<&str> = CountMap::new();
255 ///
256 /// assert_eq!(count_map.increment(&"foo"), None);
257 ///
258 /// count_map.insert_or_increment(&"foo");
259 ///
260 /// assert_eq!(count_map.increment(&"foo"), Some(2));
261 /// ```
262 pub fn increment(&mut self, element: &K) -> Option<C> {
263 self.increment_by(element, C::one())
264 }
265
266 /// Increments an existing element in the `CountMap` by the specified difference. Returns an
267 /// `Option` with the new value of the counter or `None` if the element doesn't exist.
268 ///
269 /// # Examples
270 /// ```
271 /// use countmap::CountMap;
272 ///
273 /// let mut count_map: CountMap<&str> = CountMap::new();
274 ///
275 /// assert_eq!(count_map.increment_by(&"foo", 5), None);
276 ///
277 /// count_map.insert_or_increment(&"foo");
278 ///
279 /// assert_eq!(count_map.increment_by(&"foo", 2), Some(3));
280 /// ```
281 pub fn increment_by(&mut self, element: &K, diff: C) -> Option<C> {
282 let entry = self.map.get_mut(element);
283 match entry {
284 Some(count) => {
285 // *count += diff;
286 *count = *count + diff;
287 Some(*count)
288 }
289 None => None,
290 }
291 }
292
293 /// Returns an `Option` containing the current counter value of the specified element or
294 /// `None`.
295 ///
296 /// # Examples
297 /// ```
298 /// use countmap::CountMap;
299 ///
300 /// let mut count_map: CountMap<&str> = CountMap::new();
301 ///
302 /// count_map.insert_or_increment("foo");
303 ///
304 /// assert_eq!(count_map.get_count(&"foo"), Some(1));
305 /// assert_eq!(count_map.get_count(&"bar"), None);
306 /// ```
307 pub fn get_count(&self, element: &K) -> Option<C> {
308 self.map.get(element).cloned()
309 }
310
311 /// An iterator visiting all key-value pairs in arbitrary order. The iterator element type is
312 /// (&'a K, &'a C).
313 ///
314 /// # Examples
315 /// ```
316 /// use countmap::CountMap;
317 ///
318 /// let mut map: CountMap<_, u16> = CountMap::new();
319 ///
320 /// map.insert_or_increment("foo");
321 /// map.insert_or_increment("foo");
322 /// map.insert_or_increment("bar");
323 ///
324 /// for (key, count) in map {
325 /// println!("key: {}, count: {}", key, count);
326 /// }
327 /// ```
328 pub fn iter(&self) -> Iter<K, C> {
329 self.map.iter()
330 }
331
332 /// An iterator visiting all key-value pairs in arbitrary order, with mutable references to the
333 /// values. The iterator element type is (&'a K, &'a mut V).
334 ///
335 /// # Examples
336 /// ```
337 /// use countmap::CountMap;
338 ///
339 /// let mut map: CountMap<_, u16> = CountMap::new();
340 ///
341 /// map.insert_or_increment("foo");
342 /// map.insert_or_increment("foo");
343 /// map.insert_or_increment("bar");
344 ///
345 /// for (_, count) in map.iter_mut() {
346 /// *count += 3;
347 /// }
348 ///
349 /// assert_eq!(map.get_count(&"foo"), Some(5));
350 /// assert_eq!(map.get_count(&"bar"), Some(4));
351 /// ```
352 pub fn iter_mut(&mut self) -> IterMut<K, C> {
353 self.map.iter_mut()
354 }
355
356 /// Returns the number of elements in the map.
357 ///
358 /// # Examples
359 /// ```
360 /// use countmap::CountMap;
361 ///
362 /// let mut map: CountMap<_, u16> = CountMap::new();
363 /// assert_eq!(map.len(), 0);
364 /// map.insert_or_increment("foo");
365 /// assert_eq!(map.len(), 1);
366 /// ```
367 pub fn len(&self) -> usize {
368 self.map.len()
369 }
370
371 /// Returns true if the map contains no elements.
372 ///
373 /// # Examples
374 /// ```
375 /// use countmap::CountMap;
376 ///
377 /// let mut map: CountMap<_, u16> = CountMap::new();
378 /// assert_eq!(map.is_empty(), true);
379 /// map.insert_or_increment("foo");
380 /// assert_eq!(map.is_empty(), false);
381 /// ```
382 pub fn is_empty(&self) -> bool {
383 self.map.is_empty()
384 }
385
386 /// Clears the map, returning all key-value pairs as an iterator. Keeps the allocated memory
387 /// for reuse.
388 ///
389 /// # Examples
390 /// ```
391 /// use countmap::CountMap;
392 ///
393 /// let mut map: CountMap<_, u16> = CountMap::new();
394 /// map.insert_or_increment("foo");
395 /// map.insert_or_increment("bar");
396 ///
397 /// for (k, c) in map.drain().take(1) {
398 /// assert!(k == "foo" || k == "bar");
399 /// assert_eq!(c, 1);
400 /// }
401 ///
402 /// assert!(map.is_empty());
403 /// ```
404 pub fn drain(&mut self) -> Drain<K, C> {
405 self.map.drain()
406 }
407
408 /// Clears the map, removing all key-counter pairs. Keeps the allocated memory for reuse.
409 ///
410 /// # Examples
411 /// ```
412 /// use countmap::CountMap;
413 ///
414 /// let mut map: CountMap<&str, u16> = CountMap::new();
415 /// map.insert_or_increment("foo");
416 /// map.clear();
417 /// assert!(map.is_empty())
418 /// ```
419 pub fn clear(&mut self) {
420 self.map.clear()
421 }
422
423 /// Returns true if the map contains a value for the specified key.
424 ///
425 /// # Examples
426 /// ```
427 /// use countmap::CountMap;
428 ///
429 /// let mut map: CountMap<&str, u16> = CountMap::new();
430 /// map.insert_or_increment("foo");
431 /// assert!(map.contains_key(&"foo"));
432 /// assert!(!map.contains_key(&"bar"));
433 /// ```
434 pub fn contains_key(&self, k: &K) -> bool {
435 self.map.contains_key(k)
436 }
437
438 /// Removes a key from the map, returning the value at the key if the key was previously in the
439 /// map.
440 ///
441 /// # Examples
442 /// ```
443 /// use countmap::CountMap;
444 ///
445 /// let mut map = CountMap::new();
446 /// map.insert_or_increment("foo");
447 /// assert_eq!(map.remove(&"foo"), Some(1));
448 /// assert_eq!(map.remove(&"bar"), None);
449 /// ```
450 pub fn remove(&mut self, k: &K) -> Option<C> {
451 self.map.remove(k)
452 }
453
454 /// Retains only the elements specified by the predicate.
455 ///
456 /// In other words, remove all pairs `(k, v)` such that `f(&k,&mut v)` returns `false`.
457 ///
458 /// # Examples
459 /// ```
460 /// use countmap::CountMap;
461 ///
462 /// let mut map: CountMap<_, u16> = CountMap::new();
463 /// map.insert_or_increment("foo");
464 /// map.insert_or_increment("foo");
465 /// map.insert_or_increment("foo");
466 /// map.insert_or_increment("bar");
467 ///
468 /// map.retain(|_, c| *c == 3);
469 /// assert_eq!(map.len(), 1);
470 /// ```
471 pub fn retain<F>(&mut self, f: F)
472 where
473 F: FnMut(&K, &mut C) -> bool,
474 {
475 self.map.retain(f)
476 }
477}
478
479impl<K, C> Default for CountMap<K, C>
480where
481 K: Eq + Hash,
482 // C: Unsigned,
483{
484 fn default() -> Self {
485 Self { map: HashMap::new() }
486 }
487}
488
489impl<K> PartialEq for CountMap<K>
490where
491 K: Eq + Hash,
492{
493 fn eq(&self, other: &CountMap<K>) -> bool {
494 self.map == other.map
495 }
496}
497
498impl<K> Eq for CountMap<K>
499where
500 K: Eq + Hash,
501{
502}
503
504impl<'a, K, C> IntoIterator for &'a CountMap<K, C>
505where
506 K: Eq + Hash,
507 // C: Unsigned,
508{
509 type Item = (&'a K, &'a C);
510 type IntoIter = Iter<'a, K, C>;
511
512 fn into_iter(self) -> Self::IntoIter {
513 self.map.iter()
514 }
515}
516
517impl<'a, K, C> IntoIterator for &'a mut CountMap<K, C>
518where
519 K: Eq + Hash,
520 // C: Unsigned,
521{
522 type Item = (&'a K, &'a mut C);
523 type IntoIter = IterMut<'a, K, C>;
524
525 fn into_iter(self) -> Self::IntoIter {
526 self.map.iter_mut()
527 }
528}
529
530impl<'a, K, C> IntoIterator for CountMap<K, C>
531where
532 K: Eq + Hash,
533 // C: Unsigned,
534{
535 type Item = (K, C);
536 type IntoIter = IntoIter<K, C>;
537
538 fn into_iter(self) -> Self::IntoIter {
539 self.map.into_iter()
540 }
541}
542
543impl<'a, K, C, Q> Index<&'a Q> for CountMap<K, C>
544where
545 K: Eq + Hash + Borrow<Q>,
546 // C: Unsigned,
547 Q: ?Sized + Eq + Hash,
548{
549 type Output = C;
550
551 /// # Examples
552 /// ```
553 /// use countmap::CountMap;
554 ///
555 /// let mut map: CountMap<_, u16> = CountMap::new();
556 ///
557 /// map.insert_or_increment("foo");
558 /// assert_eq!(map["foo"], 1);
559 /// ```
560 fn index(&self, index: &Q) -> &Self::Output {
561 &self.map[index]
562 }
563}
564
565impl<K, C> FromIterator<(K, C)> for CountMap<K, C>
566where
567 K: Eq + Hash,
568 C: Clone + Copy + One + Zero,
569{
570 /// Creates a `CountMap<K>` from an `Iterator<(K, C)>`.
571 ///
572 /// # Examples
573 /// ```
574 /// use countmap::CountMap;
575 /// use std::iter::FromIterator;
576 ///
577 /// let data = vec![("foo", 3), ("bar", 3), ("foo", 1)];
578 /// let map = CountMap::from_iter(data);
579 /// assert_eq!(map.get_count(&"foo"), Some(4));
580 /// assert_eq!(map.get_count(&"bar"), Some(3));
581 /// ```
582 fn from_iter<T>(iter: T) -> Self
583 where
584 T: IntoIterator<Item = (K, C)>,
585 {
586 let iter = iter.into_iter();
587 let mut map = CountMap::with_capacity(iter.size_hint().0);
588 for (k, v) in iter {
589 map.insert_or_increment_by(k, v);
590 }
591 map
592 }
593}
594
595impl<K> FromIterator<K> for CountMap<K>
596where
597 K: Eq + Hash,
598{
599 /// Creates a `CountMap<K>` from an `Iterator<K>`.
600 ///
601 /// # Examples
602 /// ```
603 /// use countmap::CountMap;
604 /// use std::iter::FromIterator;
605 ///
606 /// let data = vec!["foo", "bar", "foo"];
607 /// let map = CountMap::from_iter(data);
608 /// assert_eq!(map.get_count(&"foo"), Some(2));
609 /// assert_eq!(map.get_count(&"bar"), Some(1));
610 /// ```
611 fn from_iter<T>(iter: T) -> Self
612 where
613 T: IntoIterator<Item = K>,
614 {
615 let iter = iter.into_iter();
616 let mut map = CountMap::with_capacity(iter.size_hint().0);
617 for item in iter {
618 map.insert_or_increment(item);
619 }
620 map
621 }
622}
623
624impl<K, C> Extend<(K, C)> for CountMap<K, C>
625where
626 K: Eq + Hash,
627 C: Clone + Copy + One + Zero,
628{
629 /// Extends a `CountMap<K>` with an `Iterator<(K, C)>`.
630 ///
631 /// # Examples
632 /// ```
633 /// use countmap::CountMap;
634 ///
635 /// let data = vec![("foo", 3), ("bar", 3), ("foo", 1)];
636 /// let mut map = CountMap::new();
637 /// map.extend(data);
638 ///
639 /// assert_eq!(map.get_count(&"foo"), Some(4));
640 /// assert_eq!(map.get_count(&"bar"), Some(3));
641 /// ```
642 fn extend<T>(&mut self, iter: T)
643 where
644 T: IntoIterator<Item = (K, C)>,
645 {
646 let iter = iter.into_iter();
647 let reserve = if self.is_empty() {
648 iter.size_hint().0
649 } else {
650 (iter.size_hint().0 + 1) / 2
651 };
652 self.reserve(reserve);
653 for (k, v) in iter {
654 self.insert_or_increment_by(k, v);
655 }
656 }
657}
658
659impl<'a, K, C> Extend<(&'a K, &'a C)> for CountMap<K, C>
660where
661 K: 'a + Eq + Hash + Copy,
662 C: 'a + Clone + Copy + One + Zero,
663{
664 fn extend<T>(&mut self, iter: T)
665 where
666 T: IntoIterator<Item = (&'a K, &'a C)>,
667 {
668 self.extend(iter.into_iter().map(|(&key, &value)| (key, value)));
669 }
670}
671
672impl<K> Extend<K> for CountMap<K>
673where
674 K: Eq + Hash,
675{
676 /// Extends a `CountMap<K>` with an `Iterator<K>`.
677 ///
678 /// # Examples
679 /// ```
680 /// use countmap::CountMap;
681 ///
682 /// let data = vec!["foo", "bar", "foo"];
683 /// let mut map = CountMap::new();
684 /// map.extend(data);
685 ///
686 /// assert_eq!(map.get_count(&"foo"), Some(2));
687 /// assert_eq!(map.get_count(&"bar"), Some(1));
688 /// ```
689 fn extend<T>(&mut self, iter: T)
690 where
691 T: IntoIterator<Item = K>,
692 {
693 let iter = iter.into_iter();
694 let reserve = if self.is_empty() {
695 iter.size_hint().0
696 } else {
697 (iter.size_hint().0 + 1) / 2
698 };
699 self.reserve(reserve);
700 for k in iter {
701 self.insert_or_increment(k);
702 }
703 }
704}
705
706impl<'a, K> Extend<&'a K> for CountMap<K>
707where
708 K: 'a + Eq + Hash + Copy,
709{
710 fn extend<T>(&mut self, iter: T)
711 where
712 T: IntoIterator<Item = &'a K>,
713 {
714 self.extend(iter.into_iter().cloned());
715 }
716}