From 6b8807333aa6e254a4f73e061c502d2bd221c0df Mon Sep 17 00:00:00 2001 From: Hideki Sekine Date: Tue, 1 Oct 2019 19:16:51 +0900 Subject: [PATCH] Add .into_iter_sorted() and .drain_sorted() * `.drain_sorted()` doc change suggested by @KodrAus --- src/liballoc/collections/binary_heap.rs | 141 +++++++++++++++++++++++- src/liballoc/tests/binary_heap.rs | 81 +++++++++++++- src/liballoc/tests/lib.rs | 2 + 3 files changed, 219 insertions(+), 5 deletions(-) diff --git a/src/liballoc/collections/binary_heap.rs b/src/liballoc/collections/binary_heap.rs index 3d04f30e7bd..08d39d80949 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; @@ -672,6 +672,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 @@ -899,6 +920,47 @@ impl BinaryHeap { Drain { iter: self.data.drain(..) } } + /// Returns an iterator which retrieves elements in heap order. + /// The retrieved elements will be removed from the original heap. + /// + /// Note: + /// * Unlike other `.drain()` methods, this method removes elements *lazily*. + /// In order to remove elements in heap order, you need to retrieve elements explicitly. + /// * The remaining elements are removed on drop in arbitrary order. + /// + /// # 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); + /// + /// let len = heap.len(); + /// let removed = heap.drain_sorted() + /// .take(len).collect::>(); // removes all elements in *heap* order + /// assert_eq!(removed, vec![5, 4, 3, 2, 1]); + /// assert_eq!(heap.len(), 0); + /// + /// + /// let mut heap = BinaryHeap::from(vec![1, 2, 3, 4, 5]); + /// assert_eq!(heap.len(), 5); + /// + /// let drain_sorted = heap.drain_sorted(); + /// drop(drain_sorted); // removes all elements in *arbitrary* 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, + } + } + /// Drops all items from the binary heap. /// /// # Examples @@ -1115,6 +1177,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 +1254,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> { + inner: &'a mut BinaryHeap, +} + +#[unstable(feature = "binary_heap_drain_sorted", issue = "59278")] +impl<'a, T> Drop for DrainSorted<'a, T> { + /// Removes heap elements in arbitrary order for efficiency. + fn drop(&mut self) { + self.inner.drain(); + } +} + +#[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 b8c720264d0..f494a3e4b3d 100644 --- a/src/liballoc/tests/binary_heap.rs +++ b/src/liballoc/tests/binary_heap.rs @@ -1,5 +1,10 @@ use std::collections::BinaryHeap; use std::collections::binary_heap::{Drain, PeekMut}; +use std::panic::{self, AssertUnwindSafe}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::iter::TrustedLen; + +use rand::{thread_rng, seq::SliceRandom}; #[test] fn test_iterator() { @@ -14,7 +19,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 +29,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 +39,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 +56,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 +65,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 +270,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 676874c8b27..4508fcc8451 100644 --- a/src/liballoc/tests/lib.rs +++ b/src/liballoc/tests/lib.rs @@ -9,6 +9,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;