diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs index 8b9ffdfb49b46..9da324ba2d4f1 100644 --- a/src/liballoc/collections/btree/map.rs +++ b/src/liballoc/collections/btree/map.rs @@ -1477,6 +1477,14 @@ impl Drop for IntoIter { // 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 Drop for IntoIter { 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(); } diff --git a/src/liballoc/tests/btree/map.rs b/src/liballoc/tests/btree/map.rs index fd07a4d3926c3..d05eec19346de 100644 --- a/src/liballoc/tests/btree/map.rs +++ b/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); + } +}