Skip to content

Commit

Permalink
Make Iterator::unzip fast
Browse files Browse the repository at this point in the history
Closes rust-lang#72085

This consists of the following optimizations:

* Adds a `with_capacity` function to `Extend`. This definitely needs more thought if it's going to be stabilized, so I'm not writing an RFC yet. This takes off most of the performance gap.
* Optimizes `Vec`'s `Extend` implementation for the case where `size_hint` is 1. This shaves off the remaining performance gap.
  • Loading branch information
leo60228 committed May 13, 2020
1 parent d903a9d commit 1b90dab
Show file tree
Hide file tree
Showing 4 changed files with 24 additions and 3 deletions.
1 change: 1 addition & 0 deletions src/liballoc/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,7 @@
#![feature(alloc_layout_extra)]
#![feature(try_trait)]
#![feature(associated_type_bounds)]
#![feature(extend_with_capacity)]

// Allow testing this library

Expand Down
15 changes: 14 additions & 1 deletion src/liballoc/vec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2048,6 +2048,11 @@ impl<T> Extend<T> for Vec<T> {
fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
<Self as SpecExtend<T, I::IntoIter>>::spec_extend(self, iter.into_iter())
}

#[inline]
fn with_capacity(capacity: usize) -> Self {
Self::with_capacity(capacity)
}
}

// Specialization trait used for Vec::from_iter and Vec::extend
Expand Down Expand Up @@ -2097,7 +2102,7 @@ where
vector
}

default fn spec_extend(&mut self, iterator: I) {
default fn spec_extend(&mut self, mut iterator: I) {
// This is the case for a TrustedLen iterator.
let (low, high) = iterator.size_hint();
if let Some(high_value) = high {
Expand All @@ -2109,6 +2114,14 @@ where
);
}
if let Some(additional) = high {
if additional == 1 {
// work around inefficiencies in generic vec.extend(Some(x))
if let Some(x) = iterator.next() {
self.push(x);
}
return;
}

self.reserve(additional);
unsafe {
let mut ptr = self.as_mut_ptr().add(self.len());
Expand Down
6 changes: 6 additions & 0 deletions src/libcore/iter/traits/collect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -341,6 +341,12 @@ pub trait Extend<A> {
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
fn extend<T: IntoIterator<Item = A>>(&mut self, iter: T);

/// Creates this collection with a reserved capacity.
#[unstable(feature = "extend_with_capacity", issue = "none")]
fn with_capacity(_capacity: usize) -> Self where Self: Default {
Self::default()
}
}

#[stable(feature = "extend_for_unit", since = "1.28.0")]
Expand Down
5 changes: 3 additions & 2 deletions src/libcore/iter/traits/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2672,8 +2672,9 @@ pub trait Iterator {
}
}

let mut ts: FromA = Default::default();
let mut us: FromB = Default::default();
let cap = self.size_hint().0.saturating_add(1);
let mut ts: FromA = Extend::with_capacity(cap);
let mut us: FromB = Extend::with_capacity(cap);

self.for_each(extend(&mut ts, &mut us));

Expand Down

0 comments on commit 1b90dab

Please sign in to comment.