Auto merge of #62429 - cuviper:iter-closures, r=cramertj

Reduce the genericity of closures in the iterator traits

By default, closures inherit the generic parameters of their scope,
including `Self`. However, in most cases, the closures used to implement
iterators don't need to be generic on the iterator type, only its `Item`
type. We can reduce this genericity by redirecting such closures through
local functions.

This does make the closures more cumbersome to write, but it will
hopefully reduce duplication in their monomorphizations, as well as
their related type lengths.

Author: bors

src/libcore/iter/adapters/flatten.rs

@@ -229,16 +229,16 @@ where
                 if let elt@Some(_) = inner.next() { return elt }
             }
             match self.iter.next() {
-                None => return self.backiter.as_mut().and_then(|it| it.next()),
+                None => return self.backiter.as_mut()?.next(),
                 Some(inner) => self.frontiter = Some(inner.into_iter()),
             }
         }
     }
 
     #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let (flo, fhi) = self.frontiter.as_ref().map_or((0, Some(0)), |it| it.size_hint());
-        let (blo, bhi) = self.backiter.as_ref().map_or((0, Some(0)), |it| it.size_hint());
+        let (flo, fhi) = self.frontiter.as_ref().map_or((0, Some(0)), U::size_hint);
+        let (blo, bhi) = self.backiter.as_ref().map_or((0, Some(0)), U::size_hint);
         let lo = flo.saturating_add(blo);
         match (self.iter.size_hint(), fhi, bhi) {
             ((0, Some(0)), Some(a), Some(b)) => (lo, a.checked_add(b)),
@@ -250,20 +250,25 @@ where
     fn try_fold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R where
         Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try<Ok=Acc>
     {
+        #[inline]
+        fn flatten<'a, T: IntoIterator, Acc, R: Try<Ok = Acc>>(
+            frontiter: &'a mut Option<T::IntoIter>,
+            fold: &'a mut impl FnMut(Acc, T::Item) -> R,
+        ) -> impl FnMut(Acc, T) -> R + 'a {
+            move |acc, x| {
+                let mut mid = x.into_iter();
+                let r = mid.try_fold(acc, &mut *fold);
+                *frontiter = Some(mid);
+                r
+            }
+        }
+
         if let Some(ref mut front) = self.frontiter {
             init = front.try_fold(init, &mut fold)?;
         }
         self.frontiter = None;
 
-        {
-            let frontiter = &mut self.frontiter;
-            init = self.iter.try_fold(init, |acc, x| {
-                let mut mid = x.into_iter();
-                let r = mid.try_fold(acc, &mut fold);
-                *frontiter = Some(mid);
-                r
-            })?;
-        }
+        init = self.iter.try_fold(init, flatten(&mut self.frontiter, &mut fold))?;
         self.frontiter = None;
 
         if let Some(ref mut back) = self.backiter {
@@ -275,13 +280,20 @@ where
     }
 
     #[inline]
-    fn fold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc
+    fn fold<Acc, Fold>(self, init: Acc, ref mut fold: Fold) -> Acc
         where Fold: FnMut(Acc, Self::Item) -> Acc,
     {
+        #[inline]
+        fn flatten<U: Iterator, Acc>(
+            fold: &mut impl FnMut(Acc, U::Item) -> Acc,
+        ) -> impl FnMut(Acc, U) -> Acc + '_ {
+            move |acc, iter| iter.fold(acc, &mut *fold)
+        }
+
         self.frontiter.into_iter()
             .chain(self.iter.map(IntoIterator::into_iter))
             .chain(self.backiter)
-            .fold(init, |acc, iter| iter.fold(acc, &mut fold))
+            .fold(init, flatten(fold))
     }
 }
 
@@ -297,7 +309,7 @@ where
                 if let elt@Some(_) = inner.next_back() { return elt }
             }
             match self.iter.next_back() {
-                None => return self.frontiter.as_mut().and_then(|it| it.next_back()),
+                None => return self.frontiter.as_mut()?.next_back(),
                 next => self.backiter = next.map(IntoIterator::into_iter),
             }
         }
@@ -307,20 +319,27 @@ where
     fn try_rfold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R where
         Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try<Ok=Acc>
     {
-        if let Some(ref mut back) = self.backiter {
-            init = back.try_rfold(init, &mut fold)?;
-        }
-        self.backiter = None;
-
+        #[inline]
+        fn flatten<'a, T: IntoIterator, Acc, R: Try<Ok = Acc>>(
+            backiter: &'a mut Option<T::IntoIter>,
+            fold: &'a mut impl FnMut(Acc, T::Item) -> R,
+        ) -> impl FnMut(Acc, T) -> R + 'a where
+            T::IntoIter: DoubleEndedIterator,
         {
-            let backiter = &mut self.backiter;
-            init = self.iter.try_rfold(init, |acc, x| {
+            move |acc, x| {
                 let mut mid = x.into_iter();
-                let r = mid.try_rfold(acc, &mut fold);
+                let r = mid.try_rfold(acc, &mut *fold);
                 *backiter = Some(mid);
                 r
-            })?;
+            }
         }
