Reduce code duplication

Specialize the `Windows` struct with a `variant` field to replace the
`AxisWindows` struct

Author: jonasBoss

src/impl_methods.rs

@@ -14,6 +14,8 @@ use rawpointer::PointerExt;
 
 use crate::imp_prelude::*;
 
+use crate::iterators::AxisWindow;
+use crate::iterators::GeneralWindow;
 use crate::{arraytraits, DimMax};
 use crate::argument_traits::AssignElem;
 use crate::dimension;
@@ -34,7 +36,7 @@ use crate::zip::{IntoNdProducer, Zip};
 
 use crate::iter::{
     AxisChunksIter, AxisChunksIterMut, AxisIter, AxisIterMut, ExactChunks, ExactChunksMut,
-    IndexedIter, IndexedIterMut, Iter, IterMut, Lanes, LanesMut, Windows, AxisWindows
+    IndexedIter, IndexedIterMut, Iter, IterMut, Lanes, LanesMut, Windows
 };
 use crate::slice::{MultiSliceArg, SliceArg};
 use crate::stacking::concatenate;
@@ -1419,12 +1421,12 @@ where
     /// that fit into the array's shape.
     ///
     /// This is essentially equivalent to [`.windows_with_stride()`] with unit stride.
-    pub fn windows<E>(&self, window_size: E) -> Windows<'_, A, D>
+    pub fn windows<E>(&self, window_size: E) -> Windows<'_, A, D, GeneralWindow>
     where
         E: IntoDimension<Dim = D>,
         S: Data,
     {
-        Windows::new(self.view(), window_size)
+        Windows::new(self.view(), window_size, GeneralWindow)
     }
 
     /// Return a window producer and iterable.
@@ -1471,12 +1473,12 @@ where
     ///          ┃ a₂₀ ┃ a₂₁ ┃     │     │   │     │     ┃ a₂₂ ┃ a₂₃ ┃
     ///          ┗━━━━━┻━━━━━┹─────┴─────┘   └─────┴─────┺━━━━━┻━━━━━┛
     /// ```
-    pub fn windows_with_stride<E>(&self, window_size: E, stride: E) -> Windows<'_, A, D>
+    pub fn windows_with_stride<E>(&self, window_size: E, stride: E) -> Windows<'_, A, D, GeneralWindow>
     where
         E: IntoDimension<Dim = D>,
         S: Data,
     {
-        Windows::new_with_stride(self.view(), window_size, stride)
+        Windows::new_with_stride(self.view(), window_size, stride, GeneralWindow)
     }
 
     /// Returns a producer which traverses over all windows of a given length along an axis.
@@ -1500,7 +1502,7 @@ where
     ///     assert_eq!(window.shape(), &[4, 3, 2]);
     /// }
     /// ```
-    pub fn axis_windows(&self, axis: Axis, window_size: usize) -> AxisWindows<'_, A, D>
+    pub fn axis_windows(&self, axis: Axis, window_size: usize) -> Windows<'_, A, D, AxisWindow>
     where
         S: Data,
     {
@@ -1517,7 +1519,10 @@ where
             self.shape()
         );
 
-        AxisWindows::new(self.view(), axis, window_size)
+        let mut size = self.raw_dim();
+        size[axis_index] = window_size;
+
+        Windows::new(self.view(), size, AxisWindow{index: axis_index})
     }
 
     // Return (length, stride) for diagonal

src/iterators/iter.rs

@@ -11,5 +11,5 @@ pub use crate::indexes::{Indices, IndicesIter};
 pub use crate::iterators::{
     AxisChunksIter, AxisChunksIterMut, AxisIter, AxisIterMut, ExactChunks, ExactChunksIter,
     ExactChunksIterMut, ExactChunksMut, IndexedIter, IndexedIterMut, Iter, IterMut, Lanes,
-    LanesIter, LanesIterMut, LanesMut, Windows, AxisWindows
+    LanesIter, LanesIterMut, LanesMut, Windows
 };

src/iterators/mod.rs

@@ -26,7 +26,8 @@ use super::{Dimension, Ix, Ixs};
 
 pub use self::chunks::{ExactChunks, ExactChunksIter, ExactChunksIterMut, ExactChunksMut};
 pub use self::lanes::{Lanes, LanesMut};
-pub use self::windows::{Windows, AxisWindows};
+pub use self::windows::Windows;
+pub(crate) use self::windows::{GeneralWindow, AxisWindow};
 pub use self::into_iter::IntoIter;
 
 use std::slice::{self, Iter as SliceIter, IterMut as SliceIterMut};

src/iterators/windows.rs

@@ -5,18 +5,24 @@ use crate::Layout;
 use crate::NdProducer;
 use crate::Slice;
 
+#[derive(Clone)]
+pub struct GeneralWindow;
+#[derive(Clone)]
+pub struct AxisWindow{pub(crate) index: usize}
+
 /// Window producer and iterable
 ///
 /// See [`.windows()`](ArrayBase::windows) for more
 /// information.
-pub struct Windows<'a, A, D> {
+pub struct Windows<'a, A, D, V> {
     base: ArrayView<'a, A, D>,
     window: D,
     strides: D,
+    variant: V,
 }
 
