Auto merge of #44856 - cuviper:more-fold, r=dtolnay

Add more custom folding to `core::iter` adaptors

Many of the iterator adaptors will perform faster folds if they forward
to their inner iterator's folds, especially for inner types like `Chain`
which are optimized too.  The following types are newly specialized:

| Type        | `fold` | `rfold` |
| ----------- | ------ | ------- |
| `Enumerate` | ✓      | ✓       |
| `Filter`    | ✓      | ✓       |
| `FilterMap` | ✓      | ✓       |
| `FlatMap`   | exists | ✓       |
| `Fuse`      | ✓      | ✓       |
| `Inspect`   | ✓      | ✓       |
| `Peekable`  | ✓      | N/A¹    |
| `Skip`      | ✓      | N/A²    |
| `SkipWhile` | ✓      | N/A¹    |

¹ not a `DoubleEndedIterator`

² `Skip::next_back` doesn't pull skipped items at all, but this couldn't
be avoided if `Skip::rfold` were to call its inner iterator's `rfold`.

Benchmarks
----------

In the following results, plain `_sum` computes the sum of a million
integers -- note that `sum()` is implemented with `fold()`.  The
`_ref_sum` variants do the same on a `by_ref()` iterator, which is
limited to calling `next()` one by one, without specialized `fold`.

The `chain` variants perform the same tests on two iterators chained
together, to show a greater benefit of forwarding `fold` internally.

    test iter::bench_enumerate_chain_ref_sum  ... bench:   2,216,264 ns/iter (+/- 29,228)
    test iter::bench_enumerate_chain_sum      ... bench:     922,380 ns/iter (+/- 2,676)
    test iter::bench_enumerate_ref_sum        ... bench:     476,094 ns/iter (+/- 7,110)
    test iter::bench_enumerate_sum            ... bench:     476,438 ns/iter (+/- 3,334)

    test iter::bench_filter_chain_ref_sum     ... bench:   2,266,095 ns/iter (+/- 6,051)
    test iter::bench_filter_chain_sum         ... bench:     745,594 ns/iter (+/- 2,013)
    test iter::bench_filter_ref_sum           ... bench:     889,696 ns/iter (+/- 1,188)
    test iter::bench_filter_sum               ... bench:     667,325 ns/iter (+/- 1,894)

    test iter::bench_filter_map_chain_ref_sum ... bench:   2,259,195 ns/iter (+/- 353,440)
    test iter::bench_filter_map_chain_sum     ... bench:   1,223,280 ns/iter (+/- 1,972)
    test iter::bench_filter_map_ref_sum       ... bench:     611,607 ns/iter (+/- 2,507)
    test iter::bench_filter_map_sum           ... bench:     611,610 ns/iter (+/- 472)

    test iter::bench_fuse_chain_ref_sum       ... bench:   2,246,106 ns/iter (+/- 22,395)
    test iter::bench_fuse_chain_sum           ... bench:     634,887 ns/iter (+/- 1,341)
    test iter::bench_fuse_ref_sum             ... bench:     444,816 ns/iter (+/- 1,748)
    test iter::bench_fuse_sum                 ... bench:     316,954 ns/iter (+/- 2,616)

    test iter::bench_inspect_chain_ref_sum    ... bench:   2,245,431 ns/iter (+/- 21,371)
    test iter::bench_inspect_chain_sum        ... bench:     631,645 ns/iter (+/- 4,928)
    test iter::bench_inspect_ref_sum          ... bench:     317,437 ns/iter (+/- 702)
    test iter::bench_inspect_sum              ... bench:     315,942 ns/iter (+/- 4,320)

    test iter::bench_peekable_chain_ref_sum   ... bench:   2,243,585 ns/iter (+/- 12,186)
    test iter::bench_peekable_chain_sum       ... bench:     634,848 ns/iter (+/- 1,712)
    test iter::bench_peekable_ref_sum         ... bench:     444,808 ns/iter (+/- 480)
    test iter::bench_peekable_sum             ... bench:     317,133 ns/iter (+/- 3,309)

    test iter::bench_skip_chain_ref_sum       ... bench:   1,778,734 ns/iter (+/- 2,198)
    test iter::bench_skip_chain_sum           ... bench:     761,850 ns/iter (+/- 1,645)
    test iter::bench_skip_ref_sum             ... bench:     478,207 ns/iter (+/- 119,252)
    test iter::bench_skip_sum                 ... bench:     315,614 ns/iter (+/- 3,054)

    test iter::bench_skip_while_chain_ref_sum ... bench:   2,486,370 ns/iter (+/- 4,845)
    test iter::bench_skip_while_chain_sum     ... bench:     633,915 ns/iter (+/- 5,892)
    test iter::bench_skip_while_ref_sum       ... bench:     666,926 ns/iter (+/- 804)
    test iter::bench_skip_while_sum           ... bench:     444,405 ns/iter (+/- 571)

