Skip to content

Commit 61a7703

Browse files
committed
Customize <FlatMap as Iterator>::fold
`FlatMap` can use internal iteration for its `fold`, which shows a performance advantage in the new benchmarks: test iter::bench_flat_map_chain_ref_sum ... bench: 4,354,111 ns/iter (+/- 108,871) test iter::bench_flat_map_chain_sum ... bench: 468,167 ns/iter (+/- 2,274) test iter::bench_flat_map_ref_sum ... bench: 449,616 ns/iter (+/- 6,257) test iter::bench_flat_map_sum ... bench: 348,010 ns/iter (+/- 1,227) ... where the "ref" benches are using `by_ref()` that isn't optimized. So this change shows a decent advantage on its own, but much more when combined with a `chain` iterator that also optimizes `fold`.
1 parent 5dfc84c commit 61a7703

File tree

2 files changed

+48
-0
lines changed

2 files changed

+48
-0
lines changed

src/libcore/benches/iter.rs

+38
Original file line numberDiff line numberDiff line change
@@ -146,3 +146,41 @@ fn bench_for_each_chain_ref_fold(b: &mut Bencher) {
146146
acc
147147
});
148148
}
149+
150+
#[bench]
151+
fn bench_flat_map_sum(b: &mut Bencher) {
152+
b.iter(|| -> i64 {
153+
(0i64..1000).flat_map(|x| x..x+1000)
154+
.map(black_box)
155+
.sum()
156+
});
157+
}
158+
159+
#[bench]
160+
fn bench_flat_map_ref_sum(b: &mut Bencher) {
161+
b.iter(|| -> i64 {
162+
(0i64..1000).flat_map(|x| x..x+1000)
163+
.map(black_box)
164+
.by_ref()
165+
.sum()
166+
});
167+
}
168+
169+
#[bench]
170+
fn bench_flat_map_chain_sum(b: &mut Bencher) {
171+
b.iter(|| -> i64 {
172+
(0i64..1000000).flat_map(|x| once(x).chain(once(x)))
173+
.map(black_box)
174+
.sum()
175+
});
176+
}
177+
178+
#[bench]
179+
fn bench_flat_map_chain_ref_sum(b: &mut Bencher) {
180+
b.iter(|| -> i64 {
181+
(0i64..1000000).flat_map(|x| once(x).chain(once(x)))
182+
.map(black_box)
183+
.by_ref()
184+
.sum()
185+
});
186+
}

src/libcore/iter/mod.rs

+10
Original file line numberDiff line numberDiff line change
@@ -1902,6 +1902,16 @@ impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
19021902
_ => (lo, None)
19031903
}
19041904
}
1905+
1906+
#[inline]
1907+
fn fold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc
1908+
where Fold: FnMut(Acc, Self::Item) -> Acc,
1909+
{
1910+
self.frontiter.into_iter()
1911+
.chain(self.iter.map(self.f).map(U::into_iter))
1912+
.chain(self.backiter)
1913+
.fold(init, |acc, iter| iter.fold(acc, &mut fold))
1914+
}
19051915
}
19061916

19071917
#[stable(feature = "rust1", since = "1.0.0")]

0 commit comments

Comments
 (0)