Implement ordered/sorted iterators on BinaryHeap as per #59278

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

@@ -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;