Rollup merge of rust-lang#153107 - asder8215:btreemap_append_optimized, r=Mark-Simulacrum

Optimize BTreeMap::append() using CursorMut

Since [`BTreeMap::merge`](rust-lang#152418 (comment)) uses `CursorMut` to avoid reconstructing the map from scratch and instead inserting other `BTreeMap` at the right places or overwriting the value in self `BTreeMap` on conflict, we might as well do the same for `BTreeMap::append`. This also means that some of the code in `append.rs` can be removed; `bulk_push()` however is used by `bulk_build_from_sorted_iterator()`, which is used by the `From`/`FromIterator` trait impl on `BTreeMap`. Feels like we should rename the file or place the `bulk_push()` in an existing file.

The same additional optimization consideration that `BTreeMap::merge` has is also applied to `BTreeMap::append`.

r? @Mark-Simulacrum  since Mark has seen the `BTreeMap::merge` code already (only diff is the `Ordering::Equal` case and now one of the test assertions on a panic case has the correct value now).

Author: JonathanBrouwer

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

@@ -1219,26 +1219,8 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
         K: Ord,
         A: Clone,
     {
-        // Do we have to append anything at all?
-        if other.is_empty() {
-            return;
-        }
-
-        // We can just swap `self` and `other` if `self` is empty.
-        if self.is_empty() {
-            mem::swap(self, other);
-            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(),
-        )
+        let other = mem::replace(other, Self::new_in((*self.alloc).clone()));
+        self.merge(other, |_key, _self_val, other_val| other_val);
     }
 
     /// 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);
 }