Add shift_remove family to IndexMap/IndexSet

CHANGELOG.md

@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - Fixed unsoundness in `IndexMap:insert`.
 - Limited max size of `IndexMap` to u16::MAX + 1.
   The implementation for sizes higher than u16::MAX were unsound anyway.
+- Added `shift_remove`, `shift_remove_entry` and `shift_remove_full` to `IndexMap` and `IndexSet`.
 
 ## [v0.9.3] 2025-04-15
 

src/index_map.rs

@@ -365,7 +365,7 @@ where
     }
 
     /// SAFETY: found must be a valid entry and probe the index of the corresponding `Pos`
-    unsafe fn remove_found(&mut self, probe: usize, found: usize) -> (K, V) {
+    unsafe fn swap_remove_found(&mut self, probe: usize, found: usize) -> (K, V) {
         // index `probe` and entry `found` is to be removed
         // use swap_remove, but then we need to update the index that points
         // to the other entry that has to move
@@ -403,6 +403,33 @@ where
         (entry.key, entry.value)
     }
 
+    /// SAFETY: found must be a valid entry and probe the index of the corresponding `Pos`
+    unsafe fn shift_remove_found(&mut self, probe: usize, found: usize) -> (K, V) {
+        let old_len = self.entries.len();
+        self.indices[probe] = Pos::none();
+        // Unlike `swap_remove_found`, every entry after `found` is shifted towards the front.
+        let entry = self.entries.remove(found);
+
+        // Validate against the old length so a full-capacity `Pos` that shares the empty
+        // sentinel encoding is still updated. The removed probe is the sole empty slot.
+        for (table_index, pos) in self.indices.iter_mut().enumerate() {
+            if table_index == probe {
+                continue;
+            }
+
+            if let Some(valid) = pos.as_valid_mut::<N>(old_len) {
+                let entry_index = valid.index();
+                if entry_index > found {
+                    valid.replace(Pos::new(entry_index - 1, valid.hash()));
+                }
+            }
+        }
+
+        self.backward_shift_after_removal(probe);
+
+        (entry.key, entry.value)
+    }
+
     fn retain_in_order<F>(&mut self, mut keep: F)
     where
         F: FnMut(&mut K, &mut V) -> bool,
@@ -713,7 +740,19 @@ where
     pub fn swap_remove_entry(self) -> (K, V) {
         // SAFETY: We know that `pos` is valid from the creation of the entry
         // and that cannot have changed since we held a mutable entry to the map
-        unsafe { self.core.remove_found(self.probe, self.pos) }
+        unsafe { self.core.swap_remove_found(self.probe, self.pos) }
+    }
+
+    /// Removes this entry from the map and yields its corresponding key and value.
+    ///
+    /// Like `Vec::remove`, the pair is removed by shifting all of the elements that follow it,
+    /// preserving their relative order.
+    ///
+    /// Computes in **O(n)** time on average
+    pub fn shift_remove_entry(self) -> (K, V) {
+        // SAFETY: We know that `pos` is valid from the creation of the entry
+        // and that cannot have changed since we held a mutable reference to the map
+        unsafe { self.core.shift_remove_found(self.probe, self.pos) }
     }
 
     /// Gets a reference to the value associated with this entry
@@ -768,6 +807,16 @@ where
     pub fn swap_remove(self) -> V {
         self.swap_remove_entry().1
     }
+
+    /// Removes this entry from the map and yields its value.
+    ///
+    /// Like `Vec::remove`, the pair is removed by shifting all of the elements that follow it,
+    /// preserving their relative order.
+    ///
+    /// Computes in **O(n)** time on average
+    pub fn shift_remove(self) -> V {
+        self.shift_remove_entry().1
+    }
 }
 
 /// A view into an empty slot in the underlying map
