From 6c400e04d6008353acf34cc1eaa13576b80f3e55 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Thu, 28 Sep 2023 20:50:10 +0800 Subject: [PATCH 1/4] Implement custom fold for ZipLongest Add custom fold logic and benchmark --- benches/bench1.rs | 20 +++++++++++++++++++- src/zip_longest.rs | 25 +++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/benches/bench1.rs b/benches/bench1.rs index 4f6e1963e..7cef765a1 100644 --- a/benches/bench1.rs +++ b/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, diff --git a/src/zip_longest.rs b/src/zip_longest.rs index 98ce4e63e..e4e4f90a4 100644 --- a/src/zip_longest.rs +++ b/src/zip_longest.rs @@ -52,6 +52,31 @@ where fn size_hint(&self) -> (usize, Option) { size_hint::max(self.a.size_hint(), self.b.size_hint()) } + + #[inline] + fn fold(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 DoubleEndedIterator for ZipLongest From 55f2830ef34dff48fac5013824988cbe8d54868f Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Thu, 28 Sep 2023 22:40:44 +0800 Subject: [PATCH 2/4] Run cargo fmt to fix the CI error --- benches/bench1.rs | 2 +- src/zip_longest.rs | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/benches/bench1.rs b/benches/bench1.rs index 7cef765a1..4dd1395dc 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::{EitherOrBoth, iproduct}; use itertools::Itertools; +use itertools::{iproduct, EitherOrBoth}; use std::cmp; use std::iter::repeat; diff --git a/src/zip_longest.rs b/src/zip_longest.rs index e4e4f90a4..d89e6edef 100644 --- a/src/zip_longest.rs +++ b/src/zip_longest.rs @@ -56,7 +56,8 @@ where #[inline] fn fold(self, mut acc: B, mut f: F) -> B where - Self: Sized, F: FnMut(B, Self::Item) -> B + Self: Sized, + F: FnMut(B, Self::Item) -> B, { let ZipLongest { mut a, mut b } = self; From e53a6e0ed182affdd201bfb11ebd0654f21f7991 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Fri, 29 Sep 2023 11:29:29 +0800 Subject: [PATCH 3/4] Add ziplongest to specializations.rs --- tests/specializations.rs | 4 ++++ 1 file changed, 4 insertions(+) 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! { From a42c4de4502e74afa5c0a51a7f377ccded0f7ca2 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Fri, 29 Sep 2023 22:48:26 +0800 Subject: [PATCH 4/4] Revise benchmark test, and optimize fold logic to avoid use of for loop --- benches/bench1.rs | 4 +++- src/zip_longest.rs | 26 +++++++------------------- 2 files changed, 10 insertions(+), 20 deletions(-) diff --git a/benches/bench1.rs b/benches/bench1.rs index 4dd1395dc..8a384ee69 100644 --- a/benches/bench1.rs +++ b/benches/bench1.rs @@ -391,9 +391,11 @@ 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 = (0..768).zip_longest(0..1024); + 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, diff --git a/src/zip_longest.rs b/src/zip_longest.rs index d89e6edef..27d9f3ab6 100644 --- a/src/zip_longest.rs +++ b/src/zip_longest.rs @@ -54,29 +54,17 @@ where } #[inline] - fn fold(self, mut acc: B, mut f: F) -> B + fn fold(self, mut init: 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. - } - } + 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))) } }