ndarray/iterators/
chunks.rs

1use crate::imp_prelude::*;
2use crate::ElementsBase;
3use crate::ElementsBaseMut;
4use crate::IntoDimension;
5use crate::{Layout, NdProducer};
6
7impl_ndproducer! {
8    ['a, A, D: Dimension]
9    [Clone => 'a, A, D: Clone ]
10    ExactChunks {
11        base,
12        chunk,
13        inner_strides,
14    }
15    ExactChunks<'a, A, D> {
16        type Item = ArrayView<'a, A, D>;
17        type Dim = D;
18
19        unsafe fn item(&self, ptr) {
20            ArrayView::new_(ptr, self.chunk.clone(),
21                            self.inner_strides.clone())
22        }
23    }
24}
25
26type BaseProducerRef<'a, A, D> = ArrayView<'a, A, D>;
27type BaseProducerMut<'a, A, D> = ArrayViewMut<'a, A, D>;
28
29/// Exact chunks producer and iterable.
30///
31/// See [`.exact_chunks()`](ArrayBase::exact_chunks) for more
32/// information.
33//#[derive(Debug)]
34pub struct ExactChunks<'a, A, D> {
35    base: BaseProducerRef<'a, A, D>,
36    chunk: D,
37    inner_strides: D,
38}
39
40impl<'a, A, D: Dimension> ExactChunks<'a, A, D> {
41    /// Creates a new exact chunks producer.
42    ///
43    /// **Panics** if any chunk dimension is zero
44    pub(crate) fn new<E>(mut a: ArrayView<'a, A, D>, chunk: E) -> Self
45    where
46        E: IntoDimension<Dim = D>,
47    {
48        let chunk = chunk.into_dimension();
49        ndassert!(
50            a.ndim() == chunk.ndim(),
51            concat!(
52                "Chunk dimension {} does not match array dimension {} ",
53                "(with array of shape {:?})"
54            ),
55            chunk.ndim(),
56            a.ndim(),
57            a.shape()
58        );
59        for i in 0..a.ndim() {
60            a.dim[i] /= chunk[i];
61        }
62        let inner_strides = a.raw_strides();
63        a.strides *= &chunk;
64
65        ExactChunks {
66            base: a,
67            chunk,
68            inner_strides,
69        }
70    }
71}
72
73impl<'a, A, D> IntoIterator for ExactChunks<'a, A, D>
74where
75    D: Dimension,
76    A: 'a,
77{
78    type Item = <Self::IntoIter as Iterator>::Item;
79    type IntoIter = ExactChunksIter<'a, A, D>;
80    fn into_iter(self) -> Self::IntoIter {
81        ExactChunksIter {
82            iter: self.base.into_elements_base(),
83            chunk: self.chunk,
84            inner_strides: self.inner_strides,
85        }
86    }
87}
88
89/// Exact chunks iterator.
90///
91/// See [`.exact_chunks()`](ArrayBase::exact_chunks) for more
92/// information.
93pub struct ExactChunksIter<'a, A, D> {
94    iter: ElementsBase<'a, A, D>,
95    chunk: D,
96    inner_strides: D,
97}
98
99impl_ndproducer! {
100    ['a, A, D: Dimension]
101    [Clone => ]
102    ExactChunksMut {
103        base,
104        chunk,
105        inner_strides,
106    }
107    ExactChunksMut<'a, A, D> {
108        type Item = ArrayViewMut<'a, A, D>;
109        type Dim = D;
110
111        unsafe fn item(&self, ptr) {
112            ArrayViewMut::new_(ptr,
113                               self.chunk.clone(),
114                               self.inner_strides.clone())
115        }
116    }
117}
118
119/// Exact chunks producer and iterable.
120///
121/// See [`.exact_chunks_mut()`](ArrayBase::exact_chunks_mut)
122/// for more information.
123//#[derive(Debug)]
124pub struct ExactChunksMut<'a, A, D> {
125    base: BaseProducerMut<'a, A, D>,
126    chunk: D,
127    inner_strides: D,
128}
129
130impl<'a, A, D: Dimension> ExactChunksMut<'a, A, D> {
131    /// Creates a new exact chunks producer.
132    ///
133    /// **Panics** if any chunk dimension is zero
134    pub(crate) fn new<E>(mut a: ArrayViewMut<'a, A, D>, chunk: E) -> Self
135    where
136        E: IntoDimension<Dim = D>,
137    {
138        let chunk = chunk.into_dimension();
139        ndassert!(
140            a.ndim() == chunk.ndim(),
141            concat!(
142                "Chunk dimension {} does not match array dimension {} ",
143                "(with array of shape {:?})"
144            ),
145            chunk.ndim(),
146            a.ndim(),
147            a.shape()
148        );
149        for i in 0..a.ndim() {
150            a.dim[i] /= chunk[i];
151        }
152        let inner_strides = a.raw_strides();
153        a.strides *= &chunk;
154
155        ExactChunksMut {
156            base: a,
157            chunk,
158            inner_strides,
159        }
160    }
161}
162
163impl<'a, A, D> IntoIterator for ExactChunksMut<'a, A, D>
164where
165    D: Dimension,
166    A: 'a,
167{
168    type Item = <Self::IntoIter as Iterator>::Item;
169    type IntoIter = ExactChunksIterMut<'a, A, D>;
170    fn into_iter(self) -> Self::IntoIter {
171        ExactChunksIterMut {
172            iter: self.base.into_elements_base(),
173            chunk: self.chunk,
174            inner_strides: self.inner_strides,
175        }
176    }
177}
178
179impl_iterator! {
180    ['a, A, D: Dimension]
181    [Clone => 'a, A, D: Clone]
182    ExactChunksIter {
183        iter,
184        chunk,
185        inner_strides,
186    }
187    ExactChunksIter<'a, A, D> {
188        type Item = ArrayView<'a, A, D>;
189
190        fn item(&mut self, elt) {
191            unsafe {
192                ArrayView::new_(
193                    elt,
194                    self.chunk.clone(),
195                    self.inner_strides.clone())
196            }
197        }
198    }
199}
200
201impl_iterator! {
202    ['a, A, D: Dimension]
203    [Clone => ]
204    ExactChunksIterMut {
205        iter,
206        chunk,
207        inner_strides,
208    }
209    ExactChunksIterMut<'a, A, D> {
210        type Item = ArrayViewMut<'a, A, D>;
211
212        fn item(&mut self, elt) {
213            unsafe {
214                ArrayViewMut::new_(
215                    elt,
216                    self.chunk.clone(),
217                    self.inner_strides.clone())
218            }
219        }
220    }
221}
222
223/// Exact chunks iterator.
224///
225/// See [`.exact_chunks_mut()`](ArrayBase::exact_chunks_mut)
226/// for more information.
227pub struct ExactChunksIterMut<'a, A, D> {
228    iter: ElementsBaseMut<'a, A, D>,
229    chunk: D,
230    inner_strides: D,
231}
ndarray/iterators/chunks.rs

ndarray/iterators/
chunks.rs
