collections: add append and extend specialization for binary heap

3 files changed, 121 insertions, 5 deletions

diff --git a/src/libcollections/binary_heap.rs b/src/libcollections/binary_heap.rs
index c9dd1efb374..b01a09a2fe3 100644
--- a/src/libcollections/binary_heap.rs
+++ b/src/libcollections/binary_heap.rs
@@ -153,12 +153,15 @@
 
 use core::iter::FromIterator;
 use core::mem::swap;
+use core::mem::size_of;
 use core::ptr;
 use core::fmt;
 
 use slice;
 use vec::{self, Vec};
 
+use super::SpecExtend;
+
 /// A priority queue implemented with a binary heap.
 ///
 /// This will be a max-heap.
@@ -738,6 +741,71 @@ impl<T: Ord> BinaryHeap<T> {
     pub fn clear(&mut self) {
         self.drain();
     }
+
+    fn rebuild(&mut self) {
+        let mut n = self.len() / 2;
+        while n > 0 {
+            n -= 1;
+            self.sift_down(n);
+        }
+    }
+
+    /// Moves all the elements of `other` into `self`, leaving `other` empty.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(binary_heap_append)]
+    ///
+    /// use std::collections::BinaryHeap;
+    ///
+    /// let v = vec![-10, 1, 2, 3, 3];
+    /// let mut a = BinaryHeap::from(v);
+    ///
+    /// let v = vec![-20, 5, 43];
+    /// let mut b = BinaryHeap::from(v);
+    ///
+    /// a.append(&mut b);
+    ///
+    /// assert_eq!(a.into_sorted_vec(), [-20, -10, 1, 2, 3, 3, 5, 43]);
+    /// assert!(b.is_empty());
+    /// ```
+    #[unstable(feature = "binary_heap_append",
+           reason = "needs to be audited",
+           issue = "32526")]
+    pub fn append(&mut self, other: &mut Self) {
+        if self.len() < other.len() {
+            swap(self, other);
+        }
+
+        if other.is_empty() {
+            return;
+        }
+
+        #[inline(always)]
+        fn log2_fast(x: usize) -> usize {
+            8 * size_of::<usize>() - (x.leading_zeros() as usize) - 1
+        }
+
+        // `rebuild` takes O(len1 + len2) operations
+        // and about 2 * (len1 + len2) comparisons in the worst case
+        // while `extend` takes O(len2 * log_2(len1)) operations
+        // and about 1 * len2 * log_2(len1) comparisons in the worst case,
+        // assuming len1 >= len2.
+        #[inline]
+        fn better_to_rebuild(len1: usize, len2: usize) -> bool {
+            2 * (len1 + len2) < len2 * log2_fast(len1)
+        }
+
+        if better_to_rebuild(self.len(), other.len()) {
+            self.data.append(&mut other.data);
+            self.rebuild();
+        } else {
+            self.extend(other.drain());
+        }
+    }
 }
 
 /// Hole represents a hole in a slice i.e. an index without valid value
@@ -917,11 +985,7 @@ impl<'a, T: 'a> ExactSizeIterator for Drain<'a, T> {}
 impl<T: Ord> From<Vec<T>> for BinaryHeap<T> {
     fn from(vec: Vec<T>) -> BinaryHeap<T> {
         let mut heap = BinaryHeap { data: vec };
-        let mut n = heap.len() / 2;
-        while n > 0 {
-            n -= 1;
-            heap.sift_down(n);
-        }
+        heap.rebuild();
         heap
     }
 }
@@ -980,7 +1044,26 @@ impl<'a, T> IntoIterator for &'a BinaryHeap<T> where T: Ord {
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<T: Ord> Extend<T> for BinaryHeap<T> {
+    #[inline]
     fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
+        <Self as SpecExtend<I>>::spec_extend(self, iter);
+    }
+}
+
+impl<T: Ord, I: IntoIterator<Item = T>> SpecExtend<I> for BinaryHeap<T> {
+    default fn spec_extend(&mut self, iter: I) {
+        self.extend_desugared(iter.into_iter());
+    }
+}
+
+impl<T: Ord> SpecExtend<BinaryHeap<T>> for BinaryHeap<T> {
+    fn spec_extend(&mut self, ref mut other: BinaryHeap<T>) {
+        self.append(other);
+    }
+}
+
+impl<T: Ord> BinaryHeap<T> {
+    fn extend_desugared<I: IntoIterator<Item = T>>(&mut self, iter: I) {
         let iterator = iter.into_iter();
         let (lower, _) = iterator.size_hint();
 
diff --git a/src/libcollectionstest/binary_heap.rs b/src/libcollectionstest/binary_heap.rs
index cc4366e8ae4..58194fe75f7 100644
--- a/src/libcollectionstest/binary_heap.rs
+++ b/src/libcollectionstest/binary_heap.rs
@@ -242,3 +242,35 @@ fn test_extend_ref() {
     assert_eq!(a.len(), 5);
     assert_eq!(a.into_sorted_vec(), [1, 2, 3, 4, 5]);
 }
+
+#[test]
+fn test_append() {
+    let mut a = BinaryHeap::from(vec![-10, 1, 2, 3, 3]);
+    let mut b = BinaryHeap::from(vec![-20, 5, 43]);
+
+    a.append(&mut b);
+
+    assert_eq!(a.into_sorted_vec(), [-20, -10, 1, 2, 3, 3, 5, 43]);
+    assert!(b.is_empty());
+}
+
+#[test]
+fn test_append_to_empty() {
+    let mut a = BinaryHeap::new();
+    let mut b = BinaryHeap::from(vec![-20, 5, 43]);
+
+    a.append(&mut b);
+
+    assert_eq!(a.into_sorted_vec(), [-20, 5, 43]);
+    assert!(b.is_empty());
+}
+
+#[test]
+fn test_extend_specialization() {
+    let mut a = BinaryHeap::from(vec![-10, 1, 2, 3, 3]);
+    let b = BinaryHeap::from(vec![-20, 5, 43]);
+
+    a.extend(b);
+
+    assert_eq!(a.into_sorted_vec(), [-20, -10, 1, 2, 3, 3, 5, 43]);
+}
diff --git a/src/libcollectionstest/lib.rs b/src/libcollectionstest/lib.rs
index 211942f2294..4e08cab6db9 100644
--- a/src/libcollectionstest/lib.rs
+++ b/src/libcollectionstest/lib.rs
@@ -11,6 +11,7 @@
 #![deny(warnings)]
 
 #![feature(binary_heap_extras)]
+#![feature(binary_heap_append)]
 #![feature(box_syntax)]
 #![feature(btree_range)]
 #![feature(collections)]