Rollup merge of #151829 - max-heller:binary-heap-pop-if, r=jhpratt,joboet

Implement `BinaryHeap::pop_if()`

Implementation of https://github.com/rust-lang/rust/issues/151828

library/alloc/src/collections/binary_heap/mod.rs

@@ -649,6 +649,33 @@ impl<T: Ord, A: Allocator> BinaryHeap<T, A> {
         })
     }
 
+    /// Removes and returns the greatest item from the binary heap if the predicate
+    /// returns `true`, or [`None`] if the predicate returns false or the heap
+    /// is empty (the predicate will not be called in that case).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(binary_heap_pop_if)]
+    /// use std::collections::BinaryHeap;
+    /// let mut heap = BinaryHeap::from([1, 2]);
+    /// let pred = |x: &i32| *x % 2 == 0;
+    ///
+    /// assert_eq!(heap.pop_if(pred), Some(2));
+    /// assert_eq!(heap.as_slice(), [1]);
+    /// assert_eq!(heap.pop_if(pred), None);
+    /// assert_eq!(heap.as_slice(), [1]);
+    /// ```
+    ///
+    /// # Time complexity
+    ///
+    /// The worst case cost of `pop_if` on a heap containing *n* elements is *O*(log(*n*)).
+    #[unstable(feature = "binary_heap_pop_if", issue = "151828")]
+    pub fn pop_if(&mut self, predicate: impl FnOnce(&T) -> bool) -> Option<T> {
+        let first = self.peek()?;
+        if predicate(first) { self.pop() } else { None }
+    }
+
     /// Pushes an item onto the binary heap.
     ///
     /// # Examples

library/alloctests/tests/collections/binary_heap.rs

@@ -136,6 +136,18 @@ fn test_peek_and_pop() {
     }
 }
 
+#[test]
+fn test_pop_if() {
+    let data = vec![9, 8, 7, 6, 5, 4, 3, 2, 1, 0];
+    let mut sorted = data.clone();
+    sorted.sort();
+    let mut heap = BinaryHeap::from(data);
+    while let Some(popped) = heap.pop_if(|x| *x > 2) {
+        assert_eq!(popped, sorted.pop().unwrap());
+    }
+    assert_eq!(heap.into_sorted_vec(), vec![0, 1, 2]);
+}
+
 #[test]
 fn test_peek_mut() {
     let data = vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1];

library/alloctests/tests/lib.rs

@@ -1,4 +1,5 @@
 #![feature(allocator_api)]
+#![feature(binary_heap_pop_if)]
 #![feature(const_heap)]
 #![feature(deque_extend_front)]
 #![feature(iter_array_chunks)]