about summary refs log tree commit diff
diff options
context:
space:
mode:
authorThe8472 <git@infinite-source.de>2021-07-24 17:42:20 +0200
committerThe8472 <git@infinite-source.de>2021-07-24 19:24:11 +0200
commite015e9da717bf501ab4752b2ea045d899f804a42 (patch)
tree50978138163fc4dad124b8f7de1ff97680b6f2d0
parent18840b0719aa766a1bc49ea2eb5dc2e4cde7da3f (diff)
downloadrust-e015e9da717bf501ab4752b2ea045d899f804a42.tar.gz
rust-e015e9da717bf501ab4752b2ea045d899f804a42.zip
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)
```
-rw-r--r--library/alloc/benches/vec.rs6
-rw-r--r--library/core/src/array/iter.rs16
2 files changed, 22 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..004d1736b0f 100644
--- a/library/core/src/array/iter.rs
+++ b/library/core/src/array/iter.rs
@@ -123,6 +123,22 @@ 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;
+        (&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()
     }