vecmat/matrix/base/
iter.rs

1use crate::{Matrix, Vector};
2use core::slice;
3
4/// Iterator that iterates over flattened sequence of `IntoIter`s.
5pub struct FlatIter<I, IT, II>
6where
7    I: Iterator<Item = IT>,
8    IT: IntoIterator<IntoIter = II, Item = II::Item>,
9    II: Iterator,
10{
11    iter: I,
12    subiter: II,
13}
14
15impl<I, IT, II> FlatIter<I, IT, II>
16where
17    I: Iterator<Item = IT>,
18    IT: IntoIterator<IntoIter = II, Item = II::Item>,
19    II: Iterator,
20{
21    /// Create `FlatIter` from sequence of `IntoIter`s.
22    pub fn new(mut iter: I) -> Option<Self> {
23        iter.next().map(|a| Self {
24            iter,
25            subiter: a.into_iter(),
26        })
27    }
28}
29
30impl<I, IT, II> Iterator for FlatIter<I, IT, II>
31where
32    I: Iterator<Item = IT>,
33    IT: IntoIterator<IntoIter = II, Item = II::Item>,
34    II: Iterator,
35{
36    type Item = II::Item;
37    fn next(&mut self) -> Option<Self::Item> {
38        self.subiter.next().or_else(|| {
39            self.iter.next().and_then(|a| {
40                self.subiter = a.into_iter();
41                self.subiter.next()
42            })
43        })
44    }
45}
46
47impl<T, const M: usize, const N: usize> Matrix<T, M, N> {
48    /// Returns iterator over matrix element refrences.
49    pub fn iter(&self) -> impl Iterator<Item = &T> {
50        FlatIter::new(self.data.iter()).unwrap()
51    }
52    /// Returns iterator over matrix element mutable refrences.
53    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> {
54        FlatIter::new(self.data.iter_mut()).unwrap()
55    }
56}
57
58impl<T, const M: usize, const N: usize> IntoIterator for Matrix<T, M, N> {
59    type Item = T;
60
61    #[allow(clippy::type_complexity)]
62    type IntoIter = FlatIter<
63        <Vector<Vector<T, N>, M> as IntoIterator>::IntoIter,
64        Vector<T, N>,
65        <Vector<T, N> as IntoIterator>::IntoIter,
66    >;
67
68    fn into_iter(self) -> Self::IntoIter {
69        Self::IntoIter::new(self.data.into_iter()).unwrap()
70    }
71}
72impl<'a, T, const M: usize, const N: usize> IntoIterator for &'a Matrix<T, M, N> {
73    type Item = &'a T;
74    type IntoIter = FlatIter<slice::Iter<'a, Vector<T, N>>, &'a Vector<T, N>, slice::Iter<'a, T>>;
75    fn into_iter(self) -> Self::IntoIter {
76        Self::IntoIter::new(self.data.iter()).unwrap()
77    }
78}
79impl<'a, T, const M: usize, const N: usize> IntoIterator for &'a mut Matrix<T, M, N> {
80    type Item = &'a mut T;
81    type IntoIter =
82        FlatIter<slice::IterMut<'a, Vector<T, N>>, &'a mut Vector<T, N>, slice::IterMut<'a, T>>;
83    fn into_iter(self) -> Self::IntoIter {
84        Self::IntoIter::new(self.data.iter_mut()).unwrap()
85    }
86}
87
88impl<const M: usize, const N: usize> Matrix<(usize, usize), M, N> {
89    /// Create matrix which elements are tuples (j, i) where j and i are coordinates of the matrix cell.
90    pub fn indices() -> Self {
91        Self::try_from_iter((0..(M * N)).map(|x| (x / N, x % N))).unwrap()
92    }
93}
94
95impl<T, const M: usize, const N: usize> Matrix<T, M, N> {
96    /// Call closure for each element of the matrix passing it by value.
97    pub fn for_each<F: FnMut(T)>(self, mut f: F) {
98        self.data.for_each(|a| a.for_each(&mut f))
99    }
100    /// Map matrix elements.
101    pub fn map<U, F: FnMut(T) -> U>(self, mut f: F) -> Matrix<U, M, N> {
102        self.data.map(|a| a.map(&mut f)).into()
103    }
104    /// Zip two matrices into one.
105    pub fn zip<U>(self, other: Matrix<U, M, N>) -> Matrix<(T, U), M, N> {
106        self.data.zip(other.data).map(|(a, b)| a.zip(b)).into()
107    }
108}
109impl<T, U, const M: usize, const N: usize> Matrix<(T, U), M, N> {
110    /// Unzip matrix of tuples into two matrices.
111    pub fn unzip(self) -> (Matrix<T, M, N>, Matrix<U, M, N>) {
112        let z = self.data.map(|a| a.unzip()).unzip();
113        (z.0.into(), z.1.into())
114    }
115}
116impl<T, const M: usize, const N: usize> Matrix<T, M, N> {
117    pub fn fold<S, F: Fn(S, T) -> S>(self, s: S, mut f: F) -> S {
118        self.data.fold(s, |t, a| a.fold(t, &mut f))
119    }
120    pub fn fold_first<F: FnMut(T, T) -> T>(self, mut f: F) -> T {
121        let mut iter = self.data.into_iter();
122        let s = iter.next().unwrap().fold_first(&mut f);
123        iter.fold(s, |t, a| a.fold(t, &mut f))
124    }
125    pub fn scan<S, U, F: FnMut(&mut S, T) -> U>(self, mut s: S, mut f: F) -> Matrix<U, M, N> {
126        self.data
127            .scan(&mut s, |r, a| a.scan(r, |r, x| f(*r, x)))
128            .into()
129    }
130}