Optimize BTreeMap::append() using CursorMut

Author: asder8215

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

@@ -1,38 +1,8 @@
 use core::alloc::Allocator;
-use core::iter::FusedIterator;
 
-use super::merge_iter::MergeIterInner;
 use super::node::{self, Root};
 
 impl<K, V> Root<K, V> {
-    /// Appends 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 drops the pair from
-    /// the left iterator and appends the pair from the right iterator.
-    ///
-    /// 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 append_from_sorted_iters<I, A: Allocator + Clone>(
-        &mut self,
-        left: I,
-        right: I,
-        length: &mut usize,
-        alloc: A,
-    ) 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 = MergeIter(MergeIterInner::new(left, right));
-
-        // 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.
@@ -94,24 +64,3 @@ impl<K, V> Root<K, V> {
         self.fix_right_border_of_plentiful();
     }
 }
-
-// An iterator for merging two sorted sequences into one
-struct MergeIter<K, V, I: Iterator<Item = (K, V)>>(MergeIterInner<I>);
-
-impl<K: Ord, V, I> Iterator for MergeIter<K, V, I>
-where
-    I: Iterator<Item = (K, V)> + FusedIterator,
-{
-    type Item = (K, V);
-
-    /// If two keys are equal, returns the key from the left and the value from the right.
-    fn next(&mut self) -> Option<(K, V)> {
-        let (a_next, b_next) = self.0.nexts(|a: &(K, V), b: &(K, V)| K::cmp(&a.0, &b.0));
-        match (a_next, b_next) {
-            (Some((a_k, _)), Some((_, b_v))) => Some((a_k, b_v)),
-            (Some(a), None) => Some(a),
-            (None, Some(b)) => Some(b),
-            (None, None) => None,
-        }
-    }
-}

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

@@ -1229,15 +1229,82 @@ 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(other, Self::new_in((*self.alloc).clone())).into_iter();
-        let root = self.root.get_or_insert_with(|| Root::new((*self.alloc).clone()));
-        root.append_from_sorted_iters(
-            self_iter,
-            other_iter,
-            &mut self.length,
-            (*self.alloc).clone(),
-        )
+        // Because `other` takes a &mut Self, in order for us to own
+        // the K,V pairs, we need to use mem::replace
+        let mut other_iter = mem::replace(other, Self::new_in((*self.alloc).clone())).into_iter();
+        let (first_other_key, first_other_val) = other_iter.next().unwrap();
+
+        // find the first gap that has the smallest key greater than or equal to
+        // the first key from other
+        let mut self_cursor = self.lower_bound_mut(Bound::Included(&first_other_key));
+
+        if let Some((self_key, _)) = self_cursor.peek_next() {
+            match K::cmp(self_key, &first_other_key) {
+                Ordering::Equal => {
+                    if let Some((_, v)) = self_cursor.next() {
+                        *v = first_other_val;
+                    }
+                }
+                Ordering::Greater =>
+                // 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::Less => {
+                    unreachable!("Cursor's peek_next should return None.");
+                }
+            }
+        } else {
+            // SAFETY: reaching here means our cursor is at the end
+            // self BTreeMap so we just insert other_key here
+            unsafe {
+                self_cursor.insert_before_unchecked(first_other_key, first_other_val);
+            }
+        }
+
+        for (other_key, other_val) in other_iter {
+            loop {
+                if let Some((self_key, _)) = self_cursor.peek_next() {
+                    match K::cmp(self_key, &other_key) {
+                        Ordering::Equal => {
+                            if let Some((_, v)) = self_cursor.next() {
+                                *v = other_val;
+                            }
+                            break;
+                        }
+                        Ordering::Greater => {
+                            // SAFETY: we know our self_key's ordering is greater than other_key,
+                            // so inserting before will guarantee sorted order
+                            unsafe {
+                                self_cursor.insert_before_unchecked(other_key, other_val);
+                            }
+                            break;
+                        }
+                        Ordering::Less => {
+                            // FIXME: instead of doing a linear search here,
+                            // this can be optimized to search the tree by starting
+                            // from self_cursor and going towards the root and then
+                            // back down to the proper node -- that should probably
+                            // be a new method on Cursor*.
+                            self_cursor.next();
+                        }
+                    }
+                } else {
+                    // FIXME: If we get here, that means all of other's keys are greater than
+                    // self's keys. For performance, this should really do a bulk insertion of items
+                    // from other_iter into the end of self `BTreeMap`. Maybe this should be
+                    // a method for Cursor*?
+
+                    // SAFETY: reaching here means our cursor is at the end
+                    // self BTreeMap so we just insert other_key here
+                    unsafe {
+                        self_cursor.insert_before_unchecked(other_key, other_val);
+                    }
+                    break;
+                }
+            }
+        }
     }
 
     /// Moves all elements from `other` into `self`, leaving `other` empty.

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

@@ -2224,7 +2224,7 @@ fn test_append_drop_leak() {
 
     catch_unwind(move || left.append(&mut right)).unwrap_err();
     assert_eq!(a.dropped(), 1);
-    assert_eq!(b.dropped(), 1); // should be 2 were it not for Rust issue #47949
+    assert_eq!(b.dropped(), 2);
     assert_eq!(c.dropped(), 2);
 }