Rollup merge of #58064 - llogiq:vec-deque-try-rfold, r=scottmcm

override `VecDeque::try_rfold`, also update iterator

This keeps the slice based iteration and updates the iterator state after each slice. It also uses a loop to reduce the amount of code.

This uses unsafe code, so some thorough review would be appreciated. Cc @RalfJung

3 files changed, 117 insertions, 5 deletions

diff --git a/src/liballoc/benches/vec_deque.rs b/src/liballoc/benches/vec_deque.rs
index f7aadbdbd82..7d2d3cfa612 100644
--- a/src/liballoc/benches/vec_deque.rs
+++ b/src/liballoc/benches/vec_deque.rs
@@ -45,3 +45,10 @@ fn bench_mut_iter_1000(b: &mut Bencher) {
         black_box(sum);
     })
 }
+
+#[bench]
+fn bench_try_fold(b: &mut Bencher) {
+    let ring: VecDeque<_> = (0..1000).collect();
+
+    b.iter(|| black_box(ring.iter().try_fold(0, |a, b| Some(a + b))))
+}
diff --git a/src/liballoc/collections/vec_deque.rs b/src/liballoc/collections/vec_deque.rs
index f778c4cbfde..4e90f783ec6 100644
--- a/src/liballoc/collections/vec_deque.rs
+++ b/src/liballoc/collections/vec_deque.rs
@@ -2170,12 +2170,29 @@ impl<'a, T> Iterator for Iter<'a, T> {
         back.iter().fold(accum, &mut f)
     }
 
-    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>
+    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 (front, back) = RingSlices::ring_slices(self.ring, self.head, self.tail);
-        let accum = front.iter().try_fold(init, &mut f)?;
-        back.iter().try_fold(accum, &mut f)
+        let (mut iter, final_res);
+        if self.tail <= self.head {
+            // single slice self.ring[self.tail..self.head]
+            iter = self.ring[self.tail..self.head].iter();
+            final_res = iter.try_fold(init, &mut f);
+        } else {
+            // two slices: self.ring[self.tail..], self.ring[..self.head]
+            let (front, back) = self.ring.split_at(self.tail);
+            let mut back_iter = back.iter();
+            let res = back_iter.try_fold(init, &mut f);
+            let len = self.ring.len();
+            self.tail = (self.ring.len() - back_iter.len()) & (len - 1);
+            iter = front[..self.head].iter();
+            final_res = iter.try_fold(res?, &mut f);
+        }
+        self.tail = self.head - iter.len();
+        final_res
     }
 }
 
@@ -2197,6 +2214,30 @@ impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
         accum = back.iter().rfold(accum, &mut f);
         front.iter().rfold(accum, &mut f)
     }
+
+    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 iter, final_res);
+        if self.tail <= self.head {
+            // single slice self.ring[self.tail..self.head]
+            iter = self.ring[self.tail..self.head].iter();
+            final_res = iter.try_rfold(init, &mut f);
+        } else {
+            // two slices: self.ring[self.tail..], self.ring[..self.head]
+            let (front, back) = self.ring.split_at(self.tail);
+            let mut front_iter = front[..self.head].iter();
+            let res = front_iter.try_rfold(init, &mut f);
+            self.head = front_iter.len();
+            iter = back.iter();
+            final_res = iter.try_rfold(res?, &mut f);
+        }
+        self.head = self.tail + iter.len();
+        final_res
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
diff --git a/src/liballoc/tests/vec_deque.rs b/src/liballoc/tests/vec_deque.rs
index e0cb0e7a9e7..16ddc1444fc 100644
--- a/src/liballoc/tests/vec_deque.rs
+++ b/src/liballoc/tests/vec_deque.rs
@@ -1465,6 +1465,15 @@ fn test_try_fold_unit() {
     assert_eq!(Some(()), v.into_iter().try_fold((), |(), ()| Some(())));
 }
 
+
+#[test]
+fn test_try_fold_unit_none() {
+    let v: std::collections::VecDeque<()> = [(); 10].iter().cloned().collect();
+    let mut iter = v.into_iter();
+    assert!(iter.try_fold((), |_, _| None).is_none());
+    assert_eq!(iter.len(), 9);
+}
+
 #[test]
 fn test_try_fold_rotated() {
     let mut v: VecDeque<_> = (0..12).collect();
@@ -1477,3 +1486,58 @@ fn test_try_fold_rotated() {
         assert_eq!(Ok::<_, ()>(66), v.iter().try_fold(0, |a, b| Ok(a + b)));
     }
 }
+
+#[test]
+fn test_try_fold_moves_iter() {
+    let v: VecDeque<_> = [10, 20, 30, 40, 100, 60, 70, 80, 90].iter().collect();
+    let mut iter = v.into_iter();
+    assert_eq!(iter.try_fold(0_i8, |acc, &x| acc.checked_add(x)), None);
+    assert_eq!(iter.next(), Some(&60));
+}
+
+#[test]
+fn test_try_fold_exhaust_wrap() {
+    let mut v = VecDeque::with_capacity(7);
+    v.push_back(1);
+    v.push_back(1);
+    v.push_back(1);
+    v.pop_front();
+    v.pop_front();
+    let mut iter = v.iter();
+    let _ = iter.try_fold(0, |_, _| Some(1));
+    assert!(iter.is_empty());
+}
+
+#[test]
+fn test_try_fold_wraparound() {
+    let mut v = VecDeque::with_capacity(8);
+    v.push_back(7);
+    v.push_back(8);
+    v.push_back(9);
+    v.push_front(2);
+    v.push_front(1);
+    let mut iter = v.iter();
+    let _ = iter.find(|&&x| x == 2);
+    assert_eq!(Some(&7), iter.next());
+}
+
+#[test]
+fn test_try_rfold_rotated() {
+    let mut v: VecDeque<_> = (0..12).collect();
+    for n in 0..10 {
+        if n & 1 == 0 {
+            v.rotate_left(n);
+        } else {
+            v.rotate_right(n);
+        }
+        assert_eq!(Ok::<_, ()>(66), v.iter().try_rfold(0, |a, b| Ok(a + b)));
+    }
+}
+
+#[test]
+fn test_try_rfold_moves_iter() {
+    let v : VecDeque<_> = [10, 20, 30, 40, 100, 60, 70, 80, 90].iter().collect();
+    let mut iter = v.into_iter();
+    assert_eq!(iter.try_rfold(0_i8, |acc, &x| acc.checked_add(x)), None);
+    assert_eq!(iter.next_back(), Some(&70));
+}