Override Cycle::try_fold

This commit is contained in:
Tim Vermeulen 2019-07-17 01:16:49 +02:00
parent d36b7f6944
commit 688c11216a
2 changed files with 42 additions and 0 deletions

View File

@ -387,6 +387,36 @@ impl<I> Iterator for Cycle<I> where I: Clone + Iterator {
_ => (usize::MAX, None)
}
}
#[inline]
fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
where
F: FnMut(Acc, Self::Item) -> R,
R: Try<Ok = Acc>,
{
// fully iterate the current iterator. this is necessary because
// `self.iter` may be empty even when `self.orig` isn't
acc = self.iter.try_fold(acc, &mut f)?;
self.iter = self.orig.clone();
// complete a full cycle, keeping track of whether the cycled
// iterator is empty or not. we need to return early in case
// of an empty iterator to prevent an infinite loop
let mut is_empty = true;
acc = self.iter.try_fold(acc, |acc, x| {
is_empty = false;
f(acc, x)
})?;
if is_empty {
return Try::from_ok(acc);
}
loop {
self.iter = self.orig.clone();
acc = self.iter.try_fold(acc, &mut f)?;
}
}
}
#[stable(feature = "fused", since = "1.26.0")]

View File

@ -1015,6 +1015,18 @@ fn test_cycle() {
assert_eq!(empty::<i32>().cycle().fold(0, |acc, x| acc + x), 0);
assert_eq!(once(1).cycle().skip(1).take(4).fold(0, |acc, x| acc + x), 4);
assert_eq!((0..10).cycle().take(5).sum::<i32>(), 10);
assert_eq!((0..10).cycle().take(15).sum::<i32>(), 55);
assert_eq!((0..10).cycle().take(25).sum::<i32>(), 100);
let mut iter = (0..10).cycle();
iter.nth(14);
assert_eq!(iter.take(8).sum::<i32>(), 38);
let mut iter = (0..10).cycle();
iter.nth(9);
assert_eq!(iter.take(3).sum::<i32>(), 3);
}
#[test]