diff --git a/src/libcollections/binary_heap.rs b/src/libcollections/binary_heap.rs index c9dd1efb37435..b01a09a2fe3a5 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 BinaryHeap { 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::() - (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 From> for BinaryHeap { fn from(vec: Vec) -> BinaryHeap { 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 where T: Ord { #[stable(feature = "rust1", since = "1.0.0")] impl Extend for BinaryHeap { + #[inline] fn extend>(&mut self, iter: I) { + >::spec_extend(self, iter); + } +} + +impl> SpecExtend for BinaryHeap { + default fn spec_extend(&mut self, iter: I) { + self.extend_desugared(iter.into_iter()); + } +} + +impl SpecExtend> for BinaryHeap { + fn spec_extend(&mut self, ref mut other: BinaryHeap) { + self.append(other); + } +} + +impl BinaryHeap { + fn extend_desugared>(&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 cc4366e8ae463..58194fe75f79a 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 211942f2294ec..4e08cab6db901 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)]