+
+        if let Some(ref mut back) = self.backiter {
+            init = back.try_rfold(init, &mut fold)?;
+        }
+        self.backiter = None;
+
+        init = self.iter.try_rfold(init, flatten(&mut self.backiter, &mut fold))?;
         self.backiter = None;
 
         if let Some(ref mut front) = self.frontiter {
@@ -332,12 +351,19 @@ where
     }
 
     #[inline]
-    fn rfold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc
+    fn rfold<Acc, Fold>(self, init: Acc, ref mut fold: Fold) -> Acc
         where Fold: FnMut(Acc, Self::Item) -> Acc,
     {
+        #[inline]
+        fn flatten<U: DoubleEndedIterator, Acc>(
+            fold: &mut impl FnMut(Acc, U::Item) -> Acc,
+        ) -> impl FnMut(Acc, U) -> Acc + '_ {
+            move |acc, iter| iter.rfold(acc, &mut *fold)
+        }
+
         self.frontiter.into_iter()
             .chain(self.iter.map(IntoIterator::into_iter))
             .chain(self.backiter)
-            .rfold(init, |acc, iter| iter.rfold(acc, &mut fold))
+            .rfold(init, flatten(fold))
     }
 }

src/libcore/iter/adapters/mod.rs

@@ -94,11 +94,9 @@ impl<A, B> ZipImpl<A, B> for Zip<A, B>
 
     #[inline]
     default fn next(&mut self) -> Option<(A::Item, B::Item)> {
-        self.a.next().and_then(|x| {
-            self.b.next().and_then(|y| {
-                Some((x, y))
-            })
-        })
+        let x = self.a.next()?;
+        let y = self.b.next()?;
+        Some((x, y))
     }
 
     #[inline]

src/libcore/iter/adapters/zip.rs

@@ -394,7 +394,8 @@ impl<A, F: FnOnce() -> A> Iterator for OnceWith<F> {
 
     #[inline]
     fn next(&mut self) -> Option<A> {
-        self.gen.take().map(|f| f())
+        let f = self.gen.take()?;
+        Some(f())
     }
 
     #[inline]
@@ -608,10 +609,9 @@ impl<T, F> Iterator for Successors<T, F>
 
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        self.next.take().map(|item| {
-            self.next = (self.succ)(&item);
-            item
-        })
+        let item = self.next.take()?;
+        self.next = (self.succ)(&item);
+        Some(item)
     }
 
     #[inline]

src/libcore/iter/sources.rs

@@ -85,28 +85,28 @@ macro_rules! float_sum_product {
         #[stable(feature = "iter_arith_traits", since = "1.12.0")]
         impl Sum for $a {
             fn sum<I: Iterator<Item=$a>>(iter: I) -> $a {
-                iter.fold(0.0, |a, b| a + b)
+                iter.fold(0.0, Add::add)
             }
         }
 
         #[stable(feature = "iter_arith_traits", since = "1.12.0")]
         impl Product for $a {
             fn product<I: Iterator<Item=$a>>(iter: I) -> $a {
-                iter.fold(1.0, |a, b| a * b)
+                iter.fold(1.0, Mul::mul)
             }
         }
 
         #[stable(feature = "iter_arith_traits", since = "1.12.0")]
         impl<'a> Sum<&'a $a> for $a {
             fn sum<I: Iterator<Item=&'a $a>>(iter: I) -> $a {
-                iter.fold(0.0, |a, b| a + *b)
+                iter.fold(0.0, Add::add)
             }
         }
 
         #[stable(feature = "iter_arith_traits", since = "1.12.0")]
         impl<'a> Product<&'a $a> for $a {
             fn product<I: Iterator<Item=&'a $a>>(iter: I) -> $a {
-                iter.fold(1.0, |a, b| a * *b)
+                iter.fold(1.0, Mul::mul)
             }
         }
     )*)

src/libcore/iter/traits/accum.rs

