From 6a2f234df57f67e382550979a26951ea603f9727 Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Sun, 1 Feb 2026 23:13:55 -0500 Subject: [PATCH 1/7] feat: implement append_with() for BTreeMap --- library/alloc/src/collections/btree/append.rs | 60 ++++++++++++++++ library/alloc/src/collections/btree/map.rs | 71 +++++++++++++++++++ 2 files changed, 131 insertions(+) diff --git a/library/alloc/src/collections/btree/append.rs b/library/alloc/src/collections/btree/append.rs index 66ea22e75247c..36e46f74889f3 100644 --- a/library/alloc/src/collections/btree/append.rs +++ b/library/alloc/src/collections/btree/append.rs @@ -33,6 +33,36 @@ impl Root { self.bulk_push(iter, length, alloc) } + /// 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 constructs a pair using the + /// key from the left iterator and calls on a function `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 append_from_sorted_iters_with( + &mut self, + left: I, + right: I, + length: &mut usize, + alloc: A, + f: impl FnMut(&K, V, V) -> V, + ) where + K: Ord, + I: Iterator + 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. @@ -115,3 +145,33 @@ where } } } + +/// An iterator for merging two sorted sequences into one with +/// a callback function to return a value on conflicting keys +struct MergeIterWith> { + inner: MergeIterInner, + f: F, +} + +impl Iterator for MergeIterWith +where + F: FnMut(&K, V, V) -> V, + I: Iterator + FusedIterator, +{ + type Item = (K, V); + + /// If two keys are equal, returns the key from the left and uses `f` to return + /// a value + 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.call_mut((&a_k, a_v, b_v)); + (a_k, next_val) + }), + (Some(a), None) => Some(a), + (None, Some(b)) => Some(b), + (None, None) => None, + } + } +} diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 426be364a56b0..ec46b09df0fba 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1240,6 +1240,77 @@ impl BTreeMap { ) } + /// Moves all elements from `other` into `self`, leaving `other` empty. + /// + /// If a key from `other` is already present in `self`, then the `conflict` + /// function is used to return a value to `self` given the conflicting key, + /// the current value of `self`, and the current value of `other`. + /// An example of why one might use this function over `BTreeMap::append` + /// is to merge `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 + /// + /// # Examples + /// + /// ``` + /// 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.append_with(&mut b, |_, a_val, b_val| { + /// format!("{a_val}{b_val}") + /// }); + /// + /// assert_eq!(a.len(), 5); + /// assert_eq!(b.len(), 0); + /// + /// 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_append_with", issue = "147700")] + pub fn append_with(&mut self, other: &mut Self, 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, 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_with( + self_iter, + other_iter, + &mut self.length, + (*self.alloc).clone(), + conflict, + ) + } + /// 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). From 1611b73532dc76d3ac91e52e402ac9c3af82a139 Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Mon, 2 Feb 2026 09:04:07 -0500 Subject: [PATCH 2/7] updated comment, used () syntax for FnMut callback, and added feature block in append_with() --- library/alloc/src/collections/btree/append.rs | 4 ++-- library/alloc/src/collections/btree/map.rs | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/library/alloc/src/collections/btree/append.rs b/library/alloc/src/collections/btree/append.rs index 36e46f74889f3..7b00bbfb59533 100644 --- a/library/alloc/src/collections/btree/append.rs +++ b/library/alloc/src/collections/btree/append.rs @@ -161,12 +161,12 @@ where type Item = (K, V); /// If two keys are equal, returns the key from the left and uses `f` to return - /// a value + /// a value given knowledge of 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.call_mut((&a_k, a_v, b_v)); + let next_val = (self.f)(&a_k, a_v, b_v); (a_k, next_val) }), (Some(a), None) => Some(a), diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index ec46b09df0fba..6f866c2ee590d 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1256,6 +1256,7 @@ impl BTreeMap { /// # Examples /// /// ``` + /// #![feature(btreemap_append_with)] /// use std::collections::BTreeMap; /// /// let mut a = BTreeMap::new(); @@ -1282,7 +1283,7 @@ impl BTreeMap { /// assert_eq!(a[&4], "e"); /// assert_eq!(a[&5], "f"); /// ``` - #[unstable(feature = "btree_append_with", issue = "147700")] + #[unstable(feature = "btree_append_with", issue = "147700")] // FIXME: Change issue # to track # pub fn append_with(&mut self, other: &mut Self, conflict: impl FnMut(&K, V, V) -> V) where K: Ord, From 8136795960173edde864e7e6b5f099f8e097ede5 Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Mon, 2 Feb 2026 09:52:51 -0500 Subject: [PATCH 3/7] corrected feature doctest name --- library/alloc/src/collections/btree/map.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 6f866c2ee590d..c9786d53d95a7 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1256,7 +1256,7 @@ impl BTreeMap { /// # Examples /// /// ``` - /// #![feature(btreemap_append_with)] + /// #![feature(btree_append_with)] /// use std::collections::BTreeMap; /// /// let mut a = BTreeMap::new(); From 746641c6bce8c4bf7ff1ccd58048aeb5054ad774 Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Thu, 5 Feb 2026 00:56:05 -0500 Subject: [PATCH 4/7] renamed append_with() to merge(), changed function signature, and clarified doc comments --- library/alloc/src/collections/btree/append.rs | 8 +++---- library/alloc/src/collections/btree/map.rs | 24 +++++++++++-------- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/library/alloc/src/collections/btree/append.rs b/library/alloc/src/collections/btree/append.rs index 7b00bbfb59533..38e4542d07905 100644 --- a/library/alloc/src/collections/btree/append.rs +++ b/library/alloc/src/collections/btree/append.rs @@ -33,19 +33,19 @@ impl Root { self.bulk_push(iter, length, alloc) } - /// Appends all key-value pairs from the union of two ascending iterators, + /// 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 function `f` to return a value given + /// 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 append_from_sorted_iters_with( + pub(super) fn merge_from_sorted_iters_with( &mut self, left: I, right: I, @@ -161,7 +161,7 @@ where type Item = (K, V); /// If two keys are equal, returns the key from the left and uses `f` to return - /// a value given knowledge of conflicting key and values from left and right + /// 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) { diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index c9786d53d95a7..49686f4cb9909 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1243,20 +1243,24 @@ impl BTreeMap { /// Moves all elements from `other` into `self`, leaving `other` empty. /// /// If a key from `other` is already present in `self`, then the `conflict` - /// function is used to return a value to `self` given the conflicting key, - /// the current value of `self`, and the current value of `other`. - /// An example of why one might use this function over `BTreeMap::append` + /// closure is used to return a value to `self`. For clarity, 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 merge `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_append_with)] + /// #![feature(btree_merge)] /// use std::collections::BTreeMap; /// /// let mut a = BTreeMap::new(); @@ -1270,7 +1274,7 @@ impl BTreeMap { /// b.insert(5, String::from("f")); /// /// // concatenate a's value and b's value - /// a.append_with(&mut b, |_, a_val, b_val| { + /// a.merge(&mut b, |_, a_val, b_val| { /// format!("{a_val}{b_val}") /// }); /// @@ -1283,8 +1287,8 @@ impl BTreeMap { /// assert_eq!(a[&4], "e"); /// assert_eq!(a[&5], "f"); /// ``` - #[unstable(feature = "btree_append_with", issue = "147700")] // FIXME: Change issue # to track # - pub fn append_with(&mut self, other: &mut Self, conflict: impl FnMut(&K, V, V) -> V) + #[unstable(feature = "btree_merge", issue = "147700")] // FIXME: Change issue # to track # + pub fn merge(&mut self, mut other: Self, conflict: impl FnMut(&K, V, V) -> V) where K: Ord, A: Clone, @@ -1296,14 +1300,14 @@ impl BTreeMap { // We can just swap `self` and `other` if `self` is empty. if self.is_empty() { - mem::swap(self, other); + mem::swap(self, &mut 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 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.append_from_sorted_iters_with( + root.merge_from_sorted_iters_with( self_iter, other_iter, &mut self.length, From d3ee0fc18a059ed9edc308a223eeba47b27284fd Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Thu, 5 Feb 2026 10:53:06 -0500 Subject: [PATCH 5/7] update issue number with tracking issue # --- library/alloc/src/collections/btree/map.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 49686f4cb9909..3e8023d6e0719 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1287,7 +1287,7 @@ impl BTreeMap { /// assert_eq!(a[&4], "e"); /// assert_eq!(a[&5], "f"); /// ``` - #[unstable(feature = "btree_merge", issue = "147700")] // FIXME: Change issue # to track # + #[unstable(feature = "btree_merge", issue = "152152")] pub fn merge(&mut self, mut other: Self, conflict: impl FnMut(&K, V, V) -> V) where K: Ord, From 2f926de8f4bc9b79ca7e2366308f15db048174fa Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Thu, 5 Feb 2026 12:10:16 -0500 Subject: [PATCH 6/7] fix doctest for merge --- library/alloc/src/collections/btree/map.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 3e8023d6e0719..f1fa1584a53c0 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1274,7 +1274,7 @@ impl BTreeMap { /// b.insert(5, String::from("f")); /// /// // concatenate a's value and b's value - /// a.merge(&mut b, |_, a_val, b_val| { + /// a.merge(b, |_, a_val, b_val| { /// format!("{a_val}{b_val}") /// }); /// From 1b290821b7d71756b7c8484f4d3ee0fd39f86371 Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Fri, 6 Feb 2026 09:35:03 -0500 Subject: [PATCH 7/7] Squashed doctests and merge test commits --- library/alloc/src/collections/btree/map.rs | 7 +- .../alloc/src/collections/btree/map/tests.rs | 88 ++++++++++++++++++- library/alloctests/tests/lib.rs | 1 + 3 files changed, 91 insertions(+), 5 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index f1fa1584a53c0..b33fd21a4a06f 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1243,12 +1243,12 @@ impl BTreeMap { /// 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`. For clarity, 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 merge `self`'s value with `other`'s value when their keys conflict. + /// 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. @@ -1278,8 +1278,7 @@ impl BTreeMap { /// format!("{a_val}{b_val}") /// }); /// - /// assert_eq!(a.len(), 5); - /// assert_eq!(b.len(), 0); + /// assert_eq!(a.len(), 5); // all of b's keys in a /// /// assert_eq!(a[&1], "a"); /// assert_eq!(a[&2], "b"); diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index 938e867b85812..1b07076142019 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -1,9 +1,9 @@ use core::assert_matches; -use std::iter; use std::ops::Bound::{Excluded, Included, Unbounded}; use std::panic::{AssertUnwindSafe, catch_unwind}; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering::SeqCst; +use std::{cmp, iter}; use super::*; use crate::boxed::Box; @@ -2128,6 +2128,76 @@ create_append_test!(test_append_239, 239); #[cfg(not(miri))] // Miri is too slow create_append_test!(test_append_1700, 1700); +macro_rules! create_merge_test { + ($name:ident, $len:expr) => { + #[test] + fn $name() { + let mut a = BTreeMap::new(); + for i in 0..8 { + a.insert(i, i); + } + + let mut b = BTreeMap::new(); + for i in 5..$len { + b.insert(i, 2 * i); + } + + a.merge(b, |_, a_val, b_val| a_val + b_val); + + assert_eq!(a.len(), cmp::max($len, 8)); + + for i in 0..cmp::max($len, 8) { + if i < 5 { + assert_eq!(a[&i], i); + } else { + if i < cmp::min($len, 8) { + assert_eq!(a[&i], i + 2 * i); + } else if i >= $len { + assert_eq!(a[&i], i); + } else { + assert_eq!(a[&i], 2 * i); + } + } + } + + a.check(); + assert_eq!( + a.remove(&($len - 1)), + if $len >= 5 && $len < 8 { + Some(($len - 1) + 2 * ($len - 1)) + } else { + Some(2 * ($len - 1)) + } + ); + assert_eq!(a.insert($len - 1, 20), None); + a.check(); + } + }; +} + +// These are mostly for testing the algorithm that "fixes" the right edge after insertion. +// Single node, merge conflicting key values. +create_merge_test!(test_merge_7, 7); +// Single node. +create_merge_test!(test_merge_9, 9); +// Two leafs that don't need fixing. +create_merge_test!(test_merge_17, 17); +// Two leafs where the second one ends up underfull and needs stealing at the end. +create_merge_test!(test_merge_14, 14); +// Two leafs where the second one ends up empty because the insertion finished at the root. +create_merge_test!(test_merge_12, 12); +// Three levels; insertion finished at the root. +create_merge_test!(test_merge_144, 144); +// Three levels; insertion finished at leaf while there is an empty node on the second level. +create_merge_test!(test_merge_145, 145); +// Tests for several randomly chosen sizes. +create_merge_test!(test_merge_170, 170); +create_merge_test!(test_merge_181, 181); +#[cfg(not(miri))] // Miri is too slow +create_merge_test!(test_merge_239, 239); +#[cfg(not(miri))] // Miri is too slow +create_merge_test!(test_merge_1700, 1700); + #[test] #[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")] fn test_append_drop_leak() { @@ -2615,3 +2685,19 @@ fn test_id_based_append() { assert_eq!(lhs.pop_first().unwrap().0.name, "lhs_k".to_string()); } + +#[test] +fn test_id_based_merge() { + let mut lhs = BTreeMap::new(); + let mut rhs = BTreeMap::new(); + + lhs.insert(IdBased { id: 0, name: "lhs_k".to_string() }, "1".to_string()); + rhs.insert(IdBased { id: 0, name: "rhs_k".to_string() }, "2".to_string()); + + lhs.merge(rhs, |_, mut lhs_val, rhs_val| { + lhs_val.push_str(&rhs_val); + lhs_val + }); + + assert_eq!(lhs.pop_first().unwrap().0.name, "lhs_k".to_string()); +} diff --git a/library/alloctests/tests/lib.rs b/library/alloctests/tests/lib.rs index e15c86496cf1b..e30bd26307fbf 100644 --- a/library/alloctests/tests/lib.rs +++ b/library/alloctests/tests/lib.rs @@ -1,5 +1,6 @@ #![feature(allocator_api)] #![feature(binary_heap_pop_if)] +#![feature(btree_merge)] #![feature(const_heap)] #![feature(deque_extend_front)] #![feature(iter_array_chunks)]