Separate the implementations of get_disjoint_mut

src/map.rs

@@ -1,6 +1,7 @@
 //! [`IndexMap`] is a hash table where the iteration order of the key-value
 //! pairs is independent of the hash values of the keys.
 
+mod disjoint;
 mod entry;
 mod iter;
 mod mutable;
@@ -952,25 +953,19 @@ where
     ///
     /// ```
     /// let mut map = indexmap::IndexMap::from([(1, 'a'), (3, 'b'), (2, 'c')]);
-    /// assert_eq!(map.get_disjoint_mut([&2, &1]), [Some(&mut 'c'), Some(&mut 'a')]);
+    /// assert_eq!(
+    ///   map.get_disjoint_mut([&2, &1, &0]),
+    ///   [Some(&mut 'c'), Some(&mut 'a'), None],
+    /// );
     /// ```
     #[track_caller]
     pub fn get_disjoint_mut<Q, const N: usize>(&mut self, keys: [&Q; N]) -> [Option<&mut V>; N]
     where
         Q: ?Sized + Hash + Equivalent<K>,
     {
         let indices = keys.map(|key| self.get_index_of(key));
-        match self.as_mut_slice().get_disjoint_opt_mut(indices) {
-            Err(GetDisjointMutError::IndexOutOfBounds) => {
-                unreachable!(
-                    "Internal error: indices should never be OOB as we got them from get_index_of"
-                );
-            }
-            Err(GetDisjointMutError::OverlappingIndices) => {
-                panic!("duplicate keys found");
-            }
-            Ok(key_values) => key_values.map(|kv_opt| kv_opt.map(|kv| kv.1)),
-        }
+        disjoint::get_disjoint_opt_mut(self.as_entries_mut(), indices)
+            .map(|opt| opt.map(Bucket::value_mut))
     }
 
     /// Remove the key-value pair equivalent to `key` and return

src/map/disjoint.rs

@@ -0,0 +1,60 @@
+#![allow(unsafe_code)]
+
+use crate::GetDisjointMutError;
+
+/// Like `slice::get_disjoint_mut`, although we're not dealing with ranges (yet).
+// TODO(MSRV 1.86): remove this in favor of the standard library's method.
+pub(super) fn get_disjoint_mut<T, const N: usize>(
+    entries: &mut [T],
+    indices: [usize; N],
+) -> Result<[&mut T; N], GetDisjointMutError> {
+    // SAFETY: Can't allow duplicate indices as we would return mutable refs to the same data.
+    let len = entries.len();
+    for i in 0..N {
+        let idx = indices[i];
+        if idx >= len {
+            return Err(GetDisjointMutError::IndexOutOfBounds);
+        } else if indices[..i].contains(&idx) {
+            return Err(GetDisjointMutError::OverlappingIndices);
+        }
+    }
+
+    let entries_ptr = entries.as_mut_ptr();
+    Ok(indices.map(move |idx| {
+        // SAFETY: The base pointer is valid as it comes from a slice and the reference is always
+        // in-bounds & unique as we've already checked the indices above.
+        unsafe { &mut *(entries_ptr.add(idx)) }
+    }))
+}
+
+/// Like `slice::get_disjoint_mut` but with optional indices,
+/// allowing for absent keys from the user's original request.
+#[track_caller]
+pub(super) fn get_disjoint_opt_mut<T, const N: usize>(
+    entries: &mut [T],
+    indices: [Option<usize>; N],
+) -> [Option<&mut T>; N] {
+    // SAFETY: Can't allow duplicate indices as we would return mutable refs to the same data.
+    let len = entries.len();
+    for i in 0..N {
+        if let Some(idx) = indices[i] {
+            if idx >= len {
+                unreachable!("`get_index_of` returned an out-of-bounds index");
+            } else if indices[..i].contains(&Some(idx)) {
+                panic!("duplicate keys found");
+            }
+        }
+    }
+
+    let entries_ptr = entries.as_mut_ptr();
+    indices.map(move |idx_opt| {
+        match idx_opt {
+            Some(idx) => {
+                // SAFETY: The base pointer is valid as it comes from a slice and the reference is always
+                // in-bounds & unique as we've already checked the indices above.
+                Some(unsafe { &mut *entries_ptr.add(idx) })
+            }
+            None => None,
+        }
+    })
+}

src/map/slice.rs

@@ -350,42 +350,9 @@ impl<K, V> Slice<K, V> {
         &mut self,
         indices: [usize; N],
     ) -> Result<[(&K, &mut V); N], GetDisjointMutError> {
-        let indices = indices.map(Some);
-        let key_values = self.get_disjoint_opt_mut(indices)?;
-        Ok(key_values.map(Option::unwrap))
-    }
-
-    #[allow(unsafe_code)]
-    pub(crate) fn get_disjoint_opt_mut<const N: usize>(
-        &mut self,
-        indices: [Option<usize>; N],
-    ) -> Result<[Option<(&K, &mut V)>; N], GetDisjointMutError> {
-        // SAFETY: Can't allow duplicate indices as we would return several mutable refs to the same data.
-        let len = self.len();
-        for i in 0..N {
-            if let Some(idx) = indices[i] {
-                if idx >= len {
-                    return Err(GetDisjointMutError::IndexOutOfBounds);
-                } else if indices[..i].contains(&Some(idx)) {
-                    return Err(GetDisjointMutError::OverlappingIndices);
-                }
-            }
-        }
-
-        let entries_ptr = self.entries.as_mut_ptr();
-        let out = indices.map(|idx_opt| {
-            match idx_opt {
-                Some(idx) => {
-                    // SAFETY: The base pointer is valid as it comes from a slice and the reference is always
-                    // in-bounds & unique as we've already checked the indices above.
-                    let kv = unsafe { (*(entries_ptr.add(idx))).ref_mut() };
-                    Some(kv)
-                }
-                None => None,
-            }
-        });
-
-        Ok(out)
+        // TODO(MSRV 1.86): use the standard library's `slice::get_disjoint_mut`
+        let entries = super::disjoint::get_disjoint_mut(&mut self.entries, indices)?;
+        Ok(entries.map(Bucket::ref_mut))
     }
 }