diff --git a/src/liballoc/collections/binary_heap.rs b/src/liballoc/collections/binary_heap.rs index 3d04f30e7bde5..fda6f090fd779 100644 --- a/src/liballoc/collections/binary_heap.rs +++ b/src/liballoc/collections/binary_heap.rs @@ -146,7 +146,7 @@ #![stable(feature = "rust1", since = "1.0.0")] use core::ops::{Deref, DerefMut}; -use core::iter::{FromIterator, FusedIterator}; +use core::iter::{FromIterator, FusedIterator, TrustedLen}; use core::mem::{swap, size_of, ManuallyDrop}; use core::ptr; use core::fmt; @@ -648,6 +648,36 @@ impl BinaryHeap { self.extend(other.drain()); } } + + /// Returns an iterator which retrieves elements in heap order. + /// The retrieved elements are removed from the original heap. + /// The remaining elements will be removed on drop in heap order. + /// + /// Note: + /// * `.drain_sorted()` is O(n lg n); much slower than `.drain()`. + /// You should use the latter for most cases. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(binary_heap_drain_sorted)] + /// use std::collections::BinaryHeap; + /// + /// let mut heap = BinaryHeap::from(vec![1, 2, 3, 4, 5]); + /// assert_eq!(heap.len(), 5); + /// + /// drop(heap.drain_sorted()); // removes all elements in heap order + /// assert_eq!(heap.len(), 0); + /// ``` + #[inline] + #[unstable(feature = "binary_heap_drain_sorted", issue = "59278")] + pub fn drain_sorted(&mut self) -> DrainSorted<'_, T> { + DrainSorted { + inner: self, + } + } } impl BinaryHeap { @@ -672,6 +702,27 @@ impl BinaryHeap { Iter { iter: self.data.iter() } } + /// Returns an iterator which retrieves elements in heap order. + /// This method consumes the original heap. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(binary_heap_into_iter_sorted)] + /// use std::collections::BinaryHeap; + /// let heap = BinaryHeap::from(vec![1, 2, 3, 4, 5]); + /// + /// assert_eq!(heap.into_iter_sorted().take(2).collect::>(), vec![5, 4]); + /// ``` + #[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")] + pub fn into_iter_sorted(self) -> IntoIterSorted { + IntoIterSorted { + inner: self, + } + } + /// Returns the greatest item in the binary heap, or `None` if it is empty. /// /// # Examples @@ -1115,6 +1166,37 @@ impl ExactSizeIterator for IntoIter { #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for IntoIter {} +#[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")] +#[derive(Clone, Debug)] +pub struct IntoIterSorted { + inner: BinaryHeap, +} + +#[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")] +impl Iterator for IntoIterSorted { + type Item = T; + + #[inline] + fn next(&mut self) -> Option { + self.inner.pop() + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let exact = self.inner.len(); + (exact, Some(exact)) + } +} + +#[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")] +impl ExactSizeIterator for IntoIterSorted {} + +#[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")] +impl FusedIterator for IntoIterSorted {} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl TrustedLen for IntoIterSorted {} + /// A draining iterator over the elements of a `BinaryHeap`. /// /// This `struct` is created by the [`drain`] method on [`BinaryHeap`]. See its @@ -1161,6 +1243,52 @@ impl ExactSizeIterator for Drain<'_, T> { #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Drain<'_, T> {} +/// A draining iterator over the elements of a `BinaryHeap`. +/// +/// This `struct` is created by the [`drain_sorted`] method on [`BinaryHeap`]. See its +/// documentation for more. +/// +/// [`drain_sorted`]: struct.BinaryHeap.html#method.drain_sorted +/// [`BinaryHeap`]: struct.BinaryHeap.html +#[unstable(feature = "binary_heap_drain_sorted", issue = "59278")] +#[derive(Debug)] +pub struct DrainSorted<'a, T: Ord> { + inner: &'a mut BinaryHeap, +} + +#[unstable(feature = "binary_heap_drain_sorted", issue = "59278")] +impl<'a, T: Ord> Drop for DrainSorted<'a, T> { + /// Removes heap elements in heap order. + fn drop(&mut self) { + while let Some(_) = self.inner.pop() {} + } +} + +#[unstable(feature = "binary_heap_drain_sorted", issue = "59278")] +impl Iterator for DrainSorted<'_, T> { + type Item = T; + + #[inline] + fn next(&mut self) -> Option { + self.inner.pop() + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let exact = self.inner.len(); + (exact, Some(exact)) + } +} + +#[unstable(feature = "binary_heap_drain_sorted", issue = "59278")] +impl ExactSizeIterator for DrainSorted<'_, T> { } + +#[unstable(feature = "binary_heap_drain_sorted", issue = "59278")] +impl FusedIterator for DrainSorted<'_, T> {} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl TrustedLen for DrainSorted<'_, T> {} + #[stable(feature = "binary_heap_extras_15", since = "1.5.0")] impl From> for BinaryHeap { /// Converts a `Vec` into a `BinaryHeap`. diff --git a/src/liballoc/tests/binary_heap.rs b/src/liballoc/tests/binary_heap.rs index b8c720264d0ab..a44cf1eaf6df0 100644 --- a/src/liballoc/tests/binary_heap.rs +++ b/src/liballoc/tests/binary_heap.rs @@ -1,5 +1,6 @@ use std::collections::BinaryHeap; use std::collections::binary_heap::{Drain, PeekMut}; +use std::iter::TrustedLen; #[test] fn test_iterator() { @@ -14,7 +15,7 @@ fn test_iterator() { } #[test] -fn test_iterator_reverse() { +fn test_iter_rev_cloned_collect() { let data = vec![5, 9, 3]; let iterout = vec![3, 5, 9]; let pq = BinaryHeap::from(data); @@ -24,7 +25,7 @@ fn test_iterator_reverse() { } #[test] -fn test_move_iter() { +fn test_into_iter_collect() { let data = vec![5, 9, 3]; let iterout = vec![9, 5, 3]; let pq = BinaryHeap::from(data); @@ -34,7 +35,7 @@ fn test_move_iter() { } #[test] -fn test_move_iter_size_hint() { +fn test_into_iter_size_hint() { let data = vec![5, 9]; let pq = BinaryHeap::from(data); @@ -51,7 +52,7 @@ fn test_move_iter_size_hint() { } #[test] -fn test_move_iter_reverse() { +fn test_into_iter_rev_collect() { let data = vec![5, 9, 3]; let iterout = vec![3, 5, 9]; let pq = BinaryHeap::from(data); @@ -60,6 +61,65 @@ fn test_move_iter_reverse() { assert_eq!(v, iterout); } +#[test] +fn test_into_iter_sorted_collect() { + let heap = BinaryHeap::from(vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1]); + let it = heap.into_iter_sorted(); + let sorted = it.collect::>(); + assert_eq!(sorted, vec![10, 9, 8, 7, 6, 5, 4, 3, 2, 2, 1, 1, 0]); +} + +#[test] +fn test_drain_sorted_collect() { + let mut heap = BinaryHeap::from(vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1]); + let it = heap.drain_sorted(); + let sorted = it.collect::>(); + assert_eq!(sorted, vec![10, 9, 8, 7, 6, 5, 4, 3, 2, 2, 1, 1, 0]); +} + +fn check_exact_size_iterator(len: usize, it: I) { + let mut it = it; + + for i in 0..it.len() { + let (lower, upper) = it.size_hint(); + assert_eq!(Some(lower), upper); + assert_eq!(lower, len - i); + assert_eq!(it.len(), len - i); + it.next(); + } + assert_eq!(it.len(), 0); + assert!(it.is_empty()); +} + +#[test] +fn test_exact_size_iterator() { + let 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.clone().into_iter()); + check_exact_size_iterator(heap.len(), heap.clone().into_iter_sorted()); + check_exact_size_iterator(heap.len(), heap.clone().drain()); + check_exact_size_iterator(heap.len(), heap.clone().drain_sorted()); +} + +fn check_trusted_len(len: usize, it: I) { + let mut it = it; + for i in 0..len { + let (lower, upper) = it.size_hint(); + if upper.is_some() { + assert_eq!(Some(lower), upper); + assert_eq!(lower, len - i); + } + it.next(); + } +} + +#[test] +fn test_trusted_len() { + let heap = BinaryHeap::from(vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1]); + check_trusted_len(heap.len(), heap.clone().into_iter_sorted()); + check_trusted_len(heap.len(), heap.clone().drain_sorted()); +} + #[test] fn test_peek_and_pop() { let data = vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1]; @@ -206,6 +266,15 @@ fn test_drain() { assert!(q.is_empty()); } +#[test] +fn test_drain_sorted() { + let mut q: BinaryHeap<_> = [9, 8, 7, 6, 5, 4, 3, 2, 1].iter().cloned().collect(); + + assert_eq!(q.drain_sorted().take(5).collect::>(), vec![9, 8, 7, 6, 5]); + + assert!(q.is_empty()); +} + #[test] fn test_extend_ref() { let mut a = BinaryHeap::new(); diff --git a/src/liballoc/tests/lib.rs b/src/liballoc/tests/lib.rs index 79e5ba340b784..3273feb7b5dd4 100644 --- a/src/liballoc/tests/lib.rs +++ b/src/liballoc/tests/lib.rs @@ -8,6 +8,8 @@ #![feature(try_reserve)] #![feature(unboxed_closures)] #![feature(associated_type_bounds)] +#![feature(binary_heap_into_iter_sorted)] +#![feature(binary_heap_drain_sorted)] use std::hash::{Hash, Hasher}; use std::collections::hash_map::DefaultHasher;