1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110


use imp_prelude::*;
use super::ElementsBase;
use IntoDimension;
use Layout;
use NdProducer;

/// Window producer and iterable
///
/// See [`.windows()`](../struct.ArrayBase.html#method.windows) for more
/// information.
pub struct Windows<'a, A: 'a, D> {
    base: ArrayView<'a, A, D>,
    window: D,
    strides: D,
}

pub fn windows<A, D, E>(a: ArrayView<A, D>, window_size: E) -> Windows<A, D>
	where D: Dimension,
          E: IntoDimension<Dim=D>,
{
    let window = window_size.into_dimension();
    ndassert!(a.ndim() == window.ndim(),
        concat!("Window dimension {} does not match array dimension {} ",
        "(with array of shape {:?})"),
        window.ndim(), a.ndim(), a.shape());
    let mut size = a.dim;
    for (sz, &ws) in size.slice_mut().iter_mut().zip(window.slice())
    {
        if ws == 0 { panic!("window-size must not be zero!"); }
        // cannot use std::cmp::max(0, ..) since arithmetic underflow panics
        *sz = if *sz < ws { 0 } else { *sz - ws + 1 };
    }

    let window_strides = a.strides.clone();

    unsafe {
        Windows {
            base: ArrayView::from_shape_ptr(size.clone().strides(a.strides), a.ptr),
            window: window,
            strides: window_strides,
        }
    }
}

impl_ndproducer! {
    ['a, A, D: Dimension]
    [Clone => 'a, A, D: Clone ]
    Windows {
        base,
        window,
        strides,
    }
    Windows<'a, A, D> {
        type Item = ArrayView<'a, A, D>;
        type Dim = D;

        unsafe fn item(&self, ptr) {
            ArrayView::new_(ptr, self.window.clone(),
                            self.strides.clone())
        }
    }
}

impl<'a, A, D> IntoIterator for Windows<'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,
        }
    }
}

/// Window iterator.
///
/// See [`.windows()`](../struct.ArrayBase.html#method.windows) for more
/// information.
pub struct WindowsIter<'a, A: 'a, D> {
    iter: ElementsBase<'a, A, D>,
    window: D,
    strides: D,
}

impl_iterator!{
    ['a, A, D: Dimension]
    [Clone => 'a, A, D: Clone]
    WindowsIter {
        iter,
        window,
        strides,
    }
    WindowsIter<'a, A, D> {
        type Item = ArrayView<'a, A, D>;

        fn item(&mut self, elt) {
            unsafe {
                ArrayView::new_(
                    elt,
                    self.window.clone(),
                    self.strides.clone())
            }
        }
    }
}