Rollup merge of #72139 - nnethercote:standalone-fold, r=cuviper

Make `fold` standalone.

`fold` is currently implemented via `try_fold`, but implementing it
directly results in slightly less LLVM IR being generated, speeding up
compilation of some benchmarks.

r? @cuviper

src/libcore/iter/adapters/mod.rs

@@ -512,6 +512,9 @@ fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
             acc = self.iter.try_fold(acc, &mut f)?;
         }
     }
+
+    // No `fold` override, because `fold` doesn't make much sense for `Cycle`,
+    // and we can't do anything better than the default.
 }
 
 #[stable(feature = "fused", since = "1.26.0")]
@@ -643,6 +646,25 @@ fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item
         }
         from_fn(nth(&mut self.iter, self.step)).try_fold(acc, f)
     }
+
+    fn fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc
+    where
+        F: FnMut(Acc, Self::Item) -> Acc,
+    {
+        #[inline]
+        fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ {
+            move || iter.nth(step)
+        }
+
+        if self.first_take {
+            self.first_take = false;
+            match self.iter.next() {
+                None => return acc,
+                Some(x) => acc = f(acc, x),
+            }
+        }
+        from_fn(nth(&mut self.iter, self.step)).fold(acc, f)
+    }
 }
 
 impl<I> StepBy<I>
@@ -702,6 +724,29 @@ fn nth_back<I: DoubleEndedIterator>(
             }
         }
     }
+
+    #[inline]
+    fn rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc
+    where
+        Self: Sized,
+        F: FnMut(Acc, Self::Item) -> Acc,
+    {
+        #[inline]
+        fn nth_back<I: DoubleEndedIterator>(
+            iter: &mut I,
+            step: usize,
+        ) -> impl FnMut() -> Option<I::Item> + '_ {
+            move || iter.nth_back(step)
+        }
+
+        match self.next_back() {
+            None => init,
+            Some(x) => {
+                let acc = f(init, x);
+                from_fn(nth_back(&mut self.iter, self.step)).fold(acc, f)
+            }
+        }
+    }
 }
 
 // StepBy can only make the iterator shorter, so the len will still fit.
@@ -1767,6 +1812,20 @@ fn check<'a, T, Acc, R: Try<Ok = Acc>>(
             self.iter.try_fold(init, check(flag, p, fold)).into_try()
         }
     }
+
+    #[inline]
+    fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
+    where
+        Self: Sized,
+        Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        #[inline]
+        fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
+            move |acc, x| Ok(f(acc, x))
+        }
+
+        self.try_fold(init, ok(fold)).unwrap()
+    }
 }
 
 #[stable(feature = "fused", since = "1.26.0")]
@@ -1838,6 +1897,20 @@ fn try_fold<Acc, Fold, R>(&mut self, init: Acc, mut fold: Fold) -> R
         })
         .into_try()
     }
+
+    #[inline]
+    fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
+    where
+        Self: Sized,
+        Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        #[inline]
+        fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
+            move |acc, x| Ok(f(acc, x))
+        }
+
+        self.try_fold(init, ok(fold)).unwrap()
+    }
 }
 
 /// An iterator that skips over `n` elements of `iter`.
@@ -2006,6 +2079,18 @@ fn check<T, Acc, R: Try<Ok = Acc>>(
             self.iter.try_rfold(init, check(n, fold)).into_try()
         }
     }
+
+    fn rfold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
+    where
+        Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        #[inline]
+        fn ok<Acc, T>(mut f: impl FnMut(Acc, T) -> Acc) -> impl FnMut(Acc, T) -> Result<Acc, !> {
+            move |acc, x| Ok(f(acc, x))
+        }
+
+        self.try_rfold(init, ok(fold)).unwrap()
+    }
 }
 
 #[stable(feature = "fused", since = "1.26.0")]
@@ -2105,6 +2190,20 @@ fn check<'a, T, Acc, R: Try<Ok = Acc>>(
             self.iter.try_fold(init, check(n, fold)).into_try()
         }
     }
+
+    #[inline]
+    fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
+    where
+        Self: Sized,
+        Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        #[inline]
+        fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
+            move |acc, x| Ok(f(acc, x))
+        }
+
+        self.try_fold(init, ok(fold)).unwrap()
+    }
 }
 
 #[stable(feature = "double_ended_take_iterator", since = "1.38.0")]
