Skip to main content

compressed_intvec/fixed/
iter_mut.rs

1//! # Mutable Iterators
2//!
3//! This module provides iterators for mutable, sequential access to the
4//! elements of a [`FixedVec`].
5//!
6//! # Provided Iterators
7//!
8//! - [`IterMut`]: An iterator that yields a mutable proxy for each element.
9//! - [`ChunksMut`]: An iterator that yields non-overlapping, mutable slices.
10//!
11//! # Examples
12//!
13//! ## Mutating elements in chunks
14//!
15//! The [`ChunksMut`] iterator allows for processing a vector in mutable,
16//! non-overlapping chunks.
17//!
18//! ```rust
19//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
20//! use compressed_intvec::fixed::{FixedVec, UFixedVec, BitWidth};
21//!
22//! let data: Vec<u32> = (0..100).collect();
23//! let mut vec: UFixedVec<u32> = FixedVec::builder().bit_width(BitWidth::Explicit(8)).build(&data)?;
24//!
25//! // Process each chunk sequentially.
26//! for mut chunk in vec.chunks_mut(10) {
27//!     // Each chunk is a `FixedVecSlice` that can be mutated.
28//!     for i in 0..chunk.len() {
29//!         if let Some(mut proxy) = chunk.at_mut(i) {
30//!             *proxy *= 2;
31//!         }
32//!     }
33//! }
34//!
35//! assert_eq!(vec.get(10), Some(20));
36//! assert_eq!(vec.get(99), Some(198));
37//! # Ok(())
38//! # }
39//! ```
40
41use crate::fixed::{
42    FixedVec,
43    proxy::MutProxy,
44    slice::FixedVecSlice,
45    traits::{Storable, Word},
46};
47use dsi_bitstream::prelude::Endianness;
48use std::{cmp::min, fmt, iter::FusedIterator, marker::PhantomData};
49
50/// An iterator over non-overlapping, mutable chunks of a [`FixedVec`].
51///
52/// This struct is created by the [`chunks_mut`](super::FixedVec::chunks_mut)
53/// method.
54#[derive(Debug)]
55pub struct ChunksMut<'a, T, W, E, B>
56where
57    T: Storable<W>,
58    W: Word,
59    E: Endianness,
60    B: AsRef<[W]> + AsMut<[W]>,
61{
62    // A raw pointer to the original `FixedVec` is used to allow creating
63    // mutable slices without consuming the iterator.
64    vec_ptr: *mut FixedVec<T, W, E, B>,
65    end: usize,
66    current_pos: usize,
67    chunk_size: usize,
68    // Ensures the iterator's lifetime is tied to the original mutable borrow.
69    _phantom: PhantomData<&'a mut FixedVec<T, W, E, B>>,
70}
71
72impl<'a, T, W, E, B> ChunksMut<'a, T, W, E, B>
73where
74    T: Storable<W>,
75    W: Word,
76    E: Endianness,
77    B: AsRef<[W]> + AsMut<[W]>,
78{
79    /// Creates a new `ChunksMut` iterator.
80    pub(super) fn new(vec: &'a mut FixedVec<T, W, E, B>, chunk_size: usize) -> Self {
81        assert!(chunk_size != 0, "chunk_size cannot be zero");
82        let end = vec.len();
83        Self {
84            vec_ptr: vec as *mut _,
85            chunk_size,
86            current_pos: 0,
87            end,
88            _phantom: PhantomData,
89        }
90    }
91}
92
93impl<'a, T, W, E, B> Iterator for ChunksMut<'a, T, W, E, B>
94where
95    T: Storable<W>,
96    W: Word,
97    E: Endianness,
98    B: AsRef<[W]> + AsMut<[W]>,
99{
100    type Item = FixedVecSlice<&'a mut FixedVec<T, W, E, B>>;
101
102    fn next(&mut self) -> Option<Self::Item> {
103        if self.current_pos >= self.end {
104            return None;
105        }
106
107        let start = self.current_pos;
108        let len = min(self.chunk_size, self.end - start);
109        self.current_pos += len;
110
111        // SAFETY:
112        // 1. `self.vec_ptr` is a valid pointer to a `FixedVec` for the lifetime 'a.
113        // 2. The borrow checker ensures the original `&'a mut FixedVec` is
114        //    exclusively borrowed by this iterator for its entire lifetime.
115        // 3. Each call to `next` produces a slice for a unique, non-overlapping
116        //    range of the original vector, making the mutable borrow safe.
117        let vec_ref = unsafe { &mut *self.vec_ptr };
118        let slice = FixedVecSlice::new(vec_ref, start..start + len);
119
120        Some(slice)
121    }
122}
123
124impl<T, W, E, B> FusedIterator for ChunksMut<'_, T, W, E, B>
125where
126    T: Storable<W>,
127    W: Word,
128    E: Endianness,
129    B: AsRef<[W]> + AsMut<[W]>,
130{
131}
132
133/// A mutable iterator over the elements of a [`FixedVec`].
134///
135/// This struct is created by the [`iter_mut`](super::FixedVec::iter_mut)
136/// method. It yields a [`MutProxy`] for each element.
137pub struct IterMut<'a, T, W, E, B>
138where
139    T: Storable<W>,
140    W: Word,
141    E: Endianness,
142    B: AsRef<[W]> + AsMut<[W]>,
143{
144    // A raw pointer is used to allow creating proxies without consuming the iterator.
145    vec_ptr: *mut FixedVec<T, W, E, B>,
146    pub(super) current_index: usize,
147    pub(super) end_index: usize,
148    _phantom: PhantomData<&'a mut FixedVec<T, W, E, B>>,
149}
150
151impl<'a, T, W, E, B> IterMut<'a, T, W, E, B>
152where
153    T: Storable<W>,
154    W: Word,
155    E: Endianness,
156    B: AsRef<[W]> + AsMut<[W]>,
157{
158    /// Creates a new `IterMut` iterator.
159    pub(super) fn new(vec: &'a mut FixedVec<T, W, E, B>) -> Self {
160        let len = vec.len();
161        Self {
162            vec_ptr: vec as *mut _,
163            current_index: 0,
164            end_index: len,
165            _phantom: PhantomData,
166        }
167    }
168}
169
170impl<'a, T, W, E, B> Iterator for IterMut<'a, T, W, E, B>
171where
172    T: Storable<W>,
173    W: Word,
174    E: Endianness,
175    B: AsRef<[W]> + AsMut<[W]>,
176{
177    type Item = MutProxy<'a, T, W, E, B>;
178
179    fn next(&mut self) -> Option<Self::Item> {
180        if self.current_index >= self.end_index {
181            return None;
182        }
183
184        // SAFETY: The iterator's lifetime `'a` and bounds checking ensure
185        // that the pointer is valid and that we are creating non-overlapping
186        // mutable access proxies one at a time.
187        let proxy = unsafe { MutProxy::new(&mut *self.vec_ptr, self.current_index) };
188        self.current_index += 1;
189
190        Some(proxy)
191    }
192}
193
194impl<T, W, E, B> FusedIterator for IterMut<'_, T, W, E, B>
195where
196    T: Storable<W>,
197    W: Word,
198    E: Endianness,
199    B: AsRef<[W]> + AsMut<[W]>,
200{
201}
202
203impl<T, W, E, B> fmt::Debug for IterMut<'_, T, W, E, B>
204where
205    T: Storable<W>,
206    W: Word,
207    E: Endianness,
208    B: AsRef<[W]> + AsMut<[W]>,
209{
210    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
211        f.debug_struct("IterMut")
212            .field("remaining", &(self.end_index - self.current_index))
213            .finish()
214    }
215}
216
217/// An unchecked mutable iterator over the elements of a [`FixedVec`].
218///
219/// This struct is created by the [`iter_mut_unchecked`](super::FixedVec::iter_mut_unchecked)
220/// method. It does not perform any bounds checking.
221///
222/// # Safety
223///
224/// The iterator is safe to use only if it is guaranteed that it will not
225/// be advanced beyond the end of the vector.
226pub struct IterMutUnchecked<'a, T, W, E, B>
227where
228    T: Storable<W>,
229    W: Word,
230    E: Endianness,
231    B: AsRef<[W]> + AsMut<[W]>,
232{
233    iter: IterMut<'a, T, W, E, B>,
234}
235
236impl<'a, T, W, E, B> IterMutUnchecked<'a, T, W, E, B>
237where
238    T: Storable<W>,
239    W: Word,
240    E: Endianness,
241    B: AsRef<[W]> + AsMut<[W]>,
242{
243    /// Creates a new `IterMutUnchecked`.
244    pub(super) fn new(vec: &'a mut FixedVec<T, W, E, B>) -> Self {
245        Self {
246            iter: IterMut::new(vec),
247        }
248    }
249
250    /// Returns the next mutable proxy without bounds checking.
251    ///
252    /// # Safety
253    ///
254    /// Calling this method when the iterator is exhausted is undefined behavior.
255    #[inline]
256    pub unsafe fn next_unchecked(&mut self) -> MutProxy<'a, T, W, E, B> {
257        unsafe { self.iter.next().unwrap_unchecked() }
258    }
259}
260
261impl<T, W, E, B> fmt::Debug for IterMutUnchecked<'_, T, W, E, B>
262where
263    T: Storable<W>,
264    W: Word,
265    E: Endianness,
266    B: AsRef<[W]> + AsMut<[W]>,
267{
268    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
269        f.debug_struct("IterMutUnchecked")
270            .field("remaining", &(self.iter.end_index - self.iter.current_index))
271            .finish()
272    }
273}