From 552cde7036651f401885a631821370ca4f4e81c2 Mon Sep 17 00:00:00 2001
From: Stein Somers <git@steinsomers.be>
Date: Mon, 26 Oct 2020 13:16:18 +0100
Subject: [PATCH] BTreeMap: revert a part of #78015 and test what it broke

---
 library/alloc/src/collections/btree/map.rs    | 30 ++++++++++++++-----
 .../alloc/src/collections/btree/map/tests.rs  | 27 +++++++++++++++++
 2 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
index 2704aab322924..1d12f5d9dae4d 100644
--- a/library/alloc/src/collections/btree/map.rs
+++ b/library/alloc/src/collections/btree/map.rs
@@ -2,14 +2,13 @@ use core::borrow::Borrow;
 use core::cmp::Ordering;
 use core::fmt::{self, Debug};
 use core::hash::{Hash, Hasher};
-use core::iter::{FromIterator, FusedIterator};
+use core::iter::{FromIterator, FusedIterator, Peekable};
 use core::marker::PhantomData;
 use core::mem::{self, ManuallyDrop};
 use core::ops::{Index, RangeBounds};
 use core::ptr;
 
 use super::borrow::DormantMutRef;
-use super::merge_iter::MergeIterInner;
 use super::node::{self, marker, ForceResult::*, Handle, NodeRef};
 use super::search::{self, SearchResult::*};
 use super::unwrap_unchecked;
@@ -459,7 +458,10 @@ impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for RangeMut<'_, K, V> {
 }
 
 // An iterator for merging two sorted sequences into one
-struct MergeIter<K, V, I: Iterator<Item = (K, V)>>(MergeIterInner<I>);
+struct MergeIter<K, V, I: Iterator<Item = (K, V)>> {
+    left: Peekable<I>,
+    right: Peekable<I>,
+}
 
 impl<K: Ord, V> BTreeMap<K, V> {
     /// Makes a new empty BTreeMap.
@@ -911,7 +913,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
         // First, we merge `self` and `other` into a sorted sequence in linear time.
         let self_iter = mem::take(self).into_iter();
         let other_iter = mem::take(other).into_iter();
-        let iter = MergeIter(MergeIterInner::new(self_iter, other_iter));
+        let iter = MergeIter { left: self_iter.peekable(), right: other_iter.peekable() };
 
         // Second, we build a tree from the sorted sequence in linear time.
         self.from_sorted_iter(iter);
@@ -2224,10 +2226,24 @@ where
 {
     type Item = (K, V);
 
-    /// If two keys are equal, returns the key/value-pair from the right source.
     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));
-        b_next.or(a_next)
+        let res = match (self.left.peek(), self.right.peek()) {
+            (Some(&(ref left_key, _)), Some(&(ref right_key, _))) => left_key.cmp(right_key),
+            (Some(_), None) => Ordering::Less,
+            (None, Some(_)) => Ordering::Greater,
+            (None, None) => return None,
+        };
+
+        // Check which elements comes first and only advance the corresponding iterator.
+        // If two keys are equal, take the value from `right`.
+        match res {
+            Ordering::Less => self.left.next(),
+            Ordering::Greater => self.right.next(),
+            Ordering::Equal => {
+                self.left.next();
+                self.right.next()
+            }
+        }
     }
 }
 
diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs
index adb94972f5bb6..ca90c0b592027 100644
--- a/library/alloc/src/collections/btree/map/tests.rs
+++ b/library/alloc/src/collections/btree/map/tests.rs
@@ -1660,6 +1660,33 @@ create_append_test!(test_append_239, 239);
 #[cfg(not(miri))] // Miri is too slow
 create_append_test!(test_append_1700, 1700);
 
+#[test]
+fn test_append_drop_leak() {
+    static DROPS: AtomicUsize = AtomicUsize::new(0);
+
+    struct D;
+
+    impl Drop for D {
+        fn drop(&mut self) {
+            if DROPS.fetch_add(1, Ordering::SeqCst) == 0 {
+                panic!("panic in `drop`");
+            }
+        }
+    }
+
+    let mut left = BTreeMap::new();
+    let mut right = BTreeMap::new();
+    left.insert(0, D);
+    left.insert(1, D); // first to be dropped during append
+    left.insert(2, D);
+    right.insert(1, D);
+    right.insert(2, D);
+
+    catch_unwind(move || left.append(&mut right)).unwrap_err();
+
+    assert_eq!(DROPS.load(Ordering::SeqCst), 5);
+}
+
 fn rand_data(len: usize) -> Vec<(u32, u32)> {
     let mut rng = DeterministicRng::new();
     Vec::from_iter((0..len).map(|_| (rng.next(), rng.next())))