From 9a7278258b6b5f21e77a969198a1702762901ea1 Mon Sep 17 00:00:00 2001
From: Josh Stone <cuviper@gmail.com>
Date: Sun, 16 Aug 2020 18:23:26 -0700
Subject: [PATCH] Add const-generic par_array_windows

---
 src/slice/array.rs         | 66 +++++++++++++++++++++++++++++++++++++-
 src/slice/mod.rs           | 17 ++++++++--
 tests/clones.rs            |  1 +
 tests/debug.rs             |  1 +
 tests/producer_split_at.rs |  9 ++++++
 5 files changed, 91 insertions(+), 3 deletions(-)

diff --git a/src/slice/array.rs b/src/slice/array.rs
index be08b149b..bae9ea686 100644
--- a/src/slice/array.rs
+++ b/src/slice/array.rs
@@ -1,7 +1,7 @@
 use crate::iter::plumbing::*;
 use crate::iter::*;
 
-use super::{Iter, IterMut};
+use super::{Iter, IterMut, ParallelSlice};
 
 /// Parallel iterator over immutable non-overlapping chunks of a slice
 #[derive(Debug)]
@@ -170,3 +170,67 @@ impl<'data, T: Send + 'data, const N: usize> IndexedParallelIterator
         self.iter.with_producer(callback)
     }
 }
+
+/// Parallel iterator over immutable overlapping windows of a slice
+#[derive(Debug)]
+pub struct ArrayWindows<'data, T: Sync, const N: usize> {
+    slice: &'data [T],
+}
+
+impl<'data, T: Sync, const N: usize> ArrayWindows<'data, T, N> {
+    pub(super) fn new(slice: &'data [T]) -> Self {
+        ArrayWindows { slice }
+    }
+}
+
+impl<'data, T: Sync, const N: usize> Clone for ArrayWindows<'data, T, N> {
+    fn clone(&self) -> Self {
+        ArrayWindows { ..*self }
+    }
+}
+
+impl<'data, T: Sync + 'data, const N: usize> ParallelIterator for ArrayWindows<'data, T, N> {
+    type Item = &'data [T; N];
+
+    fn drive_unindexed<C>(self, consumer: C) -> C::Result
+    where
+        C: UnindexedConsumer<Self::Item>,
+    {
+        bridge(self, consumer)
+    }
+
+    fn opt_len(&self) -> Option<usize> {
+        Some(self.len())
+    }
+}
+
+impl<'data, T: Sync + 'data, const N: usize> IndexedParallelIterator for ArrayWindows<'data, T, N> {
+    fn drive<C>(self, consumer: C) -> C::Result
+    where
+        C: Consumer<Self::Item>,
+    {
+        bridge(self, consumer)
+    }
+
+    fn len(&self) -> usize {
+        assert!(N >= 1);
+        self.slice.len().saturating_sub(N - 1)
+    }
+
+    fn with_producer<CB>(self, callback: CB) -> CB::Output
+    where
+        CB: ProducerCallback<Self::Item>,
+    {
+        fn array<T, const N: usize>(slice: &[T]) -> &[T; N] {
+            debug_assert_eq!(slice.len(), N);
+            let ptr = slice.as_ptr() as *const [T; N];
+            unsafe { &*ptr }
+        }
+
+        // FIXME: use our own producer and the standard `array_windows`, rust-lang/rust#75027
+        self.slice
+            .par_windows(N)
+            .map(array::<T, N>)
+            .with_producer(callback)
+    }
+}
diff --git a/src/slice/mod.rs b/src/slice/mod.rs
index b64b81f15..ff38c8a05 100644
--- a/src/slice/mod.rs
+++ b/src/slice/mod.rs
@@ -13,8 +13,7 @@ mod rchunks;
 
 mod test;
 
-#[cfg(min_const_generics)]
-pub use self::array::{ArrayChunks, ArrayChunksMut};
+pub use self::array::{ArrayChunks, ArrayChunksMut, ArrayWindows};
 
 use self::mergesort::par_mergesort;
 use self::quicksort::par_quicksort;
@@ -75,6 +74,20 @@ pub trait ParallelSlice<T: Sync> {
         }
     }
 
+    /// Returns a parallel iterator over all contiguous array windows of
+    /// length `N`. The windows overlap.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use rayon::prelude::*;
+    /// let windows: Vec<_> = [1, 2, 3].par_array_windows().collect();
+    /// assert_eq!(vec![&[1, 2], &[2, 3]], windows);
+    /// ```
+    fn par_array_windows<const N: usize>(&self) -> ArrayWindows<'_, T, N> {
+        ArrayWindows::new(self.as_parallel_slice())
+    }
+
     /// Returns a parallel iterator over at most `chunk_size` elements of
     /// `self` at a time. The chunks do not overlap.
     ///
diff --git a/tests/clones.rs b/tests/clones.rs
index f79f0624a..eba309342 100644
--- a/tests/clones.rs
+++ b/tests/clones.rs
@@ -113,6 +113,7 @@ fn clone_vec() {
     check(v.par_rchunks_exact(42));
     check(v.par_array_chunks::<42>());
     check(v.par_windows(42));
+    check(v.par_array_windows::<42>());
     check(v.par_split(|x| x % 3 == 0));
     check(v.into_par_iter());
 }
diff --git a/tests/debug.rs b/tests/debug.rs
index 022f25e90..bf195f326 100644
--- a/tests/debug.rs
+++ b/tests/debug.rs
@@ -130,6 +130,7 @@ fn debug_vec() {
     check(v.par_rchunks_mut(42));
     check(v.par_rchunks_exact_mut(42));
     check(v.par_windows(42));
+    check(v.par_array_windows::<42>());
     check(v.par_split(|x| x % 3 == 0));
     check(v.par_split_mut(|x| x % 3 == 0));
     check(v.par_drain(..));
diff --git a/tests/producer_split_at.rs b/tests/producer_split_at.rs
index 592720d1f..74e6101ad 100644
--- a/tests/producer_split_at.rs
+++ b/tests/producer_split_at.rs
@@ -328,6 +328,15 @@ fn slice_windows() {
     check(&v, || s.par_windows(2));
 }
 
+#[test]
+fn slice_array_windows() {
+    use std::convert::TryInto;
+    let s: Vec<_> = (0..10).collect();
+    // FIXME: use the standard `array_windows`, rust-lang/rust#75027
+    let v: Vec<&[_; 2]> = s.windows(2).map(|s| s.try_into().unwrap()).collect();
+    check(&v, || s.par_array_windows::<2>());
+}
+
 #[test]
 fn vec() {
     let v: Vec<_> = (0..10).collect();