From 61a7703e5575795e837d16d3a0ec46551cc6b69b Mon Sep 17 00:00:00 2001
From: Josh Stone <jistone@redhat.com>
Date: Thu, 14 Sep 2017 13:51:32 -0700
Subject: [PATCH 1/2] 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 | 38 +++++++++++++++++++++++++++++++++++++
 src/libcore/iter/mod.rs     | 10 ++++++++++
 2 files changed, 48 insertions(+)

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<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")]

From 351f56a6034486c38c3b36bfbbd15a81a39ba9aa Mon Sep 17 00:00:00 2001
From: Josh Stone <jistone@redhat.com>
Date: Fri, 15 Sep 2017 10:30:56 -0700
Subject: [PATCH 2/2] Add a specific test for `FlatMap::fold`

---
 src/libcore/tests/iter.rs | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

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