Implement get_many_index_mut & and add docs

Avoid mutable aliasing by using a pointer to entries

Author: NiklasJonsson

src/map.rs

@@ -477,37 +477,35 @@ where
         }
     }
 
+    /// Return the values for `N` keys. If any key is missing a value, or there
+    /// are duplicate keys, `None` is returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut map = indexmap::IndexMap::from([(1, 'a'), (3, 'b'), (2, 'c')]);
+    /// assert_eq!(map.get_many_mut([&2, &1]), Some([&mut 'c', &mut 'a']));
+    /// ```
     pub fn get_many_mut<'a, 'b, Q: ?Sized, const N: usize>(
         &'a mut self,
         keys: [&'b Q; N],
     ) -> Option<[&'a mut V; N]>
     where
         Q: Hash + Equivalent<K>,
     {
+        let len = self.len();
         let indices = keys.map(|key| self.get_index_of(key));
-        if indices.iter().any(Option::is_none) {
-            return None;
-        }
-        let indices = indices.map(Option::unwrap);
 
-        // SAFETY: Can't allow duplicate indices as we would return several mutable refs to the same data
-        for i in 0..N {
-            let idx = indices[i];
-            if indices[i + 1..N].contains(&idx) {
-                return None;
+        // Handle out-of-bounds indices with panic as this is an internal error in get_index_of.
+        for idx in indices {
+            let idx = idx?;
+            if idx >= len {
+                panic!("Index is out of range! Got '{idx}' but length is '{len}'")
             }
         }
-
-        let entries = self.as_entries_mut();
-        let out = indices.map(|i| {
-            // SAFETY: OK to discard mutable borrow lifetime as each index is unique
-            #[allow(unsafe_code)]
-            unsafe {
-                &mut *(&mut entries[i].value as *mut V)
-            }
-        });
-
-        Some(out)
+        let indices = indices.map(Option::unwrap);
+        let entries = self.get_many_index_mut(indices)?;
+        Some(entries.map(|(_key, value)| value))
     }
 
     /// Remove the key-value pair equivalent to `key` and return
@@ -817,6 +815,45 @@ impl<K, V, S> IndexMap<K, V, S> {
         self.as_entries_mut().get_mut(index).map(Bucket::ref_mut)
     }
 
+    /// Get an array of `N` key-value pairs by `N` indices
+    ///
+    /// Valid indices are *0 <= index < self.len()* and each index needs to be unique.
+    ///
+    /// Computes in **O(1)** time.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut map = indexmap::IndexMap::from([(1, 'a'), (3, 'b'), (2, 'c')]);
+    /// assert_eq!(map.get_many_index_mut([2, 0]), Some([(&2, &mut 'c'), (&1, &mut 'a')]));
+    /// ```
+    pub fn get_many_index_mut<const N: usize>(
+        &mut self,
+        indices: [usize; N],
+    ) -> Option<[(&K, &mut V); N]> {
+        // SAFETY: Can't allow duplicate indices as we would return several mutable refs to the same data.
+        // Additionally, handle out-of-bounds indices (internal error in get_index_of) with panic.
+        let len = self.len();
+        for i in 0..N {
+            let idx = indices[i];
+            if idx >= len || indices[i + 1..N].contains(&idx) {
+                return None;
+            }
+        }
+
+        let entries_ptr = self.as_entries_mut().as_mut_ptr();
+        let out = indices.map(|i| {
+            // SAFETY: The base pointer is valid as it comes from a slice and the deref is always
+            // in-bounds as we've already checked the indices above.
+            #[allow(unsafe_code)]
+            unsafe {
+                (*(entries_ptr.add(i))).ref_mut()
+            }
+        });
+
+        Some(out)
+    }
+
     /// Returns a slice of key-value pairs in the given range of indices.
     ///
     /// Valid indices are *0 <= index < self.len()*