@@ -1409,10 +1458,97 @@ where
     {
         self.find(key).map(|(probe, found)| {
             // SAFETY: Find gives a correct bucket and the corresponding probe
-            unsafe { self.core.remove_found(probe, found) }.1
+            unsafe { self.core.swap_remove_found(probe, found) }.1
         })
     }
 
+    /// 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(1, "a").unwrap();
+    /// map.insert(2, "b").unwrap();
+    /// map.insert(3, "c").unwrap();
+    /// map.insert(4, "d").unwrap();
+    /// assert_eq!(map.shift_remove(&2), Some("b"));
+    /// assert_eq!(
+    ///     map.keys().copied().collect::<heapless::Vec<_, 8>>().as_slice(),
+    ///     &[1, 3, 4]
+    /// );
+    /// assert_eq!(map.shift_remove(&2), None);
+    /// ```
+    pub fn shift_remove<Q>(&mut self, key: &Q) -> Option<V>
+    where
+        K: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        self.find(key).map(|(probe, found)| {
+            // SAFETY: Find gives a correct bucket and the corresponding probe
+            unsafe { self.core.shift_remove_found(probe, found) }.1
+        })
+    }
+
+    /// Remove and return the key-value pair equivalent to `key`.
+    ///
+    /// 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(1, "a").unwrap();
+    /// assert_eq!(map.shift_remove_entry(&1), Some((1, "a")));
+    /// assert_eq!(map.shift_remove_entry(&1), None);
+    /// ```
+    pub fn shift_remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V)>
+    where
+        K: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        self.find(key).map(|(probe, found)| {
+            // SAFETY: Find gives a correct bucket and the corresponding probe
+            unsafe { self.core.shift_remove_found(probe, found) }
+        })
+    }
+
+    /// Remove and return the key-value pair at `index`.
+    ///
+    /// Valid indices are `0 <= index < self.len()`.
+    ///
+    /// The pair is removed by shifting all of the elements that follow it, preserving the
+    /// insertion order of the remaining elements.
+    ///
+    /// Computes in *O*(n) time (average).
+    pub fn shift_remove_index(&mut self, index: usize) -> Option<(K, V)> {
+        let probe = self.core.indices.iter().position(|pos| {
+            pos.as_valid::<N>(self.core.entries.len())
+                .is_some_and(|pos| pos.index() == index)
+        })?;
+
+        // SAFETY: The position search found a valid entry and its corresponding probe
+        Some(unsafe { self.core.shift_remove_found(probe, index) })
+    }
+
     /// Retains only the elements specified by the predicate.
     ///
     /// In other words, remove all pairs `(k, v)` for which `f(&k, &mut v)` returns `false`.
@@ -2109,6 +2245,74 @@ mod tests {
         assert_eq!(MAP_SLOTS - 1, src.len());
     }
 
