BTreeMap::merge optimized

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

@@ -1240,6 +1240,140 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
         )
     }
 
+    /// Moves all elements from `other` into `self`, leaving `other` empty.
+    ///
+    /// If a key from `other` is already present in `self`, then the `conflict`
+    /// closure is used to return a value to `self`. The `conflict`
+    /// closure takes in a borrow of `self`'s key, `self`'s value, and `other`'s value
+    /// in that order.
+    ///
+    /// An example of why one might use this method over [`append`]
+    /// is to combine `self`'s value with `other`'s value when their keys conflict.
+    ///
+    /// Similar to [`insert`], though, the key is not overwritten,
+    /// which matters for types that can be `==` without being identical.
+    ///
+    ///
+    /// [`insert`]: BTreeMap::insert
+    /// [`append`]: BTreeMap::append
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(btree_merge)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut a = BTreeMap::new();
+    /// a.insert(1, String::from("a"));
+    /// a.insert(2, String::from("b"));
+    /// a.insert(3, String::from("c")); // Note: Key (3) also present in b.
+    ///
+    /// let mut b = BTreeMap::new();
+    /// b.insert(3, String::from("d")); // Note: Key (3) also present in a.
+    /// b.insert(4, String::from("e"));
+    /// b.insert(5, String::from("f"));
+    ///
+    /// // concatenate a's value and b's value
+    /// a.merge(b, |_, a_val, b_val| {
+    ///     format!("{a_val}{b_val}")
+    /// });
+    ///
+    /// assert_eq!(a.len(), 5); // all of b's keys in a
+    ///
+    /// assert_eq!(a[&1], "a");
+    /// assert_eq!(a[&2], "b");
+    /// assert_eq!(a[&3], "cd"); // Note: "c" has been combined with "d".
+    /// assert_eq!(a[&4], "e");
+    /// assert_eq!(a[&5], "f");
+    /// ```
+    #[unstable(feature = "btree_merge", issue = "152152")]
+    pub fn merge(&mut self, mut other: Self, mut conflict: impl FnMut(&K, V, V) -> V)
+    where
+        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, &mut other);
+            return;
+        }
+
+        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 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(&first_other_key, self_key) {
+                Ordering::Equal => {
+                    self_cursor.with_next(|self_key, self_val| {
+                        conflict(self_key, self_val, first_other_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 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(&other_key, self_key) {
+                        Ordering::Equal => {
+                            self_cursor.with_next(|self_key, self_val| {
+                                conflict(self_key, self_val, other_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 => {
+                            // 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 {
+                    // 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;
+                }
+            }
+        }
+    }
+
     /// Constructs a double-ended iterator over a sub-range of elements in the map.
     /// The simplest way is to use the range syntax `min..max`, thus `range(min..max)` will
     /// yield elements from min (inclusive) to max (exclusive).
@@ -3262,6 +3396,37 @@ impl<'a, K, V, A> CursorMutKey<'a, K, V, A> {
 
 // Now the tree editing operations
 impl<'a, K: Ord, V, A: Allocator + Clone> CursorMutKey<'a, K, V, A> {
+    /// Calls a function with ownership of the next element's key and
+    /// and value and expects it to return a value to write
+    /// back to the next element's key and value. The cursor is not
+    /// advanced forward.
+    ///
+    /// If the cursor is at the end of the map then the function is not called
+    /// and this essentially does not do anything.
+    ///
+    /// # Safety
+    ///
+    /// You must ensure that the `BTreeMap` invariants are maintained.
+    /// Specifically:
+    ///
+    /// * The next element's key must be unique in the tree.
+    /// * All keys in the tree must remain in sorted order.
+    #[allow(dead_code)] /* This function exists for consistency with CursorMut */
+    pub(super) fn with_next(&mut self, f: impl FnOnce(K, V) -> (K, V)) {
+        // if `f` unwinds, the next entry is already removed leaving
+        // the tree in valid state.
+        // FIXME: Once `MaybeDangling` is implemented, we can optimize
+        // this through using a drop handler and transmutating CursorMutKey<K, V>
+        // to CursorMutKey<ManuallyDrop<K>, ManuallyDrop<V>> (see PR #152418)
+        if let Some((k, v)) = self.remove_next() {
+            // SAFETY: we remove the K, V out of the next entry,
+            // apply 'f' to get a new (K, V), and insert it back
+            // into the next entry that the cursor is pointing at
+            let (k, v) = f(k, v);
+            unsafe { self.insert_after_unchecked(k, v) };
+        }
+    }
+
     /// Inserts a new key-value pair into the map in the gap that the
     /// cursor is currently pointing to.
     ///
@@ -3467,6 +3632,29 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMutKey<'a, K, V, A> {
 }
 
 impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> {
+    /// Calls a function with a reference to the next element's key and
+    /// ownership of its value. The function is expected to return a value
+    /// to write back to the next element's value. The cursor is not
+    /// advanced forward.
+    ///
+    /// If the cursor is at the end of the map then the function is not called
+    /// and this essentially does not do anything.
+    pub(super) fn with_next(&mut self, f: impl FnOnce(&K, V) -> V) {
+        // FIXME: This can be optimized to not do all the removing/reinserting
+        // logic by using ptr::read, calling `f`, and then using ptr::write.
+        // if `f` unwinds, then we need to remove the entry while being careful to
+        // not cause UB by moving or dropping the already-dropped `V`
+        // for the entry. Some implementation ideas:
+        // https://github.com/rust-lang/rust/pull/152418#discussion_r2800232576
+        if let Some((k, v)) = self.remove_next() {
+            // SAFETY: we remove the K, V out of the next entry,
+            // apply 'f' to get a new V, and insert (K, V) back
+            // into the next entry that the cursor is pointing at
+            let v = f(&k, v);
+            unsafe { self.insert_after_unchecked(k, v) };
+        }
+    }
+
     /// Inserts a new key-value pair into the map in the gap that the
     /// cursor is currently pointing to.
     ///