Override <RangeInclusive as Iterator>::try_(r)fold

Author: sinkuu

src/libcore/benches/iter.rs

@@ -310,3 +310,39 @@ fn bench_skip_then_zip(b: &mut Bencher) {
         assert_eq!(s, 2009900);
     });
 }
+
+#[bench]
+fn bench_triples_inclusive_range(b: &mut Bencher) {
+    // Example from #45222
+    b.iter(|| {
+        (1..)
+            .flat_map(|z| {
+                (1u32..=z).flat_map(move |x| {
+                    (x..=z)
+                        .filter(move |&y| x.pow(2) + y.pow(2) == z.pow(2))
+                        .map(move |y| (x, y, z))
+                })
+            })
+            .take(20)
+            .map(|(x, y, z)| x + y + z)
+            .sum::<u32>()
+    })
+}
+
+#[bench]
+fn bench_triples_range(b: &mut Bencher) {
+    // Example from #45222
+    b.iter(|| {
+        (1..)
+            .flat_map(|z| {
+                (1u32..z+1).flat_map(move |x| {
+                    (x..z+1)
+                        .filter(move |&y| x.pow(2) + y.pow(2) == z.pow(2))
+                        .map(move |y| (x, y, z))
+                })
+            })
+            .take(20)
+            .map(|(x, y, z)| x + y + z)
+            .sum::<u32>()
+    })
+}

src/libcore/iter/range.rs

@@ -10,7 +10,7 @@
 
 use convert::TryFrom;
 use mem;
-use ops::{self, Add, Sub};
+use ops::{self, Add, Sub, Try};
 use usize;
 
 use super::{FusedIterator, TrustedLen};
@@ -406,6 +406,26 @@ impl<A: Step> Iterator for ops::RangeInclusive<A> {
     fn max(mut self) -> Option<A> {
         self.next_back()
     }
+
+    #[inline]
+    fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R where
+        Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try<Ok=B>
+    {
+        let mut accum = init;
+
+        self.compute_is_empty();
+        if self.is_empty.unwrap_or_default() {
+            return Try::from_ok(accum);
+        }
+
+        while self.start < self.end {
+            let n = self.start.add_one();
+            accum = f(accum, mem::replace(&mut self.start, n))?;
+        }
+
+        self.is_empty = Some(true);
+        f(accum, self.start.clone())
+    }
 }
 
 #[stable(feature = "inclusive_range", since = "1.26.0")]
@@ -425,6 +445,26 @@ impl<A: Step> DoubleEndedIterator for ops::RangeInclusive<A> {
             self.end.clone()
         })
     }
+
+    #[inline]
+    fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R where
+        Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try<Ok=B>
+    {
+        let mut accum = init;
+
+        self.compute_is_empty();
+        if self.is_empty.unwrap_or_default() {
+            return Try::from_ok(accum);
+        }
+
+        while self.start < self.end {
+            let n = self.end.sub_one();
+            accum = f(accum, mem::replace(&mut self.end, n))?;
+        }
+
+        self.is_empty = Some(true);
+        f(accum, self.end.clone())
+    }
 }
 
 #[stable(feature = "fused", since = "1.26.0")]