diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 109beed7de947..89e20e03afe83 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -205,16 +205,13 @@ pub struct BTreeMap< #[stable(feature = "btree_drop", since = "1.7.0")] unsafe impl<#[may_dangle] K, #[may_dangle] V, A: AllocatorClone> Drop for BTreeMap { fn drop(&mut self) { - // Skip `into_iter` for an empty map: `dying_next` is too costly to inline, so the - // empty drop isn't optimised away (see #161375). - if self.root.is_some() { - // SAFETY: `self` is not used after this and none of its fields are dropped again: - // `alloc` is `ManuallyDrop` and `root` has no drop glue. - drop(unsafe { ptr::read(self) }.into_iter()) + if let Some(root) = self.root.take() { + // SAFETY: The tree won't be used after this point. + // The allocator has not been moved out and will be moved into `drop_tree`. + unsafe { root.into_dying().drop_tree(ManuallyDrop::take(&mut self.alloc)) }; } else { - // SAFETY: With no root there are no nodes to free, so only the allocator needs - // dropping. `self` is not used after this, and `alloc` is dropped only here. - unsafe { ManuallyDrop::drop(&mut self.alloc) } + // SAFETY: The allocator has not been moved out and won't be used after this point. + unsafe { ManuallyDrop::drop(&mut self.alloc) }; } } } diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index a920254872e19..c4c9148e0abfb 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -2542,6 +2542,114 @@ fn test_into_iter_drop_leak_height_1() { } } +#[test] +fn test_drop_map_elements_multi_level() { + for size in [1, MIN_INSERTS_HEIGHT_1, MIN_INSERTS_HEIGHT_2] { + let keys = Vec::from_iter((0..size).map(CrashTestDummy::new)); + let vals = Vec::from_iter((0..size).map(CrashTestDummy::new)); + let map = { + let mut map = BTreeMap::new(); + + (0..size).map(|i| (keys[i].spawn(Panic::Never), vals[i].spawn(Panic::Never))).for_each( + |(k, v)| { + map.insert(k, v); + }, + ); + map + }; + + drop(map); + + for i in 0..size { + assert_eq!(keys[i].dropped(), 1); + assert_eq!(vals[i].dropped(), 1); + } + } +} + +#[test] +fn test_drop_empty_map_allocator() { + use core::alloc::{AllocError, Allocator, AllocatorClone, Layout}; + use core::cell::Cell; + use core::ptr::NonNull; + + struct CountedAllocator<'a>(&'a Cell); + + impl Clone for CountedAllocator<'_> { + fn clone(&self) -> Self { + Self(self.0) + } + } + + impl Drop for CountedAllocator<'_> { + fn drop(&mut self) { + self.0.set(self.0.get() + 1); + } + } + + unsafe impl Allocator for CountedAllocator<'_> { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + unsafe { Global.deallocate(ptr, layout) } + } + } + + unsafe impl AllocatorClone for CountedAllocator<'_> {} + + let dropped = Cell::new(0); + let map = BTreeMap::::new_in(CountedAllocator(&dropped)); + assert!(map.root.is_none()); + + drop(map); + + assert_eq!(dropped.get(), 1); +} + +#[test] +#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")] +fn test_drop_map_elements_multi_level_panic() { + for size in [1, MIN_INSERTS_HEIGHT_1, MIN_INSERTS_HEIGHT_2] { + for panic_point in [0, size / 2, size - 1] { + for panic_in_key in [true, false] { + let keys = Vec::from_iter((0..size).map(CrashTestDummy::new)); + let vals = Vec::from_iter((0..size).map(CrashTestDummy::new)); + let map = { + let mut map = BTreeMap::new(); + + (0..size) + .map(|i| { + let key_panic = if panic_in_key && i == panic_point { + Panic::InDrop + } else { + Panic::Never + }; + let val_panic = if !panic_in_key && i == panic_point { + Panic::InDrop + } else { + Panic::Never + }; + (keys[i].spawn(key_panic), vals[i].spawn(val_panic)) + }) + .for_each(|(k, v)| { + map.insert(k, v); + }); + map + }; + + catch_unwind(move || drop(map)).unwrap_err(); + + for i in 0..size { + assert_eq!(keys[i].dropped(), 1); + assert_eq!(vals[i].dropped(), 1); + } + } + } + } +} + #[test] fn test_into_keys() { let map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]); diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs index 28054cdeb5dd6..1f827d198468c 100644 --- a/library/alloc/src/collections/btree/navigate.rs +++ b/library/alloc/src/collections/btree/navigate.rs @@ -807,3 +807,70 @@ impl NodeRef NodeRef { + /// # Safety + /// - `self` must be the root. + /// - The tree shouldn't be accessed anymore. + pub(super) unsafe fn drop_tree(self, alloc: A) { + struct DropGuard { + node: Option>, + alloc: A, + } + + impl Drop for DropGuard { + fn drop(&mut self) { + if let Some(cur) = self.node.take() { + let mut next = + // SAFETY: When unwinding reached this guard, all kvs in cur + // have already been dropped in `drop_kvs`, while the node itself has not + // yet been deallocated. + unsafe { deallocate_and_next_postorder_node(cur, self.alloc.clone()) }; + + while let Some(mut node) = next { + // SAFETY: Every node only calls the function once + unsafe { node.drop_kvs() }; + next = + // SAFETY: `drop_kvs()` above only dropped kvs and the node is still accessible, + // and also each node is deallocated exactly once. + unsafe { deallocate_and_next_postorder_node(node, self.alloc.clone()) } + } + } + } + } + /// Deallocates `node` and returns the next node in postorder traversal + /// + /// # Safety + /// + /// The node must be initialized, all kvs in node must have already + /// been dropped or moved out and its children must have already been deallocated. + /// Also the node can't be touched anymore. + unsafe fn deallocate_and_next_postorder_node( + node: NodeRef, + alloc: A, + ) -> Option> { + // SAFETY: guaranteed by the caller + let parent = unsafe { node.deallocate_and_ascend(alloc) }?; + + Some(match parent.right_kv() { + Ok(kv) => kv.right_edge().descend().first_leaf_edge().into_node().forget_type(), + Err(last_edge) => last_edge.into_node().forget_type(), + }) + } + + let node = self.first_leaf_edge().into_node().forget_type(); + let mut guard = DropGuard { node: Some(node), alloc }; + loop { + // SAFETY: all kvs of each node are dropped once + unsafe { guard.node.as_mut().unwrap().drop_kvs() }; + // SAFETY: Nodes are deallocated in postorder, and each node is deallocated exactly once. + if let Some(next) = unsafe { + deallocate_and_next_postorder_node(guard.node.take().unwrap(), guard.alloc.clone()) + } { + guard.node = Some(next); + } else { + return; + } + } + } +} diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index 47c3fb1d79c9d..2144efc5d7d4d 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -429,6 +429,45 @@ impl NodeRef { } ret } + /// Drops all kvs in this node without deallocating the node itself + /// + /// # Safety + /// + /// - kvs in the node must be initialized and must not be accessed after this call. + /// - No kvs in the node can be accessed through any other path while this function is running. + pub(super) unsafe fn drop_kvs(&mut self) { + struct DropGuard<'a, K, V> { + node: &'a NodeRef, + idx: usize, + } + impl<'a, K, V> Drop for DropGuard<'a, K, V> { + fn drop(&mut self) { + let len = self.node.len(); + for i in self.idx..len { + // SAFETY: i < len and the kv has been initialized and + // it won't use the i-th kv more than once. + unsafe { + let kv = Handle::new_kv(ptr::read(self.node), i); + kv.drop_key_val(); + } + } + } + } + let len = self.len(); + for idx in 0..len { + // If `kv.drop_key_val()` panics, the guard would drop + // the rest and skip the panicked kv. + let guard = DropGuard { node: self, idx: idx + 1 }; + // SAFETY: idx < len and the kv has been initialized and + // it won't touch the idx-th more than once + unsafe { + let kv = Handle::new_kv(ptr::read(self), idx); + kv.drop_key_val(); + } + + core::mem::forget(guard); + } + } } impl<'a, K, V, Type> NodeRef, K, V, Type> {