@@ -219,12 +219,17 @@ pub trait DoubleEndedIterator: Iterator {
     /// ```
     #[inline]
     #[stable(feature = "iter_rfold", since = "1.27.0")]
-    fn rfold<B, F>(mut self, accum: B, mut f: F) -> B
+    fn rfold<B, F>(mut self, accum: B, f: F) -> B
     where
         Self: Sized,
         F: FnMut(B, Self::Item) -> B,
     {
-        self.try_rfold(accum, move |acc, x| Ok::<B, !>(f(acc, x))).unwrap()
+        #[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(accum, ok(f)).unwrap()
     }
 
     /// Searches for an element of an iterator from the back that satisfies a predicate.
@@ -271,15 +276,21 @@ pub trait DoubleEndedIterator: Iterator {
     /// ```
     #[inline]
     #[stable(feature = "iter_rfind", since = "1.27.0")]
-    fn rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item>
+    fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item>
     where
         Self: Sized,
         P: FnMut(&Self::Item) -> bool
     {
-        self.try_rfold((), move |(), x| {
-            if predicate(&x) { LoopState::Break(x) }
-            else { LoopState::Continue(()) }
-        }).break_value()
+        #[inline]
+        fn check<T>(
+            mut predicate: impl FnMut(&T) -> bool,
+        ) -> impl FnMut((), T) -> LoopState<(), T> {
+            move |(), x| {
+                if predicate(&x) { LoopState::Break(x) } else { LoopState::Continue(()) }
+            }
+        }
+
+        self.try_rfold((), check(predicate)).break_value()
     }
 }
 

src/libcore/iter/traits/double_ended.rs

@@ -0,0 +1,14 @@
+//! 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/libcore/iter/traits/iterator.rs

@@ -0,0 +1,10 @@
+//! 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}}

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

@@ -0,0 +1,16 @@
+// run-pass
+// only-32bit too impatient for 2⁶⁴ items
+// ignore-wasm32-bare compiled with panic=abort by default
+// compile-flags: -C debug_assertions=yes -C opt-level=3
+
+use std::panic;
+use std::usize::MAX;
+
+fn main() {
+    assert_eq!((0..MAX).by_ref().count(), MAX);
+
+    let r = panic::catch_unwind(|| {
+        (0..=MAX).by_ref().count()
+    });
+    assert!(r.is_err());
+}

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

@@ -0,0 +1,11 @@
+// run-pass
+// only-32bit too impatient for 2⁶⁴ items
+// compile-flags: -C debug_assertions=no -C opt-level=3
+
+use std::panic;
+use std::usize::MAX;
+
+fn main() {
+    assert_eq!((0..MAX).by_ref().count(), MAX);
+    assert_eq!((0..=MAX).by_ref().count(), 0);
+}

src/test/ui/iterators/iter-count-overflow-debug.rs

@@ -0,0 +1,38 @@
+// run-pass
+//! Check that type lengths don't explode with `Map` folds.
+//!
+//! The normal limit is a million, and this test used to exceed 1.5 million, but
+//! now we can survive an even tighter limit. Still seems excessive though...
+#![type_length_limit = "256000"]
+
+// Custom wrapper so Iterator methods aren't specialized.
+struct Iter<I>(I);
+
+impl<I> Iterator for Iter<I>
+where
+    I: Iterator
+{
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.0.next()
+    }
+}
+
+fn main() {
+    let c = Iter(0i32..10)
+        .map(|x| x)
+        .map(|x| x)
+        .map(|x| x)
+        .map(|x| x)
+        .map(|x| x)
+        .map(|x| x)
+        .map(|x| x)
+        .map(|x| x)
+        .map(|x| x)
+        .map(|x| x)
+        .map(|x| x)
+        .map(|x| x)
+        .count();
+    assert_eq!(c, 10);
+}

src/test/ui/iterators/iter-count-overflow-ndebug.rs

@@ -0,0 +1,22 @@
+// run-pass
+// only-32bit too impatient for 2⁶⁴ items
+// ignore-wasm32-bare compiled with panic=abort by default
+// compile-flags: -C debug_assertions=yes -C opt-level=3
+
+use std::panic;
+use std::usize::MAX;
+
+fn main() {
+    let n = MAX as u64;
+    assert_eq!((0..).by_ref().position(|i| i >= n), Some(MAX));
+
+    let r = panic::catch_unwind(|| {
+        (0..).by_ref().position(|i| i > n)
+    });
+    assert!(r.is_err());
+
+    let r = panic::catch_unwind(|| {
+        (0..=n + 1).by_ref().position(|_| false)
+    });
+    assert!(r.is_err());
+}

src/test/ui/iterators/iter-map-fold-type-length.rs

@@ -0,0 +1,13 @@
+// run-pass
+// only-32bit too impatient for 2⁶⁴ items
+// compile-flags: -C debug_assertions=no -C opt-level=3
+
+use std::panic;
+use std::usize::MAX;
+
+fn main() {
+    let n = MAX as u64;
+    assert_eq!((0..).by_ref().position(|i| i >= n), Some(MAX));
+    assert_eq!((0..).by_ref().position(|i| i > n), Some(0));
+    assert_eq!((0..=n + 1).by_ref().position(|_| false), None);
+}