diff --git a/CHANGELOG.md b/CHANGELOG.md index 3977fd9c95..b19ecc7fbe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,8 +7,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/). ## [Unreleased] -- Added `resize_with` to `Vec` -- Added `retain_back` (aka `truncate_front`) to `Deque` +- 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 diff --git a/src/binary_heap.rs b/src/binary_heap.rs index 452b84a13e..7ef85de99e 100644 --- a/src/binary_heap.rs +++ b/src/binary_heap.rs @@ -415,13 +415,45 @@ 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(&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 +503,29 @@ 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 +540,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 +945,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::>() + .as_slice(), + &[100, 36, 19, 25, 3, 2, 7, 1, 17], + ); + + unsafe { + heap.remove_unchecked(1); + } + assert_eq!( + heap.iter() + .copied() + .collect::>() + .as_slice(), + &[100, 25, 19, 17, 3, 2, 7, 1], + ); + + unsafe { + heap.remove_unchecked(2); + } + assert_eq!( + heap.iter() + .copied() + .collect::>() + .as_slice(), + &[100, 25, 7, 17, 3, 2, 1], + ); + + unsafe { + heap.remove_unchecked(3); + } + assert_eq!( + heap.iter() + .copied() + .collect::>() + .as_slice(), + &[100, 25, 7, 1, 3, 2], + ); + + unsafe { + heap.remove_unchecked(5); + } + assert_eq!( + heap.iter() + .copied() + .collect::>() + .as_slice(), + &[100, 25, 7, 1, 3], + ); + + unsafe { + heap.remove_unchecked(0); + } + assert_eq!( + heap.iter() + .copied() + .collect::>() + .as_slice(), + &[25, 3, 7, 1], + ); + } + #[test] #[cfg(feature = "zeroize")] fn test_binary_heap_zeroize() {