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    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 = if ndim == 0 {
43            len = 1;
44            stride = 1;
45            v.try_remove_axis(Axis(0))
46        } else {
47            let i = axis.index();
48            len = v.dim[i];
49            stride = v.strides[i] as isize;
50            v.try_remove_axis(axis)
51        };
52        Lanes {
53            inner_len: len,
54            inner_stride: stride,
55            base: iter_v,
56        }
57    }
58}
59
60impl_ndproducer! {
61    ['a, A, D: Dimension]
62    [Clone =>]
63    LanesMut {
64        base,
65        inner_len,
66        inner_stride,
67    }
68    LanesMut<'a, A, D> {
69        type Item = ArrayViewMut<'a, A, Ix1>;
70        type Dim = D;
71
72        unsafe fn item(&self, ptr) {
73            ArrayViewMut::new_(ptr, Ix1(self.inner_len), Ix1(self.inner_stride as Ix))
74        }
75    }
76}
77
78impl<'a, A, D> IntoIterator for Lanes<'a, A, D>
79where
80    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        LanesIter {
86            iter: self.base.into_base_iter(),
87            inner_len: self.inner_len,
88            inner_stride: self.inner_stride,
89            life: PhantomData,
90        }
91    }
92}
93
94/// See [`.lanes_mut()`](ArrayBase::lanes_mut)
95/// for more information.
96pub struct LanesMut<'a, A, D> {
97    base: ArrayViewMut<'a, A, D>,
98    inner_len: Ix,
99    inner_stride: Ixs,
100}
101
102impl<'a, A, D: Dimension> LanesMut<'a, A, D> {
103    pub(crate) fn new<Di>(v: ArrayViewMut<'a, A, Di>, axis: Axis) -> Self
104    where
105        Di: Dimension<Smaller = D>,
106    {
107        let ndim = v.ndim();
108        let len;
109        let stride;
110        let iter_v = if ndim == 0 {
111            len = 1;
112            stride = 1;
113            v.try_remove_axis(Axis(0))
114        } else {
115            let i = axis.index();
116            len = v.dim[i];
117            stride = v.strides[i] as isize;
118            v.try_remove_axis(axis)
119        };
120        LanesMut {
121            inner_len: len,
122            inner_stride: stride,
123            base: iter_v,
124        }
125    }
126}
127
128impl<'a, A, D> IntoIterator for LanesMut<'a, A, D>
129where
130    D: Dimension,
131{
132    type Item = <Self::IntoIter as Iterator>::Item;
133    type IntoIter = LanesIterMut<'a, A, D>;
134    fn into_iter(self) -> Self::IntoIter {
135        LanesIterMut {
136            iter: self.base.into_base_iter(),
137            inner_len: self.inner_len,
138            inner_stride: self.inner_stride,
139            life: PhantomData,
140        }
141    }
142}