chunked_vec/iterators/
iter_mut.rs

1use likely_stable::unlikely;
2
3use crate::ChunkedVec;
4
5/// A mutable iterator over the elements of a ChunkedVec.
6///
7/// This struct is created by the [`iter_mut`] method on [`ChunkedVec`].
8/// See its documentation for more.
9pub struct IterMut<'a, T, const N: usize> {
10    pub(crate) vec: &'a mut ChunkedVec<T, N>,
11    pub(crate) chunk_idx: usize,
12    pub(crate) offset: usize,
13    pub(crate) remaining: usize,
14}
15
16impl<T, const N: usize> ChunkedVec<T, N> {
17    /// Returns an iterator that allows modifying each element in the vector.
18    ///
19    /// The iterator yields all items from start to end.
20    ///
21    /// # Examples
22    /// ```
23    /// use chunked_vec::ChunkedVec;
24    /// let mut vec = ChunkedVec::new();
25    /// vec.push(1);
26    /// vec.push(2);
27    ///
28    /// for element in vec.iter_mut() {
29    ///     *element *= 2;
30    /// }
31    ///
32    /// assert_eq!(vec[0], 2);
33    /// assert_eq!(vec[1], 4);
34    /// ```
35    pub fn iter_mut(&mut self) -> IterMut<'_, T, N> {
36        IterMut {
37            remaining: self.len(),
38            vec: self,
39            chunk_idx: 0,
40            offset: 0,
41        }
42    }
43}
44
45impl<'a, T, const N: usize> IterMut<'a, T, N> {
46    /// Advances to the next position.
47    #[inline]
48    fn advance_position(&mut self) {
49        self.offset += 1;
50        if unlikely(self.offset == N) {
51            self.chunk_idx += 1;
52            self.offset = 0;
53        }
54        self.remaining -= 1;
55    }
56
57    /// Returns a pointer to the current element.
58    #[inline]
59    fn current_ptr(&mut self) -> *mut T {
60        self.vec.data[self.chunk_idx][self.offset].as_mut_ptr()
61    }
62}
63
64impl<'a, T, const N: usize> Iterator for IterMut<'a, T, N> {
65    type Item = &'a mut T;
66
67    fn next(&mut self) -> Option<Self::Item> {
68        if unlikely(self.remaining == 0) {
69            return None;
70        }
71
72        unsafe {
73            // 使用原始指针避免借用冲突
74            let ptr = self.current_ptr();
75            self.advance_position();
76
77            // 将原始指针转换为正确生命周期的引用
78            Some(&mut *ptr)
79        }
80    }
81
82    fn size_hint(&self) -> (usize, Option<usize>) {
83        let remaining = self.remaining;
84        (remaining, Some(remaining))
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn test_iter_mut() {
94        let mut vec = ChunkedVec::new();
95        vec.push(1);
96        vec.push(2);
97        vec.push(3);
98
99        let mut iter = vec.iter_mut();
100        assert_eq!(iter.next(), Some(&mut 1));
101        assert_eq!(iter.next(), Some(&mut 2));
102        let elem = iter.next();
103        assert_eq!(elem, Some(&mut 3));
104        *elem.unwrap() = 4;
105        assert_eq!(iter.next(), None);
106        assert_eq!(vec[2], 4);
107    }
108}
chunked_vec/iterators/iter_mut.rs

chunked_vec/iterators/
iter_mut.rs
