Fix & test leak of some BTreeMap nodes on panic during into_iter

src/liballoc/collections/btree/map.rs

@@ -1477,6 +1477,14 @@ impl<K, V> Drop for IntoIter<K, V> {
  // 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() {}
+
+ // No need to avoid the shared root, because the tree was definitely not empty.
+ unsafe {
+ let mut node = ptr::read(&self.0.front).into_node().forget_type();
+ while let Some(parent) = node.deallocate_and_ascend() {
+ node = parent.into_node().forget_type();
+ }
+ }
  }
  }
 
@@ -1491,7 +1499,8 @@ impl<K, V> Drop for IntoIter<K, V> {
  if node.is_shared_root() {
  return;
  }
-
+ // Most of the nodes have been deallocated while traversing
+ // but one pile from a leaf up to the root is left standing.
  while let Some(parent) = node.deallocate_and_ascend() {
  node = parent.into_node().forget_type();
  }

src/liballoc/tests/btree/map.rs

@@ -1021,7 +1021,7 @@ fn test_split_off_large_random_sorted() {
 }
 
 #[test]
-fn test_into_iter_drop_leak() {
+fn test_into_iter_drop_leak_1() {
  static DROPS: AtomicU32 = AtomicU32::new(0);
 
  struct D;
@@ -1045,3 +1045,27 @@ fn test_into_iter_drop_leak() {
 
  assert_eq!(DROPS.load(Ordering::SeqCst), 5);
 }
+
+#[test]
+fn test_into_iter_drop_leak_2() {
+ let size = 12; // to obtain tree with 2 levels (having edges to leaf nodes)
+ static DROPS: AtomicU32 = AtomicU32::new(0);
+ static PANIC_POINT: AtomicU32 = AtomicU32::new(0);
+
+ struct D;
+ impl Drop for D {
+ fn drop(&mut self) {
+ if DROPS.fetch_add(1, Ordering::SeqCst) == PANIC_POINT.load(Ordering::SeqCst) {
+ panic!("panic in `drop`");
+ }
+ }
+ }
+
+ for panic_point in vec![0, 1, size - 2, size - 1] {
+ DROPS.store(0, Ordering::SeqCst);
+ PANIC_POINT.store(panic_point, Ordering::SeqCst);
+ let map: BTreeMap<_, _> = (0..size).map(|i| (i, D)).collect();
+ catch_unwind(move || drop(map.into_iter())).ok();
+ assert_eq!(DROPS.load(Ordering::SeqCst), size);
+ }
+}