2 files changed, 178 insertions, 3 deletions

diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index 3c4ea974fc9..f8c6fc5c8fa 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -249,6 +249,25 @@ fn test_filter_map() {
 }
 
 #[test]
+fn test_filter_map_fold() {
+    let xs = [0, 1, 2, 3, 4, 5, 6, 7, 8];
+    let ys = [0*0, 2*2, 4*4, 6*6, 8*8];
+    let it = xs.iter().filter_map(|&x| if x % 2 == 0 { Some(x*x) } else { None });
+    let i = it.fold(0, |i, x| {
+        assert_eq!(x, ys[i]);
+        i + 1
+    });
+    assert_eq!(i, ys.len());
+
+    let it = xs.iter().filter_map(|&x| if x % 2 == 0 { Some(x*x) } else { None });
+    let i = it.rfold(ys.len(), |i, x| {
+        assert_eq!(x, ys[i - 1]);
+        i - 1
+    });
+    assert_eq!(i, 0);
+}
+
+#[test]
 fn test_iterator_enumerate() {
     let xs = [0, 1, 2, 3, 4, 5];
     let it = xs.iter().enumerate();
@@ -282,7 +301,31 @@ fn test_iterator_enumerate_nth() {
 #[test]
 fn test_iterator_enumerate_count() {
     let xs = [0, 1, 2, 3, 4, 5];
-    assert_eq!(xs.iter().count(), 6);
+    assert_eq!(xs.iter().enumerate().count(), 6);
+}
+
+#[test]
+fn test_iterator_enumerate_fold() {
+    let xs = [0, 1, 2, 3, 4, 5];
+    let mut it = xs.iter().enumerate();
+    // steal a couple to get an interesting offset
+    assert_eq!(it.next(), Some((0, &0)));
+    assert_eq!(it.next(), Some((1, &1)));
+    let i = it.fold(2, |i, (j, &x)| {
+        assert_eq!(i, j);
+        assert_eq!(x, xs[j]);
+        i + 1
+    });
+    assert_eq!(i, xs.len());
+
+    let mut it = xs.iter().enumerate();
+    assert_eq!(it.next(), Some((0, &0)));
+    let i = it.rfold(xs.len() - 1, |i, (j, &x)| {
+        assert_eq!(i, j);
+        assert_eq!(x, xs[j]);
+        i - 1
+    });
+    assert_eq!(i, 0);
 }
 
 #[test]
@@ -292,6 +335,25 @@ fn test_iterator_filter_count() {
 }
 
 #[test]
+fn test_iterator_filter_fold() {
+    let xs = [0, 1, 2, 3, 4, 5, 6, 7, 8];
+    let ys = [0, 2, 4, 6, 8];
+    let it = xs.iter().filter(|&&x| x % 2 == 0);
+    let i = it.fold(0, |i, &x| {
+        assert_eq!(x, ys[i]);
+        i + 1
+    });
+    assert_eq!(i, ys.len());
+
+    let it = xs.iter().filter(|&&x| x % 2 == 0);
+    let i = it.rfold(ys.len(), |i, &x| {
+        assert_eq!(x, ys[i - 1]);
+        i - 1
+    });
+    assert_eq!(i, 0);
+}
+
+#[test]
 fn test_iterator_peekable() {
     let xs = vec![0, 1, 2, 3, 4, 5];
     let mut it = xs.iter().cloned().peekable();
@@ -381,6 +443,18 @@ fn test_iterator_peekable_last() {
     assert_eq!(it.last(), None);
 }
 
+#[test]
+fn test_iterator_peekable_fold() {
+    let xs = [0, 1, 2, 3, 4, 5];
+    let mut it = xs.iter().peekable();
+    assert_eq!(it.peek(), Some(&&0));
+    let i = it.fold(0, |i, &x| {
+        assert_eq!(x, xs[i]);
+        i + 1
+    });
+    assert_eq!(i, xs.len());
+}
+
 /// This is an iterator that follows the Iterator contract,
 /// but it is not fused. After having returned None once, it will start
 /// producing elements if .next() is called again.
@@ -471,6 +545,26 @@ fn test_iterator_skip_while() {
 }
 
 #[test]
+fn test_iterator_skip_while_fold() {
+    let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19];
+    let ys = [15, 16, 17, 19];
+    let it = xs.iter().skip_while(|&x| *x < 15);
+    let i = it.fold(0, |i, &x| {
+        assert_eq!(x, ys[i]);
+        i + 1
+    });
+    assert_eq!(i, ys.len());
+
+    let mut it = xs.iter().skip_while(|&x| *x < 15);
+    assert_eq!(it.next(), Some(&ys[0])); // process skips before folding
+    let i = it.fold(1, |i, &x| {
+        assert_eq!(x, ys[i]);
+        i + 1
+    });
+    assert_eq!(i, ys.len());
+}
+
+#[test]
 fn test_iterator_skip() {
     let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19, 20, 30];
     let ys = [13, 15, 16, 17, 19, 20, 30];
@@ -567,6 +661,26 @@ fn test_iterator_skip_last() {
 }
 
 #[test]
+fn test_iterator_skip_fold() {
+    let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19, 20, 30];
+    let ys = [13, 15, 16, 17, 19, 20, 30];
+    let it = xs.iter().skip(5);
+    let i = it.fold(0, |i, &x| {
+        assert_eq!(x, ys[i]);
+        i + 1
+    });
+    assert_eq!(i, ys.len());
+
+    let mut it = xs.iter().skip(5);
+    assert_eq!(it.next(), Some(&ys[0])); // process skips before folding
+    let i = it.fold(1, |i, &x| {
+        assert_eq!(x, ys[i]);
+        i + 1
+    });
+    assert_eq!(i, ys.len());
+}
+
+#[test]
 fn test_iterator_take() {
     let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19];
     let ys = [0, 1, 2, 3, 5];
@@ -661,13 +775,22 @@ 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();
+    assert_eq!(it.next(), Some(0));
+    assert_eq!(it.next_back(), Some(8));
     let i = it.fold(0, |i, x| {
         assert_eq!(x, ys[i]);
         i + 1
     });
     assert_eq!(i, ys.len());
+
+    let mut it = xs.iter().flat_map(|&x| x..x+3);
+    assert_eq!(it.next(), Some(0));
+    assert_eq!(it.next_back(), Some(8));
+    let i = it.rfold(ys.len(), |i, x| {
+        assert_eq!(x, ys[i - 1]);
+        i - 1
+    });
+    assert_eq!(i, 0);
 }
 
 #[test]
