intmap/
lib.rs

1#![forbid(unsafe_code)]
2
3//! Specialized hashmap for integer based keys.
4//!
5//! For more information see the [README](https://github.com/JesperAxelsson/rust-intmap/blob/master/README.md).
6//!
7//! <div class="warning">
8//! Be aware that no effort is made against DoS attacks.
9//! </div>
10
11#[cfg(feature = "serde")]
12mod serde;
13
14mod entry;
15mod int;
16mod int_key;
17mod iter;
18
19use core::iter::{IntoIterator, Iterator};
20use int::SealedInt;
21
22pub use entry::*;
23pub use int::Int;
24pub use int_key::IntKey;
25pub use iter::*;
26
27// Test examples from the README.
28#[doc = include_str!("../README.md")]
29#[cfg(doctest)]
30pub struct ReadmeDoctests;
31
32/// A hashmap that maps an integer based `K` to `V`.
33#[derive(Clone)]
34pub struct IntMap<K, V> {
35    // The slots for the key/value pairs.
36    //
37    // The number of slots is what we call "capacity". Two or more key/value pairs occupy the same
38    // slot if they have a hash collision.
39    // The size of `cache` as binary exponent. The actual size of `cache` is `2^size`.
40    cache: Vec<Vec<(K, V)>>,
41    // The size of `cache` as binary exponent. The actual size of `cache` is `2^size`.
42    size: u32,
43    // A bit mask for calculating an index for `cache`. Must be recomputed if `size` changes.
44    mod_mask: usize,
45    // The number of stored key/value pairs.
46    count: usize,
47    // The ratio between key/value pairs and available slots that we try to ensure.
48    //
49    // Multiplied by 1000, e.g. a load factor of 90.9% will result in the value 909.
50    load_factor: usize,
51}
52
53impl<K, V> IntMap<K, V> {
54    /// Creates a new [`IntMap`].
55    ///
56    /// The [`IntMap`] is initially created with a capacity of 0, so it will not allocate until it
57    /// is first inserted into.
58    ///
59    /// The [`IntMap`] is initially created with a capacity of 0, so it will not allocate until it
60    /// is first inserted into.
61    ///
62    /// # Examples
63    ///
64    /// ```
65    /// use intmap::IntMap;
66    ///
67    /// let mut map: IntMap<u64, u64> = IntMap::new();
68    /// assert_eq!(map, IntMap::default());
69    /// ```
70    pub const fn new() -> Self {
71        Self {
72            cache: Vec::new(),
73            size: 0,
74            count: 0,
75            mod_mask: 0,
76            load_factor: 909, // 90.9%
77        }
78    }
79}
80
81impl<K: IntKey, V> IntMap<K, V> {
82    /// Creates a new [`IntMap`] with at least the given capacity.
83    ///
84    /// If the capacity is 0, the [`IntMap`] will not allocate. Otherwise the capacity is rounded
85    /// to the next power of two and space for elements is allocated accordingly.
86    ///
87    /// # Examples
88    ///
89    /// ```
90    /// use intmap::IntMap;
91    ///
92    /// let mut map: IntMap<u64, u64> = IntMap::with_capacity(20);
93    /// ```
94    pub fn with_capacity(capacity: usize) -> Self {
95        let mut map = Self::new();
96        map.reserve(capacity);
97        map
98    }
99
100    /// Sets the load factor of the [`IntMap`] rounded to the first decimal point.
101    ///
102    /// A load factor between 0.0 and 1.0 will reduce hash collisions but use more space.
103    /// A load factor above 1.0 will tolerate hash collisions and use less space.
104    ///
105    /// # Examples
106    ///
107    /// ```
108    /// use intmap::IntMap;
109    ///
110    /// let mut map: IntMap<u64, u64> = IntMap::with_capacity(20);
111    /// map.set_load_factor(0.909); // Sets load factor to 90.9%
112    /// ```
113    pub fn set_load_factor(&mut self, load_factor: f32) {
114        self.load_factor = (load_factor * 1000.) as usize;
115        self.increase_cache_if_needed();
116    }
117
118    /// Returns the current load factor.
119    pub fn get_load_factor(&self) -> f32 {
120        self.load_factor as f32 / 1000.
121    }
122
123    /// Ensures that the [`IntMap`] has space for at least `additional` more elements
124    pub fn reserve(&mut self, additional: usize) {
125        let capacity = self.count + additional;
126        while self.lim() < capacity {
127            self.increase_cache();
128        }
129    }
130
131    /// Inserts a key/value pair into the [`IntMap`].
132    ///
133    /// This function returns the previous value if any otherwise `None`.
134    ///
135    /// # Examples
136    ///
137    /// ```
138    /// use intmap::IntMap;
139    ///
140    /// let mut map: IntMap::<u64, _> = IntMap::new();
141    /// assert_eq!(map.insert(21, "Eat my shorts"), None);
142    /// assert_eq!(map.insert(21, "Ay, caramba"), Some("Eat my shorts"));
143    /// assert_eq!(map.get(21), Some(&"Ay, caramba"));
144    /// ```
145    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
146        self.increase_cache_if_needed();
147
148        let k = key.into_int();
149        let ix = k.calc_index(self.mod_mask, K::PRIME);
150
151        let vals = &mut self.cache[ix];
152        let pos = vals.iter().position(|kv| kv.0.into_int() == k);
153
154        let old = if let Some(pos) = pos {
155            Some(vals.swap_remove(pos).1)
156        } else {
157            // Only increase count if we actually add a new entry
158            self.count += 1;
159            None
160        };
161
162        vals.push((key, value));
163
164        old
165    }
166
167    /// Insert a key/value pair into the [`IntMap`] if the key is not yet inserted.
168    ///
169    /// This function returns true if key/value were inserted and false otherwise.
170    ///
171    /// # Examples
172    ///
173    /// ```
174    /// use intmap::IntMap;
175    ///
176    /// let mut map: IntMap::<u64, _> = IntMap::new();
177    /// assert!(map.insert_checked(21, "Eat my shorts"));
178    /// assert!(!map.insert_checked(21, "Ay, caramba"));
179    /// assert_eq!(map.get(21), Some(&"Eat my shorts"));
180    /// ```
181    pub fn insert_checked(&mut self, key: K, value: V) -> bool {
182        self.increase_cache_if_needed();
183
184        let k = key.into_int();
185        let ix = k.calc_index(self.mod_mask, K::PRIME);
186
187        let vals = &mut self.cache[ix];
188        if vals.iter().any(|kv| kv.0.into_int() == k) {
189            return false;
190        }
191
192        self.count += 1;
193        vals.push((key, value));
194
195        true
196    }
197
198    /// Gets the value for the given key from the [`IntMap`].
199    ///
200    /// # Examples
201    ///
202    /// ```
203    /// use intmap::IntMap;
204    ///
205    /// let mut map: IntMap<u64, u64> = IntMap::new();
206    /// map.insert(21, 42);
207    /// let val = map.get(21);
208    /// assert!(val.is_some());
209    /// assert_eq!(*val.unwrap(), 42);
210    /// assert!(map.contains_key(21));
211    /// ```
212    pub fn get(&self, key: K) -> Option<&V> {
213        if self.is_empty() {
214            return None;
215        }
216
217        let k = key.into_int();
218        let ix = k.calc_index(self.mod_mask, K::PRIME);
219
220        let vals = &self.cache[ix];
221
222        vals.iter()
223            .find_map(|kv| (kv.0.into_int() == k).then(|| &kv.1))
224    }
225
226    /// Gets the mutable value for the given key from the [`IntMap`].
227    ///
228    /// # Examples
229    ///
230    /// ```
231    /// use intmap::IntMap;
232    ///
233    /// let mut map: IntMap<u64, u64> = IntMap::new();
234    /// map.insert(21, 42);
235    ///
236    /// assert_eq!(*map.get(21).unwrap(), 42);
237    /// assert!(map.contains_key(21));
238    ///
239    /// {
240    ///     let mut val = map.get_mut(21).unwrap();
241    ///     *val+=1;
242    /// }
243    ///     assert_eq!(*map.get(21).unwrap(), 43);
244    /// ```
245    pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
246        if self.is_empty() {
247            return None;
248        }
249
250        let k = key.into_int();
251        let ix = k.calc_index(self.mod_mask, K::PRIME);
252
253        let vals = &mut self.cache[ix];
254
255        return vals
256            .iter_mut()
257            .find_map(|kv| (kv.0.into_int() == k).then(move || &mut kv.1));
258    }
259
260    /// Removes the value for given key from the [`IntMap`] and returns it.
261    ///
262    /// # Examples
263    ///
264    /// ```
265    /// use intmap::IntMap;
266    ///
267    /// let mut map: IntMap<u64, u64> = IntMap::new();
268    /// map.insert(21, 42);
269    /// let val = map.remove(21);
270    /// assert!(val.is_some());
271    /// assert_eq!(val.unwrap(), 42);
272    /// assert!(!map.contains_key(21));
273    /// ```
274    pub fn remove(&mut self, key: K) -> Option<V> {
275        if self.is_empty() {
276            return None;
277        }
278
279        let k = key.into_int();
280        let ix = k.calc_index(self.mod_mask, K::PRIME);
281
282        let vals = &mut self.cache[ix];
283
284        for i in 0..vals.len() {
285            let peek = &vals[i].0;
286
287            if peek.into_int() == k {
288                self.count -= 1;
289                let kv = vals.swap_remove(i);
290                return Some(kv.1);
291            }
292        }
293
294        None
295    }
296
297    /// Returns true if the key is present in the [`IntMap`].
298    ///
299    /// # Examples
300    ///
301    /// ```
302    /// use intmap::IntMap;
303    ///
304    /// let mut map: IntMap<u64, u64> = IntMap::new();
305    /// map.insert(21, 42);
306    /// assert!(map.contains_key(21));
307    /// ```
308    pub fn contains_key(&self, key: K) -> bool {
309        self.get(key).is_some()
310    }
311
312    /// Removes all elements from the [`IntMap`].
313    ///
314    /// # Examples
315    ///
316    /// ```
317    /// use intmap::IntMap;
318    ///
319    /// let mut map: IntMap<u64, u64> = IntMap::new();
320    /// map.insert(21, 42);
321    /// map.clear();
322    /// assert_eq!(map.len(), 0);
323    /// ```
324    pub fn clear(&mut self) {
325        for vals in &mut self.cache {
326            vals.clear();
327        }
328
329        self.count = 0;
330    }
331
332    /// Retains only the key/value pairs specified by the predicate.
333    ///
334    /// In other words, remove all elements such that `f(key, &value)` returns false.
335    ///
336    /// # Examples
337    ///
338    /// ```
339    /// use intmap::IntMap;
340    ///
341    /// let mut map: IntMap<u64, u64> = IntMap::new();
342    /// map.insert(1, 11);
343    /// map.insert(2, 12);
344    /// map.insert(4, 13);
345    ///
346    /// // retain only the odd values
347    /// map.retain(|k, v| *v % 2 == 1);
348    ///
349    /// assert_eq!(map.len(), 2);
350    /// assert!(map.contains_key(1));
351    /// assert!(map.contains_key(4));
352    /// ```
353    pub fn retain<F>(&mut self, mut f: F)
354    where
355        F: FnMut(K, &V) -> bool,
356    {
357        let mut removed = 0;
358        for vals in &mut self.cache {
359            vals.retain(|(k, v)| {
360                let keep = (f)(*k, v);
361                if !keep {
362                    removed += 1;
363                }
364                keep
365            });
366        }
367
368        self.count -= removed;
369    }
370
371    /// Returns true if the [`IntMap`] is empty
372    ///
373    /// # Examples
374    ///
375    /// ```
376    /// use intmap::IntMap;
377    ///
378    /// let mut map: IntMap<u64, u64> = IntMap::new();
379    /// map.insert(21, 42);
380    /// assert!(!map.is_empty());
381    /// map.remove(21);
382    /// assert!(map.is_empty());
383    /// ```
384    pub fn is_empty(&self) -> bool {
385        self.count == 0
386    }
387
388    //**** Iterators *****
389
390    /// Returns an [`Iterator`] over all key/value pairs.
391    pub fn iter(&self) -> Iter<K, V> {
392        Iter::new(&self.cache)
393    }
394
395    /// Returns an [`Iterator`] over all key/value pairs with mutable value.
396    pub fn iter_mut(&mut self) -> IterMut<K, V> {
397        IterMut::new(&mut self.cache)
398    }
399
400    /// Returns an [`Iterator`] over all keys.
401    pub fn keys(&self) -> Keys<K, V> {
402        Keys { inner: self.iter() }
403    }
404
405    /// Returns an [`Iterator`] over all values.
406    pub fn values(&self) -> Values<K, V> {
407        Values { inner: self.iter() }
408    }
409
410    /// Returns an [`Iterator`] over all mutable values.
411    pub fn values_mut(&mut self) -> ValuesMut<K, V> {
412        ValuesMut {
413            inner: self.iter_mut(),
414        }
415    }
416
417    /// Returns an [`Iterator`] over all key/value pairs that removes the pairs from the [`IntMap`]
418    /// during iteration.
419    ///
420    /// If the [`Iterator`] is droppend then all remaining key/value pairs will be removed from
421    /// the [`IntMap`].
422    pub fn drain(&mut self) -> Drain<K, V> {
423        Drain::new(&mut self.cache, &mut self.count)
424    }
425
426    //**** Internal hash stuff *****
427
428    #[inline(always)]
429    fn lim(&self) -> usize {
430        if self.size == 0 {
431            0
432        } else {
433            2usize.pow(self.size)
434        }
435    }
436
437    fn increase_cache(&mut self) {
438        self.size += 1;
439        let new_lim = self.lim();
440        self.mod_mask = new_lim - 1;
441
442        let mut vec: Vec<Vec<(K, V)>> = (0..new_lim).map(|_| Vec::new()).collect();
443        std::mem::swap(&mut self.cache, &mut vec);
444
445        for key in vec.into_iter().flatten() {
446            let k = key.0.into_int();
447            let ix = k.calc_index(self.mod_mask, K::PRIME);
448
449            let vals = &mut self.cache[ix];
450            vals.push(key);
451        }
452
453        debug_assert!(
454            self.cache.len() == self.lim(),
455            "cache vector the wrong length, lim: {:?} cache: {:?}",
456            self.lim(),
457            self.cache.len()
458        );
459    }
460
461    #[inline]
462    fn increase_cache_if_needed(&mut self) -> bool {
463        let initial_cache_len = self.cache.len();
464
465        // Handle empty cache to prevent division by zero.
466        if self.cache.is_empty() {
467            self.increase_cache();
468        }
469
470        // Tried using floats here but insert performance tanked.
471        while ((self.count * 1000) / self.cache.len()) > self.load_factor {
472            self.increase_cache();
473        }
474
475        initial_cache_len != self.cache.len()
476    }
477
478    //**** More public methods *****
479
480    /// Returns the number of key/value pairs in the [`IntMap`].
481    pub fn len(&self) -> usize {
482        self.count
483    }
484
485    /// Returns the number of filled slots.
486    pub fn load(&self) -> u64 {
487        self.cache.iter().filter(|vals| !vals.is_empty()).count() as u64
488    }
489
490    /// Returns the ratio between key/value pairs and available slots as percentage.
491    ///
492    /// # Examples
493    ///
494    /// ```
495    /// use intmap::IntMap;
496    ///
497    /// let mut map: IntMap<u64, u64> = IntMap::with_capacity(2);
498    /// map.set_load_factor(2.0);
499    /// assert_eq!(map.load_rate(), 0.0);
500    /// map.insert(1, 42);
501    /// assert_eq!(map.load_rate(), 50.0);
502    /// map.insert(2, 42);
503    /// assert_eq!(map.load_rate(), 100.0);
504    /// map.insert(3, 42);
505    /// assert_eq!(map.load_rate(), 150.0);
506    /// ```
507    pub fn load_rate(&self) -> f64 {
508        (self.count as f64) / (self.cache.len() as f64) * 100f64
509    }
510
511    /// Returns the total number of available slots.
512    pub fn capacity(&self) -> usize {
513        self.cache.len()
514    }
515
516    //**** Testing methods *****
517
518    /// Checks whether the actual count of key/value pairs matches [`IntMap::count`].
519    ///
520    /// Only for testing.
521    #[doc(hidden)]
522    pub fn assert_count(&self) -> bool {
523        let count = self.cache.iter().flatten().count();
524
525        self.count == count
526    }
527
528    /// Returns a new [`IntMap`] that contains only the collisions of the current [`IntMap`].
529    ///
530    /// Only for testing.
531    #[doc(hidden)]
532    pub fn collisions(&self) -> IntMap<u64, u64> {
533        let mut map = IntMap::new();
534
535        for s in self.cache.iter() {
536            let key = s.len() as u64;
537            if key > 1 {
538                if !map.contains_key(key) {
539                    map.insert(key, 1);
540                } else {
541                    let counter = map.get_mut(key).unwrap();
542                    *counter += 1;
543                }
544            }
545        }
546
547        map
548    }
549
550    //**** Entry API *****
551
552    /// Gets the [`Entry`] that corresponds to the given key.
553    ///
554    /// # Examples
555    ///
556    /// ```
557    /// use intmap::{IntMap, Entry};
558    ///
559    /// let mut counters = IntMap::new();
560    ///
561    /// for number in [10, 30, 10, 40, 50, 50, 60, 50] {
562    ///     let counter = match counters.entry(number) {
563    ///         Entry::Occupied(entry) => entry.into_mut(),
564    ///         Entry::Vacant(entry) => entry.insert(0),
565    ///     };
566    ///     *counter += 1;
567    /// }
568    ///
569    /// assert_eq!(counters.get(10), Some(&2));
570    /// assert_eq!(counters.get(20), None);
571    /// assert_eq!(counters.get(30), Some(&1));
572    /// assert_eq!(counters.get(40), Some(&1));
573    /// assert_eq!(counters.get(50), Some(&3));
574    /// assert_eq!(counters.get(60), Some(&1));
575    /// ```
576    pub fn entry(&mut self, key: K) -> Entry<K, V> {
577        Entry::new(key, self)
578    }
579}
580
581impl<K, V> Default for IntMap<K, V> {
582    fn default() -> Self {
583        Self::new()
584    }
585}
586
587// ***************** Equality *********************
588
589impl<K, V> PartialEq for IntMap<K, V>
590where
591    K: IntKey,
592    V: PartialEq,
593{
594    fn eq(&self, other: &IntMap<K, V>) -> bool {
595        self.count == other.count && self.iter().all(|(k, a)| other.get(k) == Some(a))
596    }
597}
598impl<K: IntKey, V: Eq> Eq for IntMap<K, V> {}
599
600// ***************** Debug *********************
601
602impl<K, V> std::fmt::Debug for IntMap<K, V>
603where
604    K: IntKey + std::fmt::Debug,
605    V: std::fmt::Debug,
606{
607    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
608        fmt.debug_map().entries(self.iter()).finish()
609    }
610}