Rollup merge of rust-lang#44577 - cuviper:flat_map-fold, r=alexcrichton

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`.

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()
+    });
+}

src/libcore/iter/mod.rs

@@ -1902,6 +1902,16 @@ impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
             _ => (lo, None)
         }
     }
+
+    #[inline]
+    fn fold<Acc, 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")]

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];