-impl<'a, A, D: Dimension> Windows<'a, A, D> {
-    pub(crate) fn new<E>(a: ArrayView<'a, A, D>, window_size: E) -> Self
+impl<'a, A, D: Dimension, V> Windows<'a, A, D, V> {
+    pub(crate) fn new<E>(a: ArrayView<'a, A, D>, window_size: E, variant: V) -> Self
     where
         E: IntoDimension<Dim = D>,
     {
@@ -26,10 +32,15 @@ impl<'a, A, D: Dimension> Windows<'a, A, D> {
         let mut unit_stride = D::zeros(ndim);
         unit_stride.slice_mut().fill(1);
 
-        Windows::new_with_stride(a, window, unit_stride)
+        Windows::new_with_stride(a, window, unit_stride, variant)
     }
 
-    pub(crate) fn new_with_stride<E>(a: ArrayView<'a, A, D>, window_size: E, axis_strides: E) -> Self
+    pub(crate) fn new_with_stride<E>(
+        a: ArrayView<'a, A, D>,
+        window_size: E,
+        axis_strides: E,
+        variant: V,
+    ) -> Self
     where
         E: IntoDimension<Dim = D>,
     {
@@ -77,6 +88,7 @@ impl<'a, A, D: Dimension> Windows<'a, A, D> {
             base,
             window,
             strides: window_strides,
+            variant,
         }
     }
 }
@@ -88,8 +100,9 @@ impl_ndproducer! {
         base,
         window,
         strides,
+        variant,
     }
-    Windows<'a, A, D> {
+    Windows<'a, A, D, GeneralWindow> {
         type Item = ArrayView<'a, A, D>;
         type Dim = D;
 
@@ -100,7 +113,7 @@ impl_ndproducer! {
     }
 }
 
-impl<'a, A, D> IntoIterator for Windows<'a, A, D>
+impl<'a, A, D, V> IntoIterator for Windows<'a, A, D, V>
 where
     D: Dimension,
     A: 'a,
@@ -148,55 +161,14 @@ impl_iterator! {
     }
 }
 
-/// Window producer and iterable
-///
-/// See [`.axis_windows()`](ArrayBase::axis_windows) for more
-/// information.
-pub struct AxisWindows<'a, A, D>{
-    base: ArrayView<'a, A, D>,
-    axis_idx: usize,
-    window: D,
-    strides: D,
-}
-
-impl<'a, A, D: Dimension> AxisWindows<'a, A, D> {
-    pub(crate) fn new(a: ArrayView<'a, A, D>, axis: Axis, window_size: usize) -> Self
-    {   
-        let strides = a.strides.clone();
-        let mut base = a;
-        let axis_idx = axis.index();
-        let mut window = base.raw_dim();
-        window[axis_idx] = window_size;
-
-        base.slice_each_axis_inplace(|ax_desc| {
-            let len = ax_desc.len;
-            let wsz = window[ax_desc.axis.index()];
-
-            if len < wsz {
-                Slice::new(0, Some(0), 1)
-            } else {
-                Slice::new(0, Some((len - wsz + 1) as isize), 1)
-            }
-        });
-
-        AxisWindows {
-            base,
-            axis_idx,
-            window,
-            strides,
-        }
-    }
-}
-
-
-impl<'a, A, D: Dimension> NdProducer for AxisWindows<'a, A, D> {
+impl<'a, A, D: Dimension> NdProducer for Windows<'a, A, D, AxisWindow> {
     type Item = ArrayView<'a, A, D>;
     type Dim = Ix1;
     type Ptr = *mut A;
     type Stride = isize;
 
     fn raw_dim(&self) -> Ix1 {
-        Ix1(self.base.raw_dim()[self.axis_idx])
+        Ix1(self.base.raw_dim()[self.variant.index])
     }
 
     fn layout(&self) -> Layout {
@@ -212,54 +184,38 @@ impl<'a, A, D: Dimension> NdProducer for AxisWindows<'a, A, D> {
     }
 
     unsafe fn as_ref(&self, ptr: *mut A) -> Self::Item {
-        ArrayView::new_(ptr, self.window.clone(),
-            self.strides.clone())
+        ArrayView::new_(ptr, self.window.clone(), self.strides.clone())
     }
 
     unsafe fn uget_ptr(&self, i: &Self::Dim) -> *mut A {
         let mut d = D::zeros(self.base.ndim());
-        d[self.axis_idx] = i[0];
+        d[self.variant.index] = i[0];
         self.base.uget_ptr(&d)
     }
 
     fn stride_of(&self, axis: Axis) -> isize {
         assert_eq!(axis, Axis(0));
-        self.base.stride_of(Axis(self.axis_idx))
+        self.base.stride_of(Axis(self.variant.index))
     }
 
     fn split_at(self, axis: Axis, index: usize) -> (Self, Self) {
         assert_eq!(axis, Axis(0));
-        let (a, b) = self.base.split_at(Axis(self.axis_idx), index);
-        (AxisWindows {
-            base: a,
-            axis_idx: self.axis_idx,
-            window: self.window.clone(),
-            strides: self.strides.clone()
-
-        },
-        AxisWindows {
-            base: b,
-            axis_idx: self.axis_idx,
-            window: self.window,
-            strides: self.strides,
-        })
-    }
-
-    private_impl!{}
-}
-
-impl<'a, A, D> IntoIterator for AxisWindows<'a, A, D>
-where
-    D: Dimension,
-    A: 'a,
-{
-    type Item = <Self::IntoIter as Iterator>::Item;
-    type IntoIter = WindowsIter<'a, A, D>;
-    fn into_iter(self) -> Self::IntoIter {
-        WindowsIter {
-            iter: self.base.into_elements_base(),
-            window: self.window,
-            strides: self.strides,
-        }
-    }
+        let (a, b) = self.base.split_at(Axis(self.variant.index), index);
+        (
+            Windows {
+                base: a,
+                window: self.window.clone(),
+                strides: self.strides.clone(),
+                variant: self.variant.clone(),
+            },
+            Windows {
+                base: b,
+                window: self.window,
+                strides: self.strides,
+                variant: self.variant,
+            },
+        )
+    }
+
+    private_impl! {}
 }