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