Skip to main content

mermaid/matrix/iter/
cursor.rs

1use super::{IntoIter, Iter, IterMut};
2
3pub struct Cursor<I, const N: usize>
4where
5    I: CursorIterator<N>,
6{
7    iter: I,
8    i: usize,
9    j: usize,
10}
11
12impl<I, const N: usize> Cursor<I, N>
13where
14    I: CursorIterator<N>,
15{
16    fn new(iter: I) -> Self {
17        Self { iter, i: 0, j: 0 }
18    }
19}
20
21impl<I, const N: usize> Iterator for Cursor<I, N>
22where
23    I: CursorIterator<N>,
24{
25    type Item = (usize, usize, I::Item);
26
27    fn next(&mut self) -> Option<Self::Item> {
28        if self.j == N {
29            self.i += 1;
30            self.j = 0;
31        }
32
33        let j = self.j;
34        self.j += 1;
35
36        self.iter.next().map(|x| (self.i, j, x))
37    }
38}
39
40pub trait CursorIterator<const N: usize>: Sized + Iterator {
41    fn cursor(self) -> Cursor<Self, N>;
42}
43
44impl<'a, T, const M: usize, const N: usize> CursorIterator<N> for Iter<'a, T, M, N> {
45    fn cursor(self) -> Cursor<Self, N> {
46        Cursor::new(self)
47    }
48}
49
50impl<'a, T, const M: usize, const N: usize> CursorIterator<N> for IterMut<'a, T, M, N> {
51    fn cursor(self) -> Cursor<Self, N> {
52        Cursor::new(self)
53    }
54}
55
56impl<T, const M: usize, const N: usize> CursorIterator<N> for IntoIter<T, M, N> {
57    fn cursor(self) -> Cursor<Self, N> {
58        Cursor::new(self)
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::super::super::Matrix;
65    use super::*;
66
67    #[test]
68    fn cursor() {
69        let mat = Matrix::new([[1, 2], [3, 4], [5, 6]]);
70
71        let mut iter = mat.iter().cursor();
72        assert_eq!(Some((0, 0, &1)), iter.next());
73        assert_eq!(Some((0, 1, &2)), iter.next());
74        assert_eq!(Some((1, 0, &3)), iter.next());
75        assert_eq!(Some((1, 1, &4)), iter.next());
76        assert_eq!(Some((2, 0, &5)), iter.next());
77        assert_eq!(Some((2, 1, &6)), iter.next());
78        assert_eq!(None, iter.next());
79        assert_eq!(None, iter.next());
80        assert_eq!(None, iter.next());
81    }
82}