diff --git a/library/alloc/src/collections/btree/append.rs b/library/alloc/src/collections/btree/append.rs index 66ea22e75247c..cc8d793e98e4d 100644 --- a/library/alloc/src/collections/btree/append.rs +++ b/library/alloc/src/collections/btree/append.rs @@ -1,38 +1,8 @@ use core::alloc::Allocator; -use core::iter::FusedIterator; -use super::merge_iter::MergeIterInner; use super::node::{self, Root}; impl Root { - /// Appends all key-value pairs from the union of two ascending iterators, - /// incrementing a `length` variable along the way. The latter makes it - /// easier for the caller to avoid a leak when a drop handler panicks. - /// - /// If both iterators produce the same key, this method drops the pair from - /// the left iterator and appends the pair from the right iterator. - /// - /// If you want the tree to end up in a strictly ascending order, like for - /// a `BTreeMap`, both iterators should produce keys in strictly ascending - /// order, each greater than all keys in the tree, including any keys - /// already in the tree upon entry. - pub(super) fn append_from_sorted_iters( - &mut self, - left: I, - right: I, - length: &mut usize, - alloc: A, - ) where - K: Ord, - I: Iterator + FusedIterator, - { - // We prepare to merge `left` and `right` into a sorted sequence in linear time. - let iter = MergeIter(MergeIterInner::new(left, right)); - - // Meanwhile, we build a tree from the sorted sequence in linear time. - self.bulk_push(iter, length, alloc) - } - /// Pushes all key-value pairs to the end of the tree, incrementing a /// `length` variable along the way. The latter makes it easier for the /// caller to avoid a leak when the iterator panicks. @@ -94,24 +64,3 @@ impl Root { self.fix_right_border_of_plentiful(); } } - -// An iterator for merging two sorted sequences into one -struct MergeIter>(MergeIterInner); - -impl Iterator for MergeIter -where - I: Iterator + FusedIterator, -{ - type Item = (K, V); - - /// If two keys are equal, returns the key from the left and the value from the right. - fn next(&mut self) -> Option<(K, V)> { - let (a_next, b_next) = self.0.nexts(|a: &(K, V), b: &(K, V)| K::cmp(&a.0, &b.0)); - match (a_next, b_next) { - (Some((a_k, _)), Some((_, b_v))) => Some((a_k, b_v)), - (Some(a), None) => Some(a), - (None, Some(b)) => Some(b), - (None, None) => None, - } - } -} diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index d69dad70a44e9..81273e6f5b002 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1229,15 +1229,82 @@ impl BTreeMap { return; } - let self_iter = mem::replace(self, Self::new_in((*self.alloc).clone())).into_iter(); - let other_iter = mem::replace(other, Self::new_in((*self.alloc).clone())).into_iter(); - let root = self.root.get_or_insert_with(|| Root::new((*self.alloc).clone())); - root.append_from_sorted_iters( - self_iter, - other_iter, - &mut self.length, - (*self.alloc).clone(), - ) + // Because `other` takes a &mut Self, in order for us to own + // the K,V pairs, we need to use mem::replace + let mut other_iter = mem::replace(other, Self::new_in((*self.alloc).clone())).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(self_key, &first_other_key) { + Ordering::Equal => { + if let Some((_, v)) = self_cursor.next() { + *v = first_other_val; + } + } + Ordering::Greater => + // 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::Less => { + unreachable!("Cursor's peek_next should return None."); + } + } + } else { + // SAFETY: reaching here means our cursor is at the end + // self BTreeMap so we just insert other_key here + unsafe { + self_cursor.insert_before_unchecked(first_other_key, first_other_val); + } + } + + for (other_key, other_val) in other_iter { + loop { + if let Some((self_key, _)) = self_cursor.peek_next() { + match K::cmp(self_key, &other_key) { + Ordering::Equal => { + if let Some((_, v)) = self_cursor.next() { + *v = other_val; + } + break; + } + Ordering::Greater => { + // SAFETY: we know our self_key's ordering is greater than other_key, + // so inserting before will guarantee sorted order + unsafe { + self_cursor.insert_before_unchecked(other_key, other_val); + } + break; + } + Ordering::Less => { + // FIXME: instead of doing a linear search here, + // this can be optimized to search the tree by starting + // from self_cursor and going towards the root and then + // back down to the proper node -- that should probably + // be a new method on Cursor*. + self_cursor.next(); + } + } + } else { + // FIXME: If we get here, that means all of other's keys are greater than + // self's keys. For performance, this should really do a bulk insertion of items + // from other_iter into the end of self `BTreeMap`. Maybe this should be + // a method for Cursor*? + + // SAFETY: reaching here means our cursor is at the end + // self BTreeMap so we just insert other_key here + unsafe { + self_cursor.insert_before_unchecked(other_key, other_val); + } + break; + } + } + } } /// Moves all elements from `other` into `self`, leaving `other` empty. diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index 73546caa05eac..64348745aa07d 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -2224,7 +2224,7 @@ fn test_append_drop_leak() { catch_unwind(move || left.append(&mut right)).unwrap_err(); assert_eq!(a.dropped(), 1); - assert_eq!(b.dropped(), 1); // should be 2 were it not for Rust issue #47949 + assert_eq!(b.dropped(), 2); assert_eq!(c.dropped(), 2); }