Skip to content

Commit

Permalink
Auto merge of rust-lang#87431 - the8472:array-iter-fold, r=kennytm
Browse files Browse the repository at this point in the history
implement fold() on array::IntoIter to improve flatten().collect() perf

With rust-lang#87168 flattening `array::IntoIter`s is now `TrustedLen`, the `FromIterator` implementation for `Vec` has a specialization for `TrustedLen` iterators which uses internal iteration. This implements one of the main internal iteration methods on `array::Into` to optimize the combination of those two features.

This should address the main issue in rust-lang#87411

```
# old
test vec::bench_flat_map_collect                         ... bench:   2,244,024 ns/iter (+/- 18,903)

# new
test vec::bench_flat_map_collect                         ... bench:     172,863 ns/iter (+/- 2,141)
```
  • Loading branch information
bors committed Jul 27, 2021
2 parents 998cfe5 + 2276c5e commit 99d6692
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 0 deletions.
6 changes: 6 additions & 0 deletions library/alloc/benches/vec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -726,3 +726,9 @@ fn bench_dedup_old_100000(b: &mut Bencher) {
fn bench_dedup_new_100000(b: &mut Bencher) {
bench_vec_dedup_new(b, 100000);
}

#[bench]
fn bench_flat_map_collect(b: &mut Bencher) {
let v = vec![777u32; 500000];
b.iter(|| v.iter().flat_map(|color| color.rotate_left(8).to_be_bytes()).collect::<Vec<_>>());
}
21 changes: 21 additions & 0 deletions library/core/src/array/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,27 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
(len, Some(len))
}

#[inline]
fn fold<Acc, Fold>(mut self, init: Acc, mut fold: Fold) -> Acc
where
Fold: FnMut(Acc, Self::Item) -> Acc,
{
let data = &mut self.data;
// FIXME: This uses try_fold(&mut iter) instead of fold(iter) because the latter
// would go through the blanket `impl Iterator for &mut I` implementation
// which lacks inline annotations on its methods and adding those would be a larger
// perturbation than using try_fold here.
// Whether it would be beneficial to add those annotations should be investigated separately.
(&mut self.alive)
.try_fold::<_, _, Result<_, !>>(init, |acc, idx| {
// SAFETY: idx is obtained by folding over the `alive` range, which implies the
// value is currently considered alive but as the range is being consumed each value
// we read here will only be read once and then considered dead.
Ok(fold(acc, unsafe { data.get_unchecked(idx).assume_init_read() }))
})
.unwrap()
}

fn count(self) -> usize {
self.len()
}
Expand Down

0 comments on commit 99d6692

Please sign in to comment.