about summary refs log tree commit diff
diff options
context:
space:
mode:
authorbors <bors@rust-lang.org>2021-07-27 10:38:41 +0000
committerbors <bors@rust-lang.org>2021-07-27 10:38:41 +0000
commit99d6692f6c1cebd6c56a67eb21f6aae26c12a145 (patch)
tree57e6df70c3e510b0d103d5d54429f99711333e44
parent998cfe5aad7c21eb19a4bca50f05a13354706970 (diff)
parent2276c5e3d76591561e28fed760984e75c46bc407 (diff)
downloadrust-99d6692f6c1cebd6c56a67eb21f6aae26c12a145.tar.gz
rust-99d6692f6c1cebd6c56a67eb21f6aae26c12a145.zip
Auto merge of #87431 - the8472:array-iter-fold, r=kennytm
implement fold() on array::IntoIter to improve flatten().collect() perf

With #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 #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)
```
-rw-r--r--library/alloc/benches/vec.rs6
-rw-r--r--library/core/src/array/iter.rs21
2 files changed, 27 insertions, 0 deletions
diff --git a/library/alloc/benches/vec.rs b/library/alloc/benches/vec.rs
index 91eec10d575..c93a493cadb 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::<Vec<_>>());
+}
diff --git a/library/core/src/array/iter.rs b/library/core/src/array/iter.rs
index 61ab1b1faff..f6616399610 100644
--- a/library/core/src/array/iter.rs
+++ b/library/core/src/array/iter.rs
@@ -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()
     }