Add entry API

2 files changed, 313 insertions, 45 deletions

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<V> {
-    Done,
-    Swapped { prev_value: V },
-    RobinHood { probe: usize, old_pos: Pos },
+struct Insert<V> {
+    index: usize,
+    old_value: Option<V>,
 }
 
 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<V> {
+    fn insert(&mut self, hash: HashValue, key: K, value: V) -> Insert<V> {
         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<K, V, N>,
+}
+
+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<K, V, N>,
+}
+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<K, V> { .. }
+
+    /// 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<V> {
-        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<K, V, S, N>
@@ -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<usize, usize, MAP_SLOTS> {
+        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<usize, usize, MAP_SLOTS> = 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"))]