BTreeMap::merge optimized

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

@@ -1240,6 +1240,153 @@ 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: 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_entry {
+                        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 => {
+                                // advance the cursor forward
+                                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.
     /// 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 +3409,100 @@ 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.
+    #[allow(dead_code)] /* This function exists for consistency with CursorMut */
+    pub(super) fn with_next(&mut self, f: impl FnOnce(K, V) -> (K, V)) {
+        /// Call `f` with the next entry's key and value, replacing the
+        /// next entry's key and value with the returned key and value.
+        /// if `f` unwinds, the next entry is removed.
+        /// Equivalent to a more efficient version of:
+        /// ```ignore (this is an example equivalent)
+        /// if let Some((k, v)) = self.remove_next() {
+        ///     let (k, v) = f(k, v);
+        ///     // Safety: key is unmodified
+        ///     unsafe { self.insert_after_unchecked(k, v) };
+        /// }
+        /// ```
+        struct RemoveNextOnDrop<'a, 'b, K, V, A = Global>
+        where
+            K: Ord,
+            A: Allocator + Clone,
+        {
+            cursor: &'a mut CursorMutKey<'b, K, V, A>,
+            forget_next: bool,
+        }
+
+        impl<K: Ord, V, A: Allocator + Clone> Drop for RemoveNextOnDrop<'_, '_, K, V, A> {
+            fn drop(&mut self) {
+                if self.forget_next {
+                    // Removes the next entry of the Cursor without running
+                    // running the destructor on the key and value (prevents double
+                    // free)
+                    self.cursor.forget_next_key_and_value();
+                }
+            }
+        }
+
+        let mut remove_next_on_drop = RemoveNextOnDrop {
+            cursor: self,
+            forget_next: false, // next value is not known yet
+        };
+
+        if let Some((k_mut, v_mut)) = remove_next_on_drop.cursor.peek_next() {
+            remove_next_on_drop.forget_next = true;
+            // Safety: we move the K, V out of the next entry,
+            // we marked the entry's value to be forgotten
+            // when remove_next_on_drop is dropped that
+            // way we avoid returning to the caller leaving
+            // a moved-out invalid value if `f` unwinds.
+            let k = unsafe { ptr::read(k_mut) };
+            let v = unsafe { ptr::read(v_mut) };
+            let (k, v) = f(k, v);
+            unsafe { ptr::write(k_mut, k) };
+            unsafe { ptr::write(v_mut, v) };
+
+            remove_next_on_drop.forget_next = false;
+        }
+    }
+
+    /// Forgets the next element's key and value from the `BTreeMap`.
+    ///
+    /// The element's key and value does not run its destructor. The cursor position is
+    /// unchanged (before the forgotten element).
+    fn forget_next_key_and_value(&mut self) {
+        let current = match self.current.take() {
+            Some(current) => current,
+            None => return,
+        };
+
+        if current.reborrow().next_kv().is_err() {
+            self.current = Some(current);
+            return;
+        }
+
+        let mut emptied_internal_root = false;
+        if let Ok(next_kv) = current.next_kv() {
+            let ((_, val), pos) =
+                next_kv.remove_kv_tracking(|| emptied_internal_root = true, self.alloc.clone());
+            // Don't run destructor on next element's value
+            mem::forget(val);
+            self.current = Some(pos);
+            *self.length -= 1;
+            if emptied_internal_root {
+                // SAFETY: This is safe since current does not point within the now
+                // empty root node.
+                let root = unsafe { self.root.reborrow().as_mut().unwrap() };
+                root.pop_internal_level(self.alloc.clone());
+            }
+        }
+    }
+
     /// Inserts a new key-value pair into the map in the gap that the
     /// cursor is currently pointing to.
     ///
@@ -3467,6 +3708,96 @@ 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) {
+        /// Call `f` with the next entry's key and value, replacing the
+        /// next entry's value with the returned value. if `f` unwinds,
+        /// the next entry is removed.
+        /// Equivalent to a more efficient version of:
+        /// ```ignore (this is an example equivalent)
+        /// if let Some((k, v)) = self.remove_next() {
+        ///     let v = f(&k, v);
+        ///     // Safety: key is unmodified
+        ///     unsafe { self.insert_after_unchecked(k, v) };
+        /// }
+        /// ```
+        struct RemoveNextOnDrop<'a, 'b, K, V, A = Global>
+        where
+            K: Ord,
+            A: Allocator + Clone,
+        {
+            cursor: &'a mut CursorMut<'b, K, V, A>,
+            forget_next: bool,
+        }
+
+        impl<K: Ord, V, A: Allocator + Clone> Drop for RemoveNextOnDrop<'_, '_, K, V, A> {
+            fn drop(&mut self) {
+                if self.forget_next {
+                    // Removes the next entry of the Cursor without running
+                    // running the destructor on the value (prevents double
+                    // free)
+                    self.cursor.forget_next_value();
+                }
+            }
+        }
+
+        let mut remove_next_on_drop = RemoveNextOnDrop {
+            cursor: self,
+            forget_next: false, // next value is not known yet
+        };
+
+        if let Some((k, v_mut)) = remove_next_on_drop.cursor.peek_next() {
+            remove_next_on_drop.forget_next = true;
+            // Safety: we move the V out of the next entry,
+            // we marked the entry's value to be forgotten
+            // when remove_next_on_drop is dropped that
+            // way we avoid returning to the caller leaving
+            // a moved-out invalid value if `f` unwinds.
+            let v = unsafe { ptr::read(v_mut) };
+            let v = f(k, v);
+            unsafe { ptr::write(v_mut, v) };
+            remove_next_on_drop.forget_next = false;
+        }
+    }
+
+    /// Forgets the next element's value from the `BTreeMap`.
+    ///
+    /// The element's value does not run its destructor. The cursor position is
+    /// unchanged (before the forgotten element).
+    fn forget_next_value(&mut self) {
+        let current = match self.inner.current.take() {
+            Some(current) => current,
+            None => return,
+        };
+
+        if current.reborrow().next_kv().is_err() {
+            self.inner.current = Some(current);
+            return;
+        }
+
+        let mut emptied_internal_root = false;
+        if let Ok(next_kv) = current.next_kv() {
+            let ((_, val), pos) = next_kv
+                .remove_kv_tracking(|| emptied_internal_root = true, self.inner.alloc.clone());
+            // Don't run destructor on next element's value
+            mem::forget(val);
+            self.inner.current = Some(pos);
+            *self.inner.length -= 1;
+            if emptied_internal_root {
+                // SAFETY: This is safe since current does not point within the now
+                // empty root node.
+                let root = unsafe { self.inner.root.reborrow().as_mut().unwrap() };
+                root.pop_internal_level(self.inner.alloc.clone());
+            }
+        }
+    }
+
     /// Inserts a new key-value pair into the map in the gap that the
     /// cursor is currently pointing to.
     ///