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..d9cd773e33 100644 --- a/src/binary_heap.rs +++ b/src/binary_heap.rs @@ -363,13 +363,10 @@ where /// ``` pub fn peek_mut(&mut self) -> Option> { if self.is_empty() { - None - } else { - Some(PeekMutInner { - heap: self, - sift: true, - }) + return None; } + // SAFETY: the heap has at least one item because of the previous condition. + Some(unsafe { PeekMutInner::new(self) }) } /// Removes the *top* (greatest if max-heap, smallest if min-heap) item from the binary heap and @@ -388,10 +385,10 @@ where /// ``` pub fn pop(&mut self) -> Option { if self.is_empty() { - None - } else { - Some(unsafe { self.pop_unchecked() }) + return None; } + // SAFETY: the heap has at least one item because of the previous condition. + Some(unsafe { self.pop_unchecked() }) } /// Removes the *top* (greatest if max-heap, smallest if min-heap) item from the binary heap and @@ -415,13 +412,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. @@ -441,7 +470,7 @@ where if self.data.is_full() { return Err(item); } - + // SAFETY: the heap is not full because of the previous condition. unsafe { self.push_unchecked(item) } Ok(()) } @@ -466,45 +495,87 @@ where /// ``` pub unsafe fn push_unchecked(&mut self, item: T) { let old_len = self.len(); - self.data.push_unchecked(item); - self.sift_up(0, old_len); + // SAFETY: the function precondition guarantees that the heap is not full. + unsafe { self.data.push_unchecked(item) }; + // SAFETY: `old_len` is now a valid index because a new item has been pushed. + unsafe { self.sift_up(0, old_len) }; } /* Private API */ - fn sift_down_to_bottom(&mut self, mut pos: usize) { + + /// 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 { + debug_assert!(index < self.len()); + // SAFETY: the heap is not empty because of the preconiditon. + let mut item = unsafe { self.data.pop_unchecked() }; + if let Some(item_to_remove) = self.data.get_mut(index) { + mem::swap(&mut item, item_to_remove); + // SAFETY: `index` is within the data slice because of the condition. + unsafe { 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. + /// + /// # Safety + /// + /// `pos` must be within the data slice. + unsafe fn sift_down_to_bottom(&mut self, mut pos: usize) { let end = self.len(); let start = pos; - unsafe { - let mut hole = Hole::new(self.data.as_mut_slice(), pos); + // Scope to limit the lifetime of the hole. + { + // SAFETY: `pos` is within the data slice because of the function precondition. + let mut hole = unsafe { Hole::new(self.data.as_mut_slice(), pos) }; let mut child = 2 * pos + 1; while child < end { let right = child + 1; - // compare with the greater of the two children - if right < end && hole.get(child).cmp(hole.get(right)) != K::ordering() { + // SAFETY: `child` is within the data slice because it is lower than `end`. + let child_val = unsafe { hole.get(child) }; + // Compares with the greater of the two children. + // SAFETY: `right` is within the data slice because it is lower than `end`. + if right < end && child_val.cmp(unsafe { hole.get(right) }) != K::ordering() { child = right; } - hole.move_to(child); + // SAFETY: `child` is within the data slice because it is lower than `end`. + unsafe { hole.move_to(child) }; child = 2 * hole.pos() + 1; } pos = hole.pos; } - self.sift_up(start, pos); + // Moves the element up to satisfy the order invariant. + // SAFETY: `pos` is within the data slice because `hole` keeps track of a valid index. + unsafe { self.sift_up(start, pos) }; } - fn sift_up(&mut self, start: usize, pos: usize) -> usize { - unsafe { - // Take out the value at `pos` and create a hole. - let mut hole = Hole::new(self.data.as_mut_slice(), pos); - - while hole.pos() > start { - let parent = (hole.pos() - 1) / 2; - if hole.element().cmp(hole.get(parent)) != K::ordering() { - break; - } - hole.move_to(parent); + /// Moves the element at `pos` up through the heap until either the order + /// invariant is met or the `start` position is reached. + /// + /// # Safety + /// + /// `pos` must be within the data slice. + unsafe fn sift_up(&mut self, start: usize, pos: usize) -> usize { + // Take out the value at `pos` and create a hole. + // SAFETY: `pos` is within the data slice because of the function precondition. + let mut hole = unsafe { Hole::new(self.data.as_mut_slice(), pos) }; + + while hole.pos() > start { + let parent = (hole.pos() - 1) / 2; + if hole.element().cmp(unsafe { hole.get(parent) }) != K::ordering() { + break; } - hole.pos() + // SAFETY: `parent` is within the data slice because it is lower than `pos`. + unsafe { hole.move_to(parent) }; } + hole.pos() } } @@ -522,7 +593,9 @@ struct Hole<'a, T> { impl<'a, T> Hole<'a, T> { /// Create a new Hole at index `pos`. /// - /// Unsafe because pos must be within the data slice. + /// # Safety + /// + /// `pos` must be within the data slice. #[inline] unsafe fn new(data: &'a mut [T], pos: usize) -> Self { debug_assert!(pos < data.len()); @@ -547,17 +620,22 @@ impl<'a, T> Hole<'a, T> { /// Returns a reference to the element at `index`. /// - /// Unsafe because index must be within the data slice and not equal to pos. + /// # Safety + /// + /// `index` must be within the data slice and not equal to [`Self::pos`]. #[inline] unsafe fn get(&self, index: usize) -> &T { debug_assert!(index != self.pos); debug_assert!(index < self.data.len()); - self.data.get_unchecked(index) + // SAFETY: `index` is valid because of the function precondition. + unsafe { self.data.get_unchecked(index) } } /// Move hole to new location /// - /// Unsafe because index must be within the data slice and not equal to pos. + /// # Safety + /// + /// `index` must be within the data slice and not equal to [`Self::pos`]. #[inline] unsafe fn move_to(&mut self, index: usize) { debug_assert!(index != self.pos); @@ -581,8 +659,30 @@ where K: Kind, S: VecStorage + ?Sized, { + /// Non-empty heap. heap: &'a mut BinaryHeapInner, - sift: bool, + /// Does the peek item has been removed. + /// The item can be removed using [`Self::pop`]. + removed: bool, +} +impl<'a, T, K, S> PeekMutInner<'a, T, K, S> +where + T: Ord, + K: Kind, + S: VecStorage + ?Sized, +{ + /// Creates a new instance that allows removing or mutating the first item of `heap`. + /// + /// # Safety + /// + /// The heap must not be empty. + unsafe fn new(heap: &'a mut BinaryHeapInner) -> Self { + debug_assert!(!heap.is_empty()); + Self { + heap, + removed: false, + } + } } /// Structure wrapping a mutable reference to the greatest item on a @@ -606,9 +706,12 @@ where S: VecStorage + ?Sized, { fn drop(&mut self) { - if self.sift { - self.heap.sift_down_to_bottom(0); + if self.removed { + return; } + // SAFETY: the heap has at least one item because + // `PeekMut` is only instantiated wiforth non-empty heaps and `removed` is still `false`. + unsafe { self.heap.sift_down_to_bottom(0) }; } } @@ -621,7 +724,7 @@ where type Target = T; fn deref(&self) -> &T { debug_assert!(!self.heap.is_empty()); - // SAFE: PeekMut is only instantiated for non-empty heaps + // SAFETY: `PeekMut` is only instantiated for non-empty heaps. unsafe { self.heap.data.as_slice().get_unchecked(0) } } } @@ -634,7 +737,7 @@ where { fn deref_mut(&mut self) -> &mut T { debug_assert!(!self.heap.is_empty()); - // SAFE: PeekMut is only instantiated for non-empty heaps + // SAFE: PeekMut is only instantiated for non-empty heaps. unsafe { self.heap.data.as_mut_slice().get_unchecked_mut(0) } } } @@ -647,8 +750,9 @@ where { /// Removes the peeked value from the heap and returns it. pub fn pop(mut this: Self) -> T { - let value = this.heap.pop().unwrap(); - this.sift = false; + // SAFE: PeekMut is only instantiated for non-empty heaps. + let value = unsafe { this.heap.pop_unchecked() }; + this.removed = true; value } } @@ -890,6 +994,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 private APIs. + + 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() {