Port get_disjoint_mut to ringmap

Author: cuviper

src/map.rs

@@ -838,7 +838,8 @@ 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) {
+        let (head, tail) = self.as_mut_slices();
+        match Slice::get_disjoint_opt_mut(head, tail, indices) {
             Err(GetDisjointMutError::IndexOutOfBounds) => {
                 unreachable!(
                     "Internal error: indices should never be OOB as we got them from get_index_of"
@@ -1326,7 +1327,10 @@ impl<K, V, S> RingMap<K, V, S> {
         &mut self,
         indices: [usize; N],
     ) -> Result<[(&K, &mut V); N], GetDisjointMutError> {
-        self.as_mut_slice().get_disjoint_mut(indices)
+        let indices = indices.map(Some);
+        let (head, tail) = self.as_mut_slices();
+        let key_values = Slice::get_disjoint_opt_mut(head, tail, indices)?;
+        Ok(key_values.map(Option::unwrap))
     }
 
     /// Get the first key-value pair

src/map/slice.rs

@@ -274,17 +274,21 @@ impl<K, V> Slice<K, V> {
         indices: [usize; N],
     ) -> Result<[(&K, &mut V); N], GetDisjointMutError> {
         let indices = indices.map(Some);
-        let key_values = self.get_disjoint_opt_mut(indices)?;
+        let empty_tail = Self::new_mut();
+        let key_values = Self::get_disjoint_opt_mut(self, empty_tail, indices)?;
         Ok(key_values.map(Option::unwrap))
     }
 
     #[allow(unsafe_code)]
-    pub(crate) fn get_disjoint_opt_mut<const N: usize>(
-        &mut self,
+    pub(crate) fn get_disjoint_opt_mut<'a, const N: usize>(
+        head: &mut Self,
+        tail: &mut Self,
         indices: [Option<usize>; N],
-    ) -> Result<[Option<(&K, &mut V)>; N], GetDisjointMutError> {
+    ) -> Result<[Option<(&'a K, &'a mut V)>; N], GetDisjointMutError> {
+        let mid = head.len();
+        let len = mid + tail.len();
+
         // 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 {
@@ -295,14 +299,20 @@ impl<K, V> Slice<K, V> {
             }
         }
 
-        let entries_ptr = self.entries.as_mut_ptr();
+        let head_ptr = head.entries.as_mut_ptr();
+        let tail_ptr = tail.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
+                    // SAFETY: The base pointers are valid as they come from slices 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)
+                    unsafe {
+                        let ptr = match idx.checked_sub(mid) {
+                            None => head_ptr.add(idx),
+                            Some(tidx) => tail_ptr.add(tidx),
+                        };
+                        Some((*ptr).ref_mut())
+                    }
                 }
                 None => None,
             }