collections: Make BinaryHeap panic safe in sift_up / sift_down

Use a struct called Hole that keeps track of an invalid location
in the vector and fills the hole on drop.

I include a run-pass test that the current BinaryHeap fails, and the new
one passes.

Fixes #25842

Author: Ulrik Sverdrup

src/libcollections/binary_heap.rs

@@ -153,7 +153,7 @@
 use core::prelude::*;
 
 use core::iter::{FromIterator};
-use core::mem::{zeroed, replace, swap};
+use core::mem::swap;
 use core::ptr;
 
 use slice;
@@ -484,44 +484,43 @@ impl<T: Ord> BinaryHeap<T> {
 
     // The implementations of sift_up and sift_down use unsafe blocks in
     // order to move an element out of the vector (leaving behind a
-    // zeroed element), shift along the others and move it back into the
-    // vector over the junk element. This reduces the constant factor
-    // compared to using swaps, which involves twice as many moves.
-    fn sift_up(&mut self, start: usize, mut pos: usize) {
+    // hole), shift along the others and move the removed element back into the
+    // vector at the final location of the hole.
+    // The `Hole` type is used to represent this, and make sure
+    // 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) {
         unsafe {
-            let new = replace(&mut self.data[pos], zeroed());
+            // Take out the value at `pos` and create a hole.
+            let mut hole = Hole::new(&mut self.data, pos);
 
-            while pos > start {
-                let parent = (pos - 1) >> 1;
-
-                if new <= self.data[parent] { break; }
-
-                let x = replace(&mut self.data[parent], zeroed());
-                ptr::write(&mut self.data[pos], x);
-                pos = parent;
+            while hole.pos() > start {
+                let parent = (hole.pos() - 1) >> 1;
+                if hole.removed() <= hole.get(parent) { break }
+                hole.move_to(parent);
             }
-            ptr::write(&mut self.data[pos], new);
         }
     }
 
     fn sift_down_range(&mut self, mut pos: usize, end: usize) {
         unsafe {
             let start = pos;
-            let new = replace(&mut self.data[pos], zeroed());
-
-            let mut child = 2 * pos + 1;
-            while child < end {
-                let right = child + 1;
-                if right < end && !(self.data[child] > self.data[right]) {
-                    child = right;
+            {
+                let mut hole = Hole::new(&mut self.data, pos);
+                let mut child = 2 * pos + 1;
+                while child < end {
+                    let right = child + 1;
+                    if right < end && !(hole.get(child) > hole.get(right)) {
+                        child = right;
+                    }
+                    hole.move_to(child);
+                    child = 2 * hole.pos() + 1;
                 }
-                let x = replace(&mut self.data[child], zeroed());
-                ptr::write(&mut self.data[pos], x);
-                pos = child;
-                child = 2 * pos + 1;
+
+                pos = hole.pos;
             }
 
-            ptr::write(&mut self.data[pos], new);
             self.sift_up(start, pos);
         }
     }
@@ -554,6 +553,70 @@ impl<T: Ord> BinaryHeap<T> {
     pub fn clear(&mut self) { self.drain(); }
 }
 
+/// Hole represents a hole in a slice i.e. an index without valid value
+/// (because it was moved from or duplicated).
+/// In drop, `Hole` will restore the slice by filling the hole
+/// position with the value that was originally removed.
+struct Hole<'a, T: 'a> {
+    data: &'a mut [T],
+    elt: Option<T>,
+    pos: usize,
+}
+
+impl<'a, T> Hole<'a, T> {
+    /// Create a new Hole at index `pos`.
+    pub fn new(data: &'a mut [T], pos: usize) -> Self {
+        unsafe {
+            let elt = ptr::read(&data[pos]);
+            Hole {
+                data: data,
+                elt: Some(elt),
+                pos: pos,
+            }
+        }
+    }
+
+    #[inline(always)]
+    pub fn pos(&self) -> usize { self.pos }
+
+    /// Return a reference to the element removed
+    #[inline(always)]
+    pub fn removed(&self) -> &T {
+        self.elt.as_ref().unwrap()
+    }
+
+    /// Return a reference to the element at `index`.
+    ///
+    /// Panics if the index is out of bounds.
+    ///
+    /// Unsafe because index must not equal pos.
+    #[inline(always)]
+    pub unsafe fn get(&self, index: usize) -> &T {
+        debug_assert!(index != self.pos);
+        &self.data[index]
+    }
+
+    /// Move hole to new location
+    #[inline(always)]
+    pub unsafe fn move_to(&mut self, index: usize) {
+        debug_assert!(index != self.pos);
+        let old_pos = self.pos;
+        let x = ptr::read(&mut self.data[index]);
+        ptr::write(&mut self.data[old_pos], x);
+        self.pos = index;
+    }
+}
+
+impl<'a, T> Drop for Hole<'a, T> {
+    fn drop(&mut self) {
+        // fill the hole again
+        unsafe {
+            let pos = self.pos;
+            ptr::write(&mut self.data[pos], self.elt.take().unwrap());
+        }
+    }
+}
+
 /// `BinaryHeap` iterator.
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct Iter <'a, T: 'a> {

src/test/run-pass/binary-heap-panic-safe.rs

@@ -0,0 +1,71 @@
+
+#![feature(std_misc, collections, catch_panic, rand)]
+
+use std::__rand::{thread_rng, Rng};
+use std::thread;
+
+use std::collections::BinaryHeap;
+use std::cmp::Ordering;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+// old binaryheap failed this test
+//
+// Integrity means that all elements are present after a comparison panics,
+// even if the order may not be correct.
+fn test_integrity() {
+    #[derive(Eq, PartialEq, Ord, Clone, Debug)]
+    struct PanicOrd<T>(T, bool);
+
+    impl<T: PartialOrd> PartialOrd for PanicOrd<T> {
+        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+            if self.1 || other.1 {
+                panic!("Panicking comparison");
+            }
+            self.0.partial_cmp(&other.0)
+        }
+    }
+    let mut rng = thread_rng();
+    const DATASZ: usize = 32;
+    const NTEST: usize = 10;
+
+    // don't use 0 in the data -- we want to catch the zeroed-out case.
+    let data = (1..DATASZ + 1).collect::<Vec<_>>();
+
+    // since it's a fuzzy test, run several tries.
+    for _ in 0..NTEST {
+        for i in 1..DATASZ + 1 {
+            let mut panic_ords: Vec<_> = data.iter()
+                                             .filter(|&&x| x != i)
+                                             .map(|&x| PanicOrd(x, false))
+                                             .collect();
+            let panic_item = PanicOrd(i, true);
+
+            // heapify the sane items
+            rng.shuffle(&mut panic_ords);
+            let heap = Arc::new(Mutex::new(BinaryHeap::from_vec(panic_ords)));
+            let heap_ref = heap.clone();
+
+            // push the panicking item to the heap and catch the panic
+            let thread_result = thread::catch_panic(move || {
+                heap.lock().unwrap().push(panic_item);
+            });
+            assert!(thread_result.is_err());
+
+            // now fetch the binary heap again
+            let mutex_guard = match heap_ref.lock() {
+                Ok(x) => x,
+                Err(poison) => poison.into_inner(),
+            };
+            let inner_data = mutex_guard.clone().into_vec();
+            let mut data_sorted = inner_data.into_iter().map(|p| p.0).collect::<Vec<_>>();
+            data_sorted.sort();
+            assert_eq!(data_sorted, data);
+        }
+    }
+}
+
+fn main() {
+    test_integrity();
+}
+