+    #[test]
+    fn shift_remove_preserves_order() {
+        let mut map = FnvIndexMap::<_, _, 8>::new();
+        for key in 0..5 {
+            map.insert(key, key * 10).unwrap();
+        }
+
+        assert_eq!(map.shift_remove(&2), Some(20));
+        assert_eq!(map.shift_remove(&9), None);
+        assert_eq!(
+            map.keys().copied().collect::<std::vec::Vec<_>>(),
+            [0, 1, 3, 4]
+        );
+        assert_eq!(map.get(&0), Some(&0));
+        assert_eq!(map.get(&1), Some(&10));
+        assert_eq!(map.get(&3), Some(&30));
+        assert_eq!(map.get(&4), Some(&40));
+    }
+
+    #[test]
+    fn shift_remove_entry_preserves_order() {
+        let mut map = FnvIndexMap::<_, _, 8>::new();
+        for key in 0..5 {
+            map.insert(key, key * 10).unwrap();
+        }
+
+        assert_eq!(map.shift_remove_entry(&2), Some((2, 20)));
+        assert_eq!(map.shift_remove_entry(&9), None);
+        assert_eq!(
+            map.keys().copied().collect::<std::vec::Vec<_>>(),
+            [0, 1, 3, 4]
+        );
+    }
+
+    #[test]
+    fn occupied_entry_shift_remove_methods_preserve_order() {
+        let mut map = FnvIndexMap::<_, _, 8>::new();
+        for key in 0..5 {
+            map.insert(key, key * 10).unwrap();
+        }
+
+        let Entry::Occupied(entry) = map.entry(2) else {
+            panic!("Entry not found");
+        };
+        assert_eq!(entry.shift_remove_entry(), (2, 20));
+
+        let Entry::Occupied(entry) = map.entry(1) else {
+            panic!("Entry not found");
+        };
+        assert_eq!(entry.shift_remove(), 10);
+        assert_eq!(map.keys().copied().collect::<std::vec::Vec<_>>(), [0, 3, 4]);
+    }
+
+    #[test]
+    fn shift_remove_index_preserves_order() {
+        let mut map = FnvIndexMap::<_, _, 8>::new();
+        for key in 0..5 {
+            map.insert(key, key * 10).unwrap();
+        }
+
+        assert_eq!(map.shift_remove_index(2), Some((2, 20)));
+        assert_eq!(map.shift_remove_index(4), None);
+        assert_eq!(
+            map.keys().copied().collect::<std::vec::Vec<_>>(),
+            [0, 1, 3, 4]
+        );
+    }
+
     #[test]
     fn retain() {
         let mut none = almost_filled_map();
@@ -2432,11 +2636,11 @@ mod tests {
         assert!(map.get(&ControlledHash(0xFFFF, 0)).is_none());
     }
 
-    /// Test that `remove_found` doesn't fail
+    /// Test that `swap_remove_found` doesn't fail
     ///
     /// Ensures that `probe_loop!` works correctly with the `continue` in the body
     #[test]
-    fn remove_found_loop() {
+    fn swap_remove_found_loop() {
         let mut map: FnvIndexMap<u16, u16, 4> = IndexMap::default();
         map.insert(0, 0).unwrap();
         map.insert(4, 4).unwrap();

src/index_set.rs

@@ -525,6 +525,23 @@ where
         self.map.swap_remove(value).is_some()
     }
 
+    /// Removes a value from the set. Returns `true` if the value was present in the set.
+    ///
+    /// The value is removed by shifting all of the elements that follow it, preserving the
+    /// insertion order of the remaining elements.
+    ///
+    /// The value may be any borrowed form of the set's value type, but `Hash` and `Eq` on the
+    /// borrowed form must match those for the value type.
+    ///
+    /// Computes in *O*(n) time (average).
+    pub fn shift_remove<Q>(&mut self, value: &Q) -> bool
+    where
+        T: Borrow<Q>,
+        Q: ?Sized + Eq + Hash,
+    {
+        self.map.shift_remove(value).is_some()
+    }
+
     /// Retains only the elements specified by the predicate.
     ///
     /// In other words, remove all elements `e` for which `f(&e)` returns `false`.
@@ -720,11 +737,26 @@ where
 mod tests {
     use static_assertions::assert_not_impl_any;
 
-    use super::{BuildHasherDefault, IndexSet};
+    use super::{BuildHasherDefault, FnvIndexSet, IndexSet};
 
     // Ensure a `IndexSet` containing `!Send` values stays `!Send` itself.
     assert_not_impl_any!(IndexSet<*const (), BuildHasherDefault<()>, 4>: Send);
 
+    #[test]
+    fn shift_remove_preserves_order() {
+        let mut set = FnvIndexSet::<_, 8>::new();
+        for value in 0..5 {
+            set.insert(value).unwrap();
+        }
+
+        assert!(set.shift_remove(&2));
+        assert!(!set.shift_remove(&9));
+        assert_eq!(
+            set.iter().copied().collect::<std::vec::Vec<_>>(),
+            [0, 1, 3, 4]
+        );
+    }
+
     #[test]
     #[cfg(feature = "zeroize")]
     fn test_index_set_zeroize() {