Skip to content

Commit

Permalink
Revise benchmark test, and optimize fold logic to avoid use of for loop
Browse files Browse the repository at this point in the history
  • Loading branch information
Owen-CH-Leung authored and jswrenn committed Oct 5, 2023
1 parent c177dba commit b014b06
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 20 deletions.
4 changes: 3 additions & 1 deletion benches/bench1.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
26 changes: 7 additions & 19 deletions src/zip_longest.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,29 +54,17 @@ where
}

#[inline]
fn fold<B, F>(self, mut acc: B, mut f: F) -> B
fn fold<B, F>(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)))
}
}

Expand Down

0 comments on commit b014b06

Please sign in to comment.