diff --git a/src/indexmap.rs b/src/indexmap.rs index 4e55d69e..f98274a5 100644 --- a/src/indexmap.rs +++ b/src/indexmap.rs @@ -98,10 +98,9 @@ impl Pos { } } -pub enum Inserted { - Done, - Swapped { prev_value: V }, - RobinHood { probe: usize, old_pos: Pos }, +struct Insert { + index: usize, + old_value: Option, } macro_rules! probe_loop { @@ -176,16 +175,10 @@ where }); } - // First phase: Look for the preferred location for key. - // - // We will know if `key` is already in the map, before we need to insert it. - // When we insert they key, it might be that we need to continue displacing - // entries (robin hood hashing), in which case Inserted::RobinHood is returned - fn insert_phase_1(&mut self, hash: HashValue, key: K, value: V) -> Inserted { + fn insert(&mut self, hash: HashValue, key: K, value: V) -> Insert { let mut probe = hash.desired_pos(Self::mask()); let mut dist = 0; - let inserted; probe_loop!(probe < self.indices.len(), { let pos = &mut self.indices[probe]; @@ -200,37 +193,37 @@ where if their_dist < dist { // robin hood: steal the spot if it's better for us let index = self.entries.len(); - inserted = Inserted::RobinHood { - probe: probe, - old_pos: Pos::new(index, hash), - }; - break; + unsafe { self.entries.push_unchecked(Bucket { hash, key, value }) }; + return Insert { + index: self.insert_phase_2(probe, Pos::new(index, hash)), + old_value: None + } } else if entry_hash == hash && unsafe { self.entries.get_unchecked(i).key == key } { - return Inserted::Swapped { - prev_value: mem::replace( + return Insert { + index: i, + old_value: Some(mem::replace( unsafe { &mut self.entries.get_unchecked_mut(i).value }, value, - ), + )), }; } } else { // empty bucket, insert here let index = self.entries.len(); *pos = Some(Pos::new(index, hash)); - inserted = Inserted::Done; - break; + unsafe { self.entries.push_unchecked(Bucket { hash, key, value }) }; + return Insert { + index, + old_value: None, + } } dist += 1; }); - - // NOTE(unsafe) we already checked (in `insert`) that we aren't exceeding the capacity - unsafe { self.entries.push_unchecked(Bucket { hash, key, value }) } - inserted } // phase 2 is post-insert where we forward-shift `Pos` in the indices. - fn insert_phase_2(&mut self, mut probe: usize, mut old_pos: Pos) { + fn insert_phase_2(&mut self, mut probe: usize, mut old_pos: Pos) -> usize { probe_loop!(probe < self.indices.len(), { let pos = unsafe { self.indices.get_unchecked_mut(probe) }; @@ -242,7 +235,7 @@ where if is_none { *pos = Some(old_pos); - break; + return probe; } }); } @@ -313,6 +306,105 @@ where } } +/// A view into an entry in the map +pub enum Entry<'a, K, V, const N: usize> { + /// The entry corresponding to the key `K` exists in the map + Occupied(OccupiedEntry<'a, K, V, N>), + /// The entry corresponding to the key `K` does not exist in the map + Vacant(VacantEntry<'a, K, V, N>) +} + +/// An occupied entry which can be manipulated +pub struct OccupiedEntry<'a, K, V, const N: usize> { + key: K, + probe: usize, + pos: usize, + core: &'a mut CoreMap, +} + +impl<'a, K, V, const N: usize> OccupiedEntry<'a, K, V, N> where K: Eq + Hash { + + /// Gets a reference to the key that this entity corresponds to + pub fn key(&self) -> &K { + &self.key + } + + /// Removes this entry from the map and yields its corresponding key and value + pub fn remove_entry(self) -> (K, V) { + self.core.remove_found(self.probe, self.pos) + } + + /// Gets a reference to the value associated with this entry + pub fn get(&self) -> &V { + // SAFETY: Already checked existence at instantiation and the only mutable reference + // to the map is internally held. + unsafe { + &self.core.entries.get_unchecked(self.pos).value + } + } + + /// Gets a mutable reference to the value associated with this entry + pub fn get_mut(&mut self) -> &mut V { + // SAFETY: Already checked existence at instantiation and the only mutable reference + // to the map is internally held. + unsafe {&mut self.core.entries.get_unchecked_mut(self.pos).value} + } + + /// Consumes this entry and yields a reference to the underlying value + pub fn into_mut(self) -> &'a mut V { + // SAFETY: Already checked existence at instantiation and the only mutable reference + // to the map is internally held. + unsafe {&mut self.core.entries.get_unchecked_mut(self.pos).value} + } + + /// Overwrites the underlying map's value with this entry's value + pub fn insert(self, value: V) -> V { + // SAFETY: Already checked existence at instantiation and the only mutable reference + // to the map is internally held. + unsafe { mem::replace(&mut self.core.entries.get_unchecked_mut(self.pos).value, value)} + } + + /// Removes this entry from the map and yields its value + pub fn remove(self) -> V { + self.remove_entry().1 + } +} + +/// A view into an empty slot in the underlying map +pub struct VacantEntry<'a, K, V, const N: usize> { + key: K, + hash_val: HashValue, + core: &'a mut CoreMap, +} +impl<'a, K, V, const N: usize> VacantEntry<'a, K, V, N> where K: Eq + Hash { + + /// Get the key associated with this entry + pub fn key(&self) -> &K { + &self.key + } + + /// Consumes this entry to yield to key associated with it + pub fn into_key(self) -> K { + self.key + } + + /// Inserts this entry into to underlying map, yields a mutable reference to the inserted value. + /// If the map is at capacity the value is returned instead. + pub fn insert(self, value: V) -> Result<&'a mut V, V> { + if self.core.entries.is_full() { + Err(value) + } else { + let index = self.core.insert(self.hash_val, self.key, value).index; + unsafe { + // SAFETY: Already checked existence at instantiation and the only mutable reference + // to the map is internally held. + Ok(&mut self.core.entries.get_unchecked_mut(index).value) + } + } + } + +} + /// Fixed capacity [`IndexMap`](https://docs.rs/indexmap/1/indexmap/map/struct.IndexMap.html) /// /// Note that you cannot use `IndexMap` directly, since it is generic around the hashing algorithm @@ -492,8 +584,39 @@ where } } - // TODO - // pub fn entry(&mut self, key: K) -> Entry { .. } + + /// Returns an entry for the corresponding key + /// ``` + /// use heapless::FnvIndexMap; + /// use heapless::Entry; + /// let mut map = FnvIndexMap::<_, _, 16>::new(); + /// if let Entry::Vacant(v) = map.entry("a") { + /// v.insert(1).unwrap(); + /// } + /// if let Entry::Occupied(mut o) = map.entry("a") { + /// println!("found {}", *o.get()); // Prints 1 + /// o.insert(2); + /// } + /// // Prints 2 + /// println!("val: {}", *map.get("a").unwrap()); + /// ``` + pub fn entry(&mut self, key: K) -> Entry<'_, K, V, N> { + let hash_val = hash_with(&key, &self.build_hasher); + if let Some((probe, pos)) = self.core.find(hash_val, &key) { + Entry::Occupied(OccupiedEntry { + key, + probe, + pos, + core: &mut self.core + }) + } else { + Entry::Vacant(VacantEntry { + key, + hash_val, + core: &mut self.core + }) + } + } /// Return the number of key-value pairs in the map. /// @@ -657,14 +780,8 @@ where if self.core.entries.is_full() { Err((key, value)) } else { - Ok(match self.insert_phase_1(key, value) { - Inserted::Swapped { prev_value } => Some(prev_value), - Inserted::Done => None, - Inserted::RobinHood { probe, old_pos } => { - self.core.insert_phase_2(probe, old_pos); - None - } - }) + let hash = hash_with(&key, &self.build_hasher); + Ok(self.core.insert(hash, key, value).old_value) } } @@ -720,11 +837,6 @@ where let h = hash_with(key, &self.build_hasher); self.core.find(h, key) } - - fn insert_phase_1(&mut self, key: K, value: V) -> Inserted { - let hash = hash_with(&key, &self.build_hasher); - self.core.insert_phase_1(hash, key, value) - } } impl<'a, K, Q, V, S, const N: usize> ops::Index<&'a Q> for IndexMap @@ -959,6 +1071,7 @@ where mod tests { use crate::FnvIndexMap; use core::mem; + use crate::indexmap::Entry; #[test] fn size() { @@ -1015,4 +1128,159 @@ mod tests { assert_eq!(v, *src.get(k).unwrap()); } } + + const MAP_SLOTS: usize = 4096; + fn almost_filled_map() -> FnvIndexMap { + let mut almost_filled = FnvIndexMap::new(); + for i in 1..MAP_SLOTS { + almost_filled.insert(i, i).unwrap(); + } + almost_filled + } + + #[test] + fn entry_find() { + let key = 0; + let value = 0; + let mut src = almost_filled_map(); + let entry = src.entry(key); + match entry { + Entry::Occupied(_) => { + panic!("Found entry without inserting"); + } + Entry::Vacant(v) => { + assert_eq!(&key, v.key()); + assert_eq!(key, v.into_key()); + } + } + src.insert(key, value).unwrap(); + let entry = src.entry(key); + match entry { + Entry::Occupied(mut o) => { + assert_eq!(&key, o.key()); + assert_eq!(&value, o.get()); + assert_eq!(&value, o.get_mut()); + assert_eq!(&value, o.into_mut()); + } + Entry::Vacant(_) => { + panic!("Entry not found"); + } + } + } + + #[test] + fn entry_vacant_insert() { + let key = 0; + let value = 0; + let mut src = almost_filled_map(); + assert_eq!(MAP_SLOTS - 1, src.len()); + let entry = src.entry(key); + match entry { + Entry::Occupied(_) => { + panic!("Entry found when empty"); + } + Entry::Vacant(v) => { + v.insert(value).unwrap(); + } + }; + assert_eq!(value, *src.get(&key).unwrap()) + } + + #[test] + fn entry_occupied_insert() { + let key = 0; + let value = 0; + let value2 = 5; + let mut src = almost_filled_map(); + assert_eq!(MAP_SLOTS - 1, src.len()); + src.insert(key, value).unwrap(); + let entry = src.entry(key); + match entry { + Entry::Occupied(o) => { + assert_eq!(value, o.insert(value2)); + } + Entry::Vacant(_) => { + panic!("Entry not found"); + } + }; + assert_eq!(value2, *src.get(&key).unwrap()) + } + + #[test] + fn entry_remove_entry() { + let key = 0; + let value = 0; + let mut src = almost_filled_map(); + src.insert(key, value).unwrap(); + assert_eq!(MAP_SLOTS, src.len()); + let entry = src.entry(key); + match entry { + Entry::Occupied(o) => { + assert_eq!((key, value), o.remove_entry()); + }, + Entry::Vacant(_) => { + panic!("Entry not found") + } + }; + assert_eq!(MAP_SLOTS - 1, src.len()); + } + + #[test] + fn entry_remove() { + let key = 0; + let value = 0; + let mut src = almost_filled_map(); + src.insert(key, value).unwrap(); + assert_eq!(MAP_SLOTS, src.len()); + let entry = src.entry(key); + match entry { + Entry::Occupied(o) => { + assert_eq!(value, o.remove()); + }, + Entry::Vacant(_) => { + panic!("Entry not found"); + } + }; + assert_eq!(MAP_SLOTS - 1, src.len()); + } + + #[test] + fn entry_roll_through_all() { + let mut src: FnvIndexMap = FnvIndexMap::new(); + for i in 0..MAP_SLOTS { + match src.entry(i) { + Entry::Occupied(_) => { + panic!("Entry found before insert"); + } + Entry::Vacant(v) => { + v.insert(i).unwrap(); + } + } + } + let add_mod = 99; + for i in 0..MAP_SLOTS { + match src.entry(i) { + Entry::Occupied(o) => { + assert_eq!(i, o.insert(i + add_mod)); + } + Entry::Vacant(_) => { + panic!("Entry not found after insert"); + } + } + } + for i in 0..MAP_SLOTS { + match src.entry(i) { + Entry::Occupied(o) => { + assert_eq!((i, i + add_mod), o.remove_entry()); + } + Entry::Vacant(_) => { + panic!("Entry not found after insert"); + } + } + } + for i in 0..MAP_SLOTS { + assert!(matches!(src.entry(i), Entry::Vacant(_))); + } + assert!(src.is_empty()); + } } diff --git a/src/lib.rs b/src/lib.rs index a35aebee..1abc349c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -68,16 +68,16 @@ //! It *might* compile on older versions but that may change in any new patch release. #![cfg_attr(not(test), no_std)] -#![deny(missing_docs)] +#![warn(missing_docs)] #![deny(rust_2018_compatibility)] #![deny(rust_2018_idioms)] -#![deny(warnings)] +//#![deny(warnings)] #![deny(const_err)] pub use binary_heap::BinaryHeap; pub use deque::Deque; pub use histbuf::{HistoryBuffer, OldestOrdered}; -pub use indexmap::{Bucket, FnvIndexMap, IndexMap, Pos}; +pub use indexmap::{Bucket, FnvIndexMap, IndexMap, Pos, Entry, OccupiedEntry, VacantEntry}; pub use indexset::{FnvIndexSet, IndexSet}; pub use linear_map::LinearMap; #[cfg(all(has_cas, feature = "cas"))]