WIP/RFC: shift_remove and friends

CHANGELOG.md

@@ -8,7 +8,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 ## [Unreleased]
 
 ### Added
-
+- Added `shift_remove`, `shift_remove_entry`, `shift_remove_full`, `shift_remove_index` to `IndexMap`.
 - Added `bytes::Buf` and `bytes::BufMut` implementations for `Vec`.
 - Added `format` macro.
 - Added `String::from_utf16`.

src/index_map.rs

@@ -302,43 +302,63 @@ where
     where
         F: FnMut(&mut K, &mut V) -> bool,
     {
-        const INIT: Option<Pos> = None;
-
         self.entries
             .retain_mut(|entry| keep(&mut entry.key, &mut entry.value));
 
         if self.entries.len() < self.indices.len() {
-            for index in self.indices.iter_mut() {
-                *index = INIT;
-            }
+            self.after_removal(0);
+        }
+    }
+
+    fn shift_remove_index(&mut self, index: usize) -> Option<(K, V)> {
+        if index >= self.entries.len() {
+            return None;
+        }
+
+        let bucket = self.entries.remove(index);
+
+        self.after_removal(index);
+
+        Some((bucket.key, bucket.value))
+    }
+
+    fn shift_remove_found(&mut self, _probe: usize, found: usize) -> (K, V) {
+        let entry = self.entries.remove(found);
+
+        self.after_removal(found);
+
+        (entry.key, entry.value)
+    }
 
-            for (index, entry) in self.entries.iter().enumerate() {
-                let mut probe = entry.hash.desired_pos(Self::mask());
-                let mut dist = 0;
-
-                probe_loop!(probe < self.indices.len(), {
-                    let pos = &mut self.indices[probe];
-
-                    if let Some(pos) = *pos {
-                        let entry_hash = pos.hash();
-
-                        // robin hood: steal the spot if it's better for us
-                        let their_dist = entry_hash.probe_distance(Self::mask(), probe);
-                        if their_dist < dist {
-                            Self::insert_phase_2(
-                                &mut self.indices,
-                                probe,
-                                Pos::new(index, entry.hash),
-                            );
-                            break;
-                        }
-                    } else {
-                        *pos = Some(Pos::new(index, entry.hash));
+    fn after_removal(&mut self, first_modified: usize) {
+        const INIT: Option<Pos> = None;
+
+        for index in self.indices.iter_mut() {
+            *index = INIT;
+        }
+
+        for (index, entry) in self.entries.iter().skip(first_modified).enumerate() {
+            let mut probe = entry.hash.desired_pos(Self::mask());
+            let mut dist = 0;
+
+            probe_loop!(probe < self.indices.len(), {
+                let pos = &mut self.indices[probe];
+
+                if let Some(pos) = *pos {
+                    let entry_hash = pos.hash();
+
+                    // robin hood: steal the spot if it's better for us
+                    let their_dist = entry_hash.probe_distance(Self::mask(), probe);
+                    if their_dist < dist {
+                        Self::insert_phase_2(&mut self.indices, probe, Pos::new(index, entry.hash));
                         break;
                     }
-                    dist += 1;
-                });
-            }
+                } else {
+                    *pos = Some(Pos::new(index, entry.hash));
+                    break;
+                }
+                dist += 1;
+            });
         }
     }
 
@@ -1231,6 +1251,154 @@ where
             .map(|(probe, found)| self.core.remove_found(probe, found).1)
     }
 
