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()`](ArrayBase::lanes)
27/// for more information.
28pub struct Lanes<'a, A, D>
29{
30    base: ArrayView<'a, A, D>,
31    inner_len: Ix,
32    inner_stride: Ixs,
33}
34
35impl<'a, A, D: Dimension> Lanes<'a, A, D>
36{
37    pub(crate) fn new<Di>(v: ArrayView<'a, A, Di>, axis: Axis) -> Self
38    where Di: Dimension<Smaller = D>
39    {
40        let ndim = v.ndim();
41        let len;
42        let stride;
43        let iter_v = if ndim == 0 {
44            len = 1;
45            stride = 1;
46            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            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 D: Dimension
81{
82    type Item = <Self::IntoIter as Iterator>::Item;
83    type IntoIter = LanesIter<'a, A, D>;
84    fn into_iter(self) -> Self::IntoIter
85    {
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()`](ArrayBase::lanes_mut)
96/// for more information.
97pub struct LanesMut<'a, A, D>
98{
99    base: ArrayViewMut<'a, A, D>,
100    inner_len: Ix,
101    inner_stride: Ixs,
102}
103
104impl<'a, A, D: Dimension> LanesMut<'a, A, D>
105{
106    pub(crate) fn new<Di>(v: ArrayViewMut<'a, A, Di>, axis: Axis) -> Self
107    where Di: Dimension<Smaller = D>
108    {
109        let ndim = v.ndim();
110        let len;
111        let stride;
112        let iter_v = if ndim == 0 {
113            len = 1;
114            stride = 1;
115            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            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 D: Dimension
132{
133    type Item = <Self::IntoIter as Iterator>::Item;
134    type IntoIter = LanesIterMut<'a, A, D>;
135    fn into_iter(self) -> Self::IntoIter
136    {
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}