From 310d3b0ffcf40f2f138fceb1b552bbc75a1c64e6 Mon Sep 17 00:00:00 2001 From: Victorien Elvinger Date: Wed, 8 Jul 2026 16:54:33 +0200 Subject: [PATCH 1/6] Add BinaryHeap::retain --- CHANGELOG.md | 1 + src/binary_heap.rs | 132 +++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 128 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3977fd9c95..778ae495b2 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/binary_heap.rs b/src/binary_heap.rs index 452b84a13e..6ae82ef083 100644 --- a/src/binary_heap.rs +++ b/src/binary_heap.rs @@ -415,13 +415,47 @@ 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 len = self.len(); + let mut index = 0; + while index < len { + let item = unsafe { self.data.get_unchecked(index) }; + if f(item) { + index += 1; + } else { + // SAFETY: `index` is smaller than `self.len()`. + unsafe { self.remove_unchecked(index) }; + len -= 1; + } } - item } /// Pushes an item onto the binary heap. @@ -471,6 +505,22 @@ 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(place_at_index) = self.data.get_mut(index) { + mem::swap(&mut item, place_at_index); + self.sift_down_to_bottom(index); + } + item + } + fn sift_down_to_bottom(&mut self, mut pos: usize) { let end = self.len(); let start = pos; @@ -890,6 +940,78 @@ mod tests { 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] #[cfg(feature = "zeroize")] fn test_binary_heap_zeroize() { From d8e3a5c3dfb3256be9efcf4c27f9baa49bf41ca9 Mon Sep 17 00:00:00 2001 From: Victorien Elvinger Date: Tue, 21 Jul 2026 22:10:22 +0200 Subject: [PATCH 2/6] Add tests and comments --- src/binary_heap.rs | 124 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 122 insertions(+), 2 deletions(-) diff --git a/src/binary_heap.rs b/src/binary_heap.rs index 6ae82ef083..febf66b2fe 100644 --- a/src/binary_heap.rs +++ b/src/binary_heap.rs @@ -514,16 +514,19 @@ where /// 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(place_at_index) = self.data.get_mut(index) { - mem::swap(&mut item, place_at_index); + 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; @@ -538,9 +541,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. @@ -940,6 +946,41 @@ 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(); @@ -1012,6 +1053,85 @@ mod tests { 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() { From 1ba766bfbbf6f26ed9b066c4eb1f38a4bcd43bf8 Mon Sep 17 00:00:00 2001 From: Victorien Elvinger Date: Tue, 21 Jul 2026 22:15:18 +0200 Subject: [PATCH 3/6] Add SAFETY comment --- src/binary_heap.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/binary_heap.rs b/src/binary_heap.rs index febf66b2fe..25909ccf68 100644 --- a/src/binary_heap.rs +++ b/src/binary_heap.rs @@ -447,6 +447,7 @@ where let mut len = self.len(); let mut index = 0; while index < len { + // SAFETY: `index` is smaller than `self.len()`. let item = unsafe { self.data.get_unchecked(index) }; if f(item) { index += 1; From 1535316949b720d3b0d047fbc1e53217098a5fc0 Mon Sep 17 00:00:00 2001 From: Victorien Elvinger Date: Tue, 21 Jul 2026 22:16:47 +0200 Subject: [PATCH 4/6] Simplify retain impl --- src/binary_heap.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/binary_heap.rs b/src/binary_heap.rs index 25909ccf68..0cf046269b 100644 --- a/src/binary_heap.rs +++ b/src/binary_heap.rs @@ -444,17 +444,13 @@ where where F: FnMut(&T) -> bool, { - let mut len = self.len(); let mut index = 0; - while index < len { - // SAFETY: `index` is smaller than `self.len()`. - let item = unsafe { self.data.get_unchecked(index) }; + while let Some(item) = self.data.get(index) { if f(item) { index += 1; } else { - // SAFETY: `index` is smaller than `self.len()`. + // SAFETY: `index` is valid because of the loop condition. unsafe { self.remove_unchecked(index) }; - len -= 1; } } } From a1d8043b84d6547fc187eae1fefd79d2a6850e97 Mon Sep 17 00:00:00 2001 From: Victorien Elvinger Date: Wed, 22 Jul 2026 21:39:30 +0200 Subject: [PATCH 5/6] Apply suggestions --- CHANGELOG.md | 6 +++--- src/binary_heap.rs | 2 ++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 778ae495b2..b19ecc7fbe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,9 +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 `retain` to `BinaryHeap` +- 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 0cf046269b..7ef85de99e 100644 --- a/src/binary_heap.rs +++ b/src/binary_heap.rs @@ -420,6 +420,7 @@ where } /// Retains only the elements specified by the predicate. + /// /// The elements are visited in arbitrary order. /// /// # Examples @@ -504,6 +505,7 @@ 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 From 7de39e51198d11d2dacd40adf3fbef420ec0f13e Mon Sep 17 00:00:00 2001 From: Victorien Elvinger Date: Sat, 1 Aug 2026 12:40:43 +0200 Subject: [PATCH 6/6] Add SAFETY comments and unsafe blocks --- src/binary_heap.rs | 139 ++++++++++++++++++++++++++++++--------------- 1 file changed, 94 insertions(+), 45 deletions(-) diff --git a/src/binary_heap.rs b/src/binary_heap.rs index 7ef85de99e..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 @@ -473,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(()) } @@ -498,8 +495,10 @@ 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 */ @@ -512,54 +511,71 @@ where /// /// 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(); + 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); - self.sift_down_to_bottom(index); + // 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. - fn sift_down_to_bottom(&mut self, mut pos: usize) { + /// + /// # 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; - // Moves the element down to a leaf. - 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; } // Moves the element up to satisfy the order invariant. - self.sift_up(start, pos); + // SAFETY: `pos` is within the data slice because `hole` keeps track of a valid index. + unsafe { 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. - let mut hole = Hole::new(self.data.as_mut_slice(), pos); + /// + /// # 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(hole.get(parent)) != K::ordering() { - break; - } - hole.move_to(parent); + 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() } } @@ -577,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()); @@ -602,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); @@ -636,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 @@ -661,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) }; } } @@ -676,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) } } } @@ -689,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) } } } @@ -702,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 } } @@ -1054,7 +1103,7 @@ mod tests { #[test] fn remove_unchecked() { - // This test depends on implementation details. + // This test depends on private APIs. let mut heap = BinaryHeap::<_, Max, 16>::new(); heap.push(1).unwrap();