about summary refs log tree commit diff
diff options
context:
space:
mode:
-rw-r--r--src/libcore/iter/range.rs46
-rw-r--r--src/libcore/tests/iter.rs20
2 files changed, 65 insertions, 1 deletions
diff --git a/src/libcore/iter/range.rs b/src/libcore/iter/range.rs
index 66a76a24df4..3b034efcce1 100644
--- a/src/libcore/iter/range.rs
+++ b/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};
@@ -397,6 +397,28 @@ 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;
+        if self.start <= self.end {
+            loop {
+                let (x, done) =
+                    if self.start < self.end {
+                        let n = self.start.add_one();
+                        (mem::replace(&mut self.start, n), false)
+                    } else {
+                        self.end.replace_zero();
+                        (self.start.replace_one(), true)
+                    };
+                accum = f(accum, x)?;
+                if done { break }
+            }
+        }
+        Try::from_ok(accum)
+    }
 }
 
 #[unstable(feature = "inclusive_range", reason = "recently added, follows RFC", issue = "28237")]
@@ -418,6 +440,28 @@ impl<A: Step> DoubleEndedIterator for ops::RangeInclusive<A> {
             _ => None,
         }
     }
+
+    #[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;
+        if self.start <= self.end {
+            loop {
+                let (x, done) =
+                    if self.start < self.end {
+                        let n = self.end.sub_one();
+                        (mem::replace(&mut self.end, n), false)
+                    } else {
+                        self.start.replace_one();
+                        (self.end.replace_zero(), true)
+                    };
+                accum = f(accum, x)?;
+                if done { break }
+            }
+        }
+        Try::from_ok(accum)
+    }
 }
 
 #[unstable(feature = "fused", issue = "35602")]
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index 8997cf9c6bf..e33a0b6224e 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -1398,6 +1398,26 @@ fn test_range_inclusive_min() {
 }
 
 #[test]
+fn test_range_inclusive_folds() {
+    assert_eq!((1..=10).sum::<i32>(), 55);
+    assert_eq!((1..=10).rev().sum::<i32>(), 55);
+
+    let mut it = 40..=50;
+    assert_eq!(it.try_fold(0, i8::checked_add), None);
+    assert_eq!(it, 44..=50);
+    assert_eq!(it.try_rfold(0, i8::checked_add), None);
+    assert_eq!(it, 44..=47);
+
+    let mut it = 10..=20;
+    assert_eq!(it.try_fold(0, |a,b| Some(a+b)), Some(165));
+    assert_eq!(it, 1..=0);
+
+    let mut it = 10..=20;
+    assert_eq!(it.try_rfold(0, |a,b| Some(a+b)), Some(165));
+    assert_eq!(it, 1..=0);
+}
+
+#[test]
 fn test_repeat() {
     let mut it = repeat(42);
     assert_eq!(it.next(), Some(42));