diff --git a/library/alloc/src/collections/binary_heap/mod.rs b/library/alloc/src/collections/binary_heap/mod.rs index 965fd63a52981..59e47bcd11836 100644 --- a/library/alloc/src/collections/binary_heap/mod.rs +++ b/library/alloc/src/collections/binary_heap/mod.rs @@ -359,6 +359,74 @@ impl DerefMut for PeekMut<'_, T, A> { } impl<'a, T: Ord, A: Allocator> PeekMut<'a, T, A> { + /// Sifts the current element to its new position. + /// + /// Afterwards refers to the new element. Returns if the element changed. + /// + /// ## Examples + /// + /// The condition can be used to upper bound all elements in the heap. When only few elements + /// are affected, the heap's sort ensures this is faster than a reconstruction from the raw + /// element list and requires no additional allocation. + /// + /// ``` + /// #![feature(binary_heap_peek_mut_refresh)] + /// use std::collections::BinaryHeap; + /// + /// let mut heap: BinaryHeap = (0..128).collect(); + /// let mut peek = heap.peek_mut().unwrap(); + /// + /// loop { + /// *peek = 99; + /// + /// if !peek.refresh() { + /// break; + /// } + /// } + /// + /// // Post condition, this is now an upper bound. + /// assert!(*peek < 100); + /// ``` + /// + /// When the element remains the maximum after modification, the peek remains unchanged: + /// + /// ``` + /// #![feature(binary_heap_peek_mut_refresh)] + /// use std::collections::BinaryHeap; + /// + /// let mut heap: BinaryHeap = [1, 2, 3].into(); + /// let mut peek = heap.peek_mut().unwrap(); + /// + /// assert_eq!(*peek, 3); + /// *peek = 42; + /// + /// // When we refresh, the peek is updated to the new maximum. + /// assert!(!peek.refresh(), "42 is even larger than 3"); + /// assert_eq!(*peek, 42); + /// ``` + #[unstable(feature = "binary_heap_peek_mut_refresh", issue = "138355")] + #[must_use = "is equivalent to dropping and getting a new PeekMut except for return information"] + pub fn refresh(&mut self) -> bool { + // The length of the underlying heap is unchanged by sifting down. The value stored for leak + // amplification thus remains accurate. We erase the leak amplification firstly because the + // operation is then equivalent to constructing a new PeekMut and secondly this avoids any + // future complication where original_len being non-empty would be interpreted as the heap + // having been leak amplified instead of checking the heap itself. + if let Some(original_len) = self.original_len.take() { + // SAFETY: This is how many elements were in the Vec at the time of + // the BinaryHeap::peek_mut call. + unsafe { self.heap.data.set_len(original_len.get()) }; + + // The length of the heap did not change by sifting, upholding our own invariants. + + // SAFETY: PeekMut is only instantiated for non-empty heaps. + (unsafe { self.heap.sift_down(0) }) != 0 + } else { + // The element was not modified. + false + } + } + /// Removes the peeked value from the heap and returns it. #[stable(feature = "binary_heap_peek_mut_pop", since = "1.18.0")] pub fn pop(mut this: PeekMut<'a, T, A>) -> T { @@ -670,6 +738,8 @@ impl BinaryHeap { /// # Safety /// /// The caller must guarantee that `pos < self.len()`. + /// + /// Returns the new position of the element. unsafe fn sift_up(&mut self, start: usize, pos: usize) -> usize { // Take out the value at `pos` and create a hole. // SAFETY: The caller guarantees that pos < self.len() @@ -696,10 +766,12 @@ impl BinaryHeap { /// Take an element at `pos` and move it down the heap, /// while its children are larger. /// + /// Returns the new position of the element. + /// /// # Safety /// /// The caller must guarantee that `pos < end <= self.len()`. - unsafe fn sift_down_range(&mut self, pos: usize, end: usize) { + unsafe fn sift_down_range(&mut self, pos: usize, end: usize) -> usize { // SAFETY: The caller guarantees that pos < end <= self.len(). let mut hole = unsafe { Hole::new(&mut self.data, pos) }; let mut child = 2 * hole.pos() + 1; @@ -719,7 +791,7 @@ impl BinaryHeap { // SAFETY: child is now either the old child or the old child+1 // We already proven that both are < self.len() and != hole.pos() if hole.element() >= unsafe { hole.get(child) } { - return; + return hole.pos(); } // SAFETY: same as above. @@ -734,16 +806,18 @@ impl BinaryHeap { // child == 2 * hole.pos() + 1 != hole.pos(). unsafe { hole.move_to(child) }; } + + hole.pos() } /// # Safety /// /// The caller must guarantee that `pos < self.len()`. - unsafe fn sift_down(&mut self, pos: usize) { + unsafe fn sift_down(&mut self, pos: usize) -> usize { let len = self.len(); // SAFETY: pos < len is guaranteed by the caller and // obviously len = self.len() <= self.len(). - unsafe { self.sift_down_range(pos, len) }; + unsafe { self.sift_down_range(pos, len) } } /// Take an element at `pos` and move it all the way down the heap,