Implement O(1) slice::Iter{,Mut} methods.

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);
+}