Specialize Zip::nth for TrustedRandomAccess

Makes the bench asked about on URLO 58x faster :)
This commit is contained in:
Scott McMurray 2018-03-01 01:57:25 -08:00
parent a85417f593
commit 70d5a4600b
3 changed files with 84 additions and 0 deletions

View File

@ -281,3 +281,32 @@ bench_sums! {
bench_take_while_chain_ref_sum,
(0i64..1000000).chain(1000000..).take_while(|&x| x < 1111111)
}
// Checks whether Skip<Zip<A,B>> is as fast as Zip<Skip<A>, Skip<B>>, from
// https://users.rust-lang.org/t/performance-difference-between-iterator-zip-and-skip-order/15743
#[bench]
fn bench_zip_then_skip(b: &mut Bencher) {
let v: Vec<_> = (0..100_000).collect();
let t: Vec<_> = (0..100_000).collect();
b.iter(|| {
let s = v.iter().zip(t.iter()).skip(10000)
.take_while(|t| *t.0 < 10100)
.map(|(a, b)| *a + *b)
.sum::<u64>();
assert_eq!(s, 2009900);
});
}
#[bench]
fn bench_skip_then_zip(b: &mut Bencher) {
let v: Vec<_> = (0..100_000).collect();
let t: Vec<_> = (0..100_000).collect();
b.iter(|| {
let s = v.iter().skip(10000).zip(t.iter().skip(10000))
.take_while(|t| *t.0 < 10100)
.map(|(a, b)| *a + *b)
.sum::<u64>();
assert_eq!(s, 2009900);
});
}

View File

@ -1045,6 +1045,11 @@ impl<A, B> Iterator for Zip<A, B> where A: Iterator, B: Iterator
fn size_hint(&self) -> (usize, Option<usize>) {
ZipImpl::size_hint(self)
}
#[inline]
fn nth(&mut self, n: usize) -> Option<Self::Item> {
ZipImpl::nth(self, n)
}
}
#[stable(feature = "rust1", since = "1.0.0")]
@ -1065,6 +1070,14 @@ trait ZipImpl<A, B> {
fn new(a: A, b: B) -> Self;
fn next(&mut self) -> Option<Self::Item>;
fn size_hint(&self) -> (usize, Option<usize>);
fn nth(&mut self, n: usize) -> Option<Self::Item>;
fn super_nth(&mut self, mut n: usize) -> Option<Self::Item> {
while let Some(x) = self.next() {
if n == 0 { return Some(x) }
n -= 1;
}
None
}
fn next_back(&mut self) -> Option<Self::Item>
where A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator;
@ -1094,6 +1107,12 @@ impl<A, B> ZipImpl<A, B> for Zip<A, B>
})
}
#[inline]
default fn nth(&mut self, n: usize) -> Option<Self::Item>
{
self.super_nth(n)
}
#[inline]
default fn next_back(&mut self) -> Option<(A::Item, B::Item)>
where A: DoubleEndedIterator + ExactSizeIterator,
@ -1174,6 +1193,25 @@ impl<A, B> ZipImpl<A, B> for Zip<A, B>
(len, Some(len))
}
#[inline]
fn nth(&mut self, n: usize) -> Option<Self::Item>
{
let delta = cmp::min(n, self.len - self.index);
let end = self.index + delta;
while self.index < end {
let i = self.index;
self.index += 1;
if A::may_have_side_effect() {
unsafe { self.a.get_unchecked(i); }
}
if B::may_have_side_effect() {
unsafe { self.b.get_unchecked(i); }
}
}
self.super_nth(n - delta)
}
#[inline]
fn next_back(&mut self) -> Option<(A::Item, B::Item)>
where A: DoubleEndedIterator + ExactSizeIterator,

View File

@ -144,6 +144,23 @@ fn test_iterator_chain_find() {
assert_eq!(iter.next(), None);
}
#[test]
fn test_zip_nth() {
let xs = [0, 1, 2, 4, 5];
let ys = [10, 11, 12];
let mut it = xs.iter().zip(&ys);
assert_eq!(it.nth(0), Some((&0, &10)));
assert_eq!(it.nth(1), Some((&2, &12)));
assert_eq!(it.nth(0), None);
let mut it = xs.iter().zip(&ys);
assert_eq!(it.nth(3), None);
let mut it = ys.iter().zip(&xs);
assert_eq!(it.nth(3), None);
}
#[test]
fn test_iterator_step_by() {
// Identity