vectrix/
iter.rs

1use core::fmt;
2use core::iter::{FusedIterator, Sum};
3use core::marker::PhantomData;
4use core::mem::MaybeUninit;
5use core::ops::{Add, Range};
6use core::ptr;
7
8use crate::new;
9use crate::{Column, Matrix, Row, Zero};
10
11////////////////////////////////////////////////////////////////////////////////
12// Element iteration
13////////////////////////////////////////////////////////////////////////////////
14
15/// An iterator that moves out of a matrix.
16///
17/// This `struct` is created by the `.into_iter()` method on [`Matrix`]
18/// (provided by the [`IntoIterator`] trait).
19///
20/// # Examples
21///
22/// ```
23/// # use vectrix::{matrix, IntoIter};
24/// #
25/// let m = matrix![
26///     1, 3, 5;
27///     2, 4, 6;
28/// ];
29/// let iter: IntoIter<_, 2, 3> = m.into_iter();
30/// ```
31pub struct IntoIter<T, const M: usize, const N: usize> {
32    matrix: Matrix<MaybeUninit<T>, M, N>,
33    alive: Range<usize>,
34}
35
36impl<T, const M: usize, const N: usize> fmt::Debug for IntoIter<T, M, N>
37where
38    T: fmt::Debug,
39{
40    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
41        f.debug_tuple("IntoIter").field(&self.as_slice()).finish()
42    }
43}
44
45impl<T, const M: usize, const N: usize> IntoIter<T, M, N> {
46    /// Creates a new iterator over the given matrix.
47    fn new(matrix: Matrix<T, M, N>) -> Self {
48        Self {
49            // SAFETY: we know that `T` is the same size as `MaybeUninit<T>`.
50            matrix: unsafe { new::transmute_unchecked(matrix) },
51            alive: 0..(M * N),
52        }
53    }
54
55    /// Returns the `i`-th element in the underlying matrix.
56    ///
57    /// # Safety
58    ///
59    /// The caller must make sure that `i` is only fetched once and that `i` is
60    /// in the range `alive.start <= alive.end`.
61    #[inline]
62    unsafe fn get_unchecked(&self, i: usize) -> T {
63        let ptr = unsafe { self.matrix.get_unchecked(i) }.as_ptr();
64        unsafe { ptr::read(ptr) }
65    }
66
67    /// Returns a slice of the remaining initialized elements.
68    #[inline]
69    fn as_slice(&self) -> &[T] {
70        let slice = &self.matrix.as_slice()[self.alive.clone()];
71        let ptr = slice as *const [MaybeUninit<T>] as *const [T];
72        // SAFETY: `alive` keeps track of the elements that are initialized.
73        unsafe { &*ptr }
74    }
75
76    /// Returns a mutable slice of the remaining initialized elements.
77    #[inline]
78    fn as_mut_slice(&mut self) -> &mut [T] {
79        let slice = &mut self.matrix.as_mut_slice()[self.alive.clone()];
80        let ptr = slice as *mut [MaybeUninit<T>] as *mut [T];
81        // SAFETY: `alive` keeps track of the elements that are initialized.
82        unsafe { &mut *ptr }
83    }
84}
85
86impl<T, const M: usize, const N: usize> Iterator for IntoIter<T, M, N> {
87    type Item = T;
88
89    fn next(&mut self) -> Option<Self::Item> {
90        // Get the next index from the front.
91        self.alive.next().map(|i| {
92            // SAFETY: `i` is an index into the former "alive" region of the
93            // array. This is the only time `i` will be yielded .
94            unsafe { self.get_unchecked(i) }
95        })
96    }
97
98    fn size_hint(&self) -> (usize, Option<usize>) {
99        let len = self.alive.len();
100        (len, Some(len))
101    }
102
103    fn count(self) -> usize {
104        self.alive.len()
105    }
106
107    fn last(mut self) -> Option<Self::Item> {
108        self.next_back()
109    }
110}
111
112impl<T, const M: usize, const N: usize> DoubleEndedIterator for IntoIter<T, M, N> {
113    fn next_back(&mut self) -> Option<Self::Item> {
114        // Get the next index from the back.
115        self.alive.next_back().map(|i| {
116            // SAFETY: `i` is an index into the former "alive" region of the
117            // array. This is the only time `i` will be yielded .
118            unsafe { self.get_unchecked(i) }
119        })
120    }
121}
122
123impl<T, const M: usize, const N: usize> ExactSizeIterator for IntoIter<T, M, N> {
124    fn len(&self) -> usize {
125        self.alive.len()
126    }
127}
128
129impl<T, const M: usize, const N: usize> FusedIterator for IntoIter<T, M, N> {}
130
131impl<T, const M: usize, const N: usize> IntoIterator for Matrix<T, M, N> {
132    type Item = T;
133    type IntoIter = IntoIter<T, M, N>;
134
135    #[inline]
136    fn into_iter(self) -> Self::IntoIter {
137        IntoIter::new(self)
138    }
139}
140
141impl<T, const M: usize, const N: usize> Clone for IntoIter<T, M, N>
142where
143    T: Clone,
144{
145    fn clone(&self) -> Self {
146        // Note, we don't really need to match the exact same alive range, so
147        // we can just clone into offset 0 regardless of where `self` is.
148        let mut new = Self {
149            matrix: Matrix::uninit(),
150            alive: 0..0,
151        };
152        // Clone the alive elements only.
153        for (src, dst) in self.as_slice().iter().zip(new.matrix.as_mut_slice()) {
154            // Write a clone into the new array, then update its alive range.
155            // If cloning panics, we'll correctly drop the previous items.
156            *dst = MaybeUninit::new(src.clone());
157            new.alive.end += 1;
158        }
159        new
160    }
161}
162
163impl<T, const M: usize, const N: usize> Drop for IntoIter<T, M, N> {
164    fn drop(&mut self) {
165        let slice = self.as_mut_slice();
166        // SAFETY: `slice` contains only initialized elements.
167        unsafe { ptr::drop_in_place(slice) }
168    }
169}
170
171impl<T, const M: usize, const N: usize> Sum<Matrix<T, M, N>> for Matrix<T, M, N>
172where
173    Self: Add<Output = Self>,
174    T: Copy + Zero,
175{
176    fn sum<I>(iter: I) -> Self
177    where
178        I: Iterator<Item = Self>,
179    {
180        iter.fold(Matrix::zero(), Add::add)
181    }
182}
183
184////////////////////////////////////////////////////////////////////////////////
185// Immutable row iteration
186////////////////////////////////////////////////////////////////////////////////
187
188/// An iterator over the rows in a matrix.
189pub struct IterRows<'a, T, const M: usize, const N: usize> {
190    matrix: &'a Matrix<T, M, N>,
191    alive: Range<usize>,
192}
193
194impl<'a, T, const M: usize, const N: usize> IterRows<'a, T, M, N> {
195    pub(crate) fn new(matrix: &'a Matrix<T, M, N>) -> Self {
196        Self {
197            matrix,
198            alive: 0..M,
199        }
200    }
201}
202
203impl<'a, T, const M: usize, const N: usize> Iterator for IterRows<'a, T, M, N> {
204    type Item = &'a Row<T, M, N>;
205
206    fn next(&mut self) -> Option<Self::Item> {
207        self.alive.next().map(|i| self.matrix.row(i))
208    }
209
210    fn size_hint(&self) -> (usize, Option<usize>) {
211        let len = self.alive.len();
212        (len, Some(len))
213    }
214
215    fn count(self) -> usize {
216        self.alive.len()
217    }
218
219    fn last(mut self) -> Option<Self::Item> {
220        self.next_back()
221    }
222}
223
224impl<T, const M: usize, const N: usize> DoubleEndedIterator for IterRows<'_, T, M, N> {
225    fn next_back(&mut self) -> Option<Self::Item> {
226        self.alive.next_back().map(|i| self.matrix.row(i))
227    }
228}
229
230impl<T, const M: usize, const N: usize> ExactSizeIterator for IterRows<'_, T, M, N> {
231    fn len(&self) -> usize {
232        self.alive.len()
233    }
234}
235
236impl<T, const M: usize, const N: usize> FusedIterator for IterRows<'_, T, M, N> {}
237
238////////////////////////////////////////////////////////////////////////////////
239// Immutable column iteration
240////////////////////////////////////////////////////////////////////////////////
241
242/// An iterator over the columns in a matrix.
243pub struct IterColumns<'a, T, const M: usize, const N: usize> {
244    matrix: &'a Matrix<T, M, N>,
245    alive: Range<usize>,
246}
247
248impl<'a, T, const M: usize, const N: usize> IterColumns<'a, T, M, N> {
249    pub(crate) fn new(matrix: &'a Matrix<T, M, N>) -> Self {
250        Self {
251            matrix,
252            alive: 0..N,
253        }
254    }
255}
256
257impl<'a, T, const M: usize, const N: usize> Iterator for IterColumns<'a, T, M, N> {
258    type Item = &'a Column<T, M, N>;
259
260    fn next(&mut self) -> Option<Self::Item> {
261        self.alive.next().map(|i| self.matrix.column(i))
262    }
263
264    fn size_hint(&self) -> (usize, Option<usize>) {
265        let len = self.alive.len();
266        (len, Some(len))
267    }
268
269    fn count(self) -> usize {
270        self.alive.len()
271    }
272
273    fn last(mut self) -> Option<Self::Item> {
274        self.next_back()
275    }
276}
277
278impl<T, const M: usize, const N: usize> DoubleEndedIterator for IterColumns<'_, T, M, N> {
279    fn next_back(&mut self) -> Option<Self::Item> {
280        self.alive.next_back().map(|i| self.matrix.column(i))
281    }
282}
283
284impl<T, const M: usize, const N: usize> ExactSizeIterator for IterColumns<'_, T, M, N> {
285    fn len(&self) -> usize {
286        self.alive.len()
287    }
288}
289
290impl<T, const M: usize, const N: usize> FusedIterator for IterColumns<'_, T, M, N> {}
291
292////////////////////////////////////////////////////////////////////////////////
293// Mutable row iteration
294////////////////////////////////////////////////////////////////////////////////
295
296/// A mutable iterator over the rows in a matrix.
297pub struct IterRowsMut<'a, T, const M: usize, const N: usize> {
298    // We need to use a raw pointer here because the compiler doesn't
299    // know that we are yielding mutable references to *different* data
300    // each time.
301    matrix: *mut Matrix<T, M, N>,
302    alive: Range<usize>,
303    marker: PhantomData<&'a mut Matrix<T, M, N>>,
304}
305
306impl<'a, T, const M: usize, const N: usize> IterRowsMut<'a, T, M, N> {
307    pub(crate) fn new(matrix: &'a mut Matrix<T, M, N>) -> Self {
308        Self {
309            matrix: matrix as *mut Matrix<T, M, N>,
310            alive: 0..M,
311            marker: PhantomData,
312        }
313    }
314}
315
316impl<'a, T, const M: usize, const N: usize> Iterator for IterRowsMut<'a, T, M, N> {
317    type Item = &'a mut Row<T, M, N>;
318
319    fn next(&mut self) -> Option<Self::Item> {
320        self.alive.next().map(|i| {
321            // SAFETY: we yield a different row each time and `self.matrix`'s
322            // lifetime is asserted by the `PhantomData`.
323            unsafe { (*self.matrix).row_mut(i) }
324        })
325    }
326
327    fn size_hint(&self) -> (usize, Option<usize>) {
328        let len = self.alive.len();
329        (len, Some(len))
330    }
331
332    fn count(self) -> usize {
333        self.alive.len()
334    }
335
336    fn last(mut self) -> Option<Self::Item> {
337        self.next_back()
338    }
339}
340
341impl<T, const M: usize, const N: usize> DoubleEndedIterator for IterRowsMut<'_, T, M, N> {
342    fn next_back(&mut self) -> Option<Self::Item> {
343        self.alive.next_back().map(|i| {
344            // SAFETY: we yield a different row each time and `self.matrix`'s
345            // lifetime is asserted by the `PhantomData`.
346            unsafe { (*self.matrix).row_mut(i) }
347        })
348    }
349}
350
351impl<T, const M: usize, const N: usize> ExactSizeIterator for IterRowsMut<'_, T, M, N> {
352    fn len(&self) -> usize {
353        self.alive.len()
354    }
355}
356
357impl<T, const M: usize, const N: usize> FusedIterator for IterRowsMut<'_, T, M, N> {}
358
359////////////////////////////////////////////////////////////////////////////////
360// Mutable column iteration
361////////////////////////////////////////////////////////////////////////////////
362
363/// A mutable iterator over the columns in a matrix.
364pub struct IterColumnsMut<'a, T, const M: usize, const N: usize> {
365    // We need to use a raw pointer here because the compiler doesn't
366    // know that we are yielding mutable references to *different* data
367    // each time.
368    matrix: *mut Matrix<T, M, N>,
369    alive: Range<usize>,
370    marker: PhantomData<&'a mut Matrix<T, M, N>>,
371}
372
373impl<'a, T, const M: usize, const N: usize> IterColumnsMut<'a, T, M, N> {
374    pub(crate) fn new(matrix: &'a mut Matrix<T, M, N>) -> Self {
375        Self {
376            matrix: matrix as *mut Matrix<T, M, N>,
377            alive: 0..N,
378            marker: PhantomData,
379        }
380    }
381}
382
383impl<'a, T, const M: usize, const N: usize> Iterator for IterColumnsMut<'a, T, M, N> {
384    type Item = &'a mut Column<T, M, N>;
385
386    fn next(&mut self) -> Option<Self::Item> {
387        self.alive.next().map(|i| {
388            // SAFETY: we yield a different column each time and `self.matrix`'s
389            // lifetime is asserted by the `PhantomData`.
390            unsafe { (*self.matrix).column_mut(i) }
391        })
392    }
393
394    fn size_hint(&self) -> (usize, Option<usize>) {
395        let len = self.alive.len();
396        (len, Some(len))
397    }
398
399    fn count(self) -> usize {
400        self.alive.len()
401    }
402
403    fn last(mut self) -> Option<Self::Item> {
404        self.next_back()
405    }
406}
407
408impl<T, const M: usize, const N: usize> DoubleEndedIterator for IterColumnsMut<'_, T, M, N> {
409    fn next_back(&mut self) -> Option<Self::Item> {
410        self.alive.next_back().map(|i| {
411            // SAFETY: we yield a different column each time and `self.matrix`'s
412            // lifetime is asserted by the `PhantomData`.
413            unsafe { (*self.matrix).column_mut(i) }
414        })
415    }
416}
417
418impl<T, const M: usize, const N: usize> ExactSizeIterator for IterColumnsMut<'_, T, M, N> {
419    fn len(&self) -> usize {
420        self.alive.len()
421    }
422}
423
424impl<T, const M: usize, const N: usize> FusedIterator for IterColumnsMut<'_, T, M, N> {}
vectrix/iter.rs

vectrix/
iter.rs
