From 7fdaf1757600588bd814f3eb48ff63ff2f2459d2 Mon Sep 17 00:00:00 2001
From: Chris de Claverie <chris@gamaspace.lan>
Date: Thu, 23 Feb 2023 19:18:40 +0100
Subject: [PATCH] feat: support retain operation over binary heap

---
 src/binary_heap.rs | 61 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)

diff --git a/src/binary_heap.rs b/src/binary_heap.rs
index d1fbe7d997..38dec0ed03 100644
--- a/src/binary_heap.rs
+++ b/src/binary_heap.rs
@@ -295,6 +295,41 @@ where
         }
     }
 
+    /// Removes all elements from the binary heap that do not satisfy the predicate `f`.
+    /// The elements are visited in arbitrary order.
+    ///
+    /// ```
+    /// use heapless::binary_heap::{BinaryHeap, Max};
+    ///
+    /// let mut heap: BinaryHeap<_, Max, 8> = BinaryHeap::new();
+    /// heap.push(1).unwrap();
+    /// heap.push(2).unwrap();
+    /// heap.push(3).unwrap();
+    /// heap.push(4).unwrap();
+    ///
+    /// heap.retain(|&x| x % 2 == 0);
+    ///
+    /// let mut iter = heap.iter();
+    /// assert_eq!(iter.next(), Some(&4));
+    /// assert_eq!(iter.next(), Some(&2));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn retain(&mut self, mut f: impl FnMut(&T) -> bool) {
+        let mut del = 0;
+        let len = self.len();
+        {
+            let v = self.data.as_mut_slice();
+            for i in 0..len {
+                if !f(&v[i]) {
+                    del += 1;
+                } else if del > 0 {
+                    v.swap(i - del, i);
+                }
+            }
+        }
+        self.data.truncate(len - del);
+    }
+
     /// Removes the *top* (greatest if max-heap, smallest if min-heap) item from the binary heap and
     /// returns it, without checking if the binary heap is empty.
     pub unsafe fn pop_unchecked(&mut self) -> T {
@@ -738,4 +773,30 @@ mod tests {
         assert_eq!(heap.pop(), Some(1));
         assert_eq!(heap.pop(), None);
     }
+
+    #[test]
+    fn retain() {
+        let mut heap = BinaryHeap::<_, Min, 16>::new();
+        heap.push(1).unwrap();
+        heap.push(2).unwrap();
+        heap.push(3).unwrap();
+        heap.push(17).unwrap();
+        heap.push(19).unwrap();
+        heap.push(36).unwrap();
+        heap.push(7).unwrap();
+        heap.push(25).unwrap();
+        heap.push(100).unwrap();
+
+        heap.retain(|&x| x % 2 == 0);
+
+        assert_eq!(
+            heap.iter().cloned().collect::<Vec<_>>(),
+            [2, 36, 100]
+        );
+
+        assert_eq!(heap.pop(), Some(2));
+        assert_eq!(heap.pop(), Some(36));
+        assert_eq!(heap.pop(), Some(100));
+        assert_eq!(heap.pop(), None);
+    }
 }