+    /// Remove the key-value pair at position `index` and return them.
+    ///
+    /// Like [`Vec::remove`], the pair is removed by shifting all
+    /// remaining items. This maintains the remaining elements' relative
+    /// insertion order, but is a more expensive operation
+    ///
+    /// Return `None` if `index` is not in `0..len()`.
+    ///
+    /// Computes in *O*(n) time (average).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use heapless::index_map::FnvIndexMap;
+    ///
+    /// let mut map = FnvIndexMap::<_, _, 8>::new();
+    /// map.insert(3, "a").unwrap();
+    /// map.insert(2, "b").unwrap();
+    /// map.insert(1, "c").unwrap();
+    /// let removed = map.shift_remove_index(1);
+    /// assert_eq!(removed, Some((2, "b")));
+    /// assert_eq!(map.len(), 2);
+    ///
+    /// let mut iter = map.iter();
+    /// assert_eq!(iter.next(), Some((&3, &"a")));
+    /// assert_eq!(iter.next(), Some((&1, &"c")));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn shift_remove_index(&mut self, index: usize) -> Option<(K, V)> {
+        self.core.shift_remove_index(index)
+    }
+
+    /// Remove the key-value pair equivalent to `key` and return it and
+    /// the index it had.
+    ///
+    /// Like [`Vec::remove`], the pair is removed by shifting all
+    /// remaining items. This maintains the remaining elements' relative
+    /// insertion order, but is a more expensive operation.
+    ///
+    /// Return `None` if `key` is not in map.
+    ///
+    /// Computes in *O*(n) time (average).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use heapless::index_map::FnvIndexMap;
+    ///
+    /// let mut map = FnvIndexMap::<_, _, 8>::new();
+    /// map.insert(3, "a").unwrap();
+    /// map.insert(2, "b").unwrap();
+    /// map.insert(1, "c").unwrap();
+    /// let removed = map.shift_remove_full(&2);
+    /// assert_eq!(removed, Some((1, 2, "b")));
+    /// assert_eq!(map.len(), 2);
+    /// assert_eq!(map.shift_remove_full(&2), None);
+    ///
+    /// let mut iter = map.iter();
+    /// assert_eq!(iter.next(), Some((&3, &"a")));
+    /// assert_eq!(iter.next(), Some((&1, &"c")));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn shift_remove_full<Q>(&mut self, key: &Q) -> Option<(usize, K, V)>
+    where
+        K: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        self.find(key).map(|(probe, found)| {
+            let (k, v) = self.core.shift_remove_found(probe, found);
+            (found, k, v)
+        })
+    }
+
+    /// Remove and return the key-value pair equivalent to `key`.
+    ///
+    /// Like [`Vec::remove`], the pair is removed by shifting all
+    /// remaining items. This maintains the remaining elements' relative
+    /// insertion order, but is a more expensive operation.
+    ///
+    /// Return `None` if `key` is not in map.
+    ///
+    /// Computes in *O*(n) time (average).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use heapless::index_map::FnvIndexMap;
+    ///
+    /// let mut map = FnvIndexMap::<_, _, 8>::new();
+    /// map.insert(3, "a").unwrap();
+    /// map.insert(2, "b").unwrap();
+    /// map.insert(1, "c").unwrap();
+    /// let removed = map.shift_remove_entry(&2);
+    /// assert_eq!(removed, Some((2, "b")));
+    /// assert_eq!(map.len(), 2);
+    /// assert_eq!(map.shift_remove_entry(&2), None);
+    ///
+    /// let mut iter = map.iter();
+    /// assert_eq!(iter.next(), Some((&3, &"a")));
+    /// assert_eq!(iter.next(), Some((&1, &"c")));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn shift_remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V)>
+    where
+        K: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        self.shift_remove_full(key).map(|(_idx, k, v)| (k, v))
+    }
+
+    /// Remove the key-value pair equivalent to `key` and return
+    /// its value.
+    ///
+    /// Like [`Vec::remove`], the pair is removed by shifting all of the
+    /// elements that follow it, preserving their relative order.
+    /// **This perturbs the index of all of those elements!**
+    ///
+    /// Return `None` if `key` is not in map.
+    ///
+    /// Computes in *O*(n) time (average).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use heapless::index_map::FnvIndexMap;
+    ///
+    /// let mut map = FnvIndexMap::<_, _, 8>::new();
+    /// map.insert(3, "a").unwrap();
+    /// map.insert(2, "b").unwrap();
+    /// map.insert(1, "c").unwrap();
+    /// let removed = map.shift_remove(&2);
+    /// assert_eq!(removed, Some(("b")));
+    /// assert_eq!(map.len(), 2);
+    /// assert_eq!(map.shift_remove(&2), None);
+    ///
+    /// let mut iter = map.iter();
+    /// assert_eq!(iter.next(), Some((&3, &"a")));
+    /// assert_eq!(iter.next(), Some((&1, &"c")));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn shift_remove<Q>(&mut self, key: &Q) -> Option<V>
+    where
+        K: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        self.shift_remove_full(key).map(|(_idx, _k, v)| v)
+    }
+
     /// Retains only the elements specified by the predicate.
     ///
     /// In other words, remove all pairs `(k, v)` for which `f(&k, &mut v)` returns `false`.