Override try_[r]fold for RangeInclusive

Because the last item needs special handling, it seems that LLVM has trouble canonicalizing the loops in external iteration.  With the override, it becomes obvious that the start==end case exits the loop (as opposed to the one *after* that exiting the loop in external iteration).

diff --git a/src/libcore/iter/range.rs b/src/libcore/iter/range.rs
index 66a76a24df45afc983395776ad5c69f83c651d38..3b034efcce14ccbcc972b01090d564a32a630517 100644 (file)
--- 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 @@ fn min(mut self) -> Option<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 @@ fn next_back(&mut self) -> Option<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 8997cf9c6bff9d1e50af8d14e7efeb6582b6012a..e33a0b6224e54b1022a60a2a28a9fc75b320eb4d 100644 (file)
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -1397,6 +1397,26 @@ fn test_range_inclusive_min() {
     assert_eq!(r.min(), None);
 }
 
+#[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);