Override Cycle::try_fold

Author: timvermeulen

src/libcore/iter/adapters/mod.rs

@@ -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")]

src/libcore/tests/iter.rs

@@ -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]