Skip to content

Commit

Permalink
collections: add append and extend specialization for binary heap
Browse files Browse the repository at this point in the history
  • Loading branch information
xosmig committed Apr 15, 2016
1 parent 74b3684 commit faeba73
Show file tree
Hide file tree
Showing 3 changed files with 121 additions and 5 deletions.
93 changes: 88 additions & 5 deletions src/libcollections/binary_heap.rs
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
}
}
Expand Down Expand Up @@ -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();

Expand Down
32 changes: 32 additions & 0 deletions src/libcollectionstest/binary_heap.rs
Expand Up @@ -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]);
}
1 change: 1 addition & 0 deletions src/libcollectionstest/lib.rs
Expand Up @@ -11,6 +11,7 @@
#![deny(warnings)]

#![feature(binary_heap_extras)]
#![feature(binary_heap_append)]
#![feature(box_syntax)]
#![feature(btree_range)]
#![feature(collections)]
Expand Down

0 comments on commit faeba73

Please sign in to comment.