nalgebra/base/
iter.rs

Help
1//! Matrix iterators.
2
3use std::iter::FusedIterator;
4use std::marker::PhantomData;
5use std::mem;
6
7use crate::base::dimension::{Dim, U1};
8use crate::base::storage::{RawStorage, RawStorageMut};
9use crate::base::{Matrix, MatrixSlice, MatrixSliceMut, Scalar};
10
11macro_rules! iterator {
12    (struct $Name:ident for $Storage:ident.$ptr: ident -> $Ptr:ty, $Ref:ty, $SRef: ty) => {
13        /// An iterator through a dense matrix with arbitrary strides matrix.
14        #[derive(Debug)]
15        pub struct $Name<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> {
16            ptr: $Ptr,
17            inner_ptr: $Ptr,
18            inner_end: $Ptr,
19            size: usize, // We can't use an end pointer here because a stride might be zero.
20            strides: (S::RStride, S::CStride),
21            _phantoms: PhantomData<($Ref, R, C, S)>,
22        }
23
24        // TODO: we need to specialize for the case where the matrix storage is owned (in which
25        // case the iterator is trivial because it does not have any stride).
26        impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> $Name<'a, T, R, C, S> {
27            /// Creates a new iterator for the given matrix storage.
28            pub fn new(storage: $SRef) -> $Name<'a, T, R, C, S> {
29                let shape = storage.shape();
30                let strides = storage.strides();
31                let inner_offset = shape.0.value() * strides.0.value();
32                let size = shape.0.value() * shape.1.value();
33                let ptr = storage.$ptr();
34
35                // If we have a size of 0, 'ptr' must be
36                // dangling. Howver, 'inner_offset' might
37                // not be zero if only one dimension is zero, so
38                // we don't want to call 'offset'.
39                // This pointer will never actually get used
40                // if our size is '0', so it's fine to use
41                // 'ptr' for both the start and end.
42                let inner_end = if size == 0 {
43                    ptr
44                } else {
45                    // Safety:
46                    // If 'size' is non-zero, we know that 'ptr'
47                    // is not dangling, and 'inner_offset' must lie
48                    // within the allocation
49                    unsafe { ptr.add(inner_offset) }
50                };
51
52                $Name {
53                    ptr,
54                    inner_ptr: ptr,
55                    inner_end,
56                    size: shape.0.value() * shape.1.value(),
57                    strides,
58                    _phantoms: PhantomData,
59                }
60            }
61        }
62
63        impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> Iterator for $Name<'a, T, R, C, S> {
64            type Item = $Ref;
65
66            #[inline]
67            fn next(&mut self) -> Option<$Ref> {
68                unsafe {
69                    if self.size == 0 {
70                        None
71                    } else {
72                        self.size -= 1;
73
74                        // Jump to the next outer dimension if needed.
75                        if self.ptr == self.inner_end {
76                            let stride = self.strides.1.value() as isize;
77                            // This might go past the end of the allocation,
78                            // depending on the value of 'size'. We use
79                            // `wrapping_offset` to avoid UB
80                            self.inner_end = self.ptr.wrapping_offset(stride);
81                            // This will always be in bounds, since
82                            // we're going to dereference it
83                            self.ptr = self.inner_ptr.offset(stride);
84                            self.inner_ptr = self.ptr;
85                        }
86
87                        // Go to the next element.
88                        let old = self.ptr;
89
90                        // Don't offset `self.ptr` for the last element,
91                        // as this will be out of bounds. Iteration is done
92                        // at this point (the next call to `next` will return `None`)
93                        // so this is not observable.
94                        if self.size != 0 {
95                            let stride = self.strides.0.value();
96                            self.ptr = self.ptr.add(stride);
97                        }
98
99                        // We want either `& *last` or `&mut *last` here, depending
100                        // on the mutability of `$Ref`.
101                        #[allow(clippy::transmute_ptr_to_ref)]
102                        Some(mem::transmute(old))
103                    }
104                }
105            }
106
107            #[inline]
108            fn size_hint(&self) -> (usize, Option<usize>) {
109                (self.size, Some(self.size))
110            }
111
112            #[inline]
113            fn count(self) -> usize {
114                self.size_hint().0
115            }
116        }
117
118        impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> DoubleEndedIterator
119            for $Name<'a, T, R, C, S>
120        {
121            #[inline]
122            fn next_back(&mut self) -> Option<$Ref> {
123                unsafe {
124                    if self.size == 0 {
125                        None
126                    } else {
127                        // Pre-decrement `size` such that it now counts to the
128                        // element we want to return.
129                        self.size -= 1;
130
131                        // Fetch strides
132                        let inner_stride = self.strides.0.value();
133                        let outer_stride = self.strides.1.value();
134
135                        // Compute number of rows
136                        // Division should be exact
137                        let inner_raw_size = self.inner_end.offset_from(self.inner_ptr) as usize;
138                        let inner_size = inner_raw_size / inner_stride;
139
140                        // Compute rows and cols remaining
141                        let outer_remaining = self.size / inner_size;
142                        let inner_remaining = self.size % inner_size;
143
144                        // Compute pointer to last element
145                        let last = self
146                            .ptr
147                            .add((outer_remaining * outer_stride + inner_remaining * inner_stride));
148
149                        // We want either `& *last` or `&mut *last` here, depending
150                        // on the mutability of `$Ref`.
151                        #[allow(clippy::transmute_ptr_to_ref)]
152                        Some(mem::transmute(last))
153                    }
154                }
155            }
156        }
157
158        impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> ExactSizeIterator
159            for $Name<'a, T, R, C, S>
160        {
161            #[inline]
162            fn len(&self) -> usize {
163                self.size
164            }
165        }
166
167        impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> FusedIterator
168            for $Name<'a, T, R, C, S>
169        {
170        }
171    };
172}
173
174iterator!(struct MatrixIter for RawStorage.ptr -> *const T, &'a T, &'a S);
175iterator!(struct MatrixIterMut for RawStorageMut.ptr_mut -> *mut T, &'a mut T, &'a mut S);
176
177/*
178 *
179 * Row iterators.
180 *
181 */
182#[derive(Clone, Debug)]
183/// An iterator through the rows of a matrix.
184pub struct RowIter<'a, T, R: Dim, C: Dim, S: RawStorage<T, R, C>> {
185    mat: &'a Matrix<T, R, C, S>,
186    curr: usize,
187}
188
189impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> RowIter<'a, T, R, C, S> {
190    pub(crate) fn new(mat: &'a Matrix<T, R, C, S>) -> Self {
191        RowIter { mat, curr: 0 }
192    }
193}
194
195impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> Iterator for RowIter<'a, T, R, C, S> {
196    type Item = MatrixSlice<'a, T, U1, C, S::RStride, S::CStride>;
197
198    #[inline]
199    fn next(&mut self) -> Option<Self::Item> {
200        if self.curr < self.mat.nrows() {
201            let res = self.mat.row(self.curr);
202            self.curr += 1;
203            Some(res)
204        } else {
205            None
206        }
207    }
208
209    #[inline]
210    fn size_hint(&self) -> (usize, Option<usize>) {
211        (
212            self.mat.nrows() - self.curr,
213            Some(self.mat.nrows() - self.curr),
214        )
215    }
216
217    #[inline]
218    fn count(self) -> usize {
219        self.mat.nrows() - self.curr
220    }
221}
222
223impl<'a, T: Scalar, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> ExactSizeIterator
224    for RowIter<'a, T, R, C, S>
225{
226    #[inline]
227    fn len(&self) -> usize {
228        self.mat.nrows() - self.curr
229    }
230}
231
232/// An iterator through the mutable rows of a matrix.
233#[derive(Debug)]
234pub struct RowIterMut<'a, T, R: Dim, C: Dim, S: RawStorageMut<T, R, C>> {
235    mat: *mut Matrix<T, R, C, S>,
236    curr: usize,
237    phantom: PhantomData<&'a mut Matrix<T, R, C, S>>,
238}
239
240impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> RowIterMut<'a, T, R, C, S> {
241    pub(crate) fn new(mat: &'a mut Matrix<T, R, C, S>) -> Self {
242        RowIterMut {
243            mat,
244            curr: 0,
245            phantom: PhantomData,
246        }
247    }
248
249    fn nrows(&self) -> usize {
250        unsafe { (*self.mat).nrows() }
251    }
252}
253
254impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> Iterator
255    for RowIterMut<'a, T, R, C, S>
256{
257    type Item = MatrixSliceMut<'a, T, U1, C, S::RStride, S::CStride>;
258
259    #[inline]
260    fn next(&mut self) -> Option<Self::Item> {
261        if self.curr < self.nrows() {
262            let res = unsafe { (*self.mat).row_mut(self.curr) };
263            self.curr += 1;
264            Some(res)
265        } else {
266            None
267        }
268    }
269
270    #[inline]
271    fn size_hint(&self) -> (usize, Option<usize>) {
272        (self.nrows() - self.curr, Some(self.nrows() - self.curr))
273    }
274
275    #[inline]
276    fn count(self) -> usize {
277        self.nrows() - self.curr
278    }
279}
280
281impl<'a, T: Scalar, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> ExactSizeIterator
282    for RowIterMut<'a, T, R, C, S>
283{
284    #[inline]
285    fn len(&self) -> usize {
286        self.nrows() - self.curr
287    }
288}
289
290/*
291 *
292 * Column iterators.
293 *
294 */
295#[derive(Clone, Debug)]
296/// An iterator through the columns of a matrix.
297pub struct ColumnIter<'a, T, R: Dim, C: Dim, S: RawStorage<T, R, C>> {
298    mat: &'a Matrix<T, R, C, S>,
299    curr: usize,
300}
301
302impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> ColumnIter<'a, T, R, C, S> {
303    pub(crate) fn new(mat: &'a Matrix<T, R, C, S>) -> Self {
304        ColumnIter { mat, curr: 0 }
305    }
306}
307
308impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> Iterator for ColumnIter<'a, T, R, C, S> {
309    type Item = MatrixSlice<'a, T, R, U1, S::RStride, S::CStride>;
310
311    #[inline]
312    fn next(&mut self) -> Option<Self::Item> {
313        if self.curr < self.mat.ncols() {
314            let res = self.mat.column(self.curr);
315            self.curr += 1;
316            Some(res)
317        } else {
318            None
319        }
320    }
321
322    #[inline]
323    fn size_hint(&self) -> (usize, Option<usize>) {
324        (
325            self.mat.ncols() - self.curr,
326            Some(self.mat.ncols() - self.curr),
327        )
328    }
329
330    #[inline]
331    fn count(self) -> usize {
332        self.mat.ncols() - self.curr
333    }
334}
335
336impl<'a, T: Scalar, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> ExactSizeIterator
337    for ColumnIter<'a, T, R, C, S>
338{
339    #[inline]
340    fn len(&self) -> usize {
341        self.mat.ncols() - self.curr
342    }
343}
344
345/// An iterator through the mutable columns of a matrix.
346#[derive(Debug)]
347pub struct ColumnIterMut<'a, T, R: Dim, C: Dim, S: RawStorageMut<T, R, C>> {
348    mat: *mut Matrix<T, R, C, S>,
349    curr: usize,
350    phantom: PhantomData<&'a mut Matrix<T, R, C, S>>,
351}
352
353impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> ColumnIterMut<'a, T, R, C, S> {
354    pub(crate) fn new(mat: &'a mut Matrix<T, R, C, S>) -> Self {
355        ColumnIterMut {
356            mat,
357            curr: 0,
358            phantom: PhantomData,
359        }
360    }
361
362    fn ncols(&self) -> usize {
363        unsafe { (*self.mat).ncols() }
364    }
365}
366
367impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> Iterator
368    for ColumnIterMut<'a, T, R, C, S>
369{
370    type Item = MatrixSliceMut<'a, T, R, U1, S::RStride, S::CStride>;
371
372    #[inline]
373    fn next(&mut self) -> Option<Self::Item> {
374        if self.curr < self.ncols() {
375            let res = unsafe { (*self.mat).column_mut(self.curr) };
376            self.curr += 1;
377            Some(res)
378        } else {
379            None
380        }
381    }
382
383    #[inline]
384    fn size_hint(&self) -> (usize, Option<usize>) {
385        (self.ncols() - self.curr, Some(self.ncols() - self.curr))
386    }
387
388    #[inline]
389    fn count(self) -> usize {
390        self.ncols() - self.curr
391    }
392}
393
394impl<'a, T: Scalar, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> ExactSizeIterator
395    for ColumnIterMut<'a, T, R, C, S>
396{
397    #[inline]
398    fn len(&self) -> usize {
399        self.ncols() - self.curr
400    }
401}
nalgebra/base/iter.rs

nalgebra/base/
iter.rs
