From e015e9da717bf501ab4752b2ea045d899f804a42 Mon Sep 17 00:00:00 2001 From: The8472 Date: Sat, 24 Jul 2021 17:42:20 +0200 Subject: [PATCH 1/2] implement fold() on array::IntoIter to improve flatten().collect() perf ``` # 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) ``` --- library/alloc/benches/vec.rs | 6 ++++++ library/core/src/array/iter.rs | 16 ++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/library/alloc/benches/vec.rs b/library/alloc/benches/vec.rs index 91eec10d57593..c93a493cadb0d 100644 --- a/library/alloc/benches/vec.rs +++ b/library/alloc/benches/vec.rs @@ -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::>()); +} diff --git a/library/core/src/array/iter.rs b/library/core/src/array/iter.rs index 61ab1b1faff89..004d1736b0f30 100644 --- a/library/core/src/array/iter.rs +++ b/library/core/src/array/iter.rs @@ -123,6 +123,22 @@ impl Iterator for IntoIter { (len, Some(len)) } + #[inline] + fn fold(mut self, init: Acc, mut fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + let data = &mut self.data; + (&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() } From 2276c5e3d76591561e28fed760984e75c46bc407 Mon Sep 17 00:00:00 2001 From: The8472 Date: Tue, 27 Jul 2021 00:14:19 +0200 Subject: [PATCH 2/2] from review: add a comment why try_fold was chosen instead of fold --- library/core/src/array/iter.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/library/core/src/array/iter.rs b/library/core/src/array/iter.rs index 004d1736b0f30..f6616399610a6 100644 --- a/library/core/src/array/iter.rs +++ b/library/core/src/array/iter.rs @@ -129,6 +129,11 @@ impl Iterator for IntoIter { 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