Implement custom fold for zip longest

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,

src/zip_longest.rs

@@ -52,6 +52,20 @@ where
     fn size_hint(&self) -> (usize, Option<usize>) {
         size_hint::max(self.a.size_hint(), self.b.size_hint())
     }
+
+    #[inline]
+    fn fold<B, F>(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<T, U> DoubleEndedIterator for ZipLongest<T, U>

tests/specializations.rs

@@ -89,6 +89,10 @@ quickcheck! {
         a.truncate(6);
         test_specializations(&a.iter().powerset())
     }
+
+    fn zip_longest(a: Vec<u8>, b: Vec<u8>) -> () {
+        test_specializations(&a.into_iter().zip_longest(b))
+    }
 }
 
 quickcheck! {