Specialize .zip() for efficient slice and slice iteration

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

@@ -301,7 +301,9 @@
 
 use clone::Clone;
 use cmp;
+use default::Default;
 use fmt;
+use iter_private::TrustedRandomAccess;
 use ops::FnMut;
 use option::Option::{self, Some, None};
 use usize;
@@ -622,7 +624,8 @@ 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,
+    spec: <(A, B) as ZipImplData>::Data,
 }
 
 #[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,61 @@ 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;
+}
+
+// Zip specialization data members
+#[doc(hidden)]
+trait ZipImplData {
+    type Data: 'static + Clone + Default + fmt::Debug;
+}
+
+#[doc(hidden)]
+impl<T> ZipImplData for T {
+    default type Data = ();
+}
+
+// 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,
+            spec: Default::default(), // unused
+        }
+    }
+
+    #[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,12 +722,106 @@ 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)]
+#[derive(Default, Debug, Clone)]
+struct ZipImplFields {
+    index: usize,
+    len: usize,
+}
+
+#[doc(hidden)]
+impl<A, B> ZipImplData for (A, B)
+    where A: TrustedRandomAccess, B: TrustedRandomAccess
+{
+    type Data = ZipImplFields;
+}
+
+#[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,
+            spec: ZipImplFields {
+                index: 0,
+                len: len,
+            }
+        }
+    }
+
+    #[inline]
+    fn next(&mut self) -> Option<(A::Item, B::Item)> {
+        if self.spec.index < self.spec.len {
+            let i = self.spec.index;
+            self.spec.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.spec.len - self.spec.index;
+        (len, Some(len))
+    }
+
+    #[inline]
+    fn next_back(&mut self) -> Option<(A::Item, B::Item)>
+        where A: DoubleEndedIterator + ExactSizeIterator,
+              B: DoubleEndedIterator + ExactSizeIterator
+    {
+        if self.spec.index < self.spec.len {
+            self.spec.len -= 1;
+            let i = self.spec.len;
+            unsafe {
+                Some((self.a.get_unchecked(i), self.b.get_unchecked(i)))
+            }
+        } else {
+            None
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<A, B> ExactSizeIterator for Zip<A, B>
     where A: ExactSizeIterator, B: ExactSizeIterator {}
 
+#[doc(hidden)]
+unsafe impl<A, B> TrustedRandomAccess for Zip<A, B>
+    where A: TrustedRandomAccess,
+          B: TrustedRandomAccess,
+{
+    unsafe fn get_unchecked(&mut self, i: usize) -> (A::Item, B::Item) {
+        (self.a.get_unchecked(i), self.b.get_unchecked(i))
+    }
+
+}
+
 /// An iterator that maps the values of `iter` with `f`.
 ///
 /// This `struct` is created by the [`map()`] method on [`Iterator`]. See its
@@ -982,6 +1118,15 @@ impl<I> DoubleEndedIterator for Enumerate<I> where
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<I> ExactSizeIterator for Enumerate<I> where I: ExactSizeIterator {}
 
+#[doc(hidden)]
+unsafe impl<I> TrustedRandomAccess for Enumerate<I>
+    where I: TrustedRandomAccess
+{
+    unsafe fn get_unchecked(&mut self, i: usize) -> (usize, I::Item) {
+        (self.count + i, self.iter.get_unchecked(i))
+    }
+}
+
 /// An iterator with a `peek()` that returns an optional reference to the next
 /// element.
 ///

src/libcore/iter_private.rs

@@ -0,0 +1,27 @@
+// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+
+use iter::ExactSizeIterator;
+
+/// An iterator whose items are random accessible efficiently
+///
+/// # Safety
+///
+/// The iterator's .len() and size_hint() must be exact.
+///
+/// .get_unchecked() must return distinct mutable references for distinct
+/// indices (if applicable), and must return a valid reference if index is in
+/// 0..self.len().
+#[doc(hidden)]
+pub unsafe trait TrustedRandomAccess : ExactSizeIterator {
+    unsafe fn get_unchecked(&mut self, i: usize) -> Self::Item;
+}
+

src/libcore/lib.rs

@@ -156,4 +156,5 @@ pub mod hash;
 pub mod fmt;
 
 // note: does not need to be public
+mod iter_private;
 mod tuple;

src/libcore/slice.rs

@@ -50,6 +50,7 @@ use result::Result::{Ok, Err};
 use ptr;
 use mem;
 use marker::{Copy, Send, Sync, self};
+use iter_private::TrustedRandomAccess;
 
 #[repr(C)]
 struct Repr<T> {
@@ -1942,3 +1943,17 @@ macro_rules! impl_marker_for {
 
 impl_marker_for!(BytewiseEquality,
                  u8 i8 u16 i16 u32 i32 u64 i64 usize isize char bool);
+
+#[doc(hidden)]
+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)
+    }
+}
+
+#[doc(hidden)]
+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)
+    }
+}

src/libcoretest/iter.rs

@@ -13,6 +13,7 @@ use core::{i8, i16, isize};
 use core::usize;
 
 use test::Bencher;
+use test::black_box;
 
 #[test]
 fn test_lt() {
@@ -1030,3 +1031,33 @@ fn bench_max(b: &mut Bencher) {
         it.map(scatter).max()
     })
 }
+
+pub fn copy_zip(xs: &[u8], ys: &mut [u8]) {
+    for (a, b) in ys.iter_mut().zip(xs) {
+        *a = *b;
+    }
+}
+
+pub fn add_zip(xs: &[f32], ys: &mut [f32]) {
+    for (a, b) in ys.iter_mut().zip(xs) {
+        *a += *b;
+    }
+}
+
+#[bench]
+fn bench_zip_copy(b: &mut Bencher) {
+    let source = vec![0u8; 16 * 1024];
+    let mut dst = black_box(vec![0u8; 16 * 1024]);
+    b.iter(|| {
+        copy_zip(&source, &mut dst)
+    })
+}
+
+#[bench]
+fn bench_zip_add(b: &mut Bencher) {
+    let source = vec![1.; 16 * 1024];
+    let mut dst = vec![0.; 16 * 1024];
+    b.iter(|| {
+        add_zip(&source, &mut dst)
+    });
+}

src/test/codegen/zip.rs

@@ -0,0 +1,22 @@
+// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// compile-flags: -C no-prepopulate-passes -O
+
+#![crate_type = "lib"]
+
+// CHECK-LABEL: @zip_copy
+#[no_mangle]
+pub fn zip_copy(xs: &[u8], ys: &mut [u8]) {
+// CHECK: memcpy
+    for (x, y) in xs.iter().zip(ys) {
+        *y = *x;
+    }
+}