Auto merge of #65091 - sekineh:into-iter-sorted, r=KodrAus

Implement ordered/sorted iterators on BinaryHeap as per #59278

I've implemented the ordered version of iterator on BinaryHeap as per #59278.

# Added methods:

* `.into_iter_sorted()`
  * like `.into_iter()`; but returns elements in heap order
* `.drain_sorted()`
  * like `.drain()`; but returns elements in heap order
  * It's a bit _lazy_; elements are removed on drop. (Edit: it’s similar to vec::Drain)

For `DrainSorted` struct, I implemented `Drop` trait following @scottmcm 's [suggestion](#59278 (comment))

# ~TODO~ DONE
* ~I think I need to add more tests other than doctest.~

# **Notes:**
* we renamed `_ordered` to `_sorted`, because the latter is more common in rust libs. (as suggested by @KodrAus )

Author: bors

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<T: Ord> BinaryHeap<T> {
             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<T> BinaryHeap<T> {
@@ -672,6 +702,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
@@ -1115,6 +1166,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 +1243,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: Ord> {
+    inner: &'a mut BinaryHeap<T>,
+}
+
+#[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<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>`.

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::<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];
@@ -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<_>>(), vec![9, 8, 7, 6, 5]);
+
+    assert!(q.is_empty());
+}
+
 #[test]
 fn test_extend_ref() {
     let mut a = BinaryHeap::new();

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;