Implement O(1) slice::Iter methods.

Instead of using the O(n) defaults, define O(1) shortcuts.

2 files changed, 85 insertions, 0 deletions

diff --git a/src/libcore/slice.rs b/src/libcore/slice.rs
index 1e96d761d40..e2bb881921d 100644
--- a/src/libcore/slice.rs
+++ b/src/libcore/slice.rs
@@ -666,6 +666,60 @@ 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> {
+                // could be implemented with slices, but this avoids bounds checks
+                unsafe {
+                    ::intrinsics::assume(!self.ptr.is_null());
+                    ::intrinsics::assume(!self.end.is_null());
+                    // There should be some way to use offset and optimize this to LEA but I don't
+                    // know how to do that AND detect overflow...
+                    let size = mem::size_of::<T>();
+                    if size == 0 {
+                        if let Some(new_ptr) = (self.ptr as usize).checked_add(n) {
+                            if new_ptr < (self.end as usize) {
+                                self.ptr = transmute(new_ptr + 1);
+                                return Some(&mut *(1 as *mut _))
+                            }
+                        }
+                    } else {
+                        if let Some(new_ptr) = n.checked_mul(size).and_then(|offset| {
+                            (self.ptr as usize).checked_add(offset)
+                        }) {
+                            if new_ptr < (self.end as usize) {
+                                self.ptr = transmute(new_ptr + size);
+                                return Some(transmute(new_ptr))
+                            }
+                        }
+                    }
+                    None
+                }
+            }
+
+            #[inline]
+            fn last(self) -> Option<$elem> {
+                // We could just call next_back but this avoids the memory write.
+                unsafe {
+                    ::intrinsics::assume(!self.ptr.is_null());
+                    ::intrinsics::assume(!self.end.is_null());
+                    if self.end == self.ptr {
+                        None
+                    } else {
+                        if mem::size_of::<T>() == 0 {
+                            // Use a non-null pointer value
+                            Some(&mut *(1 as *mut _))
+                        } else {
+                            Some(transmute(self.end.offset(-1)))
+                        }
+                    }
+                }
+            }
         }
 
         #[stable(feature = "rust1", since = "1.0.0")]
diff --git a/src/libcoretest/slice.rs b/src/libcoretest/slice.rs
index fe73b3b4407..7c884a73ce0 100644
--- a/src/libcoretest/slice.rs
+++ b/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);
+}