Skip to content

Commit

Permalink
Auto merge of #82162 - cuviper:flat-fold, r=Mark-Simulacrum
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
bors committed Feb 25, 2021
2 parents 1fdadbf + fc150d1 commit 63bacf1
Showing 1 changed file with 35 additions and 20 deletions.
55 changes: 35 additions & 20 deletions library/core/src/iter/adapters/flatten.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}

Expand Down Expand Up @@ -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)
}

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);
}

self.frontiter
.into_iter()
.chain(self.iter.map(IntoIterator::into_iter))
.chain(self.backiter)
.rfold(init, flatten(fold))
init
}
}

0 comments on commit 63bacf1

Please sign in to comment.