Rollup merge of #64121 - timvermeulen:iter_step_by_internal, r=scottmcm

Override `StepBy::{try_fold, try_rfold}`

Previous PR: https://github.com/rust-lang/rust/pull/51435

The previous PR was closed in favor of https://github.com/rust-lang/rust/pull/51601, which was later reverted. I don't think these implementations will make it harder to specialize `StepBy<Range<_>>` later, so we should be able to land this without any consequences.

This should fix https://github.com/rust-lang/rust/issues/57517 – in my benchmarks `iter` and `iter.step_by(1)` now perform equally well, provided internal iteration is used.

src/libcore/iter/adapters/mod.rs

@@ -5,7 +5,7 @@
 use crate::intrinsics;
 
 use super::{Iterator, DoubleEndedIterator, ExactSizeIterator, FusedIterator, TrustedLen};
-use super::LoopState;
+use super::{LoopState, from_fn};
 
 mod chain;
 mod flatten;
@@ -534,6 +534,26 @@ fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
             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 {
@@ -567,6 +587,28 @@ fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
             .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.

src/libcore/tests/iter.rs

@@ -385,6 +385,23 @@ fn nth(&mut self, n: usize) -> Option<Self::Item> {
     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);
@@ -410,6 +427,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() {