diff --git a/library/alloc/src/collections/binary_heap/mod.rs b/library/alloc/src/collections/binary_heap/mod.rs index 63828b482b9a9..ba27bd1cd6aef 100644 --- a/library/alloc/src/collections/binary_heap/mod.rs +++ b/library/alloc/src/collections/binary_heap/mod.rs @@ -776,17 +776,16 @@ impl BinaryHeap { 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; - // Loop invariant: child == 2 * hole.pos() + 1. - while child <= end.saturating_sub(2) { + // While hole's position is below parent of the first empty slot + while hole.pos() < (end - 1) / 2 { + let mut child = 2 * hole.pos() + 1; + // compare with the greater of the two children // SAFETY: child < end - 1 < self.len() and // child + 1 < end <= self.len(), so they're valid indexes. // child == 2 * hole.pos() + 1 != hole.pos() and // child + 1 == 2 * hole.pos() + 2 != hole.pos(). - // FIXME: 2 * hole.pos() + 1 or 2 * hole.pos() + 2 could overflow - // if T is a ZST child += unsafe { hole.get(child) <= hole.get(child + 1) } as usize; // if we are already in order, stop. @@ -798,17 +797,24 @@ impl BinaryHeap { // SAFETY: same as above. unsafe { hole.move_to(child) }; - child = 2 * hole.pos() + 1; } - // SAFETY: && short circuit, which means that in the - // second condition it's already true that child == end - 1 < self.len(). - if child == end - 1 && hole.element() < unsafe { hole.get(child) } { - // SAFETY: child is already proven to be a valid index and - // child == 2 * hole.pos() + 1 != hole.pos(). - unsafe { hole.move_to(child) }; + // If hole has only one child. + if hole.pos() < end / 2 { + let child = 2 * hole.pos() + 1; + debug_assert!(child == end - 1); + + // If elements are not in order, move parent + // SAFETY: The previous statements imply + // child is a valid index (child < end <= self.len()), + // and child != hole.pos(). + if hole.element() < unsafe { hole.get(child) } { + // SAFETY: same as above + unsafe { hole.move_to(child) }; + } } + // Otherwise return the same position hole.pos() } @@ -837,26 +843,30 @@ impl BinaryHeap { // SAFETY: The caller guarantees that pos < self.len(). let mut hole = unsafe { Hole::new(&mut self.data, pos) }; - let mut child = 2 * hole.pos() + 1; - // Loop invariant: child == 2 * hole.pos() + 1. - while child <= end.saturating_sub(2) { + // While hole's position is below parent of the first empty slot + while hole.pos() < (end - 1) / 2 { + let mut child = 2 * hole.pos() + 1; + // SAFETY: child < end - 1 < self.len() and // child + 1 < end <= self.len(), so they're valid indexes. // child == 2 * hole.pos() + 1 != hole.pos() and // child + 1 == 2 * hole.pos() + 2 != hole.pos(). - // FIXME: 2 * hole.pos() + 1 or 2 * hole.pos() + 2 could overflow - // if T is a ZST child += unsafe { hole.get(child) <= hole.get(child + 1) } as usize; // SAFETY: Same as above unsafe { hole.move_to(child) }; - child = 2 * hole.pos() + 1; } - if child == end - 1 { - // SAFETY: child == end - 1 < self.len(), so it's a valid index - // and child == 2 * hole.pos() + 1 != hole.pos(). + // If hole has only one child. + if hole.pos() < end / 2 { + let child = 2 * hole.pos() + 1; + debug_assert!(child == end - 1); + + // If elements are not in order, move parent + // SAFETY: The previous statements imply + // child is a valid index (child < end <= self.len()), + // and child != hole.pos(). unsafe { hole.move_to(child) }; } pos = hole.pos();