diff --git a/src/libcore/benches/iter.rs b/src/libcore/benches/iter.rs index 5b06229c21f23..827c6354c60ba 100644 --- a/src/libcore/benches/iter.rs +++ b/src/libcore/benches/iter.rs @@ -146,3 +146,41 @@ fn bench_for_each_chain_ref_fold(b: &mut Bencher) { acc }); } + +#[bench] +fn bench_flat_map_sum(b: &mut Bencher) { + b.iter(|| -> i64 { + (0i64..1000).flat_map(|x| x..x+1000) + .map(black_box) + .sum() + }); +} + +#[bench] +fn bench_flat_map_ref_sum(b: &mut Bencher) { + b.iter(|| -> i64 { + (0i64..1000).flat_map(|x| x..x+1000) + .map(black_box) + .by_ref() + .sum() + }); +} + +#[bench] +fn bench_flat_map_chain_sum(b: &mut Bencher) { + b.iter(|| -> i64 { + (0i64..1000000).flat_map(|x| once(x).chain(once(x))) + .map(black_box) + .sum() + }); +} + +#[bench] +fn bench_flat_map_chain_ref_sum(b: &mut Bencher) { + b.iter(|| -> i64 { + (0i64..1000000).flat_map(|x| once(x).chain(once(x))) + .map(black_box) + .by_ref() + .sum() + }); +} diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs index ebedfe1d743bb..a596ffd6ae8fc 100644 --- a/src/libcore/iter/mod.rs +++ b/src/libcore/iter/mod.rs @@ -1902,6 +1902,16 @@ impl Iterator for FlatMap _ => (lo, None) } } + + #[inline] + fn fold(self, init: Acc, mut fold: Fold) -> Acc + where Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.frontiter.into_iter() + .chain(self.iter.map(self.f).map(U::into_iter)) + .chain(self.backiter) + .fold(init, |acc, iter| iter.fold(acc, &mut fold)) + } } #[stable(feature = "rust1", since = "1.0.0")] diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs index ed6923929d6b0..59ae30de452c9 100644 --- a/src/libcore/tests/iter.rs +++ b/src/libcore/tests/iter.rs @@ -654,6 +654,22 @@ fn test_iterator_flat_map() { assert_eq!(i, ys.len()); } +/// Test `FlatMap::fold` with items already picked off the front and back, +/// to make sure all parts of the `FlatMap` are folded correctly. +#[test] +fn test_iterator_flat_map_fold() { + let xs = [0, 3, 6]; + let ys = [1, 2, 3, 4, 5, 6, 7]; + let mut it = xs.iter().flat_map(|&x| x..x+3); + it.next(); + it.next_back(); + let i = it.fold(0, |i, x| { + assert_eq!(x, ys[i]); + i + 1 + }); + assert_eq!(i, ys.len()); +} + #[test] fn test_inspect() { let xs = [1, 2, 3, 4];