From 78908f2e0963ceed461e9961020226d3e12c5978 Mon Sep 17 00:00:00 2001
From: Tim Vermeulen <tvermeulen@me.com>
Date: Tue, 3 Sep 2019 11:56:08 +0200
Subject: [PATCH] Override `StepBy::{try_fold, try_rfold}`

---
 src/libcore/iter/adapters/mod.rs | 44 +++++++++++++++++++++++++++++++-
 src/libcore/tests/iter.rs        | 35 +++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/src/libcore/iter/adapters/mod.rs b/src/libcore/iter/adapters/mod.rs
index f50781890ab22..7c9dde81a48fb 100644
--- a/src/libcore/iter/adapters/mod.rs
+++ b/src/libcore/iter/adapters/mod.rs
@@ -5,7 +5,7 @@ use crate::usize;
 use crate::intrinsics;
 
 use super::{Iterator, DoubleEndedIterator, ExactSizeIterator, FusedIterator, TrustedLen};
-use super::LoopState;
+use super::{LoopState, from_fn};
 
 mod chain;
 mod flatten;
@@ -541,6 +541,26 @@ impl<I> Iterator for StepBy<I> where I: Iterator {
             self.iter.nth(nth - 1);
         }
     }
+
+    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>,
+    {
+        #[inline]
+        fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ {
+            move || iter.nth(step)
+        }
+
+        if self.first_take {
+            self.first_take = false;
+            match self.iter.next() {
+                None => return Try::from_ok(acc),
+                Some(x) => acc = f(acc, x)?,
+            }
+        }
+        from_fn(nth(&mut self.iter, self.step)).try_fold(acc, f)
+    }
 }
 
 impl<I> StepBy<I> where I: ExactSizeIterator {
@@ -574,6 +594,28 @@ impl<I> DoubleEndedIterator for StepBy<I> where I: DoubleEndedIterator + ExactSi
             .saturating_add(self.next_back_index());
         self.iter.nth_back(n)
     }
+
+    fn try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R
+    where
+        F: FnMut(Acc, Self::Item) -> R,
+        R: Try<Ok = Acc>,
+    {
+        #[inline]
+        fn nth_back<I: DoubleEndedIterator>(
+            iter: &mut I,
+            step: usize,
+        ) -> impl FnMut() -> Option<I::Item> + '_ {
+            move || iter.nth_back(step)
+        }
+
+        match self.next_back() {
+            None => Try::from_ok(init),
+            Some(x) => {
+                let acc = f(init, x)?;
+                from_fn(nth_back(&mut self.iter, self.step)).try_fold(acc, f)
+            }
+        }
+    }
 }
 
 // StepBy can only make the iterator shorter, so the len will still fit.
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index 3a4f76852a0d7..f52899b500a3a 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -329,6 +329,23 @@ fn test_iterator_step_by_nth_overflow() {
     assert_eq!(it.0, (usize::MAX as Bigger) * 1);
 }
 
+#[test]
+fn test_iterator_step_by_nth_try_fold() {
+    let mut it = (0..).step_by(10);
+    assert_eq!(it.try_fold(0, i8::checked_add), None);
+    assert_eq!(it.next(), Some(60));
+    assert_eq!(it.try_fold(0, i8::checked_add), None);
+    assert_eq!(it.next(), Some(90));
+
+    let mut it = (100..).step_by(10);
+    assert_eq!(it.try_fold(50, i8::checked_add), None);
+    assert_eq!(it.next(), Some(110));
+
+    let mut it = (100..=100).step_by(10);
+    assert_eq!(it.next(), Some(100));
+    assert_eq!(it.try_fold(0, i8::checked_add), Some(0));
+}
+
 #[test]
 fn test_iterator_step_by_nth_back() {
     let mut it = (0..16).step_by(5);
@@ -354,6 +371,24 @@ fn test_iterator_step_by_nth_back() {
     assert_eq!(it().nth_back(42), None);
 }
 
+#[test]
+fn test_iterator_step_by_nth_try_rfold() {
+    let mut it = (0..100).step_by(10);
+    assert_eq!(it.try_rfold(0, i8::checked_add), None);
+    assert_eq!(it.next_back(), Some(70));
+    assert_eq!(it.next(), Some(0));
+    assert_eq!(it.try_rfold(0, i8::checked_add), None);
+    assert_eq!(it.next_back(), Some(30));
+
+    let mut it = (0..100).step_by(10);
+    assert_eq!(it.try_rfold(50, i8::checked_add), None);
+    assert_eq!(it.next_back(), Some(80));
+
+    let mut it = (100..=100).step_by(10);
+    assert_eq!(it.next_back(), Some(100));
+    assert_eq!(it.try_fold(0, i8::checked_add), Some(0));
+}
+
 #[test]
 #[should_panic]
 fn test_iterator_step_by_zero() {