Implement custom fold for ZipLongest

Add custom fold logic and benchmark

Author: Owen-CH-Leung

benches/bench1.rs

@@ -1,6 +1,6 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use itertools::free::cloned;
-use itertools::iproduct;
+use itertools::{EitherOrBoth, iproduct};
 use itertools::Itertools;
 
 use std::cmp;
@@ -390,6 +390,23 @@ fn zip_unchecked_counted_loop3(c: &mut Criterion) {
     });
 }
 
+fn ziplongest(c: &mut Criterion) {
+    c.bench_function("ziplongest", move |b| {
+        b.iter(|| {
+            let zip = (0..768).zip_longest(0..1024);
+            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 +838,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,31 @@ 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 acc: B, mut f: F) -> B
+    where
+        Self: Sized, F: FnMut(B, Self::Item) -> B
+    {
+        let ZipLongest { mut a, mut b } = self;
+
+        loop {
+            match (a.next(), b.next()) {
+                (Some(x), Some(y)) => acc = f(acc, EitherOrBoth::Both(x, y)),
+                (Some(x), None) => {
+                    acc = f(acc, EitherOrBoth::Left(x));
+                    // b is exhausted, so we can drain a.
+                    return a.fold(acc, |acc, x| f(acc, EitherOrBoth::Left(x)));
+                }
+                (None, Some(y)) => {
+                    acc = f(acc, EitherOrBoth::Right(y));
+                    // a is exhausted, so we can drain b.
+                    return b.fold(acc, |acc, y| f(acc, EitherOrBoth::Right(y)));
+                }
+                (None, None) => return acc, // Both iterators are exhausted.
+            }
+        }
+    }
 }
 
 impl<T, U> DoubleEndedIterator for ZipLongest<T, U>