ndarray/iterators/
lanes.rs

1use std::marker::PhantomData;
2
3use super::LanesIter;
4use super::LanesIterMut;
5use crate::imp_prelude::*;
6use crate::{Layout, NdProducer};
7
8impl_ndproducer! {
9    ['a, A, D: Dimension]
10    [Clone => 'a, A, D: Clone ]
11    Lanes {
12        base,
13        inner_len,
14        inner_stride,
15    }
16    Lanes<'a, A, D> {
17        type Item = ArrayView<'a, A, Ix1>;
18        type Dim = D;
19
20        unsafe fn item(&self, ptr) {
21            ArrayView::new_(ptr, Ix1(self.inner_len), Ix1(self.inner_stride as Ix))
22        }
23    }
24}
25
26/// See [`.lanes()`](../struct.ArrayBase.html#method.lanes)
27/// for more information.
28pub struct Lanes<'a, A, D> {
29    base: ArrayView<'a, A, D>,
30    inner_len: Ix,
31    inner_stride: Ixs,
32}
33
34impl<'a, A, D: Dimension> Lanes<'a, A, D> {
35    pub(crate) fn new<Di>(v: ArrayView<'a, A, Di>, axis: Axis) -> Self
36    where
37        Di: Dimension<Smaller = D>,
38    {
39        let ndim = v.ndim();
40        let len;
41        let stride;
42        let iter_v;
43        if ndim == 0 {
44            len = 1;
45            stride = 1;
46            iter_v = v.try_remove_axis(Axis(0))
47        } else {
48            let i = axis.index();
49            len = v.dim[i];
50            stride = v.strides[i] as isize;
51            iter_v = v.try_remove_axis(axis)
52        }
53        Lanes {
54            inner_len: len,
55            inner_stride: stride,
56            base: iter_v,
57        }
58    }
59}
60
61impl_ndproducer! {
62    ['a, A, D: Dimension]
63    [Clone =>]
64    LanesMut {
65        base,
66        inner_len,
67        inner_stride,
68    }
69    LanesMut<'a, A, D> {
70        type Item = ArrayViewMut<'a, A, Ix1>;
71        type Dim = D;
72
73        unsafe fn item(&self, ptr) {
74            ArrayViewMut::new_(ptr, Ix1(self.inner_len), Ix1(self.inner_stride as Ix))
75        }
76    }
77}
78
79impl<'a, A, D> IntoIterator for Lanes<'a, A, D>
80where
81    D: Dimension,
82{
83    type Item = <Self::IntoIter as Iterator>::Item;
84    type IntoIter = LanesIter<'a, A, D>;
85    fn into_iter(self) -> Self::IntoIter {
86        LanesIter {
87            iter: self.base.into_base_iter(),
88            inner_len: self.inner_len,
89            inner_stride: self.inner_stride,
90            life: PhantomData,
91        }
92    }
93}
94
95/// See [`.lanes_mut()`](../struct.ArrayBase.html#method.lanes_mut)
96/// for more information.
97pub struct LanesMut<'a, A, D> {
98    base: ArrayViewMut<'a, A, D>,
99    inner_len: Ix,
100    inner_stride: Ixs,
101}
102
103impl<'a, A, D: Dimension> LanesMut<'a, A, D> {
104    pub(crate) fn new<Di>(v: ArrayViewMut<'a, A, Di>, axis: Axis) -> Self
105    where
106        Di: Dimension<Smaller = D>,
107    {
108        let ndim = v.ndim();
109        let len;
110        let stride;
111        let iter_v;
112        if ndim == 0 {
113            len = 1;
114            stride = 1;
115            iter_v = v.try_remove_axis(Axis(0))
116        } else {
117            let i = axis.index();
118            len = v.dim[i];
119            stride = v.strides[i] as isize;
120            iter_v = v.try_remove_axis(axis)
121        }
122        LanesMut {
123            inner_len: len,
124            inner_stride: stride,
125            base: iter_v,
126        }
127    }
128}
129
130impl<'a, A, D> IntoIterator for LanesMut<'a, A, D>
131where
132    D: Dimension,
133{
134    type Item = <Self::IntoIter as Iterator>::Item;
135    type IntoIter = LanesIterMut<'a, A, D>;
136    fn into_iter(self) -> Self::IntoIter {
137        LanesIterMut {
138            iter: self.base.into_base_iter(),
139            inner_len: self.inner_len,
140            inner_stride: self.inner_stride,
141            life: PhantomData,
142        }
143    }
144}