Auto merge of #47142 - sdroege:trusted-random-access-chunks, r=kennytm

Implement TrustedRandomAccess for slice::{Chunks, ChunksMut, Windows}

As suggested by @bluss in https://github.com/rust-lang/rust/issues/47115#issuecomment-354888334

2 files changed, 73 insertions, 0 deletions

diff --git a/src/libcore/slice/mod.rs b/src/libcore/slice/mod.rs
index 6b6ec7147b3..72036d6d3a2 100644
--- a/src/libcore/slice/mod.rs
+++ b/src/libcore/slice/mod.rs
@@ -2117,6 +2117,14 @@ impl<'a, T> ExactSizeIterator for Windows<'a, T> {}
 #[unstable(feature = "fused", issue = "35602")]
 impl<'a, T> FusedIterator for Windows<'a, T> {}
 
+#[doc(hidden)]
+unsafe impl<'a, T> TrustedRandomAccess for Windows<'a, T> {
+    unsafe fn get_unchecked(&mut self, i: usize) -> &'a [T] {
+        from_raw_parts(self.v.as_ptr().offset(i as isize), self.size)
+    }
+    fn may_have_side_effect() -> bool { false }
+}
+
 /// An iterator over a slice in (non-overlapping) chunks (`chunk_size` elements at a
 /// time).
 ///
@@ -2228,6 +2236,19 @@ impl<'a, T> ExactSizeIterator for Chunks<'a, T> {}
 #[unstable(feature = "fused", issue = "35602")]
 impl<'a, T> FusedIterator for Chunks<'a, T> {}
 
+#[doc(hidden)]
+unsafe impl<'a, T> TrustedRandomAccess for Chunks<'a, T> {
+    unsafe fn get_unchecked(&mut self, i: usize) -> &'a [T] {
+        let start = i * self.chunk_size;
+        let end = match start.checked_add(self.chunk_size) {
+            None => self.v.len(),
+            Some(end) => cmp::min(end, self.v.len()),
+        };
+        from_raw_parts(self.v.as_ptr().offset(start as isize), end - start)
+    }
+    fn may_have_side_effect() -> bool { false }
+}
+
 /// An iterator over a slice in (non-overlapping) mutable chunks (`chunk_size`
 /// elements at a time). When the slice len is not evenly divided by the chunk
 /// size, the last slice of the iteration will be the remainder.
@@ -2331,6 +2352,19 @@ impl<'a, T> ExactSizeIterator for ChunksMut<'a, T> {}
 #[unstable(feature = "fused", issue = "35602")]
 impl<'a, T> FusedIterator for ChunksMut<'a, T> {}
 
+#[doc(hidden)]
+unsafe impl<'a, T> TrustedRandomAccess for ChunksMut<'a, T> {
+    unsafe fn get_unchecked(&mut self, i: usize) -> &'a mut [T] {
+        let start = i * self.chunk_size;
+        let end = match start.checked_add(self.chunk_size) {
+            None => self.v.len(),
+            Some(end) => cmp::min(end, self.v.len()),
+        };
+        from_raw_parts_mut(self.v.as_mut_ptr().offset(start as isize), end - start)
+    }
+    fn may_have_side_effect() -> bool { false }
+}
+
 //
 // Free functions
 //
diff --git a/src/libcore/tests/slice.rs b/src/libcore/tests/slice.rs
index fa4c2e9b373..d6230e93f99 100644
--- a/src/libcore/tests/slice.rs
+++ b/src/libcore/tests/slice.rs
@@ -138,6 +138,18 @@ fn test_chunks_last() {
 }
 
 #[test]
+fn test_chunks_zip() {
+    let v1: &[i32] = &[0, 1, 2, 3, 4];
+    let v2: &[i32] = &[6, 7, 8, 9, 10];
+
+    let res = v1.chunks(2)
+        .zip(v2.chunks(2))
+        .map(|(a, b)| a.iter().sum::<i32>() + b.iter().sum::<i32>())
+        .collect::<Vec<_>>();
+    assert_eq!(res, vec![14, 22, 14]);
+}
+
+#[test]
 fn test_chunks_mut_count() {
     let v: &mut [i32] = &mut [0, 1, 2, 3, 4, 5];
     let c = v.chunks_mut(3);
@@ -177,6 +189,20 @@ fn test_chunks_mut_last() {
 }
 
 #[test]
+fn test_chunks_mut_zip() {
+    let v1: &mut [i32] = &mut [0, 1, 2, 3, 4];
+    let v2: &[i32] = &[6, 7, 8, 9, 10];
+
+    for (a, b) in v1.chunks_mut(2).zip(v2.chunks(2)) {
+        let sum = b.iter().sum::<i32>();
+        for v in a {
+            *v += sum;
+        }
+    }
+    assert_eq!(v1, [13, 14, 19, 20, 14]);
+}
+
+#[test]
 fn test_windows_count() {
     let v: &[i32] = &[0, 1, 2, 3, 4, 5];
     let c = v.windows(3);
@@ -216,6 +242,19 @@ fn test_windows_last() {
 }
 
 #[test]
+fn test_windows_zip() {
+    let v1: &[i32] = &[0, 1, 2, 3, 4];
+    let v2: &[i32] = &[6, 7, 8, 9, 10];
+
+    let res = v1.windows(2)
+        .zip(v2.windows(2))
+        .map(|(a, b)| a.iter().sum::<i32>() + b.iter().sum::<i32>())
+        .collect::<Vec<_>>();
+
+    assert_eq!(res, [14, 18, 22, 26]);
+}
+
+#[test]
 fn get_range() {
     let v: &[i32] = &[0, 1, 2, 3, 4, 5];
     assert_eq!(v.get(..), Some(&[0, 1, 2, 3, 4, 5][..]));