From 6a2f234df57f67e382550979a26951ea603f9727 Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Sun, 1 Feb 2026 23:13:55 -0500 Subject: [PATCH 01/14] 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 02/14] 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 03/14] 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 04/14] 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 05/14] 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 06/14] 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 07/14] 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)] From a7fa7dd487fd34ac8e7802ac77a5dbbe918dea07 Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Mon, 9 Feb 2026 21:44:34 -0500 Subject: [PATCH 08/14] optimized version of BTreeMap::merge with CursorMut and unsafe code --- library/alloc/src/collections/btree/append.rs | 60 ----------- library/alloc/src/collections/btree/map.rs | 102 ++++++++++++++++-- 2 files changed, 91 insertions(+), 71 deletions(-) diff --git a/library/alloc/src/collections/btree/append.rs b/library/alloc/src/collections/btree/append.rs index 38e4542d07905..66ea22e75247c 100644 --- a/library/alloc/src/collections/btree/append.rs +++ b/library/alloc/src/collections/btree/append.rs @@ -33,36 +33,6 @@ impl Root { self.bulk_push(iter, length, alloc) } - /// 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 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 merge_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. @@ -145,33 +115,3 @@ 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 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) { - (Some((a_k, a_v)), Some((_, b_v))) => Some({ - let next_val = (self.f)(&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 b33fd21a4a06f..ce1ad60221889 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[&5], "f"); /// ``` #[unstable(feature = "btree_merge", issue = "152152")] - pub fn merge(&mut self, mut other: Self, conflict: impl FnMut(&K, V, V) -> V) + pub fn merge(&mut self, mut other: Self, mut conflict: impl FnMut(&K, V, V) -> V) where K: Ord, A: Clone, @@ -1303,16 +1303,96 @@ impl BTreeMap { return; } - let self_iter = mem::replace(self, 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.merge_from_sorted_iters_with( - self_iter, - other_iter, - &mut self.length, - (*self.alloc).clone(), - conflict, - ) + 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 the first key from other + let mut self_cursor = self.lower_bound_mut(Bound::Included(&first_other_key)); + + if let Some((self_key, self_val)) = self_cursor.peek_next() { + match K::cmp(&first_other_key, self_key) { + Ordering::Equal => { + // SAFETY: We read in self_val's and hand it over to our conflict function + // which will always return a value that we can use to overwrite what's + // in self_val + unsafe { + let val = ptr::read(self_val); + let next_val = (conflict)(self_key, val, first_other_val); + ptr::write(self_val, next_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 would return None in this case"); + } + } + } 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_val)) = self_entry { + match K::cmp(&other_key, self_key) { + Ordering::Equal => { + // SAFETY: We read in self_val's and hand it over to our conflict function + // which will always return a value that we can use to overwrite what's + // in self_val + unsafe { + let val = ptr::read(self_val); + let next_val = (conflict)(self_key, val, other_val); + ptr::write(self_val, next_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 => { + 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. From 54499fe5f1735f864b6851f484420658bd2376df Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Tue, 10 Feb 2026 20:30:52 -0500 Subject: [PATCH 09/14] with_next() and forget_next_*() created for Cursor* to handle panic behavior for conflict func --- library/alloc/src/collections/btree/map.rs | 224 ++++++++++++++++-- .../alloc/src/collections/btree/map/tests.rs | 39 +++ 2 files changed, 239 insertions(+), 24 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index ce1ad60221889..2b512c98caefb 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1306,29 +1306,25 @@ impl BTreeMap { 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 the first key from other + // 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_val)) = self_cursor.peek_next() { + if let Some((self_key, _)) = self_cursor.peek_next() { match K::cmp(&first_other_key, self_key) { Ordering::Equal => { - // SAFETY: We read in self_val's and hand it over to our conflict function - // which will always return a value that we can use to overwrite what's - // in self_val - unsafe { - let val = ptr::read(self_val); - let next_val = (conflict)(self_key, val, first_other_val); - ptr::write(self_val, next_val); - } + 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, + // 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 would return None in this case"); + unreachable!("Cursor's peek_next should return None."); } } } else { @@ -1345,21 +1341,16 @@ impl BTreeMap { if self_cursor.peek_next().is_some() { loop { let self_entry = self_cursor.peek_next(); - if let Some((self_key, self_val)) = self_entry { + if let Some((self_key, _)) = self_entry { match K::cmp(&other_key, self_key) { Ordering::Equal => { - // SAFETY: We read in self_val's and hand it over to our conflict function - // which will always return a value that we can use to overwrite what's - // in self_val - unsafe { - let val = ptr::read(self_val); - let next_val = (conflict)(self_key, val, other_val); - ptr::write(self_val, next_val); - } + 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, + // 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); @@ -1367,11 +1358,12 @@ impl BTreeMap { break; } Ordering::Greater => { + // advance the cursor forward self_cursor.next(); } } } else { - // SAFETY: if we reach here, that means our cursor has reached + // 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 @@ -1383,7 +1375,7 @@ impl BTreeMap { } } } else { - // SAFETY: if we reach here, that means our cursor has reached + // 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 @@ -3417,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 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. /// @@ -3622,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 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. /// diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index 1b07076142019..cb0e83395a1bc 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -2239,6 +2239,45 @@ fn test_append_ord_chaos() { map2.check(); } +#[test] +#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")] +fn test_merge_drop_leak() { + let a = CrashTestDummy::new(0); + let b = CrashTestDummy::new(1); + let c = CrashTestDummy::new(2); + let mut left = BTreeMap::new(); + let mut right = BTreeMap::new(); + left.insert(a.spawn(Panic::Never), ()); + left.insert(b.spawn(Panic::Never), ()); + left.insert(c.spawn(Panic::Never), ()); + right.insert(b.spawn(Panic::InDrop), ()); // first duplicate key, dropped during append + right.insert(c.spawn(Panic::Never), ()); + + catch_unwind(move || left.merge(right, |_, _, _| ())).unwrap_err(); + assert_eq!(a.dropped(), 1); + assert_eq!(b.dropped(), 2); + assert_eq!(c.dropped(), 2); +} + +#[test] +fn test_merge_ord_chaos() { + let mut map1 = BTreeMap::new(); + map1.insert(Cyclic3::A, ()); + map1.insert(Cyclic3::B, ()); + let mut map2 = BTreeMap::new(); + map2.insert(Cyclic3::A, ()); + map2.insert(Cyclic3::B, ()); + map2.insert(Cyclic3::C, ()); // lands first, before A + map2.insert(Cyclic3::B, ()); // lands first, before C + map1.check(); + map2.check(); // keys are not unique but still strictly ascending + assert_eq!(map1.len(), 2); + assert_eq!(map2.len(), 4); + map1.merge(map2, |_, _, _| ()); + assert_eq!(map1.len(), 5); + map1.check(); +} + fn rand_data(len: usize) -> Vec<(u32, u32)> { let mut rng = DeterministicRng::new(); Vec::from_iter((0..len).map(|_| (rng.next(), rng.next()))) From 97f554744ad5a4c5462b1fb6ad306e36fc083174 Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Tue, 10 Feb 2026 23:23:54 -0500 Subject: [PATCH 10/14] forgot to add mem::forget() on key --- library/alloc/src/collections/btree/map.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 2b512c98caefb..2247204880e6b 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1328,7 +1328,7 @@ impl BTreeMap { } } } else { - // SAFETY: if we reach here, that means our cursor has reached + // 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 @@ -3488,9 +3488,10 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMutKey<'a, K, V, A> { let mut emptied_internal_root = false; if let Ok(next_kv) = current.next_kv() { - let ((_, val), pos) = + let ((key, val), pos) = next_kv.remove_kv_tracking(|| emptied_internal_root = true, self.alloc.clone()); - // Don't run destructor on next element's value + // Don't run destructor on next element's key and value + mem::forget(key); mem::forget(val); self.current = Some(pos); *self.length -= 1; From 9b423c524433e8e9e3cdc332dc5a2be181fcc63b Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Thu, 12 Feb 2026 16:40:15 -0500 Subject: [PATCH 11/14] Changed with_next() to use remove_next() + insert_after_unchecked() instead of drop handling, refactored code a bit, and updated comments --- library/alloc/src/collections/btree/map.rs | 264 ++++-------------- .../alloc/src/collections/btree/map/tests.rs | 2 + 2 files changed, 60 insertions(+), 206 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 2247204880e6b..4349e64603d07 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1318,7 +1318,7 @@ impl BTreeMap { }); } Ordering::Less => - // Safety: we know our other_key's ordering is less than self_key, + // 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); @@ -1328,61 +1328,45 @@ impl BTreeMap { } } } 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 + // SAFETY: reaching here means our cursor is at the end + // self BTreeMap so we just insert other_key here 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(); + 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; } - } 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); + Ordering::Greater => { + // advance the cursor forward + self_cursor.next(); } - self_cursor.next(); - break; } + } else { + // SAFETY: reaching here means our cursor is at the end + // self BTreeMap so we just insert other_key here + 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(); } } } @@ -3416,91 +3400,27 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMutKey<'a, K, V, A> { /// /// 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)) { - /// 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 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) }; + // 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 + // to CursorMutKey, ManuallyDrop> (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 { 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 ((key, val), pos) = - next_kv.remove_kv_tracking(|| emptied_internal_root = true, self.alloc.clone()); - // Don't run destructor on next element's key and value - mem::forget(key); - 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()); - } + unsafe { self.insert_after_unchecked(k, v) }; } } @@ -3717,85 +3637,17 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { /// 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 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()); - } + // 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 CursorMut + // to CursorMut> (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 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) }; } } diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index cb0e83395a1bc..9e4ff7cbc57c6 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -2128,6 +2128,8 @@ create_append_test!(test_append_239, 239); #[cfg(not(miri))] // Miri is too slow create_append_test!(test_append_1700, 1700); +// On merging keys 5..8, the values are +// added together (i + 2 * i) macro_rules! create_merge_test { ($name:ident, $len:expr) => { #[test] From 7a6c065f2de24e930bc1f25f909ced85ae30bfe2 Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Thu, 12 Feb 2026 19:00:26 -0500 Subject: [PATCH 12/14] use insert_before instead of insert_after and make note on optimizing searching through tree --- library/alloc/src/collections/btree/map.rs | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 4349e64603d07..c3f982fcb8bd9 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1331,7 +1331,7 @@ impl BTreeMap { // SAFETY: reaching here means our cursor is at the end // self BTreeMap so we just insert other_key here unsafe { - self_cursor.insert_after_unchecked(first_other_key, first_other_val); + self_cursor.insert_before_unchecked(first_other_key, first_other_val); } } @@ -1354,7 +1354,11 @@ impl BTreeMap { break; } Ordering::Greater => { - // advance the cursor forward + // 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(); } } @@ -1362,9 +1366,8 @@ impl BTreeMap { // SAFETY: reaching here means our cursor is at the end // self BTreeMap so we just insert other_key here unsafe { - self_cursor.insert_after_unchecked(other_key, other_val); + self_cursor.insert_before_unchecked(other_key, other_val); } - self_cursor.next(); break; } } @@ -3637,11 +3640,12 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> { /// 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) { - // 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 CursorMut - // to CursorMut> (see PR #152418) + // 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 From f37561d5f189738f64d0879788be4ea9b1377ed7 Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Fri, 13 Feb 2026 20:24:41 -0500 Subject: [PATCH 13/14] added a test to see if KV pairs are dropped on panic within conflict function --- .../alloc/src/collections/btree/map/tests.rs | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index 9e4ff7cbc57c6..d526ff2672686 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -2261,6 +2261,45 @@ fn test_merge_drop_leak() { assert_eq!(c.dropped(), 2); } +#[test] +#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")] +fn test_merge_conflict_drop_leak() { + let a = CrashTestDummy::new(0); + let a_val_left = CrashTestDummy::new(0); + + let b = CrashTestDummy::new(1); + let b_val_left = CrashTestDummy::new(1); + let b_val_right = CrashTestDummy::new(1); + + let c = CrashTestDummy::new(2); + let c_val_left = CrashTestDummy::new(2); + let c_val_right = CrashTestDummy::new(2); + + let mut left = BTreeMap::new(); + let mut right = BTreeMap::new(); + + left.insert(a.spawn(Panic::Never), a_val_left.spawn(Panic::Never)); + left.insert(b.spawn(Panic::Never), b_val_left.spawn(Panic::Never)); + left.insert(c.spawn(Panic::Never), c_val_left.spawn(Panic::Never)); + right.insert(b.spawn(Panic::Never), b_val_right.spawn(Panic::Never)); + right.insert(c.spawn(Panic::Never), c_val_right.spawn(Panic::Never)); + + // First key that conflicts should + catch_unwind(move || { + left.merge(right, |_, _, _| panic!("Panic in conflict function")); + assert_eq!(left.len(), 1); // only 1 entry should be left + }) + .unwrap_err(); + assert_eq!(a.dropped(), 1); // should not panic + assert_eq!(a_val_left.dropped(), 1); // should not panic + assert_eq!(b.dropped(), 2); // should drop from panic (conflict) + assert_eq!(b_val_left.dropped(), 1); // should be 2 were it not for Rust issue #47949 + assert_eq!(b_val_right.dropped(), 1); // should be 2 were it not for Rust issue #47949 + assert_eq!(c.dropped(), 2); // should drop from panic (conflict) + assert_eq!(c_val_left.dropped(), 1); // should be 2 were it not for Rust issue #47949 + assert_eq!(c_val_right.dropped(), 1); // should be 2 were it not for Rust issue #47949 +} + #[test] fn test_merge_ord_chaos() { let mut map1 = BTreeMap::new(); From a67cc730c42b77d92783fb10e4457a0f4457989c Mon Sep 17 00:00:00 2001 From: Mahdi Ali-Raihan Date: Sun, 15 Feb 2026 17:20:28 -0500 Subject: [PATCH 14/14] added extra checks to test_id_based_merge test and clarified what merge_test macro was doing --- .../alloc/src/collections/btree/map/tests.rs | 28 ++++++++++++++----- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index d526ff2672686..73546caa05eac 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -2128,8 +2128,16 @@ create_append_test!(test_append_239, 239); #[cfg(not(miri))] // Miri is too slow create_append_test!(test_append_1700, 1700); -// On merging keys 5..8, the values are -// added together (i + 2 * i) +// a inserts (0, 0)..(8, 8) to its own tree +// b inserts (5, 5 * 2)..($len, 2 * $len) to its own tree +// note that between a and b, there are duplicate keys +// between 5..min($len, 8), so on merge we add the values +// of these keys together +// we check that: +// - the merged tree 'a' has a length of max(8, $len) +// - all keys in 'a' have the correct value associated +// - removing and inserting an element into the merged +// tree 'a' still keeps it in valid tree form macro_rules! create_merge_test { ($name:ident, $len:expr) => { #[test] @@ -2252,13 +2260,13 @@ fn test_merge_drop_leak() { left.insert(a.spawn(Panic::Never), ()); left.insert(b.spawn(Panic::Never), ()); left.insert(c.spawn(Panic::Never), ()); - right.insert(b.spawn(Panic::InDrop), ()); // first duplicate key, dropped during append + right.insert(b.spawn(Panic::InDrop), ()); // first duplicate key, dropped during merge right.insert(c.spawn(Panic::Never), ()); catch_unwind(move || left.merge(right, |_, _, _| ())).unwrap_err(); - assert_eq!(a.dropped(), 1); - assert_eq!(b.dropped(), 2); - assert_eq!(c.dropped(), 2); + assert_eq!(a.dropped(), 1); // this should not be dropped + assert_eq!(b.dropped(), 2); // key is dropped on panic + assert_eq!(c.dropped(), 2); // key is dropped on panic } #[test] @@ -2775,9 +2783,15 @@ fn test_id_based_merge() { rhs.insert(IdBased { id: 0, name: "rhs_k".to_string() }, "2".to_string()); lhs.merge(rhs, |_, mut lhs_val, rhs_val| { + // confirming that lhs_val comes from lhs tree, + // rhs_val comes from rhs tree + assert_eq!(lhs_val, String::from("1")); + assert_eq!(rhs_val, String::from("2")); lhs_val.push_str(&rhs_val); lhs_val }); - assert_eq!(lhs.pop_first().unwrap().0.name, "lhs_k".to_string()); + let merged_kv_pair = lhs.pop_first().unwrap(); + assert_eq!(merged_kv_pair.0.id, 0); + assert_eq!(merged_kv_pair.0.name, "lhs_k".to_string()); }