Port some rust-lang/rust performance and safety PRs

Port changes from rust-lang/rust related to binary heap performance and
use of unsafe:

- #81706: Document BinaryHeap unsafe functions
- #81127: Improve sift_down performance in BinaryHeap
- #58123: Avoid some bounds checks in binary_heap::{PeekMut,Hole}
- #72709: #[deny(unsafe_op_in_unsafe_fn)] in liballoc

Note that the following related rust-lang/rust PRs were already ported
here in earlier PRs:

- (in sekineh#28) #78857: Improve BinaryHeap performance
- (in sekineh#27) #75974: Avoid useless sift_down when
  std::collections::binary_heap::PeekMut is never mutably dereferenced

Author: clint-white

src/binary_heap.rs

@@ -155,6 +155,7 @@
 //! }
 //! ```
 
+#![deny(unsafe_op_in_unsafe_fn)]
 #![allow(clippy::needless_doctest_main)]
 #![allow(missing_docs)]
 // #![stable(feature = "rust1", since = "1.0.0")]
@@ -319,7 +320,8 @@ impl<'a, T: fmt::Debug, C: Compare<T>> fmt::Debug for PeekMut<'a, T, C> {
 impl<'a, T, C: Compare<T>> Drop for PeekMut<'a, T, C> {
     fn drop(&mut self) {
         if self.sift {
-            self.heap.sift_down(0);
+            // SAFETY: PeekMut is only instantiated for non-empty heaps.
+            unsafe { self.heap.sift_down(0) };
         }
     }
 }
@@ -328,15 +330,19 @@ impl<'a, T, C: Compare<T>> Drop for PeekMut<'a, T, C> {
 impl<'a, T, C: Compare<T>> Deref for PeekMut<'a, T, C> {
     type Target = T;
     fn deref(&self) -> &T {
-        &self.heap.data[0]
+        debug_assert!(!self.heap.is_empty());
+        // SAFE: PeekMut is only instantiated for non-empty heaps
+        unsafe { self.heap.data.get_unchecked(0) }
     }
 }
 
 // #[stable(feature = "binary_heap_peek_mut", since = "1.12.0")]
 impl<'a, T, C: Compare<T>> DerefMut for PeekMut<'a, T, C> {
     fn deref_mut(&mut self) -> &mut T {
+        debug_assert!(!self.heap.is_empty());
         self.sift = true;
-        &mut self.heap.data[0]
+        // SAFE: PeekMut is only instantiated for non-empty heaps
+        unsafe { self.heap.data.get_unchecked_mut(0) }
     }
 }
 
@@ -865,7 +871,8 @@ impl<T, C: Compare<T>> BinaryHeap<T, C> {
         self.data.pop().map(|mut item| {
             if !self.is_empty() {
                 swap(&mut item, &mut self.data[0]);
-                self.sift_down_to_bottom(0);
+                // SAFETY: !self.is_empty() means that self.len() > 0
+                unsafe { self.sift_down_to_bottom(0) };
             }
             item
         })
@@ -891,7 +898,9 @@ impl<T, C: Compare<T>> BinaryHeap<T, C> {
     pub fn push(&mut self, item: T) {
         let old_len = self.len();
         self.data.push(item);
-        self.sift_up(0, old_len);
+        // SAFETY: Since we pushed a new item it means that
+        //  old_len = self.len() - 1 < self.len()
+        unsafe { self.sift_up(0, old_len) };
     }
 
     /// Consumes the `BinaryHeap` and returns the underlying vector
@@ -946,7 +955,10 @@ impl<T, C: Compare<T>> BinaryHeap<T, C> {
                 let ptr = self.data.as_mut_ptr();
                 ptr::swap(ptr, ptr.add(end));
             }
-            self.sift_down_range(0, end);
+            // SAFETY: `end` goes from `self.len() - 1` to 1 (both included) so:
+            //  0 < 1 <= end <= self.len() - 1 < self.len()
+            //  Which means 0 < end and end < self.len().
+            unsafe { self.sift_down_range(0, end) };
         }
         self.into_vec()
     }
@@ -959,81 +971,139 @@ impl<T, C: Compare<T>> BinaryHeap<T, C> {
     // the hole is filled back at the end of its scope, even on panic.
     // Using a hole reduces the constant factor compared to using swaps,
     // which involves twice as many moves.
-    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(&mut self.data, pos);
-
-            while hole.pos() > start {
-                let parent = (hole.pos() - 1) / 2;
-                // if hole.element() <= hole.get(parent) {
-                if self.cmp.compare(hole.element(), hole.get(parent)) != Ordering::Greater {
-                    break;
-                }
-                hole.move_to(parent);
+
+    /// # Safety
+    ///
+    /// The caller must guarantee that `pos < self.len()`.
+    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()
+        let mut hole = unsafe { Hole::new(&mut self.data, pos) };
+
+        while hole.pos() > start {
+            let parent = (hole.pos() - 1) / 2;
+
+            // SAFETY: hole.pos() > start >= 0, which means hole.pos() > 0
+            //  and so hole.pos() - 1 can't underflow.
+            //  This guarantees that parent < hole.pos() so
+            //  it's a valid index and also != hole.pos().
+            if self
+                .cmp
+                .compare(hole.element(), unsafe { hole.get(parent) })
+                != Ordering::Greater
+            {
+                break;
             }
-            hole.pos()
+
+            // SAFETY: Same as above
+            unsafe { hole.move_to(parent) };
         }
+
+        hole.pos()
     }
 
     /// Take an element at `pos` and move it down the heap,
     /// while its children are larger.
-    fn sift_down_range(&mut self, pos: usize, end: usize) {
-        unsafe {
-            let mut hole = Hole::new(&mut self.data, pos);
-            let mut child = 2 * pos + 1;
-            while child < end - 1 {
-                // compare with the greater of the two children
-                // if !(hole.get(child) > hole.get(child + 1)) { child += 1 }
-                child += (self.cmp.compare(hole.get(child), hole.get(child + 1))
-                    != Ordering::Greater) as usize;
-                // if we are already in order, stop.
-                // if hole.element() >= hole.get(child) {
-                if self.cmp.compare(hole.element(), hole.get(child)) != Ordering::Less {
-                    return;
-                }
-                hole.move_to(child);
-                child = 2 * hole.pos() + 1;
-            }
-            if child == end - 1
-                && self.cmp.compare(hole.element(), hole.get(child)) == Ordering::Less
-            {
-                hole.move_to(child);
+    ///
+    /// # Safety
+    ///
+    /// The caller must guarantee that `pos < end <= self.len()`.
+    unsafe fn sift_down_range(&mut self, pos: usize, end: 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) {
+            // 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 {
+                self.cmp.compare(hole.get(child), hole.get(child + 1)) != Ordering::Greater
+            } as usize;
+
+            // if we are already in order, stop.
+            // 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 self.cmp.compare(hole.element(), unsafe { hole.get(child) }) != Ordering::Less {
+                return;
             }
+
+            // 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
+            && self.cmp.compare(hole.element(), unsafe { hole.get(child) }) == Ordering::Less
+        {
+            // SAFETY: child is already proven to be a valid index and
+            //  child == 2 * hole.pos() + 1 != hole.pos().
+            unsafe { hole.move_to(child) };
         }
     }
 
-    fn sift_down(&mut self, pos: usize) {
+    /// # Safety
+    ///
+    /// The caller must guarantee that `pos < self.len()`.
+    unsafe fn sift_down(&mut self, pos: usize) {
         let len = self.len();
-        self.sift_down_range(pos, len);
+        // SAFETY: pos < len is guaranteed by the caller and
+        //  obviously len = self.len() <= self.len().
+        unsafe { self.sift_down_range(pos, len) };
     }
 
     /// Take an element at `pos` and move it all the way down the heap,
     /// then sift it up to its position.
     ///
     /// Note: This is faster when the element is known to be large / should
     /// be closer to the bottom.
-    fn sift_down_to_bottom(&mut self, mut pos: usize) {
+    ///
+    /// # Safety
+    ///
+    /// The caller must guarantee that `pos < self.len()`.
+    unsafe fn sift_down_to_bottom(&mut self, mut pos: usize) {
         let end = self.len();
         let start = pos;
-        unsafe {
-            let mut hole = Hole::new(&mut self.data, pos);
-            let mut child = 2 * pos + 1;
-            while child < end - 1 {
-                let right = child + 1;
-                // compare with the greater of the two children
-                // if !(hole.get(child) > hole.get(right)) { child += 1 }
-                child += (self.cmp.compare(hole.get(child), hole.get(right)) != Ordering::Greater)
-                    as usize;
-                hole.move_to(child);
-                child = 2 * hole.pos() + 1;
-            }
-            if child == end - 1 {
-                hole.move_to(child);
-            }
-            pos = hole.pos;
+
+        // 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) {
+            // 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 {
+                self.cmp.compare(hole.get(child), hole.get(child + 1)) != Ordering::Greater
+            } 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().
+            unsafe { hole.move_to(child) };
         }
-        self.sift_up(start, pos);
+        pos = hole.pos();
+        drop(hole);
+
+        // SAFETY: pos is the position in the hole and was already proven
+        //  to be a valid index.
+        unsafe { self.sift_up(start, pos) };
     }
 
     /// Returns the length of the binary heap.
@@ -1129,7 +1199,10 @@ impl<T, C: Compare<T>> BinaryHeap<T, C> {
         let mut n = self.len() / 2;
         while n > 0 {
             n -= 1;
-            self.sift_down(n);
+            // SAFETY: n starts from self.len() / 2 and goes down to 0.
+            //  The only case when !(n < self.len()) is if
+            //  self.len() == 0, but it's ruled out by the loop condition.
+            unsafe { self.sift_down(n) };
         }
     }
 
@@ -1205,7 +1278,8 @@ impl<'a, T> Hole<'a, T> {
     #[inline]
     unsafe fn new(data: &'a mut [T], pos: usize) -> Self {
         debug_assert!(pos < data.len());
-        let elt = ptr::read(&data[pos]);
+        // SAFE: pos should be inside the slice
+        let elt = unsafe { ptr::read(data.get_unchecked(pos)) };
         Hole {
             data,
             elt: Some(elt),
@@ -1231,7 +1305,7 @@ impl<'a, T> Hole<'a, T> {
     unsafe fn get(&self, index: usize) -> &T {
         debug_assert!(index != self.pos);
         debug_assert!(index < self.data.len());
-        self.data.get_unchecked(index)
+        unsafe { self.data.get_unchecked(index) }
     }
 
     /// Move hole to new location
@@ -1241,9 +1315,12 @@ impl<'a, T> Hole<'a, T> {
     unsafe fn move_to(&mut self, index: usize) {
         debug_assert!(index != self.pos);
         debug_assert!(index < self.data.len());
-        let index_ptr: *const _ = self.data.get_unchecked(index);
-        let hole_ptr = self.data.get_unchecked_mut(self.pos);
-        ptr::copy_nonoverlapping(index_ptr, hole_ptr, 1);
+        unsafe {
+            let ptr = self.data.as_mut_ptr();
+            let index_ptr: *const _ = ptr.add(index);
+            let hole_ptr = ptr.add(self.pos);
+            ptr::copy_nonoverlapping(index_ptr, hole_ptr, 1);
+        }
         self.pos = index;
     }
 }