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}