specialize zip: Specialize .zip() for TrustedRandomAccess iterators

This allows common iterator compositions like a.zip(b) where a, b
are slice::{Iter, IterMut} compile to *much* better code.

Author: bluss

src/libcore/iter/iterator.rs

@@ -23,6 +23,7 @@ use super::{Chain, Cycle, Cloned, Enumerate, Filter, FilterMap, FlatMap, Fuse,
 use super::ChainState;
 use super::{DoubleEndedIterator, ExactSizeIterator, Extend, FromIterator,
             IntoIterator};
+use super::ZipImpl;
 
 fn _assert_is_object_safe(_: &Iterator<Item=()>) {}
 
@@ -383,7 +384,7 @@ pub trait Iterator {
     fn zip<U>(self, other: U) -> Zip<Self, U::IntoIter> where
         Self: Sized, U: IntoIterator
     {
-        Zip{a: self, b: other.into_iter()}
+        Zip::new(self, other.into_iter())
     }
 
     /// Takes a closure and creates an iterator which calls that closure on each

src/libcore/iter/mod.rs

@@ -302,6 +302,7 @@
 use clone::Clone;
 use cmp;
 use fmt;
+use iter_private::TrustedRandomAccess;
 use ops::FnMut;
 use option::Option::{self, Some, None};
 use usize;
@@ -622,7 +623,9 @@ impl<A, B> DoubleEndedIterator for Chain<A, B> where
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct Zip<A, B> {
     a: A,
-    b: B
+    b: B,
+    index: usize,
+    len: usize,
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -631,29 +634,13 @@ impl<A, B> Iterator for Zip<A, B> where A: Iterator, B: Iterator
     type Item = (A::Item, B::Item);
 
     #[inline]
-    fn next(&mut self) -> Option<(A::Item, B::Item)> {
-        self.a.next().and_then(|x| {
-            self.b.next().and_then(|y| {
-                Some((x, y))
-            })
-        })
+    fn next(&mut self) -> Option<Self::Item> {
+        ZipImpl::next(self)
     }
 
     #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let (a_lower, a_upper) = self.a.size_hint();
-        let (b_lower, b_upper) = self.b.size_hint();
-
-        let lower = cmp::min(a_lower, b_lower);
-
-        let upper = match (a_upper, b_upper) {
-            (Some(x), Some(y)) => Some(cmp::min(x,y)),
-            (Some(x), None) => Some(x),
-            (None, Some(y)) => Some(y),
-            (None, None) => None
-        };
-
-        (lower, upper)
+        ZipImpl::size_hint(self)
     }
 }
 
@@ -664,6 +651,51 @@ impl<A, B> DoubleEndedIterator for Zip<A, B> where
 {
     #[inline]
     fn next_back(&mut self) -> Option<(A::Item, B::Item)> {
+        ZipImpl::next_back(self)
+    }
+}
+
+// Zip specialization trait
+#[doc(hidden)]
+trait ZipImpl<A, B> {
+    type Item;
+    fn new(a: A, b: B) -> Self;
+    fn next(&mut self) -> Option<Self::Item>;
+    fn size_hint(&self) -> (usize, Option<usize>);
+    fn next_back(&mut self) -> Option<Self::Item>
+        where A: DoubleEndedIterator + ExactSizeIterator,
+              B: DoubleEndedIterator + ExactSizeIterator;
+}
+
+// General Zip impl
+#[doc(hidden)]
+impl<A, B> ZipImpl<A, B> for Zip<A, B>
+    where A: Iterator, B: Iterator
+{
+    type Item = (A::Item, B::Item);
+    default fn new(a: A, b: B) -> Self {
+        Zip {
+            a: a,
+            b: b,
+            index: 0, // not used in general case
+            len: 0,
+        }
+    }
+
+    #[inline]
+    default fn next(&mut self) -> Option<(A::Item, B::Item)> {
+        self.a.next().and_then(|x| {
+            self.b.next().and_then(|y| {
+                Some((x, y))
+            })
+        })
+    }
+
+    #[inline]
+    default fn next_back(&mut self) -> Option<(A::Item, B::Item)>
+        where A: DoubleEndedIterator + ExactSizeIterator,
+              B: DoubleEndedIterator + ExactSizeIterator
+    {
         let a_sz = self.a.len();
         let b_sz = self.b.len();
         if a_sz != b_sz {
@@ -680,6 +712,73 @@ impl<A, B> DoubleEndedIterator for Zip<A, B> where
             _ => unreachable!(),
         }
     }
+
+    #[inline]
+    default fn size_hint(&self) -> (usize, Option<usize>) {
+        let (a_lower, a_upper) = self.a.size_hint();
+        let (b_lower, b_upper) = self.b.size_hint();
+
+        let lower = cmp::min(a_lower, b_lower);
+
+        let upper = match (a_upper, b_upper) {
+            (Some(x), Some(y)) => Some(cmp::min(x,y)),
+            (Some(x), None) => Some(x),
+            (None, Some(y)) => Some(y),
+            (None, None) => None
+        };
+
+        (lower, upper)
+    }
+}
+
+#[doc(hidden)]
+impl<A, B> ZipImpl<A, B> for Zip<A, B>
+    where A: TrustedRandomAccess, B: TrustedRandomAccess
+{
+    fn new(a: A, b: B) -> Self {
+        let len = cmp::min(a.len(), b.len());
+        Zip {
+            a: a,
+            b: b,
+            index: 0,
+            len: len,
+        }
+    }
+
+    #[inline]
+    fn next(&mut self) -> Option<(A::Item, B::Item)> {
+        if self.index < self.len {
+            let i = self.index;
+            self.index += 1;
+            unsafe {
+                Some((self.a.get_unchecked(i), self.b.get_unchecked(i)))
+            }
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.len - self.index;
+        (len, Some(len))
+    }
+
+    #[inline]
+    fn next_back(&mut self) -> Option<(A::Item, B::Item)>
+        where A: DoubleEndedIterator + ExactSizeIterator,
+              B: DoubleEndedIterator + ExactSizeIterator
+    {
+        if self.index < self.len {
+            self.len -= 1;
+            let i = self.len;
+            unsafe {
+                Some((self.a.get_unchecked(i), self.b.get_unchecked(i)))
+            }
+        } else {
+            None
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]