@@ -2156,6 +2255,24 @@ fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
             }
         }
     }
+
+    #[inline]
+    fn rfold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
+    where
+        Self: Sized,
+        Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        if self.n == 0 {
+            init
+        } else {
+            let len = self.iter.len();
+            if len > self.n && self.iter.nth_back(len - self.n - 1).is_none() {
+                init
+            } else {
+                self.iter.rfold(init, fold)
+            }
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -2237,6 +2354,20 @@ fn scan<'a, T, St, B, Acc, R: Try<Ok = Acc>>(
         let f = &mut self.f;
         self.iter.try_fold(init, scan(state, f, fold)).into_try()
     }
+
+    #[inline]
+    fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
+    where
+        Self: Sized,
+        Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        #[inline]
+        fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
+            move |acc, x| Ok(f(acc, x))
+        }
+
+        self.try_fold(init, ok(fold)).unwrap()
+    }
 }
 
 /// An iterator that calls a function with a reference to each element before
@@ -2444,4 +2575,17 @@ fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
             })
             .into_try()
     }
+
+    fn fold<B, F>(mut self, init: B, fold: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        #[inline]
+        fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
+            move |acc, x| Ok(f(acc, x))
+        }
+
+        self.try_fold(init, ok(fold)).unwrap()
+    }
 }

src/libcore/iter/range.rs

@@ -658,6 +658,20 @@ fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
         Try::from_ok(accum)
     }
 
+    #[inline]
+    fn fold<B, F>(mut self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        #[inline]
+        fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
+            move |acc, x| Ok(f(acc, x))
+        }
+
+        self.try_fold(init, ok(f)).unwrap()
+    }
+
     #[inline]
     fn last(mut self) -> Option<A> {
         self.next_back()
@@ -746,6 +760,20 @@ fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
 
         Try::from_ok(accum)
     }
+
+    #[inline]
+    fn rfold<B, F>(mut self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        #[inline]
+        fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
+            move |acc, x| Ok(f(acc, x))
+        }
+
+        self.try_rfold(init, ok(f)).unwrap()
+    }
 }
 
 #[unstable(feature = "trusted_len", issue = "37572")]

src/libcore/iter/traits/double_ended.rs

@@ -221,17 +221,16 @@ fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
     /// ```
     #[inline]
     #[stable(feature = "iter_rfold", since = "1.27.0")]
-    fn rfold<B, F>(mut self, accum: B, f: F) -> B
+    fn rfold<B, F>(mut self, init: B, mut f: F) -> B
     where
         Self: Sized,
         F: FnMut(B, Self::Item) -> B,
     {
-        #[inline]
+        let mut accum = init;
-        fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
+        while let Some(x) = self.next_back() {
-            move |acc, x| Ok(f(acc, x))
+            accum = f(accum, x);
         }
-
+        accum
-        self.try_rfold(accum, ok(f)).unwrap()
     }
 
     /// Searches for an element of an iterator from the back that satisfies a predicate.

src/libcore/iter/traits/iterator.rs

@@ -1697,8 +1697,8 @@ fn extend<'a, T, B: Extend<T>>(
             mut f: impl FnMut(&T) -> bool + 'a,
             left: &'a mut B,
             right: &'a mut B,
