Reimplement TrustedRandomAccess specializations for Zip::fold and Zip::try_fold

Previously these optimizations automatically fell out of the default implementation for `fold` and `try_fold`
because `next()` was specialized. Since that is gone we now need to implement it explicitly.

Author: the8472

library/core/src/iter/adapters/zip.rs

@@ -2,6 +2,7 @@ use crate::cmp;
 use crate::fmt::{self, Debug};
 use crate::iter::{DoubleEndedIterator, ExactSizeIterator, FusedIterator, Iterator};
 use crate::iter::{InPlaceIterable, SourceIter, TrustedLen};
+use crate::ops::Try;
 
 /// An iterator that iterates two other iterators simultaneously.
 ///
@@ -95,6 +96,14 @@ where
         (lower, upper)
     }
 
+    fn fold<T, F>(self, init: T, f: F) -> T
+    where
+        Self: Sized,
+        F: FnMut(T, Self::Item) -> T,
+    {
+        self.spec_fold(init, f)
+    }
+
     #[inline]
     #[doc(hidden)]
     unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item
@@ -137,6 +146,95 @@ where
     }
 }
 
+trait ZipSpec: Iterator {
+    fn spec_fold<T, F>(self, init: T, f: F) -> T
+    where
+        Self: Sized,
+        F: FnMut(T, Self::Item) -> T;
+
+    fn spec_try_fold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Output = B>;
+
+}
+
+impl<A, B> ZipSpec for Zip<A, B>
+where
+    A: Iterator,
+    B: Iterator,
+{
+    default fn spec_fold<T, F>(mut self, init: T, mut f: F) -> T
+    where
+        Self: Sized,
+        F: FnMut(T, Self::Item) -> T,
+    {
+        let mut accum = init;
+        while let Some(x) = self.next() {
+            accum = f(accum, x);
+        }
+        accum
+    }
+
+    default fn spec_try_fold<T, F, R>(&mut self, init: T, mut f: F) -> R
+    where
+        Self: Sized,
+        F: FnMut(T, Self::Item) -> R,
+        R: Try<Output = T>,
+    {
+        let mut accum = init;
+        while let Some(x) = self.next() {
+            accum = f(accum, x)?;
+        }
+        try { accum }
+    }
+}
+
+impl<A, B> ZipSpec for Zip<A, B>
+where
+    A: Iterator,
+    B: Iterator,
+    Self: TrustedRandomAccess,
+{
+    fn spec_fold<T, F>(mut self, init: T, mut f: F) -> T
+    where
+        Self: Sized,
+        F: FnMut(T, Self::Item) -> T,
+    {
+        let _ = self.advance_by(0);
+        let len = self.size();
+        let mut accum = init;
+        for i in 0..len {
+            // SAFETY: each item is only accessed once and we run the cleanup function afterwards
+            let x = unsafe { self.__iterator_get_unchecked(i) };
+            accum = f(accum, x);
+        }
+        // FIXME drop-guard or use ForLoopDesugar
+        self.cleanup(len, true);
+        accum
+    }
+
+    fn spec_try_fold<T, F, R>(&mut self, init: T, mut f: F) -> R
+    where
+        Self: Sized,
+        F: FnMut(T, Self::Item) -> R,
+        R: Try<Output = T>,
+    {
+        let _ = self.advance_by(0);
+        let len = self.size();
+        let mut accum = init;
+        for i in 0..len {
+            // SAFETY: each item is only accessed once and we run the cleanup function afterwards
+            let x = unsafe { self.__iterator_get_unchecked(i) };
+            accum = f(accum, x)?;
+        }
+        // FIXME drop-guard or use ForLoopDesugar
+        self.cleanup(len, true);
+        try { accum }
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<A, B> ExactSizeIterator for Zip<A, B>
 where

library/core/tests/iter/adapters/zip.rs

@@ -174,19 +174,6 @@ fn test_zip_map_rev_sideffectful() {
     assert_eq!(&ys, &[1, 1, 1, 1]);
 }
 
-#[test]
-fn test_zip_nested_sideffectful() {
-    let mut xs = [0; 6];
-    let ys = [0; 4];
-
-    {
-        // test that it has the side effect nested inside enumerate
-        let it = xs.iter_mut().map(|x| *x = 1).enumerate().zip(&ys);
-        it.count();
-    }
-    assert_eq!(&xs, &[1, 1, 1, 1, 1, 0]);
-}
-
 #[test]
 fn test_zip_nth_back_side_effects_exhausted() {
     let mut a = Vec::new();