mule_map/
mule_map.rs

1use entry::{
2    Entry, OccupiedEntry, OccupiedHashMapEntry, OccupiedVecEntry, VacantEntry, VacantHashMapEntry,
3    VacantVecEntry,
4};
5use key_index::KeyIndex;
6use num_traits::AsPrimitive;
7use num_traits::PrimInt;
8use sealed::sealed;
9use std::collections::HashMap;
10use std::fmt::Debug;
11use std::num::NonZero;
12
13pub(crate) mod entry;
14mod iterators;
15mod key_index;
16
17#[sealed]
18#[doc(hidden)]
19pub trait NonZeroInt {
20    type UnderlyingType;
21    const ONE: Self;
22    #[inline]
23    fn checked_add(self, other: Self::UnderlyingType) -> Option<Self>
24    where
25        Self::UnderlyingType: bytemuck::Pod,
26        Self::UnderlyingType: std::ops::AddAssign<Self::UnderlyingType>,
27        Self: bytemuck::PodInOption,
28        Self::UnderlyingType: PrimInt,
29    {
30        let mut result = Some(self);
31        let x = bytemuck::cast_mut::<Option<Self>, Self::UnderlyingType>(&mut result);
32        *x += other;
33        result
34    }
35}
36
37#[sealed]
38impl NonZeroInt for std::num::NonZeroI128 {
39    const ONE: std::num::NonZeroI128 = const { NonZero::new(1).expect("1 is not 0") };
40    type UnderlyingType = i128;
41}
42
43#[sealed]
44impl NonZeroInt for std::num::NonZeroI16 {
45    const ONE: std::num::NonZeroI16 = const { NonZero::new(1).expect("1 is not 0") };
46    type UnderlyingType = i16;
47}
48
49#[sealed]
50impl NonZeroInt for std::num::NonZeroI32 {
51    const ONE: std::num::NonZeroI32 = const { NonZero::new(1).expect("1 is not 0") };
52    type UnderlyingType = i32;
53}
54
55#[sealed]
56impl NonZeroInt for std::num::NonZeroI64 {
57    const ONE: std::num::NonZeroI64 = const { NonZero::new(1).expect("1 is not 0") };
58    type UnderlyingType = i64;
59}
60
61#[sealed]
62impl NonZeroInt for std::num::NonZeroI8 {
63    const ONE: std::num::NonZeroI8 = const { NonZero::new(1).expect("1 is not 0") };
64    type UnderlyingType = i8;
65}
66
67#[sealed]
68impl NonZeroInt for std::num::NonZeroIsize {
69    const ONE: std::num::NonZeroIsize = const { NonZero::new(1).expect("1 is not 0") };
70    type UnderlyingType = isize;
71}
72
73#[sealed]
74impl NonZeroInt for std::num::NonZeroU128 {
75    const ONE: std::num::NonZeroU128 = const { NonZero::new(1).expect("1 is not 0") };
76    type UnderlyingType = u128;
77}
78
79#[sealed]
80impl NonZeroInt for std::num::NonZeroU16 {
81    const ONE: std::num::NonZeroU16 = const { NonZero::new(1).expect("1 is not 0") };
82    type UnderlyingType = u16;
83}
84
85#[sealed]
86impl NonZeroInt for std::num::NonZeroU32 {
87    const ONE: std::num::NonZeroU32 = const { NonZero::new(1).expect("1 is not 0") };
88    type UnderlyingType = u32;
89}
90
91#[sealed]
92impl NonZeroInt for std::num::NonZeroU64 {
93    const ONE: std::num::NonZeroU64 = const { NonZero::new(1).expect("1 is not 0") };
94    type UnderlyingType = u64;
95}
96
97#[sealed]
98impl NonZeroInt for std::num::NonZeroU8 {
99    const ONE: std::num::NonZeroU8 = const { NonZero::new(1).expect("1 is not 0") };
100    type UnderlyingType = u8;
101}
102
103#[sealed]
104impl NonZeroInt for std::num::NonZeroUsize {
105    const ONE: std::num::NonZeroUsize = const { NonZero::new(1).expect("1 is not 0") };
106    type UnderlyingType = usize;
107}
108
109/// [`MuleMap`] is a hybrid between a [`HashMap`] and a lookup table. It improves performance for frequently accessed keys
110/// in a known range. If a key (integer) is in the user specified range, then its value will be stored directly in the lookup table.
111///
112/// # Differences between [`HashMap`] and [`MuleMap`]
113///
114/// - **The key, `K`, must be an integer type.** - The key is directly mapped to the index in the lookup, so it must be
115///     an integer.
116/// - **The key, `K`, is passed by value** - Because it is a primitive integer type.
117/// - **The hash builder, `S`,  does not have a default** - You must specify your hash builder. The assumption being
118///     that if you need better performance you will likely also want to use a custom hash function.
119/// - **`TABLE_MIN_VALUE` and `TABLE_SIZE`** -  If a key is between `TABLE_MIN_VALUE` and `TABLE_MIN_VALUE + TABLE_SIZE`, then
120///     the value will be stored directly in the lookup table, instead of using the `HashMap`.
121///     **NOTE:** Currently the type of a const generic can’t depend on another generic type argument, so `TABLE_MIN_VALUE`
122///     can’t use the same type as the key. Because of this, We are using [`i128`], but that means we can’t
123///     represent values near [`u128::MAX`]. Hopefully having frequent keys near [`u128::MAX`] is extremely rare.
124///
125/// # Performance
126///
127/// Benchmarks (using random selection) start to show speed improvements when about 50% of the key accesses are in the
128/// lookup table. Performance is almost identical to `HashMap` with less than 50%.
129///
130/// ## Example
131///
132/// ```
133/// use mule_map::MuleMap;
134/// use std::num::NonZero;
135/// type Hash = fnv_rs::FnvBuildHasher;  // Use whatever hash function you prefer
136///
137/// // Using Entry API
138/// let mut mule_map = MuleMap::<u32, usize, Hash>::new();
139/// assert_eq!(mule_map.get(5), None);
140/// let entry = mule_map.entry(5);
141/// entry.or_insert(10);
142/// assert_eq!(mule_map.get(5), Some(&10));
143///
144/// // Using NonZero and bump
145/// let mut mule_map_non_zero = MuleMap::<u32, NonZero<i32>, Hash>::default();
146///
147/// mule_map_non_zero.bump_non_zero(10);
148/// mule_map_non_zero.bump_non_zero(10);
149/// mule_map_non_zero.bump_non_zero(999_999);
150/// mule_map_non_zero.bump_non_zero(999_999);
151///
152// assert_eq!(mule_map_non_zero.get(10), NonZero::<i32>::new(2).as_ref());
153// assert_eq!(mule_map_non_zero.get(999_999),NonZero::<i32>::new(2).as_ref());
154/// ```
155#[derive(Debug)]
156pub struct MuleMap<
157    K,
158    V,
159    S,
160    const TABLE_MIN_VALUE: i128 = 0,
161    const TABLE_SIZE: usize = { u8::MAX as usize },
162> {
163    hash_map: HashMap<K, V, S>,
164    table: [Option<V>; TABLE_SIZE],
165}
166
167impl<K, V, S, const TABLE_MIN_VALUE: i128, const TABLE_SIZE: usize> Default
168    for MuleMap<K, V, S, TABLE_MIN_VALUE, TABLE_SIZE>
169where
170    K: PrimInt + Eq + std::hash::Hash + KeyIndex<K, TABLE_MIN_VALUE> + TryFrom<i128> + 'static,
171    S: Default + std::hash::BuildHasher,
172    V: PartialEq + Copy,
173    i128: AsPrimitive<K>,
174    usize: AsPrimitive<K>,
175    <K as TryFrom<i128>>::Error: Debug,
176{
177    /// Creates an empty [`MuleMap`].
178    ///
179    /// # Example
180    /// ```
181    /// let mule_map = mule_map::MuleMap::<u32, usize, fnv_rs::FnvBuildHasher>::default();
182    /// ```
183    ///
184    /// See: [`MuleMap::with_capacity_and_hasher`]
185    ///
186    /// Analogous to [`HashMap::default`]
187    fn default() -> Self {
188        Self::new()
189    }
190}
191
192impl<K, V, S, const TABLE_MIN_VALUE: i128, const TABLE_SIZE: usize>
193    MuleMap<K, V, S, TABLE_MIN_VALUE, TABLE_SIZE>
194where
195    K: PrimInt + Eq + std::hash::Hash + KeyIndex<K, TABLE_MIN_VALUE> + TryFrom<i128> + 'static,
196    S: Default + std::hash::BuildHasher,
197    V: PartialEq + Copy,
198    i128: AsPrimitive<K>,
199    usize: AsPrimitive<K>,
200    <K as TryFrom<i128>>::Error: Debug,
201{
202    // Hard limit, way beyond practical lookup table size. This makes it easier to calculate the key index
203    const STATIC_ASSERT_LIMIT_SIZE_TO_I32_MAX: () =
204        assert!((TABLE_SIZE as u128) < i32::MAX as u128);
205
206    #[inline]
207    #[must_use]
208    fn use_lookup_table(key: K) -> bool {
209        // NOTE: TABLE_MIN_VALUE + TABLE_SIZE and TABLE_MIN_VALUE must fit into a key type, K
210        key < (TABLE_MIN_VALUE.as_() + TABLE_SIZE.as_()) && key >= TABLE_MIN_VALUE.as_()
211    }
212
213    /// Creates an empty [`MuleMap`].
214    ///
215    /// # Example
216    /// ```
217    /// let mule_map = mule_map::MuleMap::<u32, usize, fnv_rs::FnvBuildHasher>::new();
218    /// ```
219    ///
220    /// See: [`MuleMap::with_capacity_and_hasher`]
221    ///
222    /// Analogous to [`HashMap::new`]
223    #[must_use]
224    #[inline]
225    pub fn new() -> Self {
226        Self::with_capacity_and_hasher(0, S::default())
227    }
228
229    /// Creates an empty [`MuleMap`] with at least the provided capacity.
230    ///
231    /// # Example
232    /// ```
233    /// let mule_map = mule_map::MuleMap::<u32, usize, fnv_rs::FnvBuildHasher>::with_capacity(100);
234    /// ```
235    ///
236    /// See: [`MuleMap::with_capacity_and_hasher`]
237    ///
238    /// Analogous to [`HashMap::with_capacity`]
239    #[must_use]
240    #[inline]
241    pub fn with_capacity(capacity: usize) -> Self {
242        Self::with_capacity_and_hasher(capacity, S::default())
243    }
244
245    /// Creates an empty [`MuleMap`] using `hash_builder`.
246    ///
247    /// # Example
248    /// ```
249    /// type Hash = fnv_rs::FnvBuildHasher;
250    /// let mule_map = mule_map::MuleMap::<u32, usize, fnv_rs::FnvBuildHasher>::with_hasher(Hash::default());
251    /// ```
252    ///
253    /// See: [`MuleMap::with_capacity_and_hasher`]
254    ///
255    /// Analogous to [`HashMap::with_hasher`]
256    #[must_use]
257    #[inline]
258    pub fn with_hasher(hash_builder: S) -> Self {
259        Self::with_capacity_and_hasher(0, hash_builder)
260    }
261
262    /// Creates an empty [`MuleMap`] with at least the provided capacity and using `hash_builder`.
263    ///
264    /// # Example
265    /// ```
266    /// type Hash = fnv_rs::FnvBuildHasher;
267    /// let mule_map = mule_map::MuleMap::<u32, usize, fnv_rs::FnvBuildHasher>::with_capacity_and_hasher(100, Hash::default());
268    /// ```
269    ///
270    /// # Panics
271    ///
272    /// Panics if
273    ///  - `TABLE_MIN_VALUE` or `TABLE_MIN_VALUE + TABLE_SIZE` doesn't fit into key type, `K`.
274    ///
275    /// Analogous to [`HashMap::with_capacity_and_hasher`]
276    #[must_use]
277    #[inline]
278    pub fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> Self {
279        let () = Self::STATIC_ASSERT_LIMIT_SIZE_TO_I32_MAX;
280
281        <i128 as TryInto<K>>::try_into(TABLE_MIN_VALUE + TABLE_SIZE as i128)
282            .expect("TABLE_MIN_VALUE + TABLE_SIZE should fit into key type, K");
283        <i128 as TryInto<K>>::try_into(TABLE_MIN_VALUE)
284            .expect("TABLE_MIN_VALUE should fit into key type, K");
285
286        MuleMap::<K, V, S, TABLE_MIN_VALUE, TABLE_SIZE> {
287            hash_map: HashMap::with_capacity_and_hasher(capacity, hash_builder),
288            table: [None; TABLE_SIZE],
289        }
290    }
291
292    /// Returns capacity of the underlying hash map.
293    ///
294    /// See [`HashMap::capacity`]
295    #[must_use]
296    #[inline]
297    pub fn capacity(&self) -> usize {
298        self.hash_map.capacity()
299    }
300
301    /// Returns a reference to the value corresponding to the key.
302    ///
303    /// Analogous to [`HashMap::get`]
304    #[must_use]
305    #[inline]
306    pub fn get(&self, key: K) -> Option<&V> {
307        if Self::use_lookup_table(key) {
308            self.table[key.key_index()].as_ref()
309        } else {
310            let result = self.hash_map.get(&key);
311            result
312        }
313    }
314
315    /// Returns true if the map contains a value for the specified key.
316    ///
317    /// Analogous to [`HashMap::contains_key`]
318    #[must_use]
319    #[inline]
320    pub fn contains_key(&self, key: K) -> bool {
321        if Self::use_lookup_table(key) {
322            self.table[key.key_index()].is_some()
323        } else {
324            self.hash_map.contains_key(&key)
325        }
326    }
327
328    /// Modify the values at location `key` by calling `f` on its value. If no value present, create a new value set to
329    /// `default`.
330    #[inline]
331    pub fn modify_or_insert<F>(&mut self, key: K, f: F, default: V)
332    where
333        F: FnOnce(&mut V),
334    {
335        if Self::use_lookup_table(key) {
336            let value = &mut self.table[key.key_index()];
337            match value {
338                Some(x) => f(x),
339                None => *value = Some(default),
340            }
341        } else {
342            self.hash_map.entry(key).and_modify(f).or_insert(default);
343        }
344    }
345
346    /// Adds 1 to the value stored at location `key`. If the value is not present, the value 1 will be set at
347    /// that location.
348    ///
349    /// *NOTE:* This method can only be called with values that implement `AddAssign`, like primitives. For `NonZero<T>`
350    /// values use [`bump_non_zero`] - It uses the niche optimization for better performance.
351    ///
352    /// # Panics
353    ///
354    /// May panics if adding 1 results in overflow.
355    #[inline]
356    pub fn bump_int(&mut self, key: K)
357    where
358        V: std::ops::AddAssign<V> + num_traits::One + num_traits::Zero,
359    {
360        if Self::use_lookup_table(key) {
361            *self.table[key.key_index()].get_or_insert(V::zero()) += V::one();
362        } else {
363            self.hash_map
364                .entry(key)
365                .and_modify(|counter| *counter += V::one())
366                .or_insert(V::one());
367        }
368    }
369
370    /// Adds 1 to the value stored at location `key`. If the value is not present, the value 1 will be set at
371    /// that location. Uses the niche optimization for better performance with `Option<NonZero<T>>`.
372    ///
373    /// *NOTE:* This method can only be called with `NonZero<T>` values. For primitive values use [`bump_int`].
374    ///
375    /// # Panics
376    ///
377    /// Panics if adding 1 results in overflow.
378    #[inline]
379    pub fn bump_non_zero(&mut self, key: K)
380    where
381        V: NonZeroInt + bytemuck::PodInOption,
382        <V as NonZeroInt>::UnderlyingType: std::ops::AddAssign<V::UnderlyingType>,
383        <V as NonZeroInt>::UnderlyingType: bytemuck::Pod + PrimInt,
384    {
385        use num_traits::One;
386
387        if Self::use_lookup_table(key) {
388            *bytemuck::cast_mut::<Option<V>, V::UnderlyingType>(
389                &mut self.table[key.key_index()],
390            ) += V::UnderlyingType::one();
391        } else {
392            self.hash_map
393                .entry(key)
394                .and_modify(|counter| {
395                    *counter = counter
396                        .checked_add(V::UnderlyingType::one())
397                        .expect("Addition should not overflow");
398                })
399                .or_insert(V::ONE);
400        }
401    }
402
403    /// Gets the given key’s corresponding entry in the map for in-place manipulation.
404    ///
405    /// Analogous to [`HashMap::entry`]
406    #[must_use]
407    #[inline]
408    pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
409        if Self::use_lookup_table(key) {
410            let value: &mut Option<V> = &mut self.table[key.key_index()];
411            match value {
412                Some(_) => {
413                    Entry::<K, V>::Occupied(OccupiedEntry::Vec(OccupiedVecEntry { value, key }))
414                }
415                None => Entry::<K, V>::Vacant(VacantEntry::Vec(VacantVecEntry { value, key })),
416            }
417        } else {
418            match self.hash_map.entry(key) {
419                std::collections::hash_map::Entry::Occupied(base) => {
420                    Entry::<K, V>::Occupied(OccupiedEntry::HashMap(OccupiedHashMapEntry { base }))
421                }
422                std::collections::hash_map::Entry::Vacant(base) => {
423                    Entry::<K, V>::Vacant(VacantEntry::HashMap(VacantHashMapEntry { base }))
424                }
425            }
426        }
427    }
428}
429
430#[cfg(test)]
431mod tests {
432    use super::*;
433
434    #[test]
435    fn it_works() {
436        let mut mule_map_int = MuleMap::<u32, i32, fnv_rs::FnvBuildHasher>::default();
437        mule_map_int.bump_int(10);
438        mule_map_int.bump_int(10);
439        assert_eq!(mule_map_int.get(10), Some(&2));
440
441        let mut mule_map_non_zero = MuleMap::<u32, NonZero<i32>, fnv_rs::FnvBuildHasher>::default();
442
443        mule_map_non_zero.bump_non_zero(10);
444        mule_map_non_zero.bump_non_zero(10);
445        assert_eq!(mule_map_non_zero.get(10), NonZero::<i32>::new(2).as_ref());
446        mule_map_non_zero.bump_non_zero(999_999);
447        mule_map_non_zero.bump_non_zero(999_999);
448        assert_eq!(
449            mule_map_non_zero.get(999_999),
450            NonZero::<i32>::new(2).as_ref()
451        );
452
453        let mut mule_map = MuleMap::<u32, i32, fnv_rs::FnvBuildHasher>::default();
454
455        mule_map.modify_or_insert(100, |x| *x += 10, 1);
456        assert_eq!(mule_map.get(100), Some(&1));
457
458        mule_map.modify_or_insert(100, |x| *x += 10, 1);
459        assert_eq!(mule_map.get(100), Some(&11));
460    }
461}
mule_map/mule_map.rs

mule_map/
mule_map.rs
