diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index adde491020cf1..4b649e43371de 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -161,9 +161,7 @@ pub struct BTreeMap { #[stable(feature = "btree_drop", since = "1.7.0")] unsafe impl<#[may_dangle] K, #[may_dangle] V> Drop for BTreeMap { fn drop(&mut self) { - if let Some(root) = self.root.take() { - Dropper { front: root.into_dying().first_leaf_edge(), remaining_length: self.length }; - } + drop(unsafe { ptr::read(self) }.into_iter()) } } @@ -352,14 +350,6 @@ impl fmt::Debug for IntoIter { } } -/// A simplified version of `IntoIter` that is not double-ended and has only one -/// purpose: to drop the remainder of an `IntoIter`. Therefore it also serves to -/// drop an entire tree without the need to first look up a `back` leaf edge. -struct Dropper { - front: Handle, marker::Edge>, - remaining_length: usize, -} - /// An iterator over the keys of a `BTreeMap`. /// /// This `struct` is created by the [`keys`] method on [`BTreeMap`]. See its @@ -1458,47 +1448,57 @@ impl IntoIterator for BTreeMap { } } -impl Drop for Dropper { +#[stable(feature = "btree_drop", since = "1.7.0")] +impl Drop for IntoIter { fn drop(&mut self) { - // Similar to advancing a non-fusing iterator. - fn next_or_end( - this: &mut Dropper, - ) -> Option, marker::KV>> - { - if this.remaining_length == 0 { - unsafe { ptr::read(&this.front).deallocating_end() } - None - } else { - this.remaining_length -= 1; - Some(unsafe { this.front.deallocating_next_unchecked() }) - } - } - - struct DropGuard<'a, K, V>(&'a mut Dropper); + struct DropGuard<'a, K, V>(&'a mut IntoIter); 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(kv) = next_or_end(&mut self.0) { - kv.drop_key_val(); + while let Some(kv) = self.0.dying_next() { + // SAFETY: we consume the dying handle immediately. + unsafe { kv.drop_key_val() }; } } } - while let Some(kv) = next_or_end(self) { + while let Some(kv) = self.dying_next() { let guard = DropGuard(self); - kv.drop_key_val(); + // SAFETY: we don't touch the tree before consuming the dying handle. + unsafe { kv.drop_key_val() }; mem::forget(guard); } } } -#[stable(feature = "btree_drop", since = "1.7.0")] -impl Drop for IntoIter { - fn drop(&mut self) { - if let Some(front) = self.range.take_front() { - Dropper { front, remaining_length: self.length }; +impl IntoIter { + /// Core of a `next` method returning a dying KV handle, + /// invalidated by further calls to this function and some others. + fn dying_next( + &mut self, + ) -> Option, marker::KV>> { + if self.length == 0 { + self.range.deallocating_end(); + None + } else { + self.length -= 1; + Some(unsafe { self.range.deallocating_next_unchecked() }) + } + } + + /// Core of a `next_back` method returning a dying KV handle, + /// invalidated by further calls to this function and some others. + fn dying_next_back( + &mut self, + ) -> Option, marker::KV>> { + if self.length == 0 { + self.range.deallocating_end(); + None + } else { + self.length -= 1; + Some(unsafe { self.range.deallocating_next_back_unchecked() }) } } } @@ -1508,13 +1508,8 @@ impl Iterator for IntoIter { type Item = (K, V); fn next(&mut self) -> Option<(K, V)> { - if self.length == 0 { - None - } else { - self.length -= 1; - let kv = unsafe { self.range.deallocating_next_unchecked() }; - Some(kv.into_key_val()) - } + // SAFETY: we consume the dying handle immediately. + self.dying_next().map(unsafe { |kv| kv.into_key_val() }) } fn size_hint(&self) -> (usize, Option) { @@ -1525,13 +1520,8 @@ impl Iterator for IntoIter { #[stable(feature = "rust1", since = "1.0.0")] impl DoubleEndedIterator for IntoIter { fn next_back(&mut self) -> Option<(K, V)> { - if self.length == 0 { - None - } else { - self.length -= 1; - let kv = unsafe { self.range.deallocating_next_back_unchecked() }; - Some(kv.into_key_val()) - } + // SAFETY: we consume the dying handle immediately. + self.dying_next_back().map(unsafe { |kv| kv.into_key_val() }) } } diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs index 71edaab6ec263..7b1d4d68c4f56 100644 --- a/library/alloc/src/collections/btree/navigate.rs +++ b/library/alloc/src/collections/btree/navigate.rs @@ -167,8 +167,7 @@ impl<'a, K, V> LazyLeafRange, K, V> { } impl LazyLeafRange { - #[inline] - pub fn take_front( + fn take_front( &mut self, ) -> Option, marker::Edge>> { match self.front.take()? { @@ -194,6 +193,13 @@ impl LazyLeafRange { let back = self.init_back().unwrap(); unsafe { back.deallocating_next_back_unchecked() } } + + #[inline] + pub fn deallocating_end(&mut self) { + if let Some(front) = self.take_front() { + front.deallocating_end() + } + } } impl LazyLeafRange { @@ -488,7 +494,7 @@ impl Handle, marker::Edge> { /// both sides of the tree, and have hit the same edge. As it is intended /// only to be called when all keys and values have been returned, /// no cleanup is done on any of the keys or values. - pub fn deallocating_end(self) { + fn deallocating_end(self) { let mut edge = self.forget_node_type(); while let Some(parent_edge) = unsafe { edge.into_node().deallocate_and_ascend() } { edge = parent_edge.forget_node_type(); @@ -565,7 +571,7 @@ impl Handle, marker::Edge> { /// /// The only safe way to proceed with the updated handle is to compare it, drop it, /// or call this method or counterpart `deallocating_next_back_unchecked` again. - pub unsafe fn deallocating_next_unchecked( + unsafe fn deallocating_next_unchecked( &mut self, ) -> Handle, marker::KV> { super::mem::replace(self, |leaf_edge| unsafe { leaf_edge.deallocating_next().unwrap() }) diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index f6886b275e172..8f6a2ec9ebd1f 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -1058,7 +1058,9 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle, K, V, NodeType> impl Handle, marker::KV> { /// Extracts the key and value that the KV handle refers to. - pub fn into_key_val(mut self) -> (K, V) { + /// # Safety + /// The node that the handle refers to must not yet have been deallocated. + pub unsafe fn into_key_val(mut self) -> (K, V) { debug_assert!(self.idx < self.node.len()); let leaf = self.node.as_leaf_dying(); unsafe { @@ -1069,8 +1071,10 @@ impl Handle, marker::KV> } /// Drops the key and value that the KV handle refers to. + /// # Safety + /// The node that the handle refers to must not yet have been deallocated. #[inline] - pub fn drop_key_val(mut self) { + pub unsafe fn drop_key_val(mut self) { debug_assert!(self.idx < self.node.len()); let leaf = self.node.as_leaf_dying(); unsafe {