Expand .zip() specialization to .map() and .cloned()

Implement .zip() specialization for Map and Cloned.

The crucial thing for transparent specialization is that we want to
preserve the potential side effects.

The simplest example is that in this code snippet:

`(0..6).map(f).zip((0..4).map(g)).count()`

`f` will be called five times, and `g` four times. The last time for `f`
is when the other iterator is at its end, so this element is unused.
This side effect can be preserved without disturbing code generation for
simple uses of `.map()`.

The `Zip::next_back()` case is even more complicated, unfortunately.

3 files changed, 65 insertions, 1 deletions

diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs
index 035418570a0..9eeb2608071 100644
--- a/src/libcore/iter/mod.rs
+++ b/src/libcore/iter/mod.rs
@@ -420,6 +420,18 @@ impl<'a, I, T: 'a> FusedIterator for Cloned<I>
     where I: FusedIterator<Item=&'a T>, T: Clone
 {}
 
+#[doc(hidden)]
+unsafe impl<'a, I, T: 'a> TrustedRandomAccess for Cloned<I>
+    where I: TrustedRandomAccess<Item=&'a T>, T: Clone
+{
+    unsafe fn get_unchecked(&mut self, i: usize) -> Self::Item {
+        self.it.get_unchecked(i).clone()
+    }
+
+    #[inline]
+    fn may_have_side_effect() -> bool { true }
+}
+
 /// An iterator that repeats endlessly.
 ///
 /// This `struct` is created by the [`cycle()`] method on [`Iterator`]. See its
@@ -773,6 +785,13 @@ impl<A, B> ZipImpl<A, B> for Zip<A, B>
             unsafe {
                 Some((self.a.get_unchecked(i), self.b.get_unchecked(i)))
             }
+        } else if A::may_have_side_effect() && self.index < self.a.len() {
+            // match the base implementation's potential side effects
+            unsafe {
+                self.a.get_unchecked(self.index);
+            }
+            self.index += 1;
+            None
         } else {
             None
         }
@@ -789,6 +808,23 @@ impl<A, B> ZipImpl<A, B> for Zip<A, B>
         where A: DoubleEndedIterator + ExactSizeIterator,
               B: DoubleEndedIterator + ExactSizeIterator
     {
+        // Adjust a, b to equal length
+        if A::may_have_side_effect() {
+            let sz = self.a.len();
+            if sz > self.len {
+                for _ in 0..sz - cmp::max(self.len, self.index) {
+                    self.a.next_back();
+                }
+            }
+        }
+        if B::may_have_side_effect() {
+            let sz = self.b.len();
+            if sz > self.len {
+                for _ in 0..sz - self.len {
+                    self.b.next_back();
+                }
+            }
+        }
         if self.index < self.len {
             self.len -= 1;
             let i = self.len;
@@ -814,6 +850,9 @@ unsafe impl<A, B> TrustedRandomAccess for Zip<A, B>
         (self.a.get_unchecked(i), self.b.get_unchecked(i))
     }
 
+    fn may_have_side_effect() -> bool {
+        A::may_have_side_effect() || B::may_have_side_effect()
+    }
 }
 
 #[unstable(feature = "fused", issue = "35602")]
@@ -920,6 +959,18 @@ impl<B, I: ExactSizeIterator, F> ExactSizeIterator for Map<I, F>
 impl<B, I: FusedIterator, F> FusedIterator for Map<I, F>
     where F: FnMut(I::Item) -> B {}
 
+#[doc(hidden)]
+unsafe impl<B, I, F> TrustedRandomAccess for Map<I, F>
+    where I: TrustedRandomAccess,
+          F: FnMut(I::Item) -> B,
+{
+    unsafe fn get_unchecked(&mut self, i: usize) -> Self::Item {
+        (self.f)(self.iter.get_unchecked(i))
+    }
+    #[inline]
+    fn may_have_side_effect() -> bool { true }
+}
+
 /// An iterator that filters the elements of `iter` with `predicate`.
 ///
 /// This `struct` is created by the [`filter()`] method on [`Iterator`]. See its
@@ -1135,6 +1186,10 @@ unsafe impl<I> TrustedRandomAccess for Enumerate<I>
     unsafe fn get_unchecked(&mut self, i: usize) -> (usize, I::Item) {
         (self.count + i, self.iter.get_unchecked(i))
     }
+
+    fn may_have_side_effect() -> bool {
+        I::may_have_side_effect()
+    }
 }
 
 #[unstable(feature = "fused", issue = "35602")]
@@ -1764,6 +1819,10 @@ unsafe impl<I> TrustedRandomAccess for Fuse<I>
     unsafe fn get_unchecked(&mut self, i: usize) -> I::Item {
         self.iter.get_unchecked(i)
     }
+
+    fn may_have_side_effect() -> bool {
+        I::may_have_side_effect()
+    }
 }
 
 #[unstable(feature = "fused", issue = "35602")]
diff --git a/src/libcore/iter_private.rs b/src/libcore/iter_private.rs
index 83eeef31ab0..bc1aaa09f3d 100644
--- a/src/libcore/iter_private.rs
+++ b/src/libcore/iter_private.rs
@@ -14,6 +14,7 @@
 /// # Safety
 ///
 /// The iterator's .len() and size_hint() must be exact.
+/// `.len()` must be cheap to call.
 ///
 /// .get_unchecked() must return distinct mutable references for distinct
 /// indices (if applicable), and must return a valid reference if index is in
@@ -21,5 +22,7 @@
 #[doc(hidden)]
 pub unsafe trait TrustedRandomAccess : ExactSizeIterator {
     unsafe fn get_unchecked(&mut self, i: usize) -> Self::Item;
+    /// Return `true` if getting an iterator element may have
+    /// side effects. Remember to take inner iterators into account.
+    fn may_have_side_effect() -> bool;
 }
-
diff --git a/src/libcore/slice.rs b/src/libcore/slice.rs
index d1df56905df..31be404ba90 100644
--- a/src/libcore/slice.rs
+++ b/src/libcore/slice.rs
@@ -1968,6 +1968,7 @@ unsafe impl<'a, T> TrustedRandomAccess for Iter<'a, T> {
     unsafe fn get_unchecked(&mut self, i: usize) -> &'a T {
         &*self.ptr.offset(i as isize)
     }
+    fn may_have_side_effect() -> bool { false }
 }
 
 #[doc(hidden)]
@@ -1975,4 +1976,5 @@ unsafe impl<'a, T> TrustedRandomAccess for IterMut<'a, T> {
     unsafe fn get_unchecked(&mut self, i: usize) -> &'a mut T {
         &mut *self.ptr.offset(i as isize)
     }
+    fn may_have_side_effect() -> bool { false }
 }