about summary refs log tree commit diff
diff options
context:
space:
mode:
-rw-r--r--src/liballoc/collections/binary_heap.rs28
-rw-r--r--src/liballoc/tests/binary_heap.rs8
-rw-r--r--src/liballoc/tests/lib.rs1
3 files changed, 37 insertions, 0 deletions
diff --git a/src/liballoc/collections/binary_heap.rs b/src/liballoc/collections/binary_heap.rs
index 03c9164fb90..8e170d970bc 100644
--- a/src/liballoc/collections/binary_heap.rs
+++ b/src/liballoc/collections/binary_heap.rs
@@ -665,6 +665,34 @@ impl<T: Ord> BinaryHeap<T> {
     pub fn drain_sorted(&mut self) -> DrainSorted<'_, T> {
         DrainSorted { inner: self }
     }
+
+    /// Retains only the elements specified by the predicate.
+    ///
+    /// In other words, remove all elements `e` such that `f(&e)` returns
+    /// `false`. The elements are visited in unsorted (and unspecified) order.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(binary_heap_retain)]
+    /// use std::collections::BinaryHeap;
+    ///
+    /// let mut heap = BinaryHeap::from(vec![-10, -5, 1, 2, 4, 13]);
+    ///
+    /// heap.retain(|x| x % 2 == 0); // only keep even numbers
+    ///
+    /// assert_eq!(heap.into_sorted_vec(), [-10, 2, 4])
+    /// ```
+    #[unstable(feature = "binary_heap_retain", issue = "71503")]
+    pub fn retain<F>(&mut self, f: F)
+    where
+        F: FnMut(&T) -> bool,
+    {
+        self.data.retain(f);
+        self.rebuild();
+    }
 }
 
 impl<T> BinaryHeap<T> {
diff --git a/src/liballoc/tests/binary_heap.rs b/src/liballoc/tests/binary_heap.rs
index be5516f54f3..08ed5447772 100644
--- a/src/liballoc/tests/binary_heap.rs
+++ b/src/liballoc/tests/binary_heap.rs
@@ -372,6 +372,14 @@ fn assert_covariance() {
     }
 }
 
+#[test]
+fn test_retain() {
+    let mut a = BinaryHeap::from(vec![-10, -5, 1, 2, 4, 13]);
+    a.retain(|x| x % 2 == 0);
+
+    assert_eq!(a.into_sorted_vec(), [-10, 2, 4])
+}
+
 // old binaryheap failed this test
 //
 // Integrity means that all elements are present after a comparison panics,
diff --git a/src/liballoc/tests/lib.rs b/src/liballoc/tests/lib.rs
index ad6feaeebc6..78d49558262 100644
--- a/src/liballoc/tests/lib.rs
+++ b/src/liballoc/tests/lib.rs
@@ -14,6 +14,7 @@
 #![feature(binary_heap_drain_sorted)]
 #![feature(vec_remove_item)]
 #![feature(split_inclusive)]
+#![feature(binary_heap_retain)]
 
 use std::collections::hash_map::DefaultHasher;
 use std::hash::{Hash, Hasher};