ndarray/iterators/
chunks.rs

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

ndarray/iterators/
chunks.rs
