diff --git a/library/alloc/src/collections/binary_heap.rs b/library/alloc/src/collections/binary_heap.rs index 197e7aaaccf3d..519cf5d3582e1 100644 --- a/library/alloc/src/collections/binary_heap.rs +++ b/library/alloc/src/collections/binary_heap.rs @@ -144,7 +144,7 @@ #![stable(feature = "rust1", since = "1.0.0")] use core::fmt; -use core::iter::{FromIterator, FusedIterator, InPlaceIterable, SourceIter, TrustedLen}; +use core::iter::{FromIterator, FusedIterator, InPlaceIterable, Rev, SourceIter, TrustedLen}; use core::mem::{self, swap, ManuallyDrop}; use core::ops::{Deref, DerefMut}; use core::ptr; @@ -153,6 +153,7 @@ use crate::collections::TryReserveError; use crate::slice; use crate::vec::{self, AsVecIntoIter, Vec}; +use super::btree::borrow::DormantMutRef; use super::SpecExtend; #[cfg(test)] @@ -819,6 +820,33 @@ impl BinaryHeap { // data[0..first_removed] is untouched, so we only need to rebuild the tail: self.rebuild_tail(first_removed); } + + /// Returns a mutable iterator visiting all values in the underlying vector, in + /// arbitrary order. When `IterMut` is dropped, the binary heap is rebuilt. + /// + /// Note: If the `IterMut` value is leaked, the heap may be in an + /// inconsistent state. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(binary_heap_iter_mut)] + /// use std::collections::BinaryHeap; + /// + /// let mut heap = BinaryHeap::from([-10, -5, 1, 2, 4, 13]); + /// + /// heap.iter_mut().for_each(|x| *x += 1); + /// + /// assert_eq!(vec![-9, -4, 2, 3, 5, 14], heap.into_sorted_vec()); + /// ``` + #[unstable(feature = "binary_heap_iter_mut", issue = "98524")] + pub fn iter_mut(&mut self) -> IterMut<'_, T> { + let (this, dormant) = DormantMutRef::new(self); + let it = this.data.iter_mut().rev(); // reversed to take advantage of `rebuild_tail` on drop + IterMut { iter_mut: ManuallyDrop::new(it), heap: ManuallyDrop::new(dormant) } + } } impl BinaryHeap { @@ -1357,6 +1385,65 @@ impl ExactSizeIterator for Iter<'_, T> { #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Iter<'_, T> {} +/// A mutable iterator over the elements of a `BinaryHeap`. +/// +/// This `struct` is created by [`BinaryHeap::iter_mut()`]. See its +/// documentation for more. +/// +/// [`iter_mut`]: BinaryHeap::iter_mut +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[unstable(feature = "binary_heap_iter_mut", issue = "98524")] +pub struct IterMut<'a, T: 'a + Ord> { + iter_mut: ManuallyDrop>>, + heap: ManuallyDrop>>, +} + +#[unstable(feature = "binary_heap_iter_mut", issue = "98524")] +impl<'a, T: Ord> Iterator for IterMut<'a, T> { + type Item = &'a mut T; + + #[inline] + fn next(&mut self) -> Option<&'a mut T> { + self.iter_mut.next() + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.iter_mut.size_hint() + } +} + +#[unstable(feature = "binary_heap_iter_mut", issue = "98524")] +impl fmt::Debug for IterMut<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("IterMut").field(&self.iter_mut).finish() + } +} + +#[unstable(feature = "binary_heap_iter_mut", issue = "98524")] +impl Drop for IterMut<'_, T> { + fn drop(&mut self) { + let remaining = self.iter_mut.len(); + + // Early drop to avoid alias after recovering &mut BinaryHeap + // SAFETY: we're done using it. + unsafe { ManuallyDrop::drop(&mut self.iter_mut) }; + + // SAFETY: + // take: `self.heap` is not used again. + // awaken: `self.iter_mut` has been dropped and hence no references to the heap remain. + let heap = unsafe { ManuallyDrop::take(&mut self.heap).awaken() }; + heap.rebuild_tail(remaining); + } +} + +#[unstable(feature = "binary_heap_iter_mut", issue = "98524")] +impl ExactSizeIterator for IterMut<'_, T> { + fn is_empty(&self) -> bool { + self.iter_mut.is_empty() + } +} + /// An owning iterator over the elements of a `BinaryHeap`. /// /// This `struct` is created by [`BinaryHeap::into_iter()`] @@ -1653,6 +1740,16 @@ impl<'a, T> IntoIterator for &'a BinaryHeap { } } +#[stable(feature = "rust1", since = "1.0.0")] +impl<'a, T: Ord> IntoIterator for &'a mut BinaryHeap { + type Item = &'a mut T; + type IntoIter = IterMut<'a, T>; + + fn into_iter(self) -> IterMut<'a, T> { + self.iter_mut() + } +} + #[stable(feature = "rust1", since = "1.0.0")] impl Extend for BinaryHeap { #[inline] diff --git a/library/alloc/src/collections/binary_heap/tests.rs b/library/alloc/src/collections/binary_heap/tests.rs index 5a05215aeeddf..be2609e3e12c7 100644 --- a/library/alloc/src/collections/binary_heap/tests.rs +++ b/library/alloc/src/collections/binary_heap/tests.rs @@ -95,8 +95,9 @@ fn check_exact_size_iterator(len: usize, it: I) { #[test] fn test_exact_size_iterator() { - let heap = BinaryHeap::from(vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1]); + let mut heap = BinaryHeap::from(vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1]); check_exact_size_iterator(heap.len(), heap.iter()); + check_exact_size_iterator(heap.len(), heap.iter_mut()); check_exact_size_iterator(heap.len(), heap.clone().into_iter()); check_exact_size_iterator(heap.len(), heap.clone().into_iter_sorted()); check_exact_size_iterator(heap.len(), heap.clone().drain()); @@ -122,6 +123,28 @@ fn test_trusted_len() { check_trusted_len(heap.len(), heap.clone().drain_sorted()); } +#[test] +fn test_iter_mut() { + let mut heap = BinaryHeap::from([5, 9, 3]); + heap.iter_mut().for_each(|e| *e += 1); + assert_eq!(heap.into_vec(), vec![10, 6, 4]); + + let mut heap = BinaryHeap::from([5, 9, 3]); + *heap.iter_mut().next().unwrap() -= 10; + assert_eq!(heap.into_vec(), vec![9, 5, -7]); + + let mut heap = BinaryHeap::from([5, 9, 3]); + *heap.iter_mut().next().unwrap() += 10; + assert_eq!(heap.into_vec(), vec![13, 5, 9]); + + let mut heap = BinaryHeap::from([5, 9, 3]); + *heap.iter_mut().last().unwrap() += 10; + assert_eq!(heap.into_vec(), vec![19, 5, 3]); + + let mut heap = BinaryHeap::::new(); + assert_eq!(heap.iter_mut().next(), None); +} + #[test] fn test_peek_and_pop() { let data = vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1]; diff --git a/library/alloc/src/collections/btree/mod.rs b/library/alloc/src/collections/btree/mod.rs index 9d43ac5c5be59..0799b3cf7ef36 100644 --- a/library/alloc/src/collections/btree/mod.rs +++ b/library/alloc/src/collections/btree/mod.rs @@ -1,5 +1,5 @@ mod append; -mod borrow; +pub(super) mod borrow; mod dedup_sorted_iter; mod fix; pub mod map; diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs index 367cdcdcc061c..5eb25e547c274 100644 --- a/library/alloc/tests/lib.rs +++ b/library/alloc/tests/lib.rs @@ -23,6 +23,7 @@ #![feature(unboxed_closures)] #![feature(associated_type_bounds)] #![feature(binary_heap_into_iter_sorted)] +#![feature(binary_heap_iter_mut)] #![feature(binary_heap_drain_sorted)] #![feature(slice_ptr_get)] #![feature(binary_heap_retain)]