Move shared logic of `try_fold` and `advance_by` into `iter_try_fold`

1 files changed, 66 insertions, 52 deletions

diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs
index 15a120e35a2..0ff9510cd1b 100644
--- a/library/core/src/iter/adapters/flatten.rs
+++ b/library/core/src/iter/adapters/flatten.rs
@@ -1,6 +1,6 @@
 use crate::fmt;
 use crate::iter::{DoubleEndedIterator, Fuse, FusedIterator, Iterator, Map, TrustedLen};
-use crate::ops::Try;
+use crate::ops::{ControlFlow, Try};
 
 /// An iterator that maps each element to an iterator, and yields the elements
 /// of the produced iterators.
@@ -73,6 +73,11 @@ where
     {
         self.inner.fold(init, fold)
     }
+
+    #[inline]
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        self.inner.advance_by(n)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -214,6 +219,11 @@ where
     {
         self.inner.fold(init, fold)
     }
+
+    #[inline]
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        self.inner.advance_by(n)
+    }
 }
 
 #[stable(feature = "iterator_flatten", since = "1.29.0")]
@@ -280,6 +290,46 @@ where
     }
 }
 
+impl<I, U> FlattenCompat<I, U>
+where
+    I: Iterator<Item: IntoIterator<IntoIter = U>>,
+{
+    /// Folds over the inner iterators as long as the given function returns successfully,
+    /// always storing the most recent inner iterator in `self.frontiter`.
+    ///
+    /// Folds over the inner iterators, not over their elements. Is used by the `try_fold` and
+    /// `advance_by` methods.
+    #[inline]
+    fn iter_try_fold<Acc, Fold, R>(&mut self, mut acc: Acc, mut fold: Fold) -> R
+    where
+        Fold: FnMut(Acc, &mut U) -> R,
+        R: Try<Output = Acc>,
+    {
+        #[inline]
+        fn flatten<'a, T: IntoIterator, Acc, R: Try<Output = Acc>>(
+            frontiter: &'a mut Option<T::IntoIter>,
+            fold: &'a mut impl FnMut(Acc, &mut T::IntoIter) -> R,
+        ) -> impl FnMut(Acc, T) -> R + 'a {
+            move |acc, iter| fold(acc, frontiter.insert(iter.into_iter()))
+        }
+
+        if let Some(iter) = &mut self.frontiter {
+            acc = fold(acc, iter)?;
+        }
+        self.frontiter = None;
+
+        acc = self.iter.try_fold(acc, flatten(&mut self.frontiter, &mut fold))?;
+        self.frontiter = None;
+
+        if let Some(iter) = &mut self.backiter {
+            acc = fold(acc, iter)?;
+        }
+        self.backiter = None;
+
+        try { acc }
+    }
+}
+
 impl<I, U> Iterator for FlattenCompat<I, U>
 where
     I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
@@ -323,39 +373,20 @@ where
     }
 
     #[inline]
-    fn try_fold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R
+    fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
     where
         Self: Sized,
         Fold: FnMut(Acc, Self::Item) -> R,
         R: Try<Output = Acc>,
     {
         #[inline]
-        fn flatten<'a, T: IntoIterator, Acc, R: Try<Output = 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)?;
+        fn flatten<U: Iterator, Acc, R: Try<Output = Acc>>(
+            mut fold: impl FnMut(Acc, U::Item) -> R,
+        ) -> impl FnMut(Acc, &mut U) -> R {
+            move |acc, iter| iter.try_fold(acc, &mut fold)
         }
-        self.frontiter = None;
-
-        init = self.iter.try_fold(init, flatten(&mut self.frontiter, &mut fold))?;
-        self.frontiter = None;
 
-        if let Some(ref mut back) = self.backiter {
-            init = back.try_fold(init, &mut fold)?;
-        }
-        self.backiter = None;
-
-        try { init }
+        self.iter_try_fold(init, flatten(fold))
     }
 
     #[inline]
@@ -386,36 +417,19 @@ where
     #[inline]
     #[rustc_inherit_overflow_checks]
     fn advance_by(&mut self, n: usize) -> Result<(), usize> {
-        let mut rem = n;
-        loop {
-            if let Some(ref mut front) = self.frontiter {
-                match front.advance_by(rem) {
-                    ret @ Ok(_) => return ret,
-                    Err(advanced) => rem -= advanced,
-                }
-            }
-            self.frontiter = match self.iter.next() {
-                Some(iterable) => Some(iterable.into_iter()),
-                _ => break,
-            }
-        }
-
-        self.frontiter = None;
-
-        if let Some(ref mut back) = self.backiter {
-            match back.advance_by(rem) {
-                ret @ Ok(_) => return ret,
-                Err(advanced) => rem -= advanced,
+        #[inline]
+        #[rustc_inherit_overflow_checks]
+        fn advance<U: Iterator>(n: usize, iter: &mut U) -> ControlFlow<(), usize> {
+            match iter.advance_by(n) {
+                Ok(()) => ControlFlow::BREAK,
+                Err(advanced) => ControlFlow::Continue(n - advanced),
             }
         }
 
-        if rem > 0 {
-            return Err(n - rem);
+        match self.iter_try_fold(n, advance) {
+            ControlFlow::Continue(remaining) if remaining > 0 => Err(n - remaining),
+            _ => Ok(()),
         }
-
-        self.backiter = None;
-
-        Ok(())
     }
 }