Make get_index_mut return &K, and add get_index_mut2

Author: cuviper

src/map.rs

@@ -479,21 +479,6 @@ where
         }
     }
 
-    pub(crate) fn get_full_mut2_impl<Q: ?Sized>(
-        &mut self,
-        key: &Q,
-    ) -> Option<(usize, &mut K, &mut V)>
-    where
-        Q: Hash + Equivalent<K>,
-    {
-        if let Some(i) = self.get_index_of(key) {
-            let entry = &mut self.as_entries_mut()[i];
-            Some((i, &mut entry.key, &mut entry.value))
-        } else {
-            None
-        }
-    }
-
     /// Remove the key-value pair equivalent to `key` and return
     /// its value.
     ///
@@ -780,8 +765,8 @@ impl<K, V, S> IndexMap<K, V, S> {
     /// Valid indices are *0 <= index < self.len()*
     ///
     /// Computes in **O(1)** time.
-    pub fn get_index_mut(&mut self, index: usize) -> Option<(&mut K, &mut V)> {
-        self.as_entries_mut().get_mut(index).map(Bucket::muts)
+    pub fn get_index_mut(&mut self, index: usize) -> Option<(&K, &mut V)> {
+        self.as_entries_mut().get_mut(index).map(Bucket::ref_mut)
     }
 
     /// Get the first key-value pair

src/mutable_keys.rs

@@ -1,6 +1,6 @@
 use core::hash::{BuildHasher, Hash};
 
-use super::{Equivalent, IndexMap};
+use super::{Bucket, Entries, Equivalent, IndexMap};
 
 pub struct PrivateMarker {}
 
@@ -21,13 +21,22 @@ pub trait MutableKeys {
     type Value;
 
     /// Return item index, mutable reference to key and value
+    ///
+    /// Computes in **O(1)** time (average).
     fn get_full_mut2<Q: ?Sized>(
         &mut self,
         key: &Q,
     ) -> Option<(usize, &mut Self::Key, &mut Self::Value)>
     where
         Q: Hash + Equivalent<Self::Key>;
 
+    /// Return mutable reference to key and value at an index.
+    ///
+    /// Valid indices are *0 <= index < self.len()*
+    ///
+    /// Computes in **O(1)** time.
+    fn get_index_mut2(&mut self, index: usize) -> Option<(&mut Self::Key, &mut Self::Value)>;
+
     /// Scan through each key-value pair in the map and keep those where the
     /// closure `keep` returns `true`.
     ///
@@ -39,6 +48,7 @@ pub trait MutableKeys {
     where
         F: FnMut(&mut Self::Key, &mut Self::Value) -> bool;
 
+    #[doc(hidden)]
     /// This method is not useful in itself – it is there to “seal” the trait
     /// for external implementation, so that we can add methods without
     /// causing breaking changes.
@@ -55,11 +65,21 @@ where
 {
     type Key = K;
     type Value = V;
+
     fn get_full_mut2<Q: ?Sized>(&mut self, key: &Q) -> Option<(usize, &mut K, &mut V)>
     where
         Q: Hash + Equivalent<K>,
     {
-        self.get_full_mut2_impl(key)
+        if let Some(i) = self.get_index_of(key) {
+            let entry = &mut self.as_entries_mut()[i];
+            Some((i, &mut entry.key, &mut entry.value))
+        } else {
+            None
+        }
+    }
+
+    fn get_index_mut2(&mut self, index: usize) -> Option<(&mut K, &mut V)> {
+        self.as_entries_mut().get_mut(index).map(Bucket::muts)
     }
 
     fn retain2<F>(&mut self, keep: F)