Auto merge of #108448 - ishitatsuyuki:binary-heap, r=Mark-Simulacrum

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.
This commit is contained in:
bors 2023-04-03 10:30:20 +00:00
commit f13ccefb5b

View File

@ -154,8 +154,6 @@
use crate::slice;
use crate::vec::{self, AsVecIntoIter, Vec};
use super::SpecExtend;
#[cfg(test)]
mod tests;
@ -400,6 +398,17 @@ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
}
}
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.
///
@ -850,30 +859,19 @@ pub fn retain<F>(&mut self, mut f: F)
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);
}
}
}
}
@ -1728,7 +1726,8 @@ fn into_iter(self) -> Iter<'a, 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]
@ -1742,37 +1741,6 @@ fn extend_reserve(&mut self, additional: usize) {
}
}
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) {