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 @@ fn fold(self, init: Acc, fold: Fold) -> Acc { 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 @@ fn fold(self, init: Acc, fold: Fold) -> Acc { 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 @@ fn new(iter: I) -> FlattenCompat { } } +impl FlattenCompat +where + I: Iterator>, +{ + /// 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(&mut self, mut acc: Acc, mut fold: Fold) -> R + where + Fold: FnMut(Acc, &mut U) -> R, + R: Try, + { + #[inline] + fn flatten<'a, T: IntoIterator, Acc, R: Try>( + frontiter: &'a mut Option, + 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 Iterator for FlattenCompat where I: Iterator>, @@ -323,39 +373,20 @@ fn size_hint(&self) -> (usize, Option) { } #[inline] - fn try_fold(&mut self, mut init: Acc, mut fold: Fold) -> R + fn try_fold(&mut self, init: Acc, fold: Fold) -> R where Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try, { #[inline] - fn flatten<'a, T: IntoIterator, Acc, R: Try>( - frontiter: &'a mut Option, - 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 - } + fn flatten>( + mut fold: impl FnMut(Acc, U::Item) -> R, + ) -> impl FnMut(Acc, &mut U) -> R { + move |acc, iter| iter.try_fold(acc, &mut fold) } - if let Some(ref mut front) = self.frontiter { - init = front.try_fold(init, &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 @@ fn flatten( #[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, + #[inline] + #[rustc_inherit_overflow_checks] + fn advance(n: usize, iter: &mut U) -> ControlFlow<(), usize> { + match iter.advance_by(n) { + Ok(()) => ControlFlow::BREAK, + Err(advanced) => ControlFlow::Continue(n - advanced), } } - self.frontiter = None; - - if let Some(ref mut back) = self.backiter { - match back.advance_by(rem) { - ret @ Ok(_) => return ret, - Err(advanced) => rem -= advanced, - } + match self.iter_try_fold(n, advance) { + ControlFlow::Continue(remaining) if remaining > 0 => Err(n - remaining), + _ => Ok(()), } - - if rem > 0 { - return Err(n - rem); - } - - self.backiter = None; - - Ok(()) } }