Auto merge of rust-lang#24701 - Stebalien:slice, r=alexcrichton

Instead of using the O(n) defaults, define O(1) shortcuts. I also copied (and slightly modified) the relevant tests from the iter tests into the slice tests just in case someone comes along and changes them in the future.

Partially implements  rust-lang#24214.

Author: bors

src/libcore/slice.rs

@@ -625,6 +625,36 @@ impl<'a, T> IntoIterator for &'a mut [T] {
     }
 }
 
+#[inline(always)]
+fn size_from_ptr<T>(_: *const T) -> usize {
+    mem::size_of::<T>()
+}
+
+
+// Use macro to be generic over const/mut
+macro_rules! slice_offset {
+    ($ptr:expr, $by:expr) => {{
+        let ptr = $ptr;
+        if size_from_ptr(ptr) == 0 {
+            transmute(ptr as usize + $by)
+        } else {
+            ptr.offset($by)
+        }
+    }};
+}
+
+macro_rules! slice_ref {
+    ($ptr:expr) => {{
+        let ptr = $ptr;
+        if size_from_ptr(ptr) == 0 {
+            // Use a non-null pointer value
+            &mut *(1 as *mut _)
+        } else {
+            transmute(ptr)
+        }
+    }};
+}
+
 // The shared definition of the `Iter` and `IterMut` iterators
 macro_rules! iterator {
     (struct $name:ident -> $ptr:ty, $elem:ty) => {
@@ -641,20 +671,9 @@ macro_rules! iterator {
                     if self.ptr == self.end {
                         None
                     } else {
-                        if mem::size_of::<T>() == 0 {
-                            // purposefully don't use 'ptr.offset' because for
-                            // vectors with 0-size elements this would return the
-                            // same pointer.
-                            self.ptr = transmute(self.ptr as usize + 1);
-
-                            // Use a non-null pointer value
-                            Some(&mut *(1 as *mut _))
-                        } else {
-                            let old = self.ptr;
-                            self.ptr = self.ptr.offset(1);
-
-                            Some(transmute(old))
-                        }
+                        let old = self.ptr;
+                        self.ptr = slice_offset!(self.ptr, 1);
+                        Some(slice_ref!(old))
                     }
                 }
             }
@@ -666,6 +685,22 @@ macro_rules! iterator {
                 let exact = diff / (if size == 0 {1} else {size});
                 (exact, Some(exact))
             }
+
+            #[inline]
+            fn count(self) -> usize {
+                self.size_hint().0
+            }
+
+            #[inline]
+            fn nth(&mut self, n: usize) -> Option<$elem> {
+                // Call helper method. Can't put the definition here because mut versus const.
+                self.iter_nth(n)
+            }
+
+            #[inline]
+            fn last(mut self) -> Option<$elem> {
+                self.next_back()
+            }
         }
 
         #[stable(feature = "rust1", since = "1.0.0")]
@@ -679,17 +714,8 @@ macro_rules! iterator {
                     if self.end == self.ptr {
                         None
                     } else {
-                        if mem::size_of::<T>() == 0 {
-                            // See above for why 'ptr.offset' isn't used
-                            self.end = transmute(self.end as usize - 1);
-
-                            // Use a non-null pointer value
-                            Some(&mut *(1 as *mut _))
-                        } else {
-                            self.end = self.end.offset(-1);
-
-                            Some(transmute(self.end))
-                        }
+                        self.end = slice_offset!(self.end, -1);
+                        Some(slice_ref!(self.end))
                     }
                 }
             }
@@ -785,6 +811,20 @@ impl<'a, T> Iter<'a, T> {
     pub fn as_slice(&self) -> &'a [T] {
         make_slice!(T => &'a [T]: self.ptr, self.end)
     }
+
+    // Helper function for Iter::nth
+    fn iter_nth(&mut self, n: usize) -> Option<&'a T> {
+        match self.as_slice().get(n) {
+            Some(elem_ref) => unsafe {
+                self.ptr = slice_offset!(elem_ref as *const _, 1);
+                Some(slice_ref!(elem_ref))
+            },
+            None => {
+                self.ptr = self.end;
+                None
+            }
+        }
+    }
 }
 
 iterator!{struct Iter -> *const T, &'a T}
@@ -914,6 +954,20 @@ impl<'a, T> IterMut<'a, T> {
     pub fn into_slice(self) -> &'a mut [T] {
         make_mut_slice!(T => &'a mut [T]: self.ptr, self.end)
     }
+
+    // Helper function for IterMut::nth
+    fn iter_nth(&mut self, n: usize) -> Option<&'a mut T> {
+        match make_mut_slice!(T => &'a mut [T]: self.ptr, self.end).get_mut(n) {
+            Some(elem_ref) => unsafe {
+                self.ptr = slice_offset!(elem_ref as *mut _, 1);
+                Some(slice_ref!(elem_ref))
+            },
+            None => {
+                self.ptr = self.end;
+                None
+            }
+        }
+    }
 }
 
 iterator!{struct IterMut -> *mut T, &'a mut T}

src/libcoretest/slice.rs

@@ -82,3 +82,34 @@ fn iterator_to_slice() {
     test!([1u8,2,3]);
     test!([(),(),()]);
 }
+
+#[test]
+fn test_iterator_nth() {
+    let v: &[_] = &[0, 1, 2, 3, 4];
+    for i in 0..v.len() {
+        assert_eq!(v.iter().nth(i).unwrap(), &v[i]);
+    }
+    assert_eq!(v.iter().nth(v.len()), None);
+
+    let mut iter = v.iter();
+    assert_eq!(iter.nth(2).unwrap(), &v[2]);
+    assert_eq!(iter.nth(1).unwrap(), &v[4]);
+}
+
+#[test]
+fn test_iterator_last() {
+    let v: &[_] = &[0, 1, 2, 3, 4];
+    assert_eq!(v.iter().last().unwrap(), &4);
+    assert_eq!(v[..1].iter().last().unwrap(), &0);
+}
+
+#[test]
+fn test_iterator_count() {
+    let v: &[_] = &[0, 1, 2, 3, 4];
+    assert_eq!(v.iter().count(), 5);
+
+    let mut iter2 = v.iter();
+    iter2.next();
+    iter2.next();
+    assert_eq!(iter2.count(), 3);
+}