-        ) -> impl FnMut(T) + 'a {
+        ) -> impl FnMut((), T) + 'a {
-            move |x| {
+            move |(), x| {
                 if f(&x) {
                     left.extend(Some(x));
                 } else {
@@ -1710,7 +1710,7 @@ fn extend<'a, T, B: Extend<T>>(
         let mut left: B = Default::default();
         let mut right: B = Default::default();
 
-        self.for_each(extend(f, &mut left, &mut right));
+        self.fold((), extend(f, &mut left, &mut right));
 
         (left, right)
     }
@@ -1826,7 +1826,7 @@ fn is_partitioned<P>(mut self, mut predicate: P) -> bool
     ///
     /// # Note to Implementors
     ///
-    /// Most of the other (forward) methods have default implementations in
+    /// Several of the other (forward) methods have default implementations in
     /// terms of this one, so try to implement this explicitly if it can
     /// do something better than the default `for` loop implementation.
     ///
@@ -1944,6 +1944,15 @@ fn call<T, R>(mut f: impl FnMut(T) -> R) -> impl FnMut((), T) -> R {
     /// may not terminate for infinite iterators, even on traits for which a
     /// result is determinable in finite time.
     ///
+    /// # Note to Implementors
+    ///
+    /// Several of the other (forward) methods have default implementations in
+    /// terms of this one, so try to implement this explicitly if it can
+    /// do something better than the default `for` loop implementation.
+    ///
+    /// In particular, try to have this call `fold()` on the internal parts
+    /// from which this iterator is composed.
+    ///
     /// # Examples
     ///
     /// Basic usage:
@@ -1992,17 +2001,16 @@ fn call<T, R>(mut f: impl FnMut(T) -> R) -> impl FnMut((), T) -> R {
     /// ```
     #[inline]
     #[stable(feature = "rust1", since = "1.0.0")]
-    fn fold<B, F>(mut self, init: B, f: F) -> B
+    fn fold<B, F>(mut self, init: B, mut f: F) -> B
     where
         Self: Sized,
         F: FnMut(B, Self::Item) -> B,
     {
-        #[inline]
+        let mut accum = init;
-        fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
+        while let Some(x) = self.next() {
-            move |acc, x| Ok(f(acc, x))
+            accum = f(accum, x);
         }
-
+        accum
-        self.try_fold(init, ok(f)).unwrap()
     }
 
     /// The same as [`fold()`](#method.fold), but uses the first element in the
@@ -2273,7 +2281,7 @@ fn try_find<F, E, R>(&mut self, mut f: F) -> Result<Option<Self::Item>, E>
         F: FnMut(&Self::Item) -> R,
         R: Try<Ok = bool, Error = E>,
     {
-        self.try_for_each(move |x| match f(&x).into_result() {
+        self.try_fold((), move |(), x| match f(&x).into_result() {
             Ok(false) => LoopState::Continue(()),
             Ok(true) => LoopState::Break(Ok(x)),
             Err(x) => LoopState::Break(Err(x)),
@@ -2665,8 +2673,8 @@ fn unzip<A, B, FromA, FromB>(self) -> (FromA, FromB)
         fn extend<'a, A, B>(
             ts: &'a mut impl Extend<A>,
             us: &'a mut impl Extend<B>,
-        ) -> impl FnMut((A, B)) + 'a {
+        ) -> impl FnMut((), (A, B)) + 'a {
-            move |(t, u)| {
+            move |(), (t, u)| {
                 ts.extend(Some(t));
                 us.extend(Some(u));
             }
@@ -2675,7 +2683,7 @@ fn extend<'a, A, B>(
         let mut ts: FromA = Default::default();
         let mut us: FromB = Default::default();
 
-        self.for_each(extend(&mut ts, &mut us));
+        self.fold((), extend(&mut ts, &mut us));
 
         (ts, us)
     }

src/test/codegen/iter-fold-closure-no-dupes.rs

@@ -1,14 +0,0 @@
-//! Check that fold closures aren't duplicated for each iterator type.
-// compile-flags: -C opt-level=0
-
-fn main() {
-    (0i32..10).by_ref().count();
-    (0i32..=10).by_ref().count();
-}
-
-// `count` calls `fold`, which calls `try_fold` -- find the `fold` closure:
-// CHECK: {{^define.*Iterator::fold::.*closure}}
-//
-// Only one closure is needed for both `count` calls, even from different
-// monomorphized iterator types, as it's only generic over the item type.
-// CHECK-NOT: {{^define.*Iterator::fold::.*closure}}

src/test/codegen/iter-fold-closure-no-iterator.rs

@@ -1,10 +0,0 @@
-//! Check that fold closures aren't generic in the iterator type.
-// compile-flags: -C opt-level=0
-
-fn main() {
-    (0i32..10).by_ref().count();
-}
-
-// `count` calls `fold`, which calls `try_fold` -- that `fold` closure should
-// not be generic in the iterator type, only in the item type.
-// CHECK-NOT: {{^define.*Iterator::fold::.*closure.*Range}}