Add method to get a map entry by index

src/map.rs

@@ -32,7 +32,7 @@ use alloc::vec::Vec;
 #[cfg(feature = "std")]
 use std::collections::hash_map::RandomState;
 
-use self::core::IndexMapCore;
+use self::core::{EntryAt, IndexMapCore};
 use crate::util::{third, try_simplify_range};
 use crate::{Bucket, Entries, Equivalent, HashValue, TryReserveError};
 
@@ -423,6 +423,18 @@ where
         self.core.entry(hash, key)
     }
 
+    /// Get an entry in the map by index for in-place manipulation.
+    ///
+    /// Valid indices are *0 <= index < self.len()*
+    ///
+    /// Computes in **O(1)** time.
+    pub fn entry_at(&mut self, index: usize) -> Option<EntryAt<'_, K, V>> {
+        if index >= self.len() {
+            return None;
+        }
+        Some(EntryAt::new(&mut self.core, index))
+    }
+
     /// Return `true` if an equivalent to `key` exists in the map.
     ///
     /// Computes in **O(1)** time (average).

src/map/core.rs

@@ -754,9 +754,128 @@ impl<K: fmt::Debug, V> fmt::Debug for VacantEntry<'_, K, V> {
     }
 }
 
+/// A view into an occupied entry in a `IndexMap` obtained by index.
+///
+/// This struct is constructed from the `entry_at` method on `IndexMap`.
+// SAFETY: We have a mutable reference to the map, which keeps the index
+// valid and pointing to the correct entry.
+pub struct EntryAt<'a, K, V> {
+    map: &'a mut IndexMapCore<K, V>,
+    index: usize,
+}
+
+impl<'a, K, V> EntryAt<'a, K, V> {
+    pub(crate) fn new(map: &'a mut IndexMapCore<K, V>, index: usize) -> Self {
+        Self { map, index }
+    }
+
+    /// Gets a reference to the entry's key in the map.
+    pub fn key(&self) -> &K {
+        &self.map.entries[self.index].key
+    }
+
+    /// Gets a reference to the entry's value in the map.
+    pub fn get(&self) -> &V {
+        &self.map.entries[self.index].value
+    }
+
+    /// Gets a mutable reference to the entry's value in the map.
+    ///
+    /// If you need a reference which may outlive the destruction of the
+    /// `Entry` value, see `into_mut`.
+    pub fn get_mut(&mut self) -> &mut V {
+        &mut self.map.entries[self.index].value
+    }
+
+    /// Sets the value of the entry to `value`, and returns the entry's old value.
+    pub fn insert(&mut self, value: V) -> V {
+        mem::replace(self.get_mut(), value)
+    }
+
+    /// Return the index of the key-value pair
+    #[inline]
+    pub fn index(&self) -> usize {
+        self.index
+    }
+
+    /// Converts into a mutable reference to the entry's value in the map,
+    /// with a lifetime bound to the map itself.
+    pub fn into_mut(self) -> &'a mut V {
+        &mut self.map.entries[self.index].value
+    }
+
+    /// Remove and return the key, value pair stored in the map for this entry
+    ///
+    /// Like `Vec::swap_remove`, the pair is removed by swapping it with the
+    /// last element of the map and popping it off. **This perturbs
+    /// the position of what used to be the last element!**
+    ///
+    /// Computes in **O(1)** time (average).
+    pub fn swap_remove_entry(self) -> (K, V) {
+        self.map.swap_remove_index(self.index).unwrap()
+    }
+
+    /// Remove and return the key, value pair stored in the map for this entry
+    ///
+    /// 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!**
+    ///
+    /// Computes in **O(n)** time (average).
+    pub fn shift_remove_entry(self) -> (K, V) {
+        self.map.shift_remove_index(self.index).unwrap()
+    }
+
+    /// Remove the key, value pair stored in the map for this entry, and return the value.
+    ///
+    /// **NOTE:** This is equivalent to `.swap_remove()`.
+    pub fn remove(self) -> V {
+        self.swap_remove()
+    }
+
+    /// Remove the key, value pair stored in the map for this entry, and return the value.
+    ///
+    /// Like `Vec::swap_remove`, the pair is removed by swapping it with the
+    /// last element of the map and popping it off. **This perturbs
+    /// the position of what used to be the last element!**
+    ///
+    /// Computes in **O(1)** time (average).
+    pub fn swap_remove(self) -> V {
+        self.swap_remove_entry().1
+    }
+
+    /// Remove the key, value pair stored in the map for this entry, and return the 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!**
+    ///
+    /// Computes in **O(n)** time (average).
+    pub fn shift_remove(self) -> V {
+        self.shift_remove_entry().1
+    }
+
+    /// Remove and return the key, value pair stored in the map for this entry
+    ///
+    /// **NOTE:** This is equivalent to `.swap_remove_entry()`.
+    pub fn remove_entry(self) -> (K, V) {
+        self.swap_remove_entry()
+    }
+}
+
+impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for EntryAt<'_, K, V> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct(stringify!(OccupiedEntry))
+            .field("key", self.key())
+            .field("value", self.get())
+            .finish()
+    }
+}
+
 #[test]
 fn assert_send_sync() {
     fn assert_send_sync<T: Send + Sync>() {}
     assert_send_sync::<IndexMapCore<i32, i32>>();
     assert_send_sync::<Entry<'_, i32, i32>>();
+    assert_send_sync::<EntryAt<'_, i32, i32>>();
 }

src/map/tests.rs

@@ -359,6 +359,36 @@ fn occupied_entry_key() {
     }
 }
 
+#[test]
+fn entry_at() {
+    let mut map = IndexMap::new();
+
+    assert!(map.entry_at(0).is_none());
+
+    map.insert(0, "0");
+    map.insert(1, "1");
+    map.insert(2, "2");
+    map.insert(3, "3");
+
+    assert!(map.entry_at(4).is_none());
+
+    {
+        let e = map.entry_at(1).unwrap();
+        assert_eq!(*e.key(), 1);
+        assert_eq!(*e.get(), "1");
+        assert_eq!(e.remove(), "1");
+    }
+
+    {
+        let mut e = map.entry_at(1).unwrap();
+        assert_eq!(*e.key(), 3);
+        assert_eq!(*e.get(), "3");
+        assert_eq!(e.insert("4"), "3");
+    }
+
+    assert_eq!(*map.get(&3).unwrap(), "4");
+}
+
 #[test]
 fn keys() {
     let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')];