Skip to content

Specialize .zip() for efficient slice and slice iteration #33090

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Jun 17, 2016
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
specialize zip: Specialize .zip() for TrustedRandomAccess iterators
This allows common iterator compositions like a.zip(b) where a, b
are slice::{Iter, IterMut} compile to *much* better code.
  • Loading branch information
bluss committed Jun 14, 2016
commit a8f2e9b3597b4ff5dc6230c327b35a2b0e7122c1
3 changes: 2 additions & 1 deletion src/libcore/iter/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ use super::{Chain, Cycle, Cloned, Enumerate, Filter, FilterMap, FlatMap, Fuse,
use super::ChainState;
use super::{DoubleEndedIterator, ExactSizeIterator, Extend, FromIterator,
IntoIterator};
use super::ZipImpl;

fn _assert_is_object_safe(_: &Iterator<Item=()>) {}

Expand Down Expand Up @@ -383,7 +384,7 @@ pub trait Iterator {
fn zip<U>(self, other: U) -> Zip<Self, U::IntoIter> where
Self: Sized, U: IntoIterator
{
Zip{a: self, b: other.into_iter()}
Zip::new(self, other.into_iter())
}

/// Takes a closure and creates an iterator which calls that closure on each
Expand Down
139 changes: 119 additions & 20 deletions src/libcore/iter/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -302,6 +302,7 @@
use clone::Clone;
use cmp;
use fmt;
use iter_private::TrustedRandomAccess;
use ops::FnMut;
use option::Option::{self, Some, None};
use usize;
Expand Down Expand Up @@ -622,7 +623,9 @@ impl<A, B> DoubleEndedIterator for Chain<A, B> where
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Zip<A, B> {
a: A,
b: B
b: B,
index: usize,
len: usize,
}

#[stable(feature = "rust1", since = "1.0.0")]
Expand All @@ -631,29 +634,13 @@ impl<A, B> Iterator for Zip<A, B> where A: Iterator, B: Iterator
type Item = (A::Item, B::Item);

#[inline]
fn next(&mut self) -> Option<(A::Item, B::Item)> {
self.a.next().and_then(|x| {
self.b.next().and_then(|y| {
Some((x, y))
})
})
fn next(&mut self) -> Option<Self::Item> {
ZipImpl::next(self)
}

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let (a_lower, a_upper) = self.a.size_hint();
let (b_lower, b_upper) = self.b.size_hint();

let lower = cmp::min(a_lower, b_lower);

let upper = match (a_upper, b_upper) {
(Some(x), Some(y)) => Some(cmp::min(x,y)),
(Some(x), None) => Some(x),
(None, Some(y)) => Some(y),
(None, None) => None
};

(lower, upper)
ZipImpl::size_hint(self)
}
}

Expand All @@ -664,6 +651,51 @@ impl<A, B> DoubleEndedIterator for Zip<A, B> where
{
#[inline]
fn next_back(&mut self) -> Option<(A::Item, B::Item)> {
ZipImpl::next_back(self)
}
}

// Zip specialization trait
#[doc(hidden)]
trait ZipImpl<A, B> {
type Item;
fn new(a: A, b: B) -> Self;
fn next(&mut self) -> Option<Self::Item>;
fn size_hint(&self) -> (usize, Option<usize>);
fn next_back(&mut self) -> Option<Self::Item>
where A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator;
}

// General Zip impl
#[doc(hidden)]
impl<A, B> ZipImpl<A, B> for Zip<A, B>
where A: Iterator, B: Iterator
{
type Item = (A::Item, B::Item);
default fn new(a: A, b: B) -> Self {
Zip {
a: a,
b: b,
index: 0, // not used in general case
len: 0,
}
}

#[inline]
default fn next(&mut self) -> Option<(A::Item, B::Item)> {
self.a.next().and_then(|x| {
self.b.next().and_then(|y| {
Some((x, y))
})
})
}

#[inline]
default fn next_back(&mut self) -> Option<(A::Item, B::Item)>
where A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator
{
let a_sz = self.a.len();
let b_sz = self.b.len();
if a_sz != b_sz {
Expand All @@ -680,6 +712,73 @@ impl<A, B> DoubleEndedIterator for Zip<A, B> where
_ => unreachable!(),
}
}

#[inline]
default fn size_hint(&self) -> (usize, Option<usize>) {
let (a_lower, a_upper) = self.a.size_hint();
let (b_lower, b_upper) = self.b.size_hint();

let lower = cmp::min(a_lower, b_lower);

let upper = match (a_upper, b_upper) {
(Some(x), Some(y)) => Some(cmp::min(x,y)),
(Some(x), None) => Some(x),
(None, Some(y)) => Some(y),
(None, None) => None
};

(lower, upper)
}
}

#[doc(hidden)]
impl<A, B> ZipImpl<A, B> for Zip<A, B>
where A: TrustedRandomAccess, B: TrustedRandomAccess
{
fn new(a: A, b: B) -> Self {
let len = cmp::min(a.len(), b.len());
Zip {
a: a,
b: b,
index: 0,
len: len,
}
}

#[inline]
fn next(&mut self) -> Option<(A::Item, B::Item)> {
if self.index < self.len {
let i = self.index;
self.index += 1;
unsafe {
Some((self.a.get_unchecked(i), self.b.get_unchecked(i)))
}
} else {
None
}
}

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.len - self.index;
(len, Some(len))
}

#[inline]
fn next_back(&mut self) -> Option<(A::Item, B::Item)>
where A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator
{
if self.index < self.len {
self.len -= 1;
let i = self.len;
unsafe {
Some((self.a.get_unchecked(i), self.b.get_unchecked(i)))
}
} else {
None
}
}
}

#[stable(feature = "rust1", since = "1.0.0")]
Expand Down