create AxisWindows struct and implement it

Author: jonasBoss

src/impl_methods.rs

@@ -34,7 +34,7 @@ use crate::zip::{IntoNdProducer, Zip};
 
 use crate::iter::{
     AxisChunksIter, AxisChunksIterMut, AxisIter, AxisIterMut, ExactChunks, ExactChunksMut,
-    IndexedIter, IndexedIterMut, Iter, IterMut, Lanes, LanesMut, Windows,
+    IndexedIter, IndexedIterMut, Iter, IterMut, Lanes, LanesMut, Windows, AxisWindows
 };
 use crate::slice::{MultiSliceArg, SliceArg};
 use crate::stacking::concatenate;
@@ -1500,7 +1500,7 @@ where
     ///     assert_eq!(window.shape(), &[4, 3, 2]);
     /// }
     /// ```
-    pub fn axis_windows(&self, axis: Axis, window_size: usize) -> Windows<'_, A, D>
+    pub fn axis_windows(&self, axis: Axis, window_size: usize) -> AxisWindows<'_, A, D>
     where
         S: Data,
     {
@@ -1520,7 +1520,7 @@ where
         let mut size = self.raw_dim();
         size[axis_index] = window_size;
 
-        Windows::new(self.view(), size)
+        AxisWindows::new(self.view(), axis, window_size)
     }
 
     // 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,
+    LanesIter, LanesIterMut, LanesMut, Windows, AxisWindows
 };

src/iterators/mod.rs

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

src/iterators/windows.rs

@@ -147,3 +147,121 @@ 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>,
+    window_size: usize,
+    axis_idx: usize,
+}
+
+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 mut base = a;
+        let len = base.raw_dim()[axis.index()];
+        let indices = if len < window_size {
+            Slice::new(0, Some(0), 1)
+        } else {
+            Slice::new(0, Some((len - window_size + 1) as isize), 1)
+        };
+        base.slice_axis_inplace(axis, indices);
+
+        AxisWindows {
+            base,
+            window_size,
+            axis_idx: axis.index(),
+        }
+    }
+
+    fn window(&self) -> D{
+        let mut window = self.base.raw_dim();
+        window[self.axis_idx] = self.window_size;
+        window
+    }
+
+    fn strides_(&self) -> D{
+        let mut strides = D::zeros(self.base.ndim());
+        strides.slice_mut().fill(1);
+        strides
+    }
+}
+
+
+impl<'a, A, D: Dimension> NdProducer for AxisWindows<'a, A, D> {
+    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])
+    }
+
+    fn layout(&self) -> Layout {
+        self.base.layout()
+    }
+
+    fn as_ptr(&self) -> *mut A {
+        self.base.as_ptr() as *mut _
+    }
+
+    fn contiguous_stride(&self) -> isize {
+        self.base.contiguous_stride()
+    }
+
+    unsafe fn as_ref(&self, ptr: *mut A) -> Self::Item {
+        ArrayView::new_(ptr, self.window(),
+        self.strides_())
+    }
+
+    unsafe fn uget_ptr(&self, i: &Self::Dim) -> *mut A {
+        let mut d = D::zeros(self.base.ndim());
+        d[self.axis_idx] = 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))
+    }
+
+    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,
+            window_size: self.window_size,
+            axis_idx: self.axis_idx,
+
+        },
+        AxisWindows {
+            base: b,
+            window_size: self.window_size,
+            axis_idx: self.axis_idx,
+        })
+    }
+
+    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 {
+        let window =  self.window();
+        let strides = self.strides_();
+        WindowsIter {
+            iter: self.base.into_elements_base(),
+            window,
+            strides,
+        }
+    }
+}