Add const-generic par_array_windows

Author: cuviper

src/slice/array.rs

@@ -3,7 +3,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)]
@@ -172,3 +172,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)
+    }
+}

src/slice/mod.rs

@@ -12,7 +12,7 @@ mod quicksort;
 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;
@@ -71,6 +71,21 @@ 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);
+    /// ```
+    #[cfg(min_const_generics)]
+    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.
     ///

tests/clones.rs

@@ -104,6 +104,7 @@ fn clone_vec() {
     check(v.par_chunks_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());
 }

tests/debug.rs

@@ -126,6 +126,7 @@ fn debug_vec() {
     check(v.par_chunks_exact_mut(42));
     check(v.par_array_chunks_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(..));

tests/producer_split_at.rs

@@ -284,6 +284,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();