vec_vec/
lending_iter_mut.rs

1use lending_iterator::prelude::*;
2use nougat::gat;
3
4/// A [lending iterator] over the mutable references to the elements of a `Vec<Vec<T>>`.
5///
6/// [lending iterator]: https://blog.rust-lang.org/2022/10/28/gats-stabilization.html#what-are-gats
7pub struct LendingIterMut<'a, T> {
8    container: &'a mut Vec<Vec<T>>,
9    outer_idx: usize,
10    inner_idx: usize,
11    is_finished: bool,
12}
13
14impl<'a, T> LendingIterMut<'a, T> {
15    pub(crate) fn new(container: &'a mut Vec<Vec<T>>) -> Self {
16        let mut oi = 0;
17        let mut is_finished = false;
18        loop {
19            if oi >= container.len() {
20                is_finished = true;
21                break;
22            };
23            if !container[oi].is_empty() {
24                break;
25            };
26            oi += 1;
27        }
28        Self {
29            container,
30            outer_idx: oi,
31            inner_idx: 0,
32            is_finished,
33        }
34    }
35
36    /// Returns a shared reference to the `Vec<Vec<T>>`.
37    pub fn container(&self) -> &'_ Vec<Vec<T>> {
38        self.container
39    }
40
41    // /// Returns a mutable reference to the `Vec<Vec<T>>`.
42    // pub fn container_mut(&mut self) -> &'_ mut Vec<Vec<T>> {
43    //     self.container
44    // }
45}
46
47#[gat]
48impl<'a, T> LendingIterator for LendingIterMut<'a, T> {
49    type Item<'next>
50    where
51        Self: 'next,
52    = &'next mut T;
53
54    fn next<'next>(self: &'next mut LendingIterMut<'a, T>) -> Option<&'next mut T> {
55        let Self {
56            container: c,
57            outer_idx: oi,
58            inner_idx: ii,
59            is_finished,
60        } = self;
61        if *is_finished {
62            return None;
63        };
64        let old_oi = *oi;
65        let old_ii = *ii;
66        *ii = match ii.checked_add(1) {
67            // safety: at the beginning, *oi is 0,
68            // And *oi gets modified only in the `_` case, where it is checked to be less than the length of `c`.
69            Some(next_ii) if next_ii < unsafe { c.get_unchecked(*oi) }.len() => next_ii,
70            _ => {
71                loop {
72                    *oi = oi.checked_add(1)?;
73                    if *oi >= c.len() {
74                        *is_finished = true;
75                        // safety: `oi` is less than the length of `c` and `ii` is less than the length of `c[*oi]`
76                        return Some(unsafe {
77                            c.get_unchecked_mut(old_oi).get_unchecked_mut(old_ii)
78                        });
79                    };
80                    if !unsafe { c.get_unchecked(*oi) }.is_empty() {
81                        break;
82                    }
83                }
84                0
85            }
86        };
87
88        // safety: `oi` is less than the length of `c` and `ii` is less than the length of `c[*oi]`
89        Some(unsafe { c.get_unchecked_mut(old_oi).get_unchecked_mut(old_ii) })
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use crate::VecVecExt;
96    use lending_iterator::prelude::*;
97
98    #[test]
99    fn general_case() {
100        let mut v = vec![vec![2, 3, 5], vec![], vec![7, 11, 13]];
101        let mut iter = v.lending_iter_mut();
102
103        while let Some(x) = iter.next() {
104            *x += 1;
105        }
106
107        assert_eq!(v, vec![vec![3, 4, 6], vec![], vec![8, 12, 14]]);
108    }
109
110    #[test]
111    fn for_empty_outer() {
112        let mut v = Vec::<Vec<i32>>::new();
113        let mut iter = v.lending_iter_mut();
114
115        while let Some(x) = iter.next() {
116            *x += 1;
117        }
118
119        assert_eq!(v, Vec::<Vec<i32>>::new());
120    }
121
122    #[test]
123    fn for_empty_inner() {
124        let mut v: Vec<Vec<i32>> = vec![vec![], vec![], vec![]];
125        let mut iter = v.lending_iter_mut();
126
127        while let Some(x) = iter.next() {
128            *x += 1;
129        }
130
131        assert_eq!(v, vec![vec![], vec![], vec![]]);
132    }
133}
vec_vec/lending_iter_mut.rs

vec_vec/
lending_iter_mut.rs
