Skip to content

Commit

Permalink
Auto merge of rust-lang#108448 - ishitatsuyuki:binary-heap, r=Mark-Si…
Browse files Browse the repository at this point in the history
…mulacrum

binary_heap: Optimize Extend implementation.

This PR makes the `Extend` implementation for `BinaryHeap` no longer rely on specialization, so that it always use the bulk rebuild optimization that was previously only available for the `Vec` specialization.
  • Loading branch information
bors committed Apr 3, 2023
2 parents 932c173 + 116bb4d commit f13ccef
Showing 1 changed file with 19 additions and 51 deletions.
70 changes: 19 additions & 51 deletions library/alloc/src/collections/binary_heap/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -154,8 +154,6 @@ use crate::collections::TryReserveError;
use crate::slice;
use crate::vec::{self, AsVecIntoIter, Vec};

use super::SpecExtend;

#[cfg(test)]
mod tests;

Expand Down Expand Up @@ -400,6 +398,17 @@ impl<T: fmt::Debug> fmt::Debug for BinaryHeap<T> {
}
}

struct RebuildOnDrop<'a, T: Ord> {
heap: &'a mut BinaryHeap<T>,
rebuild_from: usize,
}

impl<'a, T: Ord> Drop for RebuildOnDrop<'a, T> {
fn drop(&mut self) {
self.heap.rebuild_tail(self.rebuild_from);
}
}

impl<T: Ord> BinaryHeap<T> {
/// Creates an empty `BinaryHeap` as a max-heap.
///
Expand Down Expand Up @@ -850,30 +859,19 @@ impl<T: Ord> BinaryHeap<T> {
where
F: FnMut(&T) -> bool,
{
struct RebuildOnDrop<'a, T: Ord> {
heap: &'a mut BinaryHeap<T>,
first_removed: usize,
}

let mut guard = RebuildOnDrop { first_removed: self.len(), heap: self };

// rebuild_start will be updated to the first touched element below, and the rebuild will
// only be done for the tail.
let mut guard = RebuildOnDrop { rebuild_from: self.len(), heap: self };
let mut i = 0;

guard.heap.data.retain(|e| {
let keep = f(e);
if !keep && i < guard.first_removed {
guard.first_removed = i;
if !keep && i < guard.rebuild_from {
guard.rebuild_from = i;
}
i += 1;
keep
});

impl<'a, T: Ord> Drop for RebuildOnDrop<'a, T> {
fn drop(&mut self) {
// data[..first_removed] is untouched, so we only need to
// rebuild the tail:
self.heap.rebuild_tail(self.first_removed);
}
}
}
}

Expand Down Expand Up @@ -1728,7 +1726,8 @@ impl<'a, T> IntoIterator for &'a BinaryHeap<T> {
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);
let guard = RebuildOnDrop { rebuild_from: self.len(), heap: self };
guard.heap.data.extend(iter);
}

#[inline]
Expand All @@ -1742,37 +1741,6 @@ impl<T: Ord> Extend<T> for BinaryHeap<T> {
}
}

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<Vec<T>> for BinaryHeap<T> {
fn spec_extend(&mut self, ref mut other: Vec<T>) {
let start = self.data.len();
self.data.append(other);
self.rebuild_tail(start);
}
}

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();

self.reserve(lower);

iterator.for_each(move |elem| self.push(elem));
}
}

#[stable(feature = "extend_ref", since = "1.2.0")]
impl<'a, T: 'a + Ord + Copy> Extend<&'a T> for BinaryHeap<T> {
fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
Expand Down

0 comments on commit f13ccef

Please sign in to comment.