From 688c11216aca1d7449b07b3ebbcee3ba114d0d51 Mon Sep 17 00:00:00 2001
From: Tim Vermeulen <tvermeulen@me.com>
Date: Wed, 17 Jul 2019 01:16:49 +0200
Subject: [PATCH] Override Cycle::try_fold

---
 src/libcore/iter/adapters/mod.rs | 30 ++++++++++++++++++++++++++++++
 src/libcore/tests/iter.rs        | 12 ++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/src/libcore/iter/adapters/mod.rs b/src/libcore/iter/adapters/mod.rs
index c2edcd22f95..07526e58fd7 100644
--- a/src/libcore/iter/adapters/mod.rs
+++ b/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")]
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index b7b0849e212..19497566106 100644
--- a/src/libcore/tests/iter.rs
+++ b/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]