optimized version of BTreeMap::merge with CursorMut and unsafe code

Author: asder8215

library/alloc/src/collections/btree/append.rs

@@ -33,36 +33,6 @@ impl<K, V> Root<K, V> {
         self.bulk_push(iter, length, alloc)
     }
 
-    /// Merges all key-value pairs from the union of two ascending iterators,
-    /// incrementing a `length` variable along the way. The latter makes it
-    /// easier for the caller to avoid a leak when a drop handler panicks.
-    ///
-    /// If both iterators produce the same key, this method constructs a pair using the
-    /// key from the left iterator and calls on a closure `f` to return a value given
-    /// the conflicting key and value from left and right iterators.
-    ///
-    /// If you want the tree to end up in a strictly ascending order, like for
-    /// a `BTreeMap`, both iterators should produce keys in strictly ascending
-    /// order, each greater than all keys in the tree, including any keys
-    /// already in the tree upon entry.
-    pub(super) fn merge_from_sorted_iters_with<I, A: Allocator + Clone>(
-        &mut self,
-        left: I,
-        right: I,
-        length: &mut usize,
-        alloc: A,
-        f: impl FnMut(&K, V, V) -> V,
-    ) where
-        K: Ord,
-        I: Iterator<Item = (K, V)> + FusedIterator,
-    {
-        // We prepare to merge `left` and `right` into a sorted sequence in linear time.
-        let iter = MergeIterWith { inner: MergeIterInner::new(left, right), f };
-
-        // Meanwhile, we build a tree from the sorted sequence in linear time.
-        self.bulk_push(iter, length, alloc)
-    }
-
     /// Pushes all key-value pairs to the end of the tree, incrementing a
     /// `length` variable along the way. The latter makes it easier for the
     /// caller to avoid a leak when the iterator panicks.
@@ -145,33 +115,3 @@ where
         }
     }
 }
-
-/// An iterator for merging two sorted sequences into one with
-/// a callback function to return a value on conflicting keys
-struct MergeIterWith<F, K, V, I: Iterator<Item = (K, V)>> {
-    inner: MergeIterInner<I>,
-    f: F,
-}
-
-impl<F, K: Ord, V, I> Iterator for MergeIterWith<F, K, V, I>
-where
-    F: FnMut(&K, V, V) -> V,
-    I: Iterator<Item = (K, V)> + FusedIterator,
-{
-    type Item = (K, V);
-
-    /// If two keys are equal, returns the key from the left and uses `f` to return
-    /// a value given the conflicting key and values from left and right
-    fn next(&mut self) -> Option<(K, V)> {
-        let (a_next, b_next) = self.inner.nexts(|a: &(K, V), b: &(K, V)| K::cmp(&a.0, &b.0));
-        match (a_next, b_next) {
-            (Some((a_k, a_v)), Some((_, b_v))) => Some({
-                let next_val = (self.f)(&a_k, a_v, b_v);
-                (a_k, next_val)
-            }),
-            (Some(a), None) => Some(a),
-            (None, Some(b)) => Some(b),
-            (None, None) => None,
-        }
-    }
-}

library/alloc/src/collections/btree/map.rs

@@ -1287,7 +1287,7 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
     /// assert_eq!(a[&5], "f");
     /// ```
     #[unstable(feature = "btree_merge", issue = "152152")]
-    pub fn merge(&mut self, mut other: Self, conflict: impl FnMut(&K, V, V) -> V)
+    pub fn merge(&mut self, mut other: Self, mut conflict: impl FnMut(&K, V, V) -> V)
     where
         K: Ord,
         A: Clone,
@@ -1303,16 +1303,96 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
             return;
         }
 
-        let self_iter = mem::replace(self, Self::new_in((*self.alloc).clone())).into_iter();
-        let other_iter = mem::replace(&mut other, Self::new_in((*self.alloc).clone())).into_iter();
-        let root = self.root.get_or_insert_with(|| Root::new((*self.alloc).clone()));
-        root.merge_from_sorted_iters_with(
-            self_iter,
-            other_iter,
-            &mut self.length,
-            (*self.alloc).clone(),
-            conflict,
-        )
+        let mut other_iter = other.into_iter();
+        let (first_other_key, first_other_val) = other_iter.next().unwrap();
+
+        // find the first gap that has the smallest key greater than the first key from other
+        let mut self_cursor = self.lower_bound_mut(Bound::Included(&first_other_key));
+
+        if let Some((self_key, self_val)) = self_cursor.peek_next() {
+            match K::cmp(&first_other_key, self_key) {
+                Ordering::Equal => {
+                    // SAFETY: We read in self_val's and hand it over to our conflict function
+                    // which will always return a value that we can use to overwrite what's
+                    // in self_val
+                    unsafe {
+                        let val = ptr::read(self_val);
+                        let next_val = (conflict)(self_key, val, first_other_val);
+                        ptr::write(self_val, next_val);
+                    }
+                }
+                Ordering::Less =>
+                // SAFETY: we know our other_key's ordering is less than self_key,
+                // so inserting before will guarantee sorted order
+                unsafe {
+                    self_cursor.insert_before_unchecked(first_other_key, first_other_val);
+                },
+                Ordering::Greater => {
+                    unreachable!("cursor's peek_next would return None in this case");
+                }
+            }
+        } else {
+            // SAFETY: if we reach here, that means our cursor has reached
+            // the end of self BTreeMap, (other_key is greater than all the
+            // previous self BTreeMap keys) so we just insert other_key here
+            // at the end of the CursorMut
+            unsafe {
+                self_cursor.insert_after_unchecked(first_other_key, first_other_val);
+            }
+        }
+
+        for (other_key, other_val) in other_iter {
+            if self_cursor.peek_next().is_some() {
+                loop {
+                    let self_entry = self_cursor.peek_next();
+                    if let Some((self_key, self_val)) = self_entry {
+                        match K::cmp(&other_key, self_key) {
+                            Ordering::Equal => {
+                                // SAFETY: We read in self_val's and hand it over to our conflict function
+                                // which will always return a value that we can use to overwrite what's
+                                // in self_val
+                                unsafe {
+                                    let val = ptr::read(self_val);
+                                    let next_val = (conflict)(self_key, val, other_val);
+                                    ptr::write(self_val, next_val);
+                                }
+                                break;
+                            }
+                            Ordering::Less => {
+                                // SAFETY: we know our other_key's ordering is less than self_key,
+                                // so inserting before will guarantee sorted order
+                                unsafe {
+                                    self_cursor.insert_before_unchecked(other_key, other_val);
+                                }
+                                break;
+                            }
+                            Ordering::Greater => {
+                                self_cursor.next();
+                            }
+                        }
+                    } else {
+                        // SAFETY: if we reach here, that means our cursor has reached
+                        // the end of self BTreeMap, (other_key is greater than all the
+                        // previous self BTreeMap keys) so we just insert other_key here
+                        // at the end of the Cursor
+                        unsafe {
+                            self_cursor.insert_after_unchecked(other_key, other_val);
+                        }
+                        self_cursor.next();
+                        break;
+                    }
+                }
+            } else {
+                // SAFETY: if we reach here, that means our cursor has reached
+                // the end of self BTreeMap, (other_key is greater than all the
+                // previous self BTreeMap keys) so we just insert the rest of
+                // other_keys here at the end of CursorMut
+                unsafe {
+                    self_cursor.insert_after_unchecked(other_key, other_val);
+                }
+                self_cursor.next();
+            }
+        }
     }
 
     /// Constructs a double-ended iterator over a sub-range of elements in the map.