From fc150d17b5e705132f0f70595dfc213aa6b564d3 Mon Sep 17 00:00:00 2001
From: Josh Stone <jistone@redhat.com>
Date: Mon, 15 Feb 2021 11:52:28 -0800
Subject: [PATCH] Expand FlattenCompat folds

The former `chain`+`chain`+`fold` implementation looked nice from a
functional-programming perspective, but it introduced unnecessary layers
of abstraction on every `flat_map`/`flatten` fold. It's straightforward
to just fold each part in turn, and this makes it look like a simplified
version of the existing `try_fold` implementation.

For the `iter::bench_flat_map*` benchmarks, I get a large improvement in
`bench_flat_map_chain_sum`, from 1,598,473 ns/iter to 499,889 ns/iter,
and the rest are unchanged.
---
 library/core/src/iter/adapters/flatten.rs | 55 ++++++++++++++---------
 1 file changed, 35 insertions(+), 20 deletions(-)

diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs
index 081f282edcf..0114d7af4f4 100644
--- a/library/core/src/iter/adapters/flatten.rs
+++ b/library/core/src/iter/adapters/flatten.rs
@@ -325,22 +325,28 @@ where
     }
 
     #[inline]
-    fn fold<Acc, Fold>(self, init: Acc, ref mut fold: Fold) -> Acc
+    fn fold<Acc, Fold>(self, mut init: Acc, 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)
+        fn flatten<T: IntoIterator, Acc>(
+            fold: &mut impl FnMut(Acc, T::Item) -> Acc,
+        ) -> impl FnMut(Acc, T) -> Acc + '_ {
+            move |acc, x| x.into_iter().fold(acc, &mut *fold)
         }
 
-        self.frontiter
-            .into_iter()
-            .chain(self.iter.map(IntoIterator::into_iter))
-            .chain(self.backiter)
-            .fold(init, flatten(fold))
+        if let Some(front) = self.frontiter {
+            init = front.fold(init, &mut fold);
+        }
+
+        init = self.iter.fold(init, flatten(&mut fold));
+
+        if let Some(back) = self.backiter {
+            init = back.fold(init, &mut fold);
+        }
+
+        init
     }
 }
 
@@ -411,21 +417,30 @@ where
     }
 
     #[inline]
-    fn rfold<Acc, Fold>(self, init: Acc, ref mut fold: Fold) -> Acc
+    fn rfold<Acc, Fold>(self, mut init: Acc, 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)
+        fn flatten<T: IntoIterator, Acc>(
+            fold: &mut impl FnMut(Acc, T::Item) -> Acc,
+        ) -> impl FnMut(Acc, T) -> Acc + '_
+        where
+            T::IntoIter: DoubleEndedIterator,
+        {
+            move |acc, x| x.into_iter().rfold(acc, &mut *fold)
         }
 
-        self.frontiter
-            .into_iter()
-            .chain(self.iter.map(IntoIterator::into_iter))
-            .chain(self.backiter)
-            .rfold(init, flatten(fold))
+        if let Some(back) = self.backiter {
+            init = back.rfold(init, &mut fold);
+        }
+
+        init = self.iter.rfold(init, flatten(&mut fold));
+
+        if let Some(front) = self.frontiter {
+            init = front.rfold(init, &mut fold);
+        }
+
+        init
     }
 }