Add .into_iter_sorted() and .drain_sorted()

* `.drain_sorted()` doc change suggested by @KodrAus

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<T> BinaryHeap<T> {
         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<_>>(), vec![5, 4]);
+    /// ```
+    #[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
+    pub fn into_iter_sorted(self) -> IntoIterSorted<T> {
+        IntoIterSorted {
+            inner: self,
+        }
+    }
+
     /// Returns the greatest item in the binary heap, or `None` if it is empty.
     ///
     /// # Examples
@@ -899,6 +920,47 @@ impl<T> BinaryHeap<T> {
         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::<Vec<_>>(); // 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<T> ExactSizeIterator for IntoIter<T> {
 #[stable(feature = "fused", since = "1.26.0")]
 impl<T> FusedIterator for IntoIter<T> {}
 
+#[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
+#[derive(Clone, Debug)]
+pub struct IntoIterSorted<T> {
+    inner: BinaryHeap<T>,
+}
+
+#[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
+impl<T: Ord> Iterator for IntoIterSorted<T> {
+    type Item = T;
+
+    #[inline]
+    fn next(&mut self) -> Option<T> {
+        self.inner.pop()
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let exact = self.inner.len();
+        (exact, Some(exact))
+    }
+}
+
+#[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
+impl<T: Ord> ExactSizeIterator for IntoIterSorted<T> {}
+
+#[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
+impl<T: Ord> FusedIterator for IntoIterSorted<T> {}
+
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<T: Ord> TrustedLen for IntoIterSorted<T> {}
+
 /// 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<T> ExactSizeIterator for Drain<'_, T> {
 #[stable(feature = "fused", since = "1.26.0")]
 impl<T> 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<T>,
+}
+
+#[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<T: Ord> Iterator for DrainSorted<'_, T> {
+    type Item = T;
+
+    #[inline]
+    fn next(&mut self) -> Option<T> {
+        self.inner.pop()
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let exact = self.inner.len();
+        (exact, Some(exact))
+    }
+}
+
+#[unstable(feature = "binary_heap_drain_sorted", issue = "59278")]
+impl<T: Ord> ExactSizeIterator for DrainSorted<'_, T> { }
+
+#[unstable(feature = "binary_heap_drain_sorted", issue = "59278")]
+impl<T: Ord> FusedIterator for DrainSorted<'_, T> {}
+
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<T: Ord> TrustedLen for DrainSorted<'_, T> {}
+
 #[stable(feature = "binary_heap_extras_15", since = "1.5.0")]
 impl<T: Ord> From<Vec<T>> for BinaryHeap<T> {
     /// Converts a `Vec<T>` into a `BinaryHeap<T>`.
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);
@@ -61,6 +66,65 @@ fn test_move_iter_reverse() {
 }
 
 #[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::<Vec<_>>();
+    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::<Vec<_>>();
+    assert_eq!(sorted, vec![10, 9, 8, 7, 6, 5, 4, 3, 2, 2, 1, 1, 0]);
+}
+
+fn check_exact_size_iterator<I: ExactSizeIterator>(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<I: TrustedLen>(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];
     let mut sorted = data.clone();
@@ -207,6 +271,15 @@ fn test_drain() {
 }
 
 #[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<_>>(), vec![9, 8, 7, 6, 5]);
+
+    assert!(q.is_empty());
+}
+
+#[test]
 fn test_extend_ref() {
     let mut a = BinaryHeap::new();
     a.push(1);
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;