Fix leak in btree_map::IntoIter when drop panics

Author: jonas-schievink

src/liballoc/collections/btree/map.rs

@@ -1381,7 +1381,22 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {
 #[stable(feature = "btree_drop", since = "1.7.0")]
 impl<K, V> Drop for IntoIter<K, V> {
     fn drop(&mut self) {
-        self.for_each(drop);
+        struct DropGuard<'a, K, V>(&'a mut IntoIter<K, V>);
+
+        impl<'a, K, V> Drop for DropGuard<'a, K, V> {
+            fn drop(&mut self) {
+                // Continue the same loop we perform below. This only runs when unwinding, so we
+                // don't have to care about panics this time (they'll abort).
+                while let Some(_) = self.0.next() {}
+            }
+        }
+
+        while let Some(pair) = self.next() {
+            let guard = DropGuard(self);
+            drop(pair);
+            mem::forget(guard);
+        }
+
         unsafe {
             let leaf_node = ptr::read(&self.front).into_node();
             if leaf_node.is_shared_root() {

src/liballoc/tests/btree/map.rs

@@ -5,7 +5,9 @@ use std::fmt::Debug;
 use std::iter::FromIterator;
 use std::ops::Bound::{self, Excluded, Included, Unbounded};
 use std::ops::RangeBounds;
+use std::panic::catch_unwind;
 use std::rc::Rc;
+use std::sync::atomic::{AtomicU32, Ordering};
 
 use super::DeterministicRng;
 
@@ -980,3 +982,29 @@ fn test_split_off_large_random_sorted() {
     assert!(map.into_iter().eq(data.clone().into_iter().filter(|x| x.0 < key)));
     assert!(right.into_iter().eq(data.into_iter().filter(|x| x.0 >= key)));
 }
+
+#[test]
+fn test_into_iter_drop_leak() {
+    static DROPS: AtomicU32 = AtomicU32::new(0);
+
+    struct D;
+
+    impl Drop for D {
+        fn drop(&mut self) {
+            if DROPS.fetch_add(1, Ordering::SeqCst) == 3 {
+                panic!("panic in `drop`");
+            }
+        }
+    }
+
+    let mut map = BTreeMap::new();
+    map.insert("a", D);
+    map.insert("b", D);
+    map.insert("c", D);
+    map.insert("d", D);
+    map.insert("e", D);
+
+    catch_unwind(move || drop(map.into_iter())).ok();
+
+    assert_eq!(DROPS.load(Ordering::SeqCst), 5);
+}