Implement get_disjoint_mut for arrays of keys

Author: NiklasJonsson

src/map.rs

@@ -790,6 +790,37 @@ 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_disjoint_mut([&2, &1]), Some([&mut 'c', &mut 'a']));
+    /// ```
+    pub fn get_disjoint_mut<Q, const N: usize>(&mut self, keys: [&Q; N]) -> Option<[&mut V; N]>
+    where
+        Q: Hash + Equivalent<K> + ?Sized,
+    {
+        let len = self.len();
+        let indices = keys.map(|key| self.get_index_of(key));
+
+        // Handle out-of-bounds indices with panic as this is an internal error in get_index_of.
+        for idx in indices {
+            let idx = idx?;
+            debug_assert!(
+                idx < len,
+                "Index is out of range! Got '{}' but length is '{}'",
+                idx,
+                len
+            );
+        }
+        let indices = indices.map(Option::unwrap);
+        let entries = self.get_disjoint_indices_mut(indices)?;
+        Some(entries.map(|(_key, value)| value))
+    }
+
     /// Remove the key-value pair equivalent to `key` and return
     /// its value.
     ///
@@ -1196,6 +1227,44 @@ impl<K, V, S> IndexMap<K, V, S> {
         Some(IndexedEntry::new(&mut self.core, index))
     }
 
+    /// 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_disjoint_indices_mut([2, 0]), Some([(&2, &mut 'c'), (&1, &mut 'a')]));
+    /// ```
+    pub fn get_disjoint_indices_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.
+        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()`.

src/map/tests.rs

@@ -828,3 +828,94 @@ move_index_oob!(test_move_index_out_of_bounds_0_10, 0, 10);
 move_index_oob!(test_move_index_out_of_bounds_0_max, 0, usize::MAX);
 move_index_oob!(test_move_index_out_of_bounds_10_0, 10, 0);
 move_index_oob!(test_move_index_out_of_bounds_max_0, usize::MAX, 0);
+
+#[test]
+fn disjoint_mut_empty_map() {
+    let mut map: IndexMap<u32, u32> = IndexMap::default();
+    assert!(map.get_disjoint_mut([&0, &1, &2, &3]).is_none());
+}
+
+#[test]
+fn disjoint_mut_empty_param() {
+    let mut map: IndexMap<u32, u32> = IndexMap::default();
+    map.insert(1, 10);
+    assert!(map.get_disjoint_mut([] as [&u32; 0]).is_some());
+}
+
+#[test]
+fn disjoint_mut_single_fail() {
+    let mut map: IndexMap<u32, u32> = IndexMap::default();
+    map.insert(1, 10);
+    assert!(map.get_disjoint_mut([&0]).is_none());
+}
+
+#[test]
+fn disjoint_mut_single_success() {
+    let mut map: IndexMap<u32, u32> = IndexMap::default();
+    map.insert(1, 10);
+    assert_eq!(map.get_disjoint_mut([&1]), Some([&mut 10]));
+}
+
+#[test]
+fn disjoint_mut_multi_success() {
+    let mut map: IndexMap<u32, u32> = IndexMap::default();
+    map.insert(1, 100);
+    map.insert(2, 200);
+    map.insert(3, 300);
+    map.insert(4, 400);
+    assert_eq!(map.get_disjoint_mut([&1, &2]), Some([&mut 100, &mut 200]));
+    assert_eq!(map.get_disjoint_mut([&1, &3]), Some([&mut 100, &mut 300]));
+    assert_eq!(
+        map.get_disjoint_mut([&3, &1, &4, &2]),
+        Some([&mut 300, &mut 100, &mut 400, &mut 200])
+    );
+}
+
+#[test]
+fn disjoint_mut_multi_success_unsized_key() {
+    let mut map: IndexMap<&'static str, u32> = IndexMap::default();
+    map.insert("1", 100);
+    map.insert("2", 200);
+    map.insert("3", 300);
+    map.insert("4", 400);
+    assert_eq!(map.get_disjoint_mut(["1", "2"]), Some([&mut 100, &mut 200]));
+    assert_eq!(map.get_disjoint_mut(["1", "3"]), Some([&mut 100, &mut 300]));
+    assert_eq!(
+        map.get_disjoint_mut(["3", "1", "4", "2"]),
+        Some([&mut 300, &mut 100, &mut 400, &mut 200])
+    );
+}
+
+#[test]
+fn disjoint_mut_multi_fail_missing() {
+    let mut map: IndexMap<u32, u32> = IndexMap::default();
+    map.insert(1, 10);
+    map.insert(1123, 100);
+    map.insert(321, 20);
+    map.insert(1337, 30);
+    assert_eq!(map.get_disjoint_mut([&121, &1123]), None);
+    assert_eq!(map.get_disjoint_mut([&1, &1337, &56]), None);
+    assert_eq!(map.get_disjoint_mut([&1337, &123, &321, &1, &1123]), None);
+}
+
+#[test]
+fn disjoint_mut_multi_fail_duplicate() {
+    let mut map: IndexMap<u32, u32> = IndexMap::default();
+    map.insert(1, 10);
+    map.insert(1123, 100);
+    map.insert(321, 20);
+    map.insert(1337, 30);
+    assert_eq!(map.get_disjoint_mut([&1, &1]), None);
+    assert_eq!(
+        map.get_disjoint_mut([&1337, &123, &321, &1337, &1, &1123]),
+        None
+    );
+}
+
+#[test]
+fn many_index_mut_fail_oob() {
+    let mut map: IndexMap<u32, u32> = IndexMap::default();
+    map.insert(1, 10);
+    map.insert(321, 20);
+    assert_eq!(map.get_disjoint_indices_mut([1, 3]), None);
+}