Auto merge of #156654 - AngelicosPhosphoros:angelicos_phosphoros/improve_range_inclusive_fold, r=<try>

Call visitor only in one place in bodies of folds for RangeInclusive

Author: bors

library/core/src/iter/range.rs

@@ -1325,18 +1325,33 @@ impl<T: TrustedStep> RangeInclusiveIteratorImpl for ops::RangeInclusive<T> {
         }
 
         let mut accum = init;
+        let mut i: T = self.start;
+        // We use loop with postcondition because
+        // it allows us to finish loop even when `self.end` is `T::MAX`.
+        // Precondition is checked before loop at the beginning of the function.
+        loop {
+            let current = i;
+            let is_last: bool = i == self.end;
+
+            // Need to update both fields before calling function
+            // because we can eagerly exit due to panic or ? operator.
+
+            // Writing to `self.exhausted` at every iteration is cheaper than branching.
+            // Especially because we need to write neighbouring field _anyway_.
+            self.exhausted = is_last;
+            // This should compile to conditional move for simple types.
+            self.start = if is_last {
+                current
+            } else {
+                // SAFETY: just checked precondition
+                i = unsafe { Step::forward_unchecked(i, 1) };
+                i
+            };
 
-        while self.start < self.end {
-            // SAFETY: just checked precondition
-            let n = unsafe { Step::forward_unchecked(self.start, 1) };
-            let n = mem::replace(&mut self.start, n);
-            accum = f(accum, n)?;
-        }
-
-        self.exhausted = true;
-
-        if self.start == self.end {
-            accum = f(accum, self.start)?;
+            accum = f(accum, current)?;
+            if is_last {
+                break;
+            }
         }
 
         try { accum }
@@ -1365,23 +1380,33 @@ impl<T: TrustedStep> RangeInclusiveIteratorImpl for ops::RangeInclusive<T> {
         F: FnMut(B, T) -> R,
         R: Try<Output = B>,
     {
+        // Code of this implementation is almost the same as `spec_try_fold`
+        // except it updates self.end and decrements counter.
+        // See it for comments for justification of code structure
+
         if self.is_empty() {
             return try { init };
         }
 
         let mut accum = init;
+        let mut i: T = self.end;
+        loop {
+            let current = i;
+            let is_last: bool = i == self.start;
+
+            self.exhausted = is_last;
+            self.end = if is_last {
+                current
+            } else {
+                // SAFETY: just checked precondition
+                i = unsafe { Step::backward_unchecked(i, 1) };
+                i
+            };
 
-        while self.start < self.end {
-            // SAFETY: just checked precondition
-            let n = unsafe { Step::backward_unchecked(self.end, 1) };
-            let n = mem::replace(&mut self.end, n);
-            accum = f(accum, n)?;
-        }
-
-        self.exhausted = true;
-
-        if self.start == self.end {
-            accum = f(accum, self.start)?;
+            accum = f(accum, current)?;
+            if is_last {
+                break;
+            }
         }
 
         try { accum }

library/coretests/tests/ascii_char.rs

@@ -38,3 +38,46 @@ fn test_extend() {
     s.extend(Char::CapitalA..=Char::CapitalC);
     assert_eq!(s, String::from("abcABC"));
 }
+
+#[test]
+fn test_range_inclusive_try_fold() {
+    // Visit every value once.
+    let mut it = Char::MIN..=Char::MAX;
+    let mut expected: u32 = 0;
+    it.try_fold((), |_, x| {
+        assert!(expected <= u32::from(Char::MAX));
+        assert_eq!(expected, u32::from(x));
+        expected += 1;
+        Some(())
+    });
+
+    let mut it = Char::MIN..=Char::MAX;
+    it.try_fold((), |_, x| (x < Char::CapitalA).then_some(()));
+    assert!(!it.is_empty());
+    assert_eq!(it.next(), Some(Char::CapitalB));
+
+    let mut it = Char::MIN..=Char::MAX;
+    it.try_fold((), |_, x| (x != Char::MAX).then_some(()));
+    assert!(it.is_empty());
+    assert_eq!(it.next(), None);
+
+    // Visit every value once.
+    let mut it = Char::MIN..=Char::MAX;
+    let mut expected: i32 = u8::from(Char::MAX).into();
+    it.try_rfold((), |_, x| {
+        assert!(expected >= 0);
+        assert_eq!(expected, i32::from(u8::from(x)));
+        expected -= 1;
+        Some(())
+    });
+
+    let mut it = Char::MIN..=Char::MAX;
+    it.try_rfold((), |_, x| (x > Char::CapitalB).then_some(()));
+    assert!(!it.is_empty());
+    assert_eq!(it.next_back(), Some(Char::CapitalA));
+
+    let mut it = Char::MIN..=Char::MAX;
+    it.try_rfold((), |_, x| (x != Char::MIN).then_some(()));
+    assert!(it.is_empty());
+    assert_eq!(it.next(), None);
+}

library/coretests/tests/iter/range.rs

@@ -422,6 +422,64 @@ fn test_range_inclusive_folds() {
     assert!(it.is_empty());
     assert_eq!(it.try_rfold(0, |a, b| Some(a + b)), Some(0));
     assert!(it.is_empty());
+
+    // No endless loop.
+    let mut it = 0u8..=u8::MAX;
+    assert_eq!(it.try_fold(0u32, |a, b| Some(a + u32::from(b))), Some(32640));
+    assert!(it.is_empty());
+    assert_eq!(it.try_fold(0, |a, b| Some(a + b)), Some(0));
+    assert!(it.is_empty());
+
+    // No endless loop.
+    let mut it = i8::MIN..=i8::MAX;
+    assert_eq!(it.try_rfold(0i32, |a, b| Some(a + i32::from(b))), Some(-128));
+    assert!(it.is_empty());
+    assert_eq!(it.try_rfold(0, |a, b| Some(a + b)), Some(0));
+    assert!(it.is_empty());
+
+    // Visit every value exactly once in order.
+    let mut expected: Option<u8> = Some(0);
+    let mut it = 0u8..=u8::MAX;
+    it.try_fold((), |_, x| {
+        assert!(expected.is_some());
+        assert_eq!(Some(x), expected);
+        expected = expected.and_then(|e| e.checked_add(1));
+        Some(())
+    });
+
+    // Visit every value exactly once in reverse order.
+    let mut expected: Option<u8> = Some(u8::MAX);
+    let mut it = 0u8..=u8::MAX;
+    it.try_rfold((), |_, x| {
+        assert!(expected.is_some());
+        assert_eq!(Some(x), expected);
+        expected = expected.and_then(|e| e.checked_sub(1));
+        Some(())
+    });
+
+    // Early stop updates state correctly.
+    let mut it = 0..=100;
+    it.try_fold((), |_, x| (x < 10).then_some(()));
+    assert!(!it.is_empty());
+    assert_eq!(it.next(), Some(11));
+
+    // Early stop updates state correctly at the end.
+    let mut it = 0..=100;
+    it.try_fold((), |_, x| (x < 100).then_some(()));
+    assert!(it.is_empty());
+    assert_eq!(it.next(), None);
+
+    // Early stop updates state correctly.
+    let mut it = 0..=100;
+    it.try_rfold((), |_, x| (x > 10).then_some(()));
+    assert!(!it.is_empty());
+    assert_eq!(it.next_back(), Some(9));
+
+    // Early stop updates state correctly at the end.
+    let mut it = 0..=100;
+    it.try_rfold((), |_, x| (x != 0).then_some(()));
+    assert!(it.is_empty());
+    assert_eq!(it.next(), None);
 }
 
 #[test]