diff --git a/library/alloc/src/collections/binary_heap.rs b/library/alloc/src/collections/binary_heap.rs index 4583bc9a158ef..a420d9b4af9aa 100644 --- a/library/alloc/src/collections/binary_heap.rs +++ b/library/alloc/src/collections/binary_heap.rs @@ -398,6 +398,9 @@ impl BinaryHeap { /// Returns a mutable reference to the greatest item in the binary heap, or /// `None` if it is empty. /// + /// # Logic + /// This method assumes that [`BinaryHeap`]'s heap invariant is upheld. + /// /// Note: If the `PeekMut` value is leaked, the heap may be in an /// inconsistent state. /// @@ -432,6 +435,8 @@ impl BinaryHeap { /// Removes the greatest item from the binary heap and returns it, or `None` if it /// is empty. /// + /// # Logic + /// This method assumes that [`BinaryHeap`]'s heap invariant is upheld. /// # Examples /// /// Basic usage: @@ -462,6 +467,8 @@ impl BinaryHeap { /// Pushes an item onto the binary heap. /// + /// # Logic + /// This method assumes that [`BinaryHeap`]'s heap invariant is upheld. /// # Examples /// /// Basic usage: @@ -504,6 +511,8 @@ impl BinaryHeap { /// Consumes the `BinaryHeap` and returns a vector in sorted /// (ascending) order. /// + /// # Logic + /// This method assumes that `BinaryHeap`'s heap invariant is upheld. /// # Examples /// /// Basic usage: @@ -710,7 +719,42 @@ impl BinaryHeap { } } - fn rebuild(&mut self) { + /// Rebuilds the [`BinaryHeap`] into a valid state. + /// This conversion happens in-place, and has *O*(*n*) time complexity. + /// Commonly used after operations such as [`BinaryHeap::drain_filter`] that violate [`BinaryHeap`]'s heap invariant, + /// or methods that allow interior mutability such as [`Cell`][core::cell::Cell]. + /// + /// # Examples + /// Mutates the heap interiorly using [`Cell`][core::cell::Cell] and restores it afterwards with `rebuild`. + /// + /// ``` + /// #![feature(binary_heap_rebuild)] + /// #![feature(binary_heap_into_iter_sorted)] + /// use std::{ + /// collections::BinaryHeap, + /// cell::Cell, + /// }; + /// + /// let mut a = BinaryHeap::from(vec![Cell::new(0), Cell::new(1), Cell::new(2), Cell::new(3)]); + /// + /// let sorted_values = |heap: &BinaryHeap>| { + /// heap.clone() + /// // this method assumes the heap is in a valid state. + /// .into_iter_sorted() + /// .map(|x| x.get()) + /// .collect::>() + /// }; + /// + /// // internal mutation invalidates the heap order and so the sort fails + /// a.peek().unwrap().set(0); + /// assert_eq!(sorted_values(&a), [0, 2, 1, 0]); + /// + /// // the heap is rebuilt to a valid state and so the sort works + /// a.rebuild(); + /// assert_eq!(sorted_values(&a), [2, 1, 0, 0]); + /// ``` + #[unstable(feature = "binary_heap_rebuild", reason = "recently added", issue = "none")] + pub fn rebuild(&mut self) { let mut n = self.len() / 2; while n > 0 { n -= 1; @@ -723,6 +767,8 @@ impl BinaryHeap { /// Moves all the elements of `other` into `self`, leaving `other` empty. /// + /// # Logic + /// This method assumes that [`BinaryHeap`]'s heap invariant is upheld. /// # Examples /// /// Basic usage: @@ -762,6 +808,8 @@ impl BinaryHeap { /// * `.drain_sorted()` is *O*(*n* \* log(*n*)); much slower than `.drain()`. /// You should use the latter for most cases. /// + /// # Logic + /// This method assumes that [`BinaryHeap`]’s heap invariant is upheld. /// # Examples /// /// Basic usage: @@ -787,6 +835,8 @@ impl BinaryHeap { /// In other words, remove all elements `e` for which `f(&e)` returns /// `false`. The elements are visited in unsorted (and unspecified) order. /// + /// # Logic + /// This method assumes that [`BinaryHeap`]'s heap invariant is upheld. /// # Examples /// /// Basic usage: @@ -846,6 +896,8 @@ impl BinaryHeap { /// Returns an iterator which retrieves elements in heap order. /// This method consumes the original heap. /// + /// # Logic + /// This method assumes that [`BinaryHeap`]’s heap invariant is upheld. /// # Examples /// /// Basic usage: @@ -864,6 +916,8 @@ impl BinaryHeap { /// Returns the greatest item in the binary heap, or `None` if it is empty. /// + /// # Logic + /// This method assumes that [`BinaryHeap`]’s heap invariant is upheld. /// # Examples /// /// Basic usage: @@ -1220,6 +1274,34 @@ impl BinaryHeap { pub fn clear(&mut self) { self.drain(); } + + /// Creates an iterator which uses a closure to determine if an element should be removed from the underlying vec. + /// If the closure returns true, then the element is removed and yielded. If the closure returns false, the element will remain in the binary heap and will not be yielded by the iterator. + /// # Logic + /// This operation violates [`BinaryHeap`]'s heap invariant. + /// It is necessary to call [`BinaryHeap::rebuild`] afterwards if you plan to use methods that assume a heap invariant. + /// # Examples + /// Splitting the binary heap into evens and odds, reusing the original allocation: + /// + /// ``` + /// #![feature(binary_heap_drain_filter)] + /// #![feature(binary_heap_rebuild)] + /// use std::collections::BinaryHeap; + /// let mut a = BinaryHeap::from(vec![1, 2, 3, 4, 5]); + /// let mut evens = a.drain_filter(|x| *x % 2 == 0).collect::>(); + /// evens.sort(); + /// a.rebuild(); //restores heap + /// let odds = a.into_sorted_vec(); + /// assert_eq!(evens, vec![2, 4]); + /// assert_eq!(odds, vec![1, 3, 5]); + /// ``` + #[unstable(feature = "binary_heap_drain_filter", reason = "recently added", issue = "42849")] + pub fn drain_filter(&mut self, filter: F) -> DrainFilter<'_, T, F> + where + F: FnMut(&mut T) -> bool, + { + DrainFilter { inner: self.data.drain_filter(filter) } + } } /// Hole represents a hole in a slice i.e., an index without valid value @@ -1570,6 +1652,36 @@ impl FusedIterator for DrainSorted<'_, T> {} #[unstable(feature = "trusted_len", issue = "37572")] unsafe impl TrustedLen for DrainSorted<'_, T> {} +/// An iterator which uses a closure to determine if an element should be removed from the underlying vec. +/// +/// This `struct` is created by [`BinaryHeap::drain_filter()`]. See its +/// documentation for more. +/// +/// [`drain_filter`]: BinaryHeap::drain_filter +#[unstable(feature = "binary_heap_drain_filter", reason = "recently added", issue = "42849")] +#[derive(Debug)] +pub struct DrainFilter<'a, T, F> +where + F: FnMut(&mut T) -> bool, +{ + inner: crate::vec::DrainFilter<'a, T, F>, +} + +#[unstable(feature = "binary_heap_drain_filter", reason = "recently added", issue = "42849")] +impl Iterator for DrainFilter<'_, T, F> +where + F: FnMut(&mut T) -> bool, +{ + type Item = T; + fn next(&mut self) -> Option { + self.inner.next() + } + + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() + } +} + #[stable(feature = "binary_heap_extras_15", since = "1.5.0")] impl From> for BinaryHeap { /// Converts a `Vec` into a `BinaryHeap`. diff --git a/library/alloc/src/collections/binary_heap/tests.rs b/library/alloc/src/collections/binary_heap/tests.rs index 5a05215aeeddf..b487dbbb22be6 100644 --- a/library/alloc/src/collections/binary_heap/tests.rs +++ b/library/alloc/src/collections/binary_heap/tests.rs @@ -1,5 +1,6 @@ use super::*; use crate::boxed::Box; +use core::cell::Cell; use std::iter::TrustedLen; use std::panic::{catch_unwind, AssertUnwindSafe}; use std::sync::atomic::{AtomicU32, Ordering}; @@ -405,6 +406,37 @@ fn test_retain() { assert!(a.is_empty()); } +#[test] +fn test_drain_filter() { + let mut a = BinaryHeap::from(vec![1, 2, 3, 4, 5]); + let mut evens = a.drain_filter(|x| *x % 2 == 0).collect::>(); + evens.sort(); + a.rebuild(); + let odds = a.into_sorted_vec(); + assert_eq!(evens, vec![2, 4]); + assert_eq!(odds, vec![1, 3, 5]); +} + +#[test] +fn test_rebuild() { + let mut a = BinaryHeap::from(vec![Cell::new(0), Cell::new(1), Cell::new(2), Cell::new(3)]); + + let sorted_values = |heap: &BinaryHeap>| { + heap.clone() + // this method assumes the heap is in a valid state. + .into_iter_sorted() + .map(|x| x.get()) + .collect::>() + }; + // internal mutation invalidates the heap order and so the sort fails + a.peek().unwrap().set(0); + assert_eq!(sorted_values(&a), [0, 2, 1, 0]); + + // the heap is rebuilt to a valid state and so the sort works + a.rebuild(); + assert_eq!(sorted_values(&a), [2, 1, 0, 0]); +} + // old binaryheap failed this test // // Integrity means that all elements are present after a comparison panics,