Add .into_iter_sorted() and .drain_sorted()

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

2 files changed, 79 insertions, 4 deletions

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;