Add BinaryHeap::retain

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 - Added `resize_with` to `Vec`
 - Added `retain_back` (aka `truncate_front`) to `Deque`
+- Added `retain` to `BinaryHeap`
 - Fixed unsoundness in `Vec::IntoIter::drop` and `HistoryBuf::write`in the context of panicking drop implementations.
 
 ## [v0.9.3] 2025-04-15

src/binary_heap.rs

@@ -415,13 +415,44 @@ where
     /// # Ok::<(), u8>(())
     /// ```
     pub unsafe fn pop_unchecked(&mut self) -> T {
-        let mut item = self.data.pop_unchecked();
+        // SAFETY: the binary heap is not empty, thus `0` is smaller than `self.len()`.
+        unsafe { self.remove_unchecked(0) }
+    }
 
-        if !self.is_empty() {
-            mem::swap(&mut item, self.data.as_mut_slice().get_unchecked_mut(0));
-            self.sift_down_to_bottom(0);
+    /// Retains only the elements specified by the predicate.
+    /// The elements are visited in arbitrary order.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// 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<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&T) -> bool,
+    {
+        let mut index = 0;
+        while let Some(item) = self.data.get(index) {
+            if f(item) {
+                index += 1;
+            } else {
+                // SAFETY: `index` is valid because of the loop condition.
+                unsafe { self.remove_unchecked(index) };
+            }
         }
-        item
     }
 
     /// Pushes an item onto the binary heap.
@@ -471,9 +502,28 @@ where
     }
 
     /* Private API */
+
+    /// Removes and returns the element at position `index` within the inner vec.
+    /// The elements are shifted to preserve the invariants of the binary heap.
+    ///
+    /// # Safety
+    ///
+    /// The length of the heap must be larger than `index`.
+    unsafe fn remove_unchecked(&mut self, index: usize) -> T {
+        let mut item = self.data.pop_unchecked();
+        if let Some(item_to_remove) = self.data.get_mut(index) {
+            mem::swap(&mut item, item_to_remove);
+            self.sift_down_to_bottom(index);
+        }
+        item
+    }
+
+    /// Moves the element from `pos` down through the heap until it becomes a leaf,
+    /// then sift up the element to meet the order invariant.
     fn sift_down_to_bottom(&mut self, mut pos: usize) {
         let end = self.len();
         let start = pos;
+        // Moves the element down to a leaf.
         unsafe {
             let mut hole = Hole::new(self.data.as_mut_slice(), pos);
             let mut child = 2 * pos + 1;
@@ -488,9 +538,12 @@ where
             }
             pos = hole.pos;
         }
+        // Moves the element up to satisfy the order invariant.
         self.sift_up(start, pos);
     }
 
