O(1) count,nth,last for slice::Windows,Chunks(Mut)

Implemented count, nth, and last in constant time for Windows, Chunks,
and ChunksMut created from a slice.

Included checks for overflow in the implementation of nth().

Also added a test for each implemented method to libcoretest.

Addresses #24214

2 files changed, 215 insertions, 0 deletions

diff --git a/src/libcore/slice.rs b/src/libcore/slice.rs
index f765cdc54d8..6513c49a06e 100644
--- a/src/libcore/slice.rs
+++ b/src/libcore/slice.rs
@@ -51,6 +51,7 @@ use ptr;
 use mem;
 use mem::size_of;
 use marker::{Send, Sync, self};
+use num::wrapping::OverflowingOps;
 use raw::Repr;
 // Avoid conflicts with *both* the Slice trait (buggy) and the `slice::raw` module.
 use raw::Slice as RawSlice;
@@ -1183,6 +1184,34 @@ impl<'a, T> Iterator for Windows<'a, T> {
             (size, Some(size))
         }
     }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.size_hint().0
+    }
+
+    #[inline]
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        let (end, overflow) = self.size.overflowing_add(n);
+        if end > self.v.len() || overflow {
+            self.v = &[];
+            None
+        } else {
+            let nth = &self.v[n..end];
+            self.v = &self.v[n+1..];
+            Some(nth)
+        }
+    }
+
+    #[inline]
+    fn last(self) -> Option<Self::Item> {
+        if self.size > self.v.len() {
+            None
+        } else {
+            let start = self.v.len() - self.size;
+            Some(&self.v[start..])
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -1269,6 +1298,38 @@ impl<'a, T> Iterator for Chunks<'a, T> {
             (n, Some(n))
         }
     }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.size_hint().0
+    }
+
+    #[inline]
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        let (start, overflow) = n.overflowing_mul(self.size);
+        if start >= self.v.len() || overflow {
+            self.v = &[];
+            None
+        } else {
+            let end = match start.checked_add(self.size) {
+                Some(sum) => cmp::min(self.v.len(), sum),
+                None => self.v.len(),
+            };
+            let nth = &self.v[start..end];
+            self.v = &self.v[end..];
+            Some(nth)
+        }
+    }
+
+    #[inline]
+    fn last(self) -> Option<Self::Item> {
+        if self.v.is_empty() {
+            None
+        } else {
+            let start = (self.v.len() - 1) / self.size * self.size;
+            Some(&self.v[start..])
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -1349,6 +1410,40 @@ impl<'a, T> Iterator for ChunksMut<'a, T> {
             (n, Some(n))
         }
     }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.size_hint().0
+    }
+
+    #[inline]
+    fn nth(&mut self, n: usize) -> Option<&'a mut [T]> {
+        let (start, overflow) = n.overflowing_mul(self.chunk_size);
+        if start >= self.v.len() || overflow {
+            self.v = &mut [];
+            None
+        } else {
+            let end = match start.checked_add(self.chunk_size) {
+                Some(sum) => cmp::min(self.v.len(), sum),
+                None => self.v.len(),
+            };
+            let tmp = mem::replace(&mut self.v, &mut []);
+            let (head, tail) = tmp.split_at_mut(end);
+            let (_, nth) =  head.split_at_mut(start);
+            self.v = tail;
+            Some(nth)
+        }
+    }
+
+    #[inline]
+    fn last(self) -> Option<Self::Item> {
+        if self.v.is_empty() {
+            None
+        } else {
+            let start = (self.v.len() - 1) / self.chunk_size * self.chunk_size;
+            Some(&mut self.v[start..])
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
diff --git a/src/libcoretest/slice.rs b/src/libcoretest/slice.rs
index 5a1166aeb50..d60eeb76ccd 100644
--- a/src/libcoretest/slice.rs
+++ b/src/libcoretest/slice.rs
@@ -64,3 +64,123 @@ fn test_iterator_count() {
     iter2.next();
     assert_eq!(iter2.count(), 3);
 }
+
+#[test]
+fn test_chunks_count() {
+    let v: &[i32] = &[0, 1, 2, 3, 4, 5];
+    let c = v.chunks(3);
+    assert_eq!(c.count(), 2);
+
+    let v2: &[i32] = &[0, 1, 2, 3, 4];
+    let c2 = v2.chunks(2);
+    assert_eq!(c2.count(), 3);
+
+    let v3: &[i32] = &[];
+    let c3 = v3.chunks(2);
+    assert_eq!(c3.count(), 0);
+}
+
+#[test]
+fn test_chunks_nth() {
+    let v: &[i32] = &[0, 1, 2, 3, 4, 5];
+    let mut c = v.chunks(2);
+    assert_eq!(c.nth(1).unwrap()[1], 3);
+    assert_eq!(c.next().unwrap()[0], 4);
+
+    let v2: &[i32] = &[0, 1, 2, 3, 4];
+    let mut c2 = v2.chunks(3);
+    assert_eq!(c2.nth(1).unwrap()[1], 4);
+    assert_eq!(c2.next(), None);
+}
+
+#[test]
+fn test_chunks_last() {
+    let v: &[i32] = &[0, 1, 2, 3, 4, 5];
+    let c = v.chunks(2);
+    assert_eq!(c.last().unwrap()[1], 5);
+
+    let v2: &[i32] = &[0, 1, 2, 3, 4];
+    let c2 = v2.chunks(2);
+    assert_eq!(c2.last().unwrap()[0], 4);
+}
+
+#[test]
+fn test_chunks_mut_count() {
+    let mut v: &mut [i32] = &mut [0, 1, 2, 3, 4, 5];
+    let c = v.chunks_mut(3);
+    assert_eq!(c.count(), 2);
+
+    let mut v2: &mut [i32] = &mut [0, 1, 2, 3, 4];
+    let c2 = v2.chunks_mut(2);
+    assert_eq!(c2.count(), 3);
+
+    let mut v3: &mut [i32] = &mut [];
+    let c3 = v3.chunks_mut(2);
+    assert_eq!(c3.count(), 0);
+}
+
+#[test]
+fn test_chunks_mut_nth() {
+    let mut v: &mut [i32] = &mut [0, 1, 2, 3, 4, 5];
+    let mut c = v.chunks_mut(2);
+    assert_eq!(c.nth(1).unwrap()[1], 3);
+    assert_eq!(c.next().unwrap()[0], 4);
+
+    let mut v2: &mut [i32] = &mut [0, 1, 2, 3, 4];
+    let mut c2 = v2.chunks_mut(3);
+    assert_eq!(c2.nth(1).unwrap()[1], 4);
+    assert_eq!(c2.next(), None);
+}
+
+#[test]
+fn test_chunks_mut_last() {
+    let v: &mut [i32] = &mut [0, 1, 2, 3, 4, 5];
+    let c = v.chunks_mut(2);
+    assert_eq!(c.last().unwrap()[1], 5);
+
+    let v2: &mut [i32] = &mut [0, 1, 2, 3, 4];
+    let c2 = v2.chunks_mut(2);
+    assert_eq!(c2.last().unwrap()[0], 4);
+}
+
+
+
+
+#[test]
+fn test_windows_count() {
+    let v: &[i32] = &[0, 1, 2, 3, 4, 5];
+    let c = v.windows(3);
+    assert_eq!(c.count(), 4);
+
+    let v2: &[i32] = &[0, 1, 2, 3, 4];
+    let c2 = v2.windows(6);
+    assert_eq!(c2.count(), 0);
+
+    let v3: &[i32] = &[];
+    let c3 = v3.windows(2);
+    assert_eq!(c3.count(), 0);
+}
+
+#[test]
+fn test_windows_nth() {
+    let v: &[i32] = &[0, 1, 2, 3, 4, 5];
+    let mut c = v.windows(2);
+    assert_eq!(c.nth(2).unwrap()[1], 3);
+    assert_eq!(c.next().unwrap()[0], 3);
+
+    let v2: &[i32] = &[0, 1, 2, 3, 4];
+    let mut c2 = v2.windows(4);
+    assert_eq!(c2.nth(1).unwrap()[1], 2);
+    assert_eq!(c2.next(), None);
+}
+
+#[test]
+fn test_windows_last() {
+    let v: &[i32] = &[0, 1, 2, 3, 4, 5];
+    let c = v.windows(2);
+    assert_eq!(c.last().unwrap()[1], 5);
+
+    let v2: &[i32] = &[0, 1, 2, 3, 4];
+    let c2 = v2.windows(2);
+    assert_eq!(c2.last().unwrap()[0], 3);
+}