diff --git a/benches/bench1.rs b/benches/bench1.rs index 4f6e1963e..8a384ee69 100644 --- a/benches/bench1.rs +++ b/benches/bench1.rs @@ -1,7 +1,7 @@ use criterion::{black_box, criterion_group, criterion_main, Criterion}; use itertools::free::cloned; -use itertools::iproduct; use itertools::Itertools; +use itertools::{iproduct, EitherOrBoth}; use std::cmp; use std::iter::repeat; @@ -390,6 +390,25 @@ fn zip_unchecked_counted_loop3(c: &mut Criterion) { }); } +fn ziplongest(c: &mut Criterion) { + let v1 = black_box((0..768).collect_vec()); + let v2 = black_box((0..1024).collect_vec()); + c.bench_function("ziplongest", move |b| { + b.iter(|| { + let zip = v1.iter().zip_longest(v2.iter()); + let sum = zip.fold(0u32, |mut acc, val| { + match val { + EitherOrBoth::Both(x, y) => acc += x * y, + EitherOrBoth::Left(x) => acc += x, + EitherOrBoth::Right(y) => acc += y, + } + acc + }); + sum + }) + }); +} + fn group_by_lazy_1(c: &mut Criterion) { let mut data = vec![0; 1024]; for (index, elt) in data.iter_mut().enumerate() { @@ -821,6 +840,7 @@ criterion_group!( zipdot_i32_unchecked_counted_loop, zipdot_f32_unchecked_counted_loop, zip_unchecked_counted_loop3, + ziplongest, group_by_lazy_1, group_by_lazy_2, slice_chunks, diff --git a/src/zip_longest.rs b/src/zip_longest.rs index 98ce4e63e..27d9f3ab6 100644 --- a/src/zip_longest.rs +++ b/src/zip_longest.rs @@ -52,6 +52,20 @@ where fn size_hint(&self) -> (usize, Option) { size_hint::max(self.a.size_hint(), self.b.size_hint()) } + + #[inline] + fn fold(self, mut init: B, mut f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + let ZipLongest { a, mut b } = self; + init = a.fold(init, |init, a| match b.next() { + Some(b) => f(init, EitherOrBoth::Both(a, b)), + None => f(init, EitherOrBoth::Left(a)), + }); + b.fold(init, |init, b| f(init, EitherOrBoth::Right(b))) + } } impl DoubleEndedIterator for ZipLongest diff --git a/tests/specializations.rs b/tests/specializations.rs index 85b46e48b..f4463a1b2 100644 --- a/tests/specializations.rs +++ b/tests/specializations.rs @@ -89,6 +89,10 @@ quickcheck! { a.truncate(6); test_specializations(&a.iter().powerset()) } + + fn zip_longest(a: Vec, b: Vec) -> () { + test_specializations(&a.into_iter().zip_longest(b)) + } } quickcheck! {