+    /// Moves the element at `pos` up through the heap until either the order
+    /// invariant is met or the `start` position is reached.
     fn sift_up(&mut self, start: usize, pos: usize) -> usize {
         unsafe {
             // Take out the value at `pos` and create a hole.
@@ -890,6 +943,192 @@ mod tests {
         assert_eq!(heap.pop(), None);
     }
 
+    #[test]
+    fn peek_mut() {
+        let mut heap = BinaryHeap::<_, Max, 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();
+
+        {
+            let mut val = heap.peek_mut().unwrap();
+            *val = 22;
+        }
+
+        {
+            let mut val = heap.peek_mut().unwrap();
+            *val = 9;
+        }
+
+        assert_eq!(heap.pop(), Some(25));
+        assert_eq!(heap.pop(), Some(22));
+        assert_eq!(heap.pop(), Some(19));
+        assert_eq!(heap.pop(), Some(17));
+        assert_eq!(heap.pop(), Some(9));
+        assert_eq!(heap.pop(), Some(7));
+        assert_eq!(heap.pop(), Some(3));
+        assert_eq!(heap.pop(), Some(2));
+        assert_eq!(heap.pop(), Some(1));
+        assert_eq!(heap.pop(), None);
+    }
+
+    #[test]
+    fn retain() {
+        let mut heap = BinaryHeap::<_, Max, 8>::new();
+        heap.retain(|_: &i32| true);
+        heap.retain(|_: &i32| false);
+        assert_eq!(heap.len(), 0);
+
+        let mut heap = BinaryHeap::<_, Max, 8>::new();
+        heap.push(1).unwrap();
+        heap.push(2).unwrap();
+        heap.push(3).unwrap();
+        heap.retain(|&e| e != 1);
+        assert_eq!(heap.pop(), Some(3));
+        assert_eq!(heap.pop(), Some(2));
+        assert_eq!(heap.pop(), None);
+
+        let mut heap = BinaryHeap::<_, Max, 8>::new();
+        heap.push(1).unwrap();
+        heap.push(2).unwrap();
+        heap.push(3).unwrap();
+        heap.retain(|&e| e != 3);
+        assert_eq!(heap.pop(), Some(2));
+        assert_eq!(heap.pop(), Some(1));
+        assert_eq!(heap.pop(), None);
+
+        let mut heap = BinaryHeap::<_, Max, 8>::new();
+        heap.push(1).unwrap();
+        heap.push(2).unwrap();
+        heap.push(3).unwrap();
+        heap.retain(|&e| e != 2);
+        assert_eq!(heap.pop(), Some(3));
+        assert_eq!(heap.pop(), Some(1));
+        assert_eq!(heap.pop(), None);
+
+        let mut heap = BinaryHeap::<_, Max, 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.pop(), Some(100));
+        assert_eq!(heap.pop(), Some(36));
+        assert_eq!(heap.pop(), Some(2));
+        assert_eq!(heap.pop(), None);
+
+        let mut heap = BinaryHeap::<_, Max, 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.pop(), Some(25));
+        assert_eq!(heap.pop(), Some(19));
+        assert_eq!(heap.pop(), Some(17));
+        assert_eq!(heap.pop(), Some(7));
+        assert_eq!(heap.pop(), Some(3));
+        assert_eq!(heap.pop(), Some(1));
+        assert_eq!(heap.pop(), None);
+    }
+
+    #[test]
+    fn remove_unchecked() {
+        // This test depends on implementation details.
+
+        let mut heap = BinaryHeap::<_, Max, 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();
+
+        assert_eq!(
+            heap.iter()
+                .copied()
+                .collect::<crate::vec::Vec::<i32, 16>>()
+                .as_slice(),
+            &[100, 36, 19, 25, 3, 2, 7, 1, 17],
+        );
+
+        unsafe {
+            heap.remove_unchecked(1);
+        }
+        assert_eq!(
+            heap.iter()
+                .copied()
+                .collect::<crate::vec::Vec::<i32, 16>>()
+                .as_slice(),
+            &[100, 25, 19, 17, 3, 2, 7, 1],
+        );
+
+        unsafe {
+            heap.remove_unchecked(2);
+        }
+        assert_eq!(
+            heap.iter()
+                .copied()
+                .collect::<crate::vec::Vec::<i32, 16>>()
+                .as_slice(),
+            &[100, 25, 7, 17, 3, 2, 1],
+        );
+
+        unsafe {
+            heap.remove_unchecked(3);
+        }
+        assert_eq!(
+            heap.iter()
+                .copied()
+                .collect::<crate::vec::Vec::<i32, 16>>()
+                .as_slice(),
+            &[100, 25, 7, 1, 3, 2],
+        );
+
+        unsafe {
+            heap.remove_unchecked(5);
+        }
+        assert_eq!(
+            heap.iter()
+                .copied()
+                .collect::<crate::vec::Vec::<i32, 16>>()
+                .as_slice(),
+            &[100, 25, 7, 1, 3],
+        );
+
+        unsafe {
+            heap.remove_unchecked(0);
+        }
+        assert_eq!(
+            heap.iter()
+                .copied()
+                .collect::<crate::vec::Vec::<i32, 16>>()
+                .as_slice(),
+            &[25, 3, 7, 1],
+        );
+    }
+
     #[test]
     #[cfg(feature = "zeroize")]
     fn test_binary_heap_zeroize() {