Auto merge of #37315 - bluss:fold-more, r=alexcrichton

Implement Iterator::fold for .chain(), .cloned(), .map() and the VecDeque iterators.

Chain can do something interesting here where it passes on the fold
into its inner iterators.

The lets the underlying iterator's custom fold() be used, and skips the
regular chain logic in next.

Also implement .fold() specifically for .map() and .cloned() so that any
inner fold improvements are available through map and cloned.

The same way, a VecDeque iterator fold can be turned into two slice folds.

These changes lend the power of the slice iterator's loop codegen to
VecDeque, and to chains of slice iterators, and so on.
It's an improvement for .sum() and .product(), and other uses of fold.

3 files changed, 98 insertions, 20 deletions

diff --git a/src/libcollections/vec_deque.rs b/src/libcollections/vec_deque.rs
index cfed647f5d8..5397193cab4 100644
--- a/src/libcollections/vec_deque.rs
+++ b/src/libcollections/vec_deque.rs
@@ -743,16 +743,8 @@ impl<T> VecDeque<T> {
     #[stable(feature = "deque_extras_15", since = "1.5.0")]
     pub fn as_slices(&self) -> (&[T], &[T]) {
         unsafe {
-            let contiguous = self.is_contiguous();
             let buf = self.buffer_as_slice();
-            if contiguous {
-                let (empty, buf) = buf.split_at(0);
-                (&buf[self.tail..self.head], empty)
-            } else {
-                let (mid, right) = buf.split_at(self.tail);
-                let (left, _) = mid.split_at(self.head);
-                (right, left)
-            }
+            RingSlices::ring_slices(buf, self.head, self.tail)
         }
     }
 
@@ -780,20 +772,10 @@ impl<T> VecDeque<T> {
     #[stable(feature = "deque_extras_15", since = "1.5.0")]
     pub fn as_mut_slices(&mut self) -> (&mut [T], &mut [T]) {
         unsafe {
-            let contiguous = self.is_contiguous();
             let head = self.head;
             let tail = self.tail;
             let buf = self.buffer_as_mut_slice();
-
-            if contiguous {
-                let (empty, buf) = buf.split_at_mut(0);
-                (&mut buf[tail..head], empty)
-            } else {
-                let (mid, right) = buf.split_at_mut(tail);
-                let (left, _) = mid.split_at_mut(head);
-
-                (right, left)
-            }
+            RingSlices::ring_slices(buf, head, tail)
         }
     }
 
@@ -1829,6 +1811,42 @@ fn wrap_index(index: usize, size: usize) -> usize {
     index & (size - 1)
 }
 
+/// Returns the two slices that cover the VecDeque's valid range
+trait RingSlices : Sized {
+    fn slice(self, from: usize, to: usize) -> Self;
+    fn split_at(self, i: usize) -> (Self, Self);
+
+    fn ring_slices(buf: Self, head: usize, tail: usize) -> (Self, Self) {
+        let contiguous = tail <= head;
+        if contiguous {
+            let (empty, buf) = buf.split_at(0);
+            (buf.slice(tail, head), empty)
+        } else {
+            let (mid, right) = buf.split_at(tail);
+            let (left, _) = mid.split_at(head);
+            (right, left)
+        }
+    }
+}
+
+impl<'a, T> RingSlices for &'a [T] {
+    fn slice(self, from: usize, to: usize) -> Self {
+        &self[from..to]
+    }
+    fn split_at(self, i: usize) -> (Self, Self) {
+        (*self).split_at(i)
+    }
+}
+
+impl<'a, T> RingSlices for &'a mut [T] {
+    fn slice(self, from: usize, to: usize) -> Self {
+        &mut self[from..to]
+    }
+    fn split_at(self, i: usize) -> (Self, Self) {
+        (*self).split_at_mut(i)
+    }
+}
+
 /// Calculate the number of elements left to be read in the buffer
 #[inline]
 fn count(tail: usize, head: usize, size: usize) -> usize {
@@ -1875,6 +1893,14 @@ impl<'a, T> Iterator for Iter<'a, T> {
         let len = count(self.tail, self.head, self.ring.len());
         (len, Some(len))
     }
+
+    fn fold<Acc, F>(self, mut accum: Acc, mut f: F) -> Acc
+        where F: FnMut(Acc, Self::Item) -> Acc,
+    {
+        let (front, back) = RingSlices::ring_slices(self.ring, self.head, self.tail);
+        accum = front.iter().fold(accum, &mut f);
+        back.iter().fold(accum, &mut f)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -1927,6 +1953,14 @@ impl<'a, T> Iterator for IterMut<'a, T> {
         let len = count(self.tail, self.head, self.ring.len());
         (len, Some(len))
     }
+
+    fn fold<Acc, F>(self, mut accum: Acc, mut f: F) -> Acc
+        where F: FnMut(Acc, Self::Item) -> Acc,
+    {
+        let (front, back) = RingSlices::ring_slices(self.ring, self.head, self.tail);
+        accum = front.iter_mut().fold(accum, &mut f);
+        back.iter_mut().fold(accum, &mut f)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs
index 9eeb2608071..df4f5e5c576 100644
--- a/src/libcore/iter/mod.rs
+++ b/src/libcore/iter/mod.rs
@@ -399,6 +399,12 @@ impl<'a, I, T: 'a> Iterator for Cloned<I>
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.it.size_hint()
     }
+
+    fn fold<Acc, F>(self, init: Acc, mut f: F) -> Acc
+        where F: FnMut(Acc, Self::Item) -> Acc,
+    {
+        self.it.fold(init, move |acc, elt| f(acc, elt.clone()))
+    }
 }
 
 #[stable(feature = "iter_cloned", since = "1.1.0")]
@@ -544,6 +550,25 @@ impl<A, B> Iterator for Chain<A, B> where
         }
     }
 
+    fn fold<Acc, F>(self, init: Acc, mut f: F) -> Acc
+        where F: FnMut(Acc, Self::Item) -> Acc,
+    {
+        let mut accum = init;
+        match self.state {
+            ChainState::Both | ChainState::Front => {
+                accum = self.a.fold(accum, &mut f);
+            }
+            _ => { }
+        }
+        match self.state {
+            ChainState::Both | ChainState::Back => {
+                accum = self.b.fold(accum, &mut f);
+            }
+            _ => { }
+        }
+        accum
+    }
+
     #[inline]
     fn nth(&mut self, mut n: usize) -> Option<A::Item> {
         match self.state {
@@ -939,6 +964,13 @@ impl<B, I: Iterator, F> Iterator for Map<I, F> where F: FnMut(I::Item) -> B {
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.iter.size_hint()
     }
+
+    fn fold<Acc, G>(self, init: Acc, mut g: G) -> Acc
+        where G: FnMut(Acc, Self::Item) -> Acc,
+    {
+        let mut f = self.f;
+        self.iter.fold(init, move |acc, elt| g(acc, f(elt)))
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
diff --git a/src/libcoretest/iter.rs b/src/libcoretest/iter.rs
index 27eb25537f3..58b6444ef88 100644
--- a/src/libcoretest/iter.rs
+++ b/src/libcoretest/iter.rs
@@ -985,6 +985,18 @@ fn test_empty() {
     assert_eq!(it.next(), None);
 }
 
+#[test]
+fn test_chain_fold() {
+    let xs = [1, 2, 3];
+    let ys = [1, 2, 0];
+
+    let mut iter = xs.iter().chain(&ys);
+    iter.next();
+    let mut result = Vec::new();
+    iter.fold((), |(), &elt| result.push(elt));
+    assert_eq!(&[2, 3, 1, 2, 0], &result[..]);
+}
+
 #[bench]
 fn bench_rposition(b: &mut Bencher) {
     let it: Vec<usize> = (0..300).collect();