From 8ff8d0527956aefc2da09d93fe14a3e87fd6ca1a Mon Sep 17 00:00:00 2001 From: Tim Vermeulen Date: Wed, 20 Jul 2022 16:33:39 +0200 Subject: [PATCH 1/4] Move shared logic of `try_fold` and `advance_by` into `iter_try_fold` --- library/core/src/iter/adapters/flatten.rs | 118 ++++++++++++---------- 1 file changed, 66 insertions(+), 52 deletions(-) diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs index 15a120e35a2fa..0ff9510cd1b5e 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 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 @@ where } #[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 - } - } - - if let Some(ref mut front) = self.frontiter { - init = front.try_fold(init, &mut fold)?; + 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) } - 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(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(()) } } From cbc5f627827a82f4aa570f60cfbefe8357ba805b Mon Sep 17 00:00:00 2001 From: Tim Vermeulen Date: Wed, 20 Jul 2022 16:33:50 +0200 Subject: [PATCH 2/4] Move shared logic of `try_rfold` and `advance_back_by` into `iter_try_rfold` --- library/core/src/iter/adapters/flatten.rs | 119 ++++++++++++---------- 1 file changed, 65 insertions(+), 54 deletions(-) diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs index 0ff9510cd1b5e..67f9d7fcc94db 100644 --- a/library/core/src/iter/adapters/flatten.rs +++ b/library/core/src/iter/adapters/flatten.rs @@ -108,6 +108,11 @@ where { self.inner.rfold(init, fold) } + + #[inline] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + self.inner.advance_back_by(n) + } } #[stable(feature = "fused", since = "1.26.0")] @@ -254,6 +259,11 @@ where { self.inner.rfold(init, fold) } + + #[inline] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + self.inner.advance_back_by(n) + } } #[stable(feature = "iterator_flatten", since = "1.29.0")] @@ -330,6 +340,46 @@ where } } +impl FlattenCompat +where + I: DoubleEndedIterator>, +{ + /// Folds over the inner iterators in reverse order as long as the given function returns + /// successfully, always storing the most recent inner iterator in `self.backiter`. + /// + /// Folds over the inner iterators, not over their elements. Is used by the `try_rfold` and + /// `advance_back_by` methods. + #[inline] + fn iter_try_rfold(&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>( + backiter: &'a mut Option, + fold: &'a mut impl FnMut(Acc, &mut T::IntoIter) -> R, + ) -> impl FnMut(Acc, T) -> R + 'a { + move |acc, iter| fold(acc, backiter.insert(iter.into_iter())) + } + + if let Some(iter) = &mut self.backiter { + acc = fold(acc, iter)?; + } + self.backiter = None; + + acc = self.iter.try_rfold(acc, flatten(&mut self.backiter, &mut fold))?; + self.backiter = None; + + if let Some(iter) = &mut self.frontiter { + acc = fold(acc, iter)?; + } + self.frontiter = None; + + try { acc } + } +} + impl Iterator for FlattenCompat where I: Iterator>, @@ -452,42 +502,20 @@ where } #[inline] - fn try_rfold(&mut self, mut init: Acc, mut fold: Fold) -> R + fn try_rfold(&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>( - backiter: &'a mut Option, - fold: &'a mut impl FnMut(Acc, T::Item) -> R, - ) -> impl FnMut(Acc, T) -> R + 'a - where - T::IntoIter: DoubleEndedIterator, - { - move |acc, x| { - let mut mid = x.into_iter(); - 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 { - init = front.try_rfold(init, &mut fold)?; + fn flatten>( + mut fold: impl FnMut(Acc, U::Item) -> R, + ) -> impl FnMut(Acc, &mut U) -> R { + move |acc, iter| iter.try_rfold(acc, &mut fold) } - self.frontiter = None; - try { init } + self.iter_try_rfold(init, flatten(fold)) } #[inline] @@ -521,36 +549,19 @@ where #[inline] #[rustc_inherit_overflow_checks] fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { - let mut rem = n; - loop { - if let Some(ref mut back) = self.backiter { - match back.advance_back_by(rem) { - ret @ Ok(_) => return ret, - Err(advanced) => rem -= advanced, - } - } - match self.iter.next_back() { - Some(iterable) => self.backiter = Some(iterable.into_iter()), - _ => break, - } - } - - self.backiter = None; - - if let Some(ref mut front) = self.frontiter { - match front.advance_back_by(rem) { - ret @ Ok(_) => return ret, - Err(advanced) => rem -= advanced, + #[inline] + #[rustc_inherit_overflow_checks] + fn advance(n: usize, iter: &mut U) -> ControlFlow<(), usize> { + match iter.advance_back_by(n) { + Ok(()) => ControlFlow::BREAK, + Err(advanced) => ControlFlow::Continue(n - advanced), } } - if rem > 0 { - return Err(n - rem); + match self.iter_try_rfold(n, advance) { + ControlFlow::Continue(remaining) if remaining > 0 => Err(n - remaining), + _ => Ok(()), } - - self.frontiter = None; - - Ok(()) } } From 3f7004920c9ae9179873039b4101f31f8c40f3db Mon Sep 17 00:00:00 2001 From: Tim Vermeulen Date: Wed, 20 Jul 2022 16:33:53 +0200 Subject: [PATCH 3/4] Move `fold` logic to `iter_fold` method and reuse it in `count` and `last` --- library/core/src/iter/adapters/flatten.rs | 92 +++++++++++++++++---- library/core/tests/iter/adapters/flatten.rs | 42 ++++++++++ 2 files changed, 118 insertions(+), 16 deletions(-) diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs index 67f9d7fcc94db..ba9767e5c4a0a 100644 --- a/library/core/src/iter/adapters/flatten.rs +++ b/library/core/src/iter/adapters/flatten.rs @@ -78,6 +78,16 @@ where fn advance_by(&mut self, n: usize) -> Result<(), usize> { self.inner.advance_by(n) } + + #[inline] + fn count(self) -> usize { + self.inner.count() + } + + #[inline] + fn last(self) -> Option { + self.inner.last() + } } #[stable(feature = "rust1", since = "1.0.0")] @@ -229,6 +239,16 @@ where fn advance_by(&mut self, n: usize) -> Result<(), usize> { self.inner.advance_by(n) } + + #[inline] + fn count(self) -> usize { + self.inner.count() + } + + #[inline] + fn last(self) -> Option { + self.inner.last() + } } #[stable(feature = "iterator_flatten", since = "1.29.0")] @@ -304,6 +324,35 @@ impl FlattenCompat where I: Iterator>, { + /// Folds the inner iterators into an accumulator by applying an operation. + /// + /// Folds over the inner iterators, not over their elements. Is used by the `fold`, `count`, + /// and `last` methods. + #[inline] + fn iter_fold(self, mut acc: Acc, mut fold: Fold) -> Acc + where + Fold: FnMut(Acc, U) -> Acc, + { + #[inline] + fn flatten( + fold: &mut impl FnMut(Acc, T::IntoIter) -> Acc, + ) -> impl FnMut(Acc, T) -> Acc + '_ { + move |acc, iter| fold(acc, iter.into_iter()) + } + + if let Some(iter) = self.frontiter { + acc = fold(acc, iter); + } + + acc = self.iter.fold(acc, flatten(&mut fold)); + + if let Some(iter) = self.backiter { + acc = fold(acc, iter); + } + + acc + } + /// Folds over the inner iterators as long as the given function returns successfully, /// always storing the most recent inner iterator in `self.frontiter`. /// @@ -440,28 +489,18 @@ where } #[inline] - fn fold(self, mut init: Acc, mut fold: Fold) -> Acc + fn fold(self, init: Acc, fold: Fold) -> Acc where Fold: FnMut(Acc, Self::Item) -> Acc, { #[inline] - fn flatten( - fold: &mut impl FnMut(Acc, T::Item) -> Acc, - ) -> impl FnMut(Acc, T) -> Acc + '_ { - move |acc, x| x.into_iter().fold(acc, &mut *fold) + fn flatten( + mut fold: impl FnMut(Acc, U::Item) -> Acc, + ) -> impl FnMut(Acc, U) -> Acc { + move |acc, iter| iter.fold(acc, &mut 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 + self.iter_fold(init, flatten(fold)) } #[inline] @@ -481,6 +520,27 @@ where _ => Ok(()), } } + + #[inline] + fn count(self) -> usize { + #[inline] + #[rustc_inherit_overflow_checks] + fn count(acc: usize, iter: U) -> usize { + acc + iter.count() + } + + self.iter_fold(0, count) + } + + #[inline] + fn last(self) -> Option { + #[inline] + fn last(last: Option, iter: U) -> Option { + iter.last().or(last) + } + + self.iter_fold(None, last) + } } impl DoubleEndedIterator for FlattenCompat diff --git a/library/core/tests/iter/adapters/flatten.rs b/library/core/tests/iter/adapters/flatten.rs index f8ab8c9d4447b..690fd0c21974b 100644 --- a/library/core/tests/iter/adapters/flatten.rs +++ b/library/core/tests/iter/adapters/flatten.rs @@ -168,3 +168,45 @@ fn test_trusted_len_flatten() { assert_trusted_len(&iter); assert_eq!(iter.size_hint(), (20, Some(20))); } + +#[test] +fn test_flatten_count() { + let mut it = once(0..10).chain(once(10..30)).chain(once(30..40)).flatten(); + + assert_eq!(it.clone().count(), 40); + it.advance_by(5).unwrap(); + assert_eq!(it.clone().count(), 35); + it.advance_back_by(5).unwrap(); + assert_eq!(it.clone().count(), 30); + it.advance_by(10).unwrap(); + assert_eq!(it.clone().count(), 20); + it.advance_back_by(8).unwrap(); + assert_eq!(it.clone().count(), 12); + it.advance_by(4).unwrap(); + assert_eq!(it.clone().count(), 8); + it.advance_back_by(5).unwrap(); + assert_eq!(it.clone().count(), 3); + it.advance_by(3).unwrap(); + assert_eq!(it.clone().count(), 0); +} + +#[test] +fn test_flatten_last() { + let mut it = once(0..10).chain(once(10..30)).chain(once(30..40)).flatten(); + + assert_eq!(it.clone().last(), Some(39)); + it.advance_by(5).unwrap(); // 5..40 + assert_eq!(it.clone().last(), Some(39)); + it.advance_back_by(5).unwrap(); // 5..35 + assert_eq!(it.clone().last(), Some(34)); + it.advance_by(10).unwrap(); // 15..35 + assert_eq!(it.clone().last(), Some(34)); + it.advance_back_by(8).unwrap(); // 15..27 + assert_eq!(it.clone().last(), Some(26)); + it.advance_by(4).unwrap(); // 19..27 + assert_eq!(it.clone().last(), Some(26)); + it.advance_back_by(5).unwrap(); // 19..22 + assert_eq!(it.clone().last(), Some(21)); + it.advance_by(3).unwrap(); // 22..22 + assert_eq!(it.clone().last(), None); +} From 38bb0b173e370fccc35c7bc02e29525c9d82c767 Mon Sep 17 00:00:00 2001 From: Tim Vermeulen Date: Wed, 20 Jul 2022 16:33:57 +0200 Subject: [PATCH 4/4] Move `rfold` logic into `iter_rfold` --- library/core/src/iter/adapters/flatten.rs | 54 +++++++++++++++-------- 1 file changed, 35 insertions(+), 19 deletions(-) diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs index ba9767e5c4a0a..307016c269099 100644 --- a/library/core/src/iter/adapters/flatten.rs +++ b/library/core/src/iter/adapters/flatten.rs @@ -393,6 +393,35 @@ impl FlattenCompat where I: DoubleEndedIterator>, { + /// Folds the inner iterators into an accumulator by applying an operation, starting form the + /// back. + /// + /// Folds over the inner iterators, not over their elements. Is used by the `rfold` method. + #[inline] + fn iter_rfold(self, mut acc: Acc, mut fold: Fold) -> Acc + where + Fold: FnMut(Acc, U) -> Acc, + { + #[inline] + fn flatten( + fold: &mut impl FnMut(Acc, T::IntoIter) -> Acc, + ) -> impl FnMut(Acc, T) -> Acc + '_ { + move |acc, iter| fold(acc, iter.into_iter()) + } + + if let Some(iter) = self.backiter { + acc = fold(acc, iter); + } + + acc = self.iter.rfold(acc, flatten(&mut fold)); + + if let Some(iter) = self.frontiter { + acc = fold(acc, iter); + } + + acc + } + /// Folds over the inner iterators in reverse order as long as the given function returns /// successfully, always storing the most recent inner iterator in `self.backiter`. /// @@ -579,31 +608,18 @@ where } #[inline] - fn rfold(self, mut init: Acc, mut fold: Fold) -> Acc + fn rfold(self, init: Acc, fold: Fold) -> Acc where Fold: FnMut(Acc, Self::Item) -> Acc, { #[inline] - fn flatten( - 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) - } - - 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); + fn flatten( + mut fold: impl FnMut(Acc, U::Item) -> Acc, + ) -> impl FnMut(Acc, U) -> Acc { + move |acc, iter| iter.rfold(acc, &mut fold) } - init + self.iter_rfold(init, flatten(fold)) } #[inline]