@@ -685,6 +808,32 @@ fn test_inspect() {
 }
 
 #[test]
+fn test_inspect_fold() {
+    let xs = [1, 2, 3, 4];
+    let mut n = 0;
+    {
+        let it = xs.iter().inspect(|_| n += 1);
+        let i = it.fold(0, |i, &x| {
+            assert_eq!(x, xs[i]);
+            i + 1
+        });
+        assert_eq!(i, xs.len());
+    }
+    assert_eq!(n, xs.len());
+
+    let mut n = 0;
+    {
+        let it = xs.iter().inspect(|_| n += 1);
+        let i = it.rfold(xs.len(), |i, &x| {
+            assert_eq!(x, xs[i - 1]);
+            i - 1
+        });
+        assert_eq!(i, 0);
+    }
+    assert_eq!(n, xs.len());
+}
+
+#[test]
 fn test_cycle() {
     let cycle_len = 3;
     let it = (0..).step_by(1).take(cycle_len).cycle();
@@ -1242,6 +1391,31 @@ fn test_fuse_count() {
 }
 
 #[test]
+fn test_fuse_fold() {
+    let xs = [0, 1, 2];
+    let it = xs.iter(); // `FusedIterator`
+    let i = it.fuse().fold(0, |i, &x| {
+        assert_eq!(x, xs[i]);
+        i + 1
+    });
+    assert_eq!(i, xs.len());
+
+    let it = xs.iter(); // `FusedIterator`
+    let i = it.fuse().rfold(xs.len(), |i, &x| {
+        assert_eq!(x, xs[i - 1]);
+        i - 1
+    });
+    assert_eq!(i, 0);
+
+    let it = xs.iter().scan((), |_, &x| Some(x)); // `!FusedIterator`
+    let i = it.fuse().fold(0, |i, x| {
+        assert_eq!(x, xs[i]);
+        i + 1
+    });
+    assert_eq!(i, xs.len());
+}
+
+#[test]
 fn test_once() {
     let mut it = once(42);
     assert_eq!(it.next(), Some(42));
diff --git a/src/libcore/tests/lib.rs b/src/libcore/tests/lib.rs
index 47995597a0a..d928276ad2f 100644
--- a/src/libcore/tests/lib.rs
+++ b/src/libcore/tests/lib.rs
@@ -25,6 +25,7 @@
 #![feature(inclusive_range)]
 #![feature(inclusive_range_syntax)]
 #![feature(iter_rfind)]
+#![feature(iter_rfold)]
 #![feature(nonzero)]
 #![feature(rand)]
 #![feature(raw)]