heapless/
indexmap.rs

Help
1use core::{borrow::Borrow, fmt, iter::FromIterator, mem, num::NonZeroU32, ops, slice};
2
3use hash32::{BuildHasher, BuildHasherDefault, FnvHasher, Hash, Hasher};
4
5use crate::Vec;
6
7/// A [`heapless::IndexMap`](./struct.IndexMap.html) using the default FNV hasher
8///
9/// A list of all Methods and Traits available for `FnvIndexMap` can be found in
10/// the [`heapless::IndexMap`](./struct.IndexMap.html) documentation.
11///
12/// # Examples
13/// ```
14/// use heapless::FnvIndexMap;
15///
16/// // A hash map with a capacity of 16 key-value pairs allocated on the stack
17/// let mut book_reviews = FnvIndexMap::<_, _, 16>::new();
18///
19/// // review some books.
20/// book_reviews.insert("Adventures of Huckleberry Finn",    "My favorite book.").unwrap();
21/// book_reviews.insert("Grimms' Fairy Tales",               "Masterpiece.").unwrap();
22/// book_reviews.insert("Pride and Prejudice",               "Very enjoyable.").unwrap();
23/// book_reviews.insert("The Adventures of Sherlock Holmes", "Eye lyked it alot.").unwrap();
24///
25/// // check for a specific one.
26/// if !book_reviews.contains_key("Les Misérables") {
27///     println!("We've got {} reviews, but Les Misérables ain't one.",
28///              book_reviews.len());
29/// }
30///
31/// // oops, this review has a lot of spelling mistakes, let's delete it.
32/// book_reviews.remove("The Adventures of Sherlock Holmes");
33///
34/// // look up the values associated with some keys.
35/// let to_find = ["Pride and Prejudice", "Alice's Adventure in Wonderland"];
36/// for book in &to_find {
37///     match book_reviews.get(book) {
38///         Some(review) => println!("{}: {}", book, review),
39///         None => println!("{} is unreviewed.", book)
40///     }
41/// }
42///
43/// // iterate over everything.
44/// for (book, review) in &book_reviews {
45///     println!("{}: \"{}\"", book, review);
46/// }
47/// ```
48pub type FnvIndexMap<K, V, const N: usize> = IndexMap<K, V, BuildHasherDefault<FnvHasher>, N>;
49
50#[derive(Clone, Copy, Eq, PartialEq)]
51struct HashValue(u16);
52
53impl HashValue {
54    fn desired_pos(&self, mask: usize) -> usize {
55        usize::from(self.0) & mask
56    }
57
58    fn probe_distance(&self, mask: usize, current: usize) -> usize {
59        current.wrapping_sub(self.desired_pos(mask) as usize) & mask
60    }
61}
62
63#[doc(hidden)]
64#[derive(Clone)]
65pub struct Bucket<K, V> {
66    hash: HashValue,
67    key: K,
68    value: V,
69}
70
71#[doc(hidden)]
72#[derive(Clone, Copy, PartialEq)]
73pub struct Pos {
74    // compact representation of `{ hash_value: u16, index: u16 }`
75    // To get the most from `NonZero` we store the *value minus 1*. This way `None::Option<Pos>`
76    // is equivalent to the very unlikely value of  `{ hash_value: 0xffff, index: 0xffff }` instead
77    // the more likely of `{ hash_value: 0x00, index: 0x00 }`
78    nz: NonZeroU32,
79}
80
81impl Pos {
82    fn new(index: usize, hash: HashValue) -> Self {
83        Pos {
84            nz: unsafe {
85                NonZeroU32::new_unchecked(
86                    ((u32::from(hash.0) << 16) + index as u32).wrapping_add(1),
87                )
88            },
89        }
90    }
91
92    fn hash(&self) -> HashValue {
93        HashValue((self.nz.get().wrapping_sub(1) >> 16) as u16)
94    }
95
96    fn index(&self) -> usize {
97        self.nz.get().wrapping_sub(1) as u16 as usize
98    }
99}
100
101enum Insert<K, V> {
102    Success(Inserted<V>),
103    Full((K, V)),
104}
105struct Inserted<V> {
106    index: usize,
107    old_value: Option<V>,
108}
109
110macro_rules! probe_loop {
111    ($probe_var: ident < $len: expr, $body: expr) => {
112        loop {
113            if $probe_var < $len {
114                $body
115                    $probe_var += 1;
116            } else {
117                $probe_var = 0;
118            }
119        }
120    }
121}
122
123struct CoreMap<K, V, const N: usize> {
124    entries: Vec<Bucket<K, V>, N>,
125    indices: [Option<Pos>; N],
126}
127
128impl<K, V, const N: usize> CoreMap<K, V, N> {
129    const fn new() -> Self {
130        const INIT: Option<Pos> = None;
131
132        CoreMap {
133            entries: Vec::new(),
134            indices: [INIT; N],
135        }
136    }
137}
138
139impl<K, V, const N: usize> CoreMap<K, V, N>
140where
141    K: Eq + Hash,
142{
143    fn capacity() -> usize {
144        N
145    }
146
147    fn mask() -> usize {
148        Self::capacity() - 1
149    }
150
151    fn find<Q>(&self, hash: HashValue, query: &Q) -> Option<(usize, usize)>
152    where
153        K: Borrow<Q>,
154        Q: ?Sized + Eq,
155    {
156        let mut probe = hash.desired_pos(Self::mask());
157        let mut dist = 0;
158
159        probe_loop!(probe < self.indices.len(), {
160            if let Some(pos) = self.indices[probe] {
161                let entry_hash = pos.hash();
162                // NOTE(i) we use unchecked indexing below
163                let i = pos.index();
164                debug_assert!(i < self.entries.len());
165
166                if dist > entry_hash.probe_distance(Self::mask(), probe) {
167                    // give up when probe distance is too long
168                    return None;
169                } else if entry_hash == hash
170                    && unsafe { self.entries.get_unchecked(i).key.borrow() == query }
171                {
172                    return Some((probe, i));
173                }
174            } else {
175                return None;
176            }
177
178            dist += 1;
179        });
180    }
181
182    fn insert(&mut self, hash: HashValue, key: K, value: V) -> Insert<K, V> {
183        let mut probe = hash.desired_pos(Self::mask());
184        let mut dist = 0;
185
186        probe_loop!(probe < self.indices.len(), {
187            let pos = &mut self.indices[probe];
188
189            if let Some(pos) = *pos {
190                let entry_hash = pos.hash();
191                // NOTE(i) we use unchecked indexing below
192                let i = pos.index();
193                debug_assert!(i < self.entries.len());
194
195                let their_dist = entry_hash.probe_distance(Self::mask(), probe);
196
197                if their_dist < dist {
198                    if self.entries.is_full() {
199                        return Insert::Full((key, value));
200                    }
201                    // robin hood: steal the spot if it's better for us
202                    let index = self.entries.len();
203                    unsafe { self.entries.push_unchecked(Bucket { hash, key, value }) };
204                    return Insert::Success(Inserted {
205                        index: self.insert_phase_2(probe, Pos::new(index, hash)),
206                        old_value: None,
207                    });
208                } else if entry_hash == hash && unsafe { self.entries.get_unchecked(i).key == key }
209                {
210                    return Insert::Success(Inserted {
211                        index: i,
212                        old_value: Some(mem::replace(
213                            unsafe { &mut self.entries.get_unchecked_mut(i).value },
214                            value,
215                        )),
216                    });
217                }
218            } else {
219                if self.entries.is_full() {
220                    return Insert::Full((key, value));
221                }
222                // empty bucket, insert here
223                let index = self.entries.len();
224                *pos = Some(Pos::new(index, hash));
225                unsafe { self.entries.push_unchecked(Bucket { hash, key, value }) };
226                return Insert::Success(Inserted {
227                    index,
228                    old_value: None,
229                });
230            }
231            dist += 1;
232        });
233    }
234
235    // phase 2 is post-insert where we forward-shift `Pos` in the indices.
236    fn insert_phase_2(&mut self, mut probe: usize, mut old_pos: Pos) -> usize {
237        probe_loop!(probe < self.indices.len(), {
238            let pos = unsafe { self.indices.get_unchecked_mut(probe) };
239
240            let mut is_none = true; // work around lack of NLL
241            if let Some(pos) = pos.as_mut() {
242                old_pos = mem::replace(pos, old_pos);
243                is_none = false;
244            }
245
246            if is_none {
247                *pos = Some(old_pos);
248                return probe;
249            }
250        });
251    }
252
253    fn remove_found(&mut self, probe: usize, found: usize) -> (K, V) {
254        // index `probe` and entry `found` is to be removed
255        // use swap_remove, but then we need to update the index that points
256        // to the other entry that has to move
257        self.indices[probe] = None;
258        let entry = unsafe { self.entries.swap_remove_unchecked(found) };
259
260        // correct index that points to the entry that had to swap places
261        if let Some(entry) = self.entries.get(found) {
262            // was not last element
263            // examine new element in `found` and find it in indices
264            let mut probe = entry.hash.desired_pos(Self::mask());
265
266            probe_loop!(probe < self.indices.len(), {
267                if let Some(pos) = self.indices[probe] {
268                    if pos.index() >= self.entries.len() {
269                        // found it
270                        self.indices[probe] = Some(Pos::new(found, entry.hash));
271                        break;
272                    }
273                }
274            });
275        }
276
277        self.backward_shift_after_removal(probe);
278
279        (entry.key, entry.value)
280    }
281
282    fn backward_shift_after_removal(&mut self, probe_at_remove: usize) {
283        // backward shift deletion in self.indices
284        // after probe, shift all non-ideally placed indices backward
285        let mut last_probe = probe_at_remove;
286        let mut probe = probe_at_remove + 1;
287
288        probe_loop!(probe < self.indices.len(), {
289            if let Some(pos) = self.indices[probe] {
290                let entry_hash = pos.hash();
291
292                if entry_hash.probe_distance(Self::mask(), probe) > 0 {
293                    unsafe { *self.indices.get_unchecked_mut(last_probe) = self.indices[probe] }
294                    self.indices[probe] = None;
295                } else {
296                    break;
297                }
298            } else {
299                break;
300            }
301            last_probe = probe;
302        });
303    }
304}
305
306impl<K, V, const N: usize> Clone for CoreMap<K, V, N>
307where
308    K: Eq + Hash + Clone,
309    V: Clone,
310{
311    fn clone(&self) -> Self {
312        Self {
313            entries: self.entries.clone(),
314            indices: self.indices.clone(),
315        }
316    }
317}
318
319/// A view into an entry in the map
320pub enum Entry<'a, K, V, const N: usize> {
321    /// The entry corresponding to the key `K` exists in the map
322    Occupied(OccupiedEntry<'a, K, V, N>),
323    /// The entry corresponding to the key `K` does not exist in the map
324    Vacant(VacantEntry<'a, K, V, N>),
325}
326
327/// An occupied entry which can be manipulated
328pub struct OccupiedEntry<'a, K, V, const N: usize> {
329    key: K,
330    probe: usize,
331    pos: usize,
332    core: &'a mut CoreMap<K, V, N>,
333}
334
335impl<'a, K, V, const N: usize> OccupiedEntry<'a, K, V, N>
336where
337    K: Eq + Hash,
338{
339    /// Gets a reference to the key that this entity corresponds to
340    pub fn key(&self) -> &K {
341        &self.key
342    }
343
344    /// Removes this entry from the map and yields its corresponding key and value
345    pub fn remove_entry(self) -> (K, V) {
346        self.core.remove_found(self.probe, self.pos)
347    }
348
349    /// Gets a reference to the value associated with this entry
350    pub fn get(&self) -> &V {
351        // SAFETY: Already checked existence at instantiation and the only mutable reference
352        // to the map is internally held.
353        unsafe { &self.core.entries.get_unchecked(self.pos).value }
354    }
355
356    /// Gets a mutable reference to the value associated with this entry
357    pub fn get_mut(&mut self) -> &mut V {
358        // SAFETY: Already checked existence at instantiation and the only mutable reference
359        // to the map is internally held.
360        unsafe { &mut self.core.entries.get_unchecked_mut(self.pos).value }
361    }
362
363    /// Consumes this entry and yields a reference to the underlying value
364    pub fn into_mut(self) -> &'a mut V {
365        // SAFETY: Already checked existence at instantiation and the only mutable reference
366        // to the map is internally held.
367        unsafe { &mut self.core.entries.get_unchecked_mut(self.pos).value }
368    }
369
370    /// Overwrites the underlying map's value with this entry's value
371    pub fn insert(self, value: V) -> V {
372        // SAFETY: Already checked existence at instantiation and the only mutable reference
373        // to the map is internally held.
374        unsafe {
375            mem::replace(
376                &mut self.core.entries.get_unchecked_mut(self.pos).value,
377                value,
378            )
379        }
380    }
381
382    /// Removes this entry from the map and yields its value
383    pub fn remove(self) -> V {
384        self.remove_entry().1
385    }
386}
387
388/// A view into an empty slot in the underlying map
389pub struct VacantEntry<'a, K, V, const N: usize> {
390    key: K,
391    hash_val: HashValue,
392    core: &'a mut CoreMap<K, V, N>,
393}
394impl<'a, K, V, const N: usize> VacantEntry<'a, K, V, N>
395where
396    K: Eq + Hash,
397{
398    /// Get the key associated with this entry
399    pub fn key(&self) -> &K {
400        &self.key
401    }
402
403    /// Consumes this entry to yield to key associated with it
404    pub fn into_key(self) -> K {
405        self.key
406    }
407
408    /// Inserts this entry into to underlying map, yields a mutable reference to the inserted value.
409    /// If the map is at capacity the value is returned instead.
410    pub fn insert(self, value: V) -> Result<&'a mut V, V> {
411        if self.core.entries.is_full() {
412            Err(value)
413        } else {
414            match self.core.insert(self.hash_val, self.key, value) {
415                Insert::Success(inserted) => {
416                    unsafe {
417                        // SAFETY: Already checked existence at instantiation and the only mutable reference
418                        // to the map is internally held.
419                        Ok(&mut (*self.core.entries.as_mut_ptr().add(inserted.index)).value)
420                    }
421                }
422                Insert::Full((_, v)) => Err(v),
423            }
424        }
425    }
426}
427
428/// Fixed capacity [`IndexMap`](https://docs.rs/indexmap/1/indexmap/map/struct.IndexMap.html)
429///
430/// Note that you cannot use `IndexMap` directly, since it is generic around the hashing algorithm
431/// in use. Pick a concrete instantiation like [`FnvIndexMap`](./type.FnvIndexMap.html) instead
432/// or create your own.
433///
434/// Note that the capacity of the `IndexMap` must be a power of 2.
435///
436/// # Examples
437/// Since `IndexMap` cannot be used directly, we're using its `FnvIndexMap` instantiation
438/// for this example.
439///
440/// ```
441/// use heapless::FnvIndexMap;
442///
443/// // A hash map with a capacity of 16 key-value pairs allocated on the stack
444/// let mut book_reviews = FnvIndexMap::<_, _, 16>::new();
445///
446/// // review some books.
447/// book_reviews.insert("Adventures of Huckleberry Finn",    "My favorite book.").unwrap();
448/// book_reviews.insert("Grimms' Fairy Tales",               "Masterpiece.").unwrap();
449/// book_reviews.insert("Pride and Prejudice",               "Very enjoyable.").unwrap();
450/// book_reviews.insert("The Adventures of Sherlock Holmes", "Eye lyked it alot.").unwrap();
451///
452/// // check for a specific one.
453/// if !book_reviews.contains_key("Les Misérables") {
454///     println!("We've got {} reviews, but Les Misérables ain't one.",
455///              book_reviews.len());
456/// }
457///
458/// // oops, this review has a lot of spelling mistakes, let's delete it.
459/// book_reviews.remove("The Adventures of Sherlock Holmes");
460///
461/// // look up the values associated with some keys.
462/// let to_find = ["Pride and Prejudice", "Alice's Adventure in Wonderland"];
463/// for book in &to_find {
464///     match book_reviews.get(book) {
465///         Some(review) => println!("{}: {}", book, review),
466///         None => println!("{} is unreviewed.", book)
467///     }
468/// }
469///
470/// // iterate over everything.
471/// for (book, review) in &book_reviews {
472///     println!("{}: \"{}\"", book, review);
473/// }
474/// ```
475pub struct IndexMap<K, V, S, const N: usize> {
476    core: CoreMap<K, V, N>,
477    build_hasher: S,
478}
479
480impl<K, V, S, const N: usize> IndexMap<K, V, BuildHasherDefault<S>, N> {
481    /// Creates an empty `IndexMap`.
482    pub const fn new() -> Self {
483        // Const assert
484        crate::sealed::greater_than_1::<N>();
485        crate::sealed::power_of_two::<N>();
486
487        IndexMap {
488            build_hasher: BuildHasherDefault::new(),
489            core: CoreMap::new(),
490        }
491    }
492}
493
494impl<K, V, S, const N: usize> IndexMap<K, V, S, N>
495where
496    K: Eq + Hash,
497    S: BuildHasher,
498{
499    /* Public API */
500    /// Returns the number of elements the map can hold
501    pub fn capacity(&self) -> usize {
502        N
503    }
504
505    /// Return an iterator over the keys of the map, in their order
506    ///
507    /// ```
508    /// use heapless::FnvIndexMap;
509    ///
510    /// let mut map = FnvIndexMap::<_, _, 16>::new();
511    /// map.insert("a", 1).unwrap();
512    /// map.insert("b", 2).unwrap();
513    /// map.insert("c", 3).unwrap();
514    ///
515    /// for key in map.keys() {
516    ///     println!("{}", key);
517    /// }
518    /// ```
519    pub fn keys(&self) -> impl Iterator<Item = &K> {
520        self.core.entries.iter().map(|bucket| &bucket.key)
521    }
522
523    /// Return an iterator over the values of the map, in their order
524    ///
525    /// ```
526    /// use heapless::FnvIndexMap;
527    ///
528    /// let mut map = FnvIndexMap::<_, _, 16>::new();
529    /// map.insert("a", 1).unwrap();
530    /// map.insert("b", 2).unwrap();
531    /// map.insert("c", 3).unwrap();
532    ///
533    /// for val in map.values() {
534    ///     println!("{}", val);
535    /// }
536    /// ```
537    pub fn values(&self) -> impl Iterator<Item = &V> {
538        self.core.entries.iter().map(|bucket| &bucket.value)
539    }
540
541    /// Return an iterator over mutable references to the the values of the map, in their order
542    ///
543    /// ```
544    /// use heapless::FnvIndexMap;
545    ///
546    /// let mut map = FnvIndexMap::<_, _, 16>::new();
547    /// map.insert("a", 1).unwrap();
548    /// map.insert("b", 2).unwrap();
549    /// map.insert("c", 3).unwrap();
550    ///
551    /// for val in map.values_mut() {
552    ///     *val += 10;
553    /// }
554    ///
555    /// for val in map.values() {
556    ///     println!("{}", val);
557    /// }
558    /// ```
559    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> {
560        self.core.entries.iter_mut().map(|bucket| &mut bucket.value)
561    }
562
563    /// Return an iterator over the key-value pairs of the map, in their order
564    ///
565    /// ```
566    /// use heapless::FnvIndexMap;
567    ///
568    /// let mut map = FnvIndexMap::<_, _, 16>::new();
569    /// map.insert("a", 1).unwrap();
570    /// map.insert("b", 2).unwrap();
571    /// map.insert("c", 3).unwrap();
572    ///
573    /// for (key, val) in map.iter() {
574    ///     println!("key: {} val: {}", key, val);
575    /// }
576    /// ```
577    pub fn iter(&self) -> Iter<'_, K, V> {
578        Iter {
579            iter: self.core.entries.iter(),
580        }
581    }
582
583    /// Return an iterator over the key-value pairs of the map, in their order
584    ///
585    /// ```
586    /// use heapless::FnvIndexMap;
587    ///
588    /// let mut map = FnvIndexMap::<_, _, 16>::new();
589    /// map.insert("a", 1).unwrap();
590    /// map.insert("b", 2).unwrap();
591    /// map.insert("c", 3).unwrap();
592    ///
593    /// for (_, val) in map.iter_mut() {
594    ///     *val = 2;
595    /// }
596    ///
597    /// for (key, val) in &map {
598    ///     println!("key: {} val: {}", key, val);
599    /// }
600    /// ```
601    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
602        IterMut {
603            iter: self.core.entries.iter_mut(),
604        }
605    }
606
607    /// Get the first key-value pair
608    ///
609    /// Computes in **O(1)** time
610    pub fn first(&self) -> Option<(&K, &V)> {
611        self.core
612            .entries
613            .first()
614            .map(|bucket| (&bucket.key, &bucket.value))
615    }
616
617    /// Get the first key-value pair, with mutable access to the value
618    ///
619    /// Computes in **O(1)** time
620    pub fn first_mut(&mut self) -> Option<(&K, &mut V)> {
621        self.core
622            .entries
623            .first_mut()
624            .map(|bucket| (&bucket.key, &mut bucket.value))
625    }
626
627    /// Get the last key-value pair
628    ///
629    /// Computes in **O(1)** time
630    pub fn last(&self) -> Option<(&K, &V)> {
631        self.core
632            .entries
633            .last()
634            .map(|bucket| (&bucket.key, &bucket.value))
635    }
636
637    /// Get the last key-value pair, with mutable access to the value
638    ///
639    /// Computes in **O(1)** time
640    pub fn last_mut(&mut self) -> Option<(&K, &mut V)> {
641        self.core
642            .entries
643            .last_mut()
644            .map(|bucket| (&bucket.key, &mut bucket.value))
645    }
646
647    /// Returns an entry for the corresponding key
648    /// ```
649    /// use heapless::FnvIndexMap;
650    /// use heapless::Entry;
651    /// let mut map = FnvIndexMap::<_, _, 16>::new();
652    /// if let Entry::Vacant(v) = map.entry("a") {
653    ///     v.insert(1).unwrap();
654    /// }
655    /// if let Entry::Occupied(mut o) = map.entry("a") {
656    ///     println!("found {}", *o.get()); // Prints 1
657    ///     o.insert(2);
658    /// }
659    /// // Prints 2
660    /// println!("val: {}", *map.get("a").unwrap());
661    /// ```
662    pub fn entry(&mut self, key: K) -> Entry<'_, K, V, N> {
663        let hash_val = hash_with(&key, &self.build_hasher);
664        if let Some((probe, pos)) = self.core.find(hash_val, &key) {
665            Entry::Occupied(OccupiedEntry {
666                key,
667                probe,
668                pos,
669                core: &mut self.core,
670            })
671        } else {
672            Entry::Vacant(VacantEntry {
673                key,
674                hash_val,
675                core: &mut self.core,
676            })
677        }
678    }
679
680    /// Return the number of key-value pairs in the map.
681    ///
682    /// Computes in **O(1)** time.
683    ///
684    /// ```
685    /// use heapless::FnvIndexMap;
686    ///
687    /// let mut a = FnvIndexMap::<_, _, 16>::new();
688    /// assert_eq!(a.len(), 0);
689    /// a.insert(1, "a").unwrap();
690    /// assert_eq!(a.len(), 1);
691    /// ```
692    pub fn len(&self) -> usize {
693        self.core.entries.len()
694    }
695
696    /// Returns true if the map contains no elements.
697    ///
698    /// Computes in **O(1)** time.
699    ///
700    /// ```
701    /// use heapless::FnvIndexMap;
702    ///
703    /// let mut a = FnvIndexMap::<_, _, 16>::new();
704    /// assert!(a.is_empty());
705    /// a.insert(1, "a");
706    /// assert!(!a.is_empty());
707    /// ```
708    pub fn is_empty(&self) -> bool {
709        self.len() == 0
710    }
711
712    /// Remove all key-value pairs in the map, while preserving its capacity.
713    ///
714    /// Computes in **O(n)** time.
715    ///
716    /// ```
717    /// use heapless::FnvIndexMap;
718    ///
719    /// let mut a = FnvIndexMap::<_, _, 16>::new();
720    /// a.insert(1, "a");
721    /// a.clear();
722    /// assert!(a.is_empty());
723    /// ```
724    pub fn clear(&mut self) {
725        self.core.entries.clear();
726        for pos in self.core.indices.iter_mut() {
727            *pos = None;
728        }
729    }
730
731    /// Returns a reference to the value corresponding to the key.
732    ///
733    /// The key may be any borrowed form of the map's key type, but `Hash` and `Eq` on the borrowed
734    /// form *must* match those for the key type.
735    ///
736    /// Computes in **O(1)** time (average).
737    ///
738    /// ```
739    /// use heapless::FnvIndexMap;
740    ///
741    /// let mut map = FnvIndexMap::<_, _, 16>::new();
742    /// map.insert(1, "a").unwrap();
743    /// assert_eq!(map.get(&1), Some(&"a"));
744    /// assert_eq!(map.get(&2), None);
745    /// ```
746    pub fn get<Q>(&self, key: &Q) -> Option<&V>
747    where
748        K: Borrow<Q>,
749        Q: ?Sized + Hash + Eq,
750    {
751        self.find(key)
752            .map(|(_, found)| unsafe { &self.core.entries.get_unchecked(found).value })
753    }
754
755    /// Returns true if the map contains a value for the specified key.
756    ///
757    /// The key may be any borrowed form of the map's key type, but `Hash` and `Eq` on the borrowed
758    /// form *must* match those for the key type.
759    ///
760    /// Computes in **O(1)** time (average).
761    ///
762    /// # Examples
763    ///
764    /// ```
765    /// use heapless::FnvIndexMap;
766    ///
767    /// let mut map = FnvIndexMap::<_, _, 8>::new();
768    /// map.insert(1, "a").unwrap();
769    /// assert_eq!(map.contains_key(&1), true);
770    /// assert_eq!(map.contains_key(&2), false);
771    /// ```
772    pub fn contains_key<Q>(&self, key: &Q) -> bool
773    where
774        K: Borrow<Q>,
775        Q: ?Sized + Eq + Hash,
776    {
777        self.find(key).is_some()
778    }
779
780    /// Returns a mutable reference to the value corresponding to the key.
781    ///
782    /// The key may be any borrowed form of the map's key type, but `Hash` and `Eq` on the borrowed
783    /// form *must* match those for the key type.
784    ///
785    /// Computes in **O(1)** time (average).
786    ///
787    /// # Examples
788    ///
789    /// ```
790    /// use heapless::FnvIndexMap;
791    ///
792    /// let mut map = FnvIndexMap::<_, _, 8>::new();
793    /// map.insert(1, "a").unwrap();
794    /// if let Some(x) = map.get_mut(&1) {
795    ///     *x = "b";
796    /// }
797    /// assert_eq!(map[&1], "b");
798    /// ```
799    pub fn get_mut<'v, Q>(&'v mut self, key: &Q) -> Option<&'v mut V>
800    where
801        K: Borrow<Q>,
802        Q: ?Sized + Hash + Eq,
803    {
804        if let Some((_, found)) = self.find(key) {
805            Some(unsafe { &mut self.core.entries.get_unchecked_mut(found).value })
806        } else {
807            None
808        }
809    }
810
811    /// Inserts a key-value pair into the map.
812    ///
813    /// If an equivalent key already exists in the map: the key remains and retains in its place in
814    /// the order, its corresponding value is updated with `value` and the older value is returned
815    /// inside `Some(_)`.
816    ///
817    /// If no equivalent key existed in the map: the new key-value pair is inserted, last in order,
818    /// and `None` is returned.
819    ///
820    /// Computes in **O(1)** time (average).
821    ///
822    /// See also entry if you you want to insert or modify or if you need to get the index of the
823    /// corresponding key-value pair.
824    ///
825    /// # Examples
826    ///
827    /// ```
828    /// use heapless::FnvIndexMap;
829    ///
830    /// let mut map = FnvIndexMap::<_, _, 8>::new();
831    /// assert_eq!(map.insert(37, "a"), Ok(None));
832    /// assert_eq!(map.is_empty(), false);
833    ///
834    /// map.insert(37, "b");
835    /// assert_eq!(map.insert(37, "c"), Ok(Some("b")));
836    /// assert_eq!(map[&37], "c");
837    /// ```
838    pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>, (K, V)> {
839        let hash = hash_with(&key, &self.build_hasher);
840        match self.core.insert(hash, key, value) {
841            Insert::Success(inserted) => Ok(inserted.old_value),
842            Insert::Full((k, v)) => Err((k, v)),
843        }
844    }
845
846    /// Same as [`swap_remove`](struct.IndexMap.html#method.swap_remove)
847    ///
848    /// Computes in **O(1)** time (average).
849    ///
850    /// # Examples
851    ///
852    /// ```
853    /// use heapless::FnvIndexMap;
854    ///
855    /// let mut map = FnvIndexMap::<_, _, 8>::new();
856    /// map.insert(1, "a").unwrap();
857    /// assert_eq!(map.remove(&1), Some("a"));
858    /// assert_eq!(map.remove(&1), None);
859    /// ```
860    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
861    where
862        K: Borrow<Q>,
863        Q: ?Sized + Hash + Eq,
864    {
865        self.swap_remove(key)
866    }
867
868    /// Remove the key-value pair equivalent to `key` and return its value.
869    ///
870    /// Like `Vec::swap_remove`, the pair is removed by swapping it with the last element of the map
871    /// and popping it off. **This perturbs the postion of what used to be the last element!**
872    ///
873    /// Return `None` if `key` is not in map.
874    ///
875    /// Computes in **O(1)** time (average).
876    pub fn swap_remove<Q>(&mut self, key: &Q) -> Option<V>
877    where
878        K: Borrow<Q>,
879        Q: ?Sized + Hash + Eq,
880    {
881        self.find(key)
882            .map(|(probe, found)| self.core.remove_found(probe, found).1)
883    }
884
885    /* Private API */
886    /// Return probe (indices) and position (entries)
887    fn find<Q>(&self, key: &Q) -> Option<(usize, usize)>
888    where
889        K: Borrow<Q>,
890        Q: ?Sized + Hash + Eq,
891    {
892        if self.len() == 0 {
893            return None;
894        }
895        let h = hash_with(key, &self.build_hasher);
896        self.core.find(h, key)
897    }
898}
899
900impl<'a, K, Q, V, S, const N: usize> ops::Index<&'a Q> for IndexMap<K, V, S, N>
901where
902    K: Eq + Hash + Borrow<Q>,
903    Q: ?Sized + Eq + Hash,
904    S: BuildHasher,
905{
906    type Output = V;
907
908    fn index(&self, key: &Q) -> &V {
909        self.get(key).expect("key not found")
910    }
911}
912
913impl<'a, K, Q, V, S, const N: usize> ops::IndexMut<&'a Q> for IndexMap<K, V, S, N>
914where
915    K: Eq + Hash + Borrow<Q>,
916    Q: ?Sized + Eq + Hash,
917    S: BuildHasher,
918{
919    fn index_mut(&mut self, key: &Q) -> &mut V {
920        self.get_mut(key).expect("key not found")
921    }
922}
923
924impl<K, V, S, const N: usize> Clone for IndexMap<K, V, S, N>
925where
926    K: Eq + Hash + Clone,
927    V: Clone,
928    S: Clone,
929{
930    fn clone(&self) -> Self {
931        Self {
932            core: self.core.clone(),
933            build_hasher: self.build_hasher.clone(),
934        }
935    }
936}
937
938impl<K, V, S, const N: usize> fmt::Debug for IndexMap<K, V, S, N>
939where
940    K: Eq + Hash + fmt::Debug,
941    V: fmt::Debug,
942    S: BuildHasher,
943{
944    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
945        f.debug_map().entries(self.iter()).finish()
946    }
947}
948
949impl<K, V, S, const N: usize> Default for IndexMap<K, V, S, N>
950where
951    K: Eq + Hash,
952    S: BuildHasher + Default,
953{
954    fn default() -> Self {
955        // Const assert
956        crate::sealed::greater_than_1::<N>();
957        crate::sealed::power_of_two::<N>();
958
959        IndexMap {
960            build_hasher: <_>::default(),
961            core: CoreMap::new(),
962        }
963    }
964}
965
966impl<K, V, S, S2, const N: usize, const N2: usize> PartialEq<IndexMap<K, V, S2, N2>>
967    for IndexMap<K, V, S, N>
968where
969    K: Eq + Hash,
970    V: Eq,
971    S: BuildHasher,
972    S2: BuildHasher,
973{
974    fn eq(&self, other: &IndexMap<K, V, S2, N2>) -> bool {
975        self.len() == other.len()
976            && self
977                .iter()
978                .all(|(key, value)| other.get(key).map_or(false, |v| *value == *v))
979    }
980}
981
982impl<K, V, S, const N: usize> Eq for IndexMap<K, V, S, N>
983where
984    K: Eq + Hash,
985    V: Eq,
986    S: BuildHasher,
987{
988}
989
990impl<K, V, S, const N: usize> Extend<(K, V)> for IndexMap<K, V, S, N>
991where
992    K: Eq + Hash,
993    S: BuildHasher,
994{
995    fn extend<I>(&mut self, iterable: I)
996    where
997        I: IntoIterator<Item = (K, V)>,
998    {
999        for (k, v) in iterable {
1000            self.insert(k, v).ok().unwrap();
1001        }
1002    }
1003}
1004
1005impl<'a, K, V, S, const N: usize> Extend<(&'a K, &'a V)> for IndexMap<K, V, S, N>
1006where
1007    K: Eq + Hash + Copy,
1008    V: Copy,
1009    S: BuildHasher,
1010{
1011    fn extend<I>(&mut self, iterable: I)
1012    where
1013        I: IntoIterator<Item = (&'a K, &'a V)>,
1014    {
1015        self.extend(iterable.into_iter().map(|(&key, &value)| (key, value)))
1016    }
1017}
1018
1019impl<K, V, S, const N: usize> FromIterator<(K, V)> for IndexMap<K, V, S, N>
1020where
1021    K: Eq + Hash,
1022    S: BuildHasher + Default,
1023{
1024    fn from_iter<I>(iterable: I) -> Self
1025    where
1026        I: IntoIterator<Item = (K, V)>,
1027    {
1028        let mut map = IndexMap::default();
1029        map.extend(iterable);
1030        map
1031    }
1032}
1033
1034#[derive(Clone)]
1035pub struct IntoIter<K, V, const N: usize> {
1036    entries: Vec<Bucket<K, V>, N>,
1037}
1038
1039impl<K, V, const N: usize> Iterator for IntoIter<K, V, N> {
1040    type Item = (K, V);
1041
1042    fn next(&mut self) -> Option<Self::Item> {
1043        self.entries.pop().map(|bucket| (bucket.key, bucket.value))
1044    }
1045}
1046
1047impl<K, V, S, const N: usize> IntoIterator for IndexMap<K, V, S, N>
1048where
1049    K: Eq + Hash,
1050    S: BuildHasher,
1051{
1052    type Item = (K, V);
1053    type IntoIter = IntoIter<K, V, N>;
1054
1055    fn into_iter(self) -> Self::IntoIter {
1056        IntoIter {
1057            entries: self.core.entries,
1058        }
1059    }
1060}
1061
1062impl<'a, K, V, S, const N: usize> IntoIterator for &'a IndexMap<K, V, S, N>
1063where
1064    K: Eq + Hash,
1065    S: BuildHasher,
1066{
1067    type Item = (&'a K, &'a V);
1068    type IntoIter = Iter<'a, K, V>;
1069
1070    fn into_iter(self) -> Self::IntoIter {
1071        self.iter()
1072    }
1073}
1074
1075impl<'a, K, V, S, const N: usize> IntoIterator for &'a mut IndexMap<K, V, S, N>
1076where
1077    K: Eq + Hash,
1078    S: BuildHasher,
1079{
1080    type Item = (&'a K, &'a mut V);
1081    type IntoIter = IterMut<'a, K, V>;
1082
1083    fn into_iter(self) -> Self::IntoIter {
1084        self.iter_mut()
1085    }
1086}
1087
1088pub struct Iter<'a, K, V> {
1089    iter: slice::Iter<'a, Bucket<K, V>>,
1090}
1091
1092impl<'a, K, V> Iterator for Iter<'a, K, V> {
1093    type Item = (&'a K, &'a V);
1094
1095    fn next(&mut self) -> Option<Self::Item> {
1096        self.iter.next().map(|bucket| (&bucket.key, &bucket.value))
1097    }
1098}
1099
1100impl<'a, K, V> Clone for Iter<'a, K, V> {
1101    fn clone(&self) -> Self {
1102        Self {
1103            iter: self.iter.clone(),
1104        }
1105    }
1106}
1107
1108pub struct IterMut<'a, K, V> {
1109    iter: slice::IterMut<'a, Bucket<K, V>>,
1110}
1111
1112impl<'a, K, V> Iterator for IterMut<'a, K, V> {
1113    type Item = (&'a K, &'a mut V);
1114
1115    fn next(&mut self) -> Option<Self::Item> {
1116        self.iter
1117            .next()
1118            .map(|bucket| (&bucket.key, &mut bucket.value))
1119    }
1120}
1121
1122fn hash_with<K, S>(key: &K, build_hasher: &S) -> HashValue
1123where
1124    K: ?Sized + Hash,
1125    S: BuildHasher,
1126{
1127    let mut h = build_hasher.build_hasher();
1128    key.hash(&mut h);
1129    HashValue(h.finish() as u16)
1130}
1131
1132#[cfg(test)]
1133mod tests {
1134    use crate::{indexmap::Entry, FnvIndexMap};
1135
1136    use core::mem;
1137
1138    #[test]
1139    fn size() {
1140        const CAP: usize = 4;
1141        assert_eq!(
1142            mem::size_of::<FnvIndexMap<i16, u16, CAP>>(),
1143            CAP * mem::size_of::<u32>() + // indices
1144                CAP * (mem::size_of::<i16>() + // key
1145                     mem::size_of::<u16>() + // value
1146                     mem::size_of::<u16>() // hash
1147                ) + // buckets
1148                mem::size_of::<usize>() // entries.length
1149        )
1150    }
1151
1152    #[test]
1153    fn partial_eq() {
1154        {
1155            let mut a: FnvIndexMap<_, _, 4> = FnvIndexMap::new();
1156            a.insert("k1", "v1").unwrap();
1157
1158            let mut b: FnvIndexMap<_, _, 4> = FnvIndexMap::new();
1159            b.insert("k1", "v1").unwrap();
1160
1161            assert!(a == b);
1162
1163            b.insert("k2", "v2").unwrap();
1164
1165            assert!(a != b);
1166        }
1167
1168        {
1169            let mut a: FnvIndexMap<_, _, 4> = FnvIndexMap::new();
1170            a.insert("k1", "v1").unwrap();
1171            a.insert("k2", "v2").unwrap();
1172
1173            let mut b: FnvIndexMap<_, _, 4> = FnvIndexMap::new();
1174            b.insert("k2", "v2").unwrap();
1175            b.insert("k1", "v1").unwrap();
1176
1177            assert!(a == b);
1178        }
1179    }
1180
1181    #[test]
1182    fn into_iter() {
1183        let mut src: FnvIndexMap<_, _, 4> = FnvIndexMap::new();
1184        src.insert("k1", "v1").unwrap();
1185        src.insert("k2", "v2").unwrap();
1186        src.insert("k3", "v3").unwrap();
1187        src.insert("k4", "v4").unwrap();
1188        let clone = src.clone();
1189        for (k, v) in clone.into_iter() {
1190            assert_eq!(v, *src.get(k).unwrap());
1191        }
1192    }
1193
1194    #[test]
1195    fn insert_replaces_on_full_map() {
1196        let mut a: FnvIndexMap<_, _, 2> = FnvIndexMap::new();
1197        a.insert("k1", "v1").unwrap();
1198        a.insert("k2", "v2").unwrap();
1199        a.insert("k1", "v2").unwrap();
1200        assert_eq!(a.get("k1"), a.get("k2"));
1201    }
1202
1203    const MAP_SLOTS: usize = 4096;
1204    fn almost_filled_map() -> FnvIndexMap<usize, usize, MAP_SLOTS> {
1205        let mut almost_filled = FnvIndexMap::new();
1206        for i in 1..MAP_SLOTS {
1207            almost_filled.insert(i, i).unwrap();
1208        }
1209        almost_filled
1210    }
1211
1212    #[test]
1213    fn entry_find() {
1214        let key = 0;
1215        let value = 0;
1216        let mut src = almost_filled_map();
1217        let entry = src.entry(key);
1218        match entry {
1219            Entry::Occupied(_) => {
1220                panic!("Found entry without inserting");
1221            }
1222            Entry::Vacant(v) => {
1223                assert_eq!(&key, v.key());
1224                assert_eq!(key, v.into_key());
1225            }
1226        }
1227        src.insert(key, value).unwrap();
1228        let entry = src.entry(key);
1229        match entry {
1230            Entry::Occupied(mut o) => {
1231                assert_eq!(&key, o.key());
1232                assert_eq!(&value, o.get());
1233                assert_eq!(&value, o.get_mut());
1234                assert_eq!(&value, o.into_mut());
1235            }
1236            Entry::Vacant(_) => {
1237                panic!("Entry not found");
1238            }
1239        }
1240    }
1241
1242    #[test]
1243    fn entry_vacant_insert() {
1244        let key = 0;
1245        let value = 0;
1246        let mut src = almost_filled_map();
1247        assert_eq!(MAP_SLOTS - 1, src.len());
1248        let entry = src.entry(key);
1249        match entry {
1250            Entry::Occupied(_) => {
1251                panic!("Entry found when empty");
1252            }
1253            Entry::Vacant(v) => {
1254                v.insert(value).unwrap();
1255            }
1256        };
1257        assert_eq!(value, *src.get(&key).unwrap())
1258    }
1259
1260    #[test]
1261    fn entry_occupied_insert() {
1262        let key = 0;
1263        let value = 0;
1264        let value2 = 5;
1265        let mut src = almost_filled_map();
1266        assert_eq!(MAP_SLOTS - 1, src.len());
1267        src.insert(key, value).unwrap();
1268        let entry = src.entry(key);
1269        match entry {
1270            Entry::Occupied(o) => {
1271                assert_eq!(value, o.insert(value2));
1272            }
1273            Entry::Vacant(_) => {
1274                panic!("Entry not found");
1275            }
1276        };
1277        assert_eq!(value2, *src.get(&key).unwrap())
1278    }
1279
1280    #[test]
1281    fn entry_remove_entry() {
1282        let key = 0;
1283        let value = 0;
1284        let mut src = almost_filled_map();
1285        src.insert(key, value).unwrap();
1286        assert_eq!(MAP_SLOTS, src.len());
1287        let entry = src.entry(key);
1288        match entry {
1289            Entry::Occupied(o) => {
1290                assert_eq!((key, value), o.remove_entry());
1291            }
1292            Entry::Vacant(_) => {
1293                panic!("Entry not found")
1294            }
1295        };
1296        assert_eq!(MAP_SLOTS - 1, src.len());
1297    }
1298
1299    #[test]
1300    fn entry_remove() {
1301        let key = 0;
1302        let value = 0;
1303        let mut src = almost_filled_map();
1304        src.insert(key, value).unwrap();
1305        assert_eq!(MAP_SLOTS, src.len());
1306        let entry = src.entry(key);
1307        match entry {
1308            Entry::Occupied(o) => {
1309                assert_eq!(value, o.remove());
1310            }
1311            Entry::Vacant(_) => {
1312                panic!("Entry not found");
1313            }
1314        };
1315        assert_eq!(MAP_SLOTS - 1, src.len());
1316    }
1317
1318    #[test]
1319    fn entry_roll_through_all() {
1320        let mut src: FnvIndexMap<usize, usize, MAP_SLOTS> = FnvIndexMap::new();
1321        for i in 0..MAP_SLOTS {
1322            match src.entry(i) {
1323                Entry::Occupied(_) => {
1324                    panic!("Entry found before insert");
1325                }
1326                Entry::Vacant(v) => {
1327                    v.insert(i).unwrap();
1328                }
1329            }
1330        }
1331        let add_mod = 99;
1332        for i in 0..MAP_SLOTS {
1333            match src.entry(i) {
1334                Entry::Occupied(o) => {
1335                    assert_eq!(i, o.insert(i + add_mod));
1336                }
1337                Entry::Vacant(_) => {
1338                    panic!("Entry not found after insert");
1339                }
1340            }
1341        }
1342        for i in 0..MAP_SLOTS {
1343            match src.entry(i) {
1344                Entry::Occupied(o) => {
1345                    assert_eq!((i, i + add_mod), o.remove_entry());
1346                }
1347                Entry::Vacant(_) => {
1348                    panic!("Entry not found after insert");
1349                }
1350            }
1351        }
1352        for i in 0..MAP_SLOTS {
1353            assert!(matches!(src.entry(i), Entry::Vacant(_)));
1354        }
1355        assert!(src.is_empty());
1356    }
1357
1358    #[test]
1359    fn first_last() {
1360        let mut map = FnvIndexMap::<_, _, 4>::new();
1361
1362        assert_eq!(None, map.first());
1363        assert_eq!(None, map.last());
1364
1365        map.insert(0, 0).unwrap();
1366        map.insert(2, 2).unwrap();
1367
1368        assert_eq!(Some((&0, &0)), map.first());
1369        assert_eq!(Some((&2, &2)), map.last());
1370
1371        map.insert(1, 1).unwrap();
1372
1373        assert_eq!(Some((&1, &1)), map.last());
1374
1375        *map.first_mut().unwrap().1 += 1;
1376        *map.last_mut().unwrap().1 += 1;
1377
1378        assert_eq!(Some((&0, &1)), map.first());
1379        assert_eq!(Some((&1, &2)), map.last());
1380    }
1381}
heapless/indexmap.rs

heapless/
indexmap.rs
