block_grid/
iters.rs

1//! A bunch of custom 2D iterators.
2//!
3//! You probably won't need to interact with this module unless you need to name one of the
4//! iterator types explicitly.
5
6use core::iter::FusedIterator;
7use core::marker::PhantomData;
8use core::ptr::NonNull;
9use core::slice::{ChunksExact, ChunksExactMut, Iter, IterMut};
10
11use crate::{Block, BlockDim, BlockGrid, BlockMut, Coords};
12
13/// Provides an interface for iterators that can also yield 2D coordinates.
14///
15/// Note that this trait is sealed, meaning it cannot be implemented by downstream crates. This
16/// pattern is as described [here][sealing].
17///
18/// [`coords`]: Self::coords
19/// [sealing]: https://rust-lang.github.io/api-guidelines/future-proofing.html#sealed-traits-protect-against-downstream-implementations-c-sealed
20pub trait CoordsIterator: Iterator + private::Sealed {
21    /// Returns the coordinates of the item *to be* yielded next.
22    ///
23    /// This is really just an implementation detail of [`WithCoordsIter`], so it probably shouldn't
24    /// be used. Just use [`coords`][Self::coords] to get the coordinates.
25    #[doc(hidden)]
26    fn current_coords(&self) -> Coords;
27
28    /// Returns an iterator adapter that also gives coordinates as well as the next value.
29    ///
30    /// The iterator returned yields 2-tuples `(coords, elem)`, where `coords` is the coordinates
31    /// of the next element `elem`. This is essentially a 2D version of [`Iterator::enumerate`].
32    fn coords(self) -> WithCoordsIter<Self>
33    where
34        Self: Sized,
35    {
36        WithCoordsIter { iter: self }
37    }
38}
39
40/// Immutable iterator in memory order.
41///
42/// Created by the [`BlockGrid::each_iter`] method.
43#[derive(Clone, Debug)]
44pub struct EachIter<'a, T, B: BlockDim> {
45    row: usize,
46    col: usize,
47    cols: usize,
48    iter: Iter<'a, T>,
49    _phantom: PhantomData<B>,
50}
51
52/// Mutable iterator in memory order.
53///
54/// Created by the [`BlockGrid::each_iter_mut`] method.
55#[derive(Debug)]
56pub struct EachIterMut<'a, T, B: BlockDim> {
57    row: usize,
58    col: usize,
59    cols: usize,
60    iter: IterMut<'a, T>,
61    _phantom: PhantomData<B>,
62}
63
64/// Immutable iterator over entire blocks.
65///
66/// Created by the [`BlockGrid::block_iter`] method.
67#[derive(Clone, Debug)]
68pub struct BlockIter<'a, T, B: BlockDim> {
69    block_row: usize,
70    block_col: usize,
71    col_blocks: usize,
72    chunks: ChunksExact<'a, T>,
73    _phantom: PhantomData<B>,
74}
75
76/// Mutable iterator over entire blocks.
77///
78/// Created by the [`BlockGrid::block_iter_mut`] method.
79#[derive(Debug)]
80pub struct BlockIterMut<'a, T, B: BlockDim> {
81    block_row: usize,
82    block_col: usize,
83    col_blocks: usize,
84    chunks: ChunksExactMut<'a, T>,
85    _phantom: PhantomData<B>,
86}
87
88/// Immutable iterator in row-major order.
89///
90/// Created by the [`BlockGrid::row_major_iter`] method.
91#[derive(Clone, Debug)]
92pub struct RowMajorIter<'a, T, B: BlockDim> {
93    row: usize,
94    col: usize,
95    grid: &'a BlockGrid<T, B>,
96}
97
98/// Mutable iterator in row-major order.
99///
100/// Created by the [`BlockGrid::row_major_iter_mut`] method.
101#[derive(Debug)]
102pub struct RowMajorIterMut<'a, T, B: BlockDim> {
103    row: usize,
104    col: usize,
105    grid: NonNull<BlockGrid<T, B>>,
106    _phantom: PhantomData<&'a mut BlockGrid<T, B>>,
107}
108
109/// An iterator adapter that yields the coordinates and the element.
110///
111/// This is created by the [`CoordsIterator::coords`] method on all the iterator types that
112/// implement the trait. See its documentation for more info.
113#[derive(Clone, Debug)]
114pub struct WithCoordsIter<I> {
115    iter: I,
116}
117
118impl<'a, T, B: BlockDim> EachIter<'a, T, B> {
119    pub(crate) fn new(grid: &'a BlockGrid<T, B>) -> Self {
120        Self {
121            row: 0,
122            col: 0,
123            cols: grid.cols(),
124            iter: grid.raw().iter(),
125            _phantom: PhantomData,
126        }
127    }
128}
129
130impl<T, B: BlockDim> CoordsIterator for EachIter<'_, T, B> {
131    #[inline]
132    fn current_coords(&self) -> Coords {
133        (self.row, self.col)
134    }
135}
136
137impl<'a, T, B: BlockDim> Iterator for EachIter<'a, T, B> {
138    type Item = &'a T;
139
140    #[inline]
141    fn next(&mut self) -> Option<Self::Item> {
142        // TODO: Try out bitwise ops for potential speedup?
143        self.col += 1;
144        if self.col % B::WIDTH == 0 {
145            self.row += 1;
146            if self.row % B::WIDTH == 0 {
147                if self.col == self.cols {
148                    self.col = 0;
149                } else {
150                    self.row -= B::WIDTH;
151                }
152            } else {
153                self.col -= B::WIDTH;
154            }
155        }
156        self.iter.next()
157    }
158
159    #[inline]
160    fn size_hint(&self) -> (usize, Option<usize>) {
161        self.iter.size_hint()
162    }
163
164    #[inline]
165    fn count(self) -> usize {
166        self.iter.count()
167    }
168}
169
170impl<T, B: BlockDim> ExactSizeIterator for EachIter<'_, T, B> {
171    #[inline]
172    fn len(&self) -> usize {
173        self.iter.len()
174    }
175}
176
177impl<T, B: BlockDim> FusedIterator for EachIter<'_, T, B> {}
178
179impl<'a, T, B: BlockDim> EachIterMut<'a, T, B> {
180    pub(crate) fn new(grid: &'a mut BlockGrid<T, B>) -> Self {
181        Self {
182            row: 0,
183            col: 0,
184            cols: grid.cols(),
185            iter: grid.raw_mut().iter_mut(),
186            _phantom: PhantomData,
187        }
188    }
189}
190
191impl<T, B: BlockDim> CoordsIterator for EachIterMut<'_, T, B> {
192    #[inline]
193    fn current_coords(&self) -> Coords {
194        (self.row, self.col)
195    }
196}
197
198impl<'a, T, B: BlockDim> Iterator for EachIterMut<'a, T, B> {
199    type Item = &'a mut T;
200
201    #[inline]
202    fn next(&mut self) -> Option<Self::Item> {
203        self.col += 1;
204        if self.col % B::WIDTH == 0 {
205            self.row += 1;
206            if self.row % B::WIDTH == 0 {
207                if self.col == self.cols {
208                    self.col = 0;
209                } else {
210                    self.row -= B::WIDTH;
211                }
212            } else {
213                self.col -= B::WIDTH;
214            }
215        }
216        self.iter.next()
217    }
218
219    #[inline]
220    fn size_hint(&self) -> (usize, Option<usize>) {
221        self.iter.size_hint()
222    }
223
224    #[inline]
225    fn count(self) -> usize {
226        self.iter.count()
227    }
228}
229
230impl<T, B: BlockDim> ExactSizeIterator for EachIterMut<'_, T, B> {
231    #[inline]
232    fn len(&self) -> usize {
233        self.iter.len()
234    }
235}
236
237impl<T, B: BlockDim> FusedIterator for EachIterMut<'_, T, B> {}
238
239impl<'a, T, B: BlockDim> BlockIter<'a, T, B> {
240    pub(crate) fn new(grid: &'a BlockGrid<T, B>) -> Self {
241        Self {
242            block_row: 0,
243            block_col: 0,
244            col_blocks: grid.col_blocks(),
245            chunks: grid.raw().chunks_exact(B::AREA),
246            _phantom: PhantomData,
247        }
248    }
249}
250
251impl<T, B: BlockDim> CoordsIterator for BlockIter<'_, T, B> {
252    #[inline]
253    fn current_coords(&self) -> Coords {
254        (self.block_row, self.block_col)
255    }
256}
257
258impl<'a, T, B: BlockDim> Iterator for BlockIter<'a, T, B> {
259    type Item = Block<'a, T, B>;
260
261    #[inline]
262    fn next(&mut self) -> Option<Self::Item> {
263        let chunk = self.chunks.next()?;
264        // SAFETY: `self.chunks` gives slices of exactly `B::AREA` length
265        let block = unsafe { Block::new(self.current_coords(), chunk) };
266        self.block_col += 1;
267        if self.block_col == self.col_blocks {
268            self.block_row += 1;
269            self.block_col = 0;
270        }
271        Some(block)
272    }
273
274    #[inline]
275    fn size_hint(&self) -> (usize, Option<usize>) {
276        self.chunks.size_hint()
277    }
278
279    #[inline]
280    fn count(self) -> usize {
281        self.chunks.count()
282    }
283}
284
285impl<T, B: BlockDim> ExactSizeIterator for BlockIter<'_, T, B> {
286    #[inline]
287    fn len(&self) -> usize {
288        self.chunks.len()
289    }
290}
291
292impl<T, B: BlockDim> FusedIterator for BlockIter<'_, T, B> {}
293
294impl<'a, T, B: BlockDim> BlockIterMut<'a, T, B> {
295    pub(crate) fn new(grid: &'a mut BlockGrid<T, B>) -> Self {
296        Self {
297            block_row: 0,
298            block_col: 0,
299            col_blocks: grid.col_blocks(),
300            chunks: grid.raw_mut().chunks_exact_mut(B::AREA),
301            _phantom: PhantomData,
302        }
303    }
304}
305
306impl<T, B: BlockDim> CoordsIterator for BlockIterMut<'_, T, B> {
307    #[inline]
308    fn current_coords(&self) -> Coords {
309        (self.block_row, self.block_col)
310    }
311}
312
313impl<'a, T, B: BlockDim> Iterator for BlockIterMut<'a, T, B> {
314    type Item = BlockMut<'a, T, B>;
315
316    #[inline]
317    fn next(&mut self) -> Option<Self::Item> {
318        let chunk = self.chunks.next()?;
319        // SAFETY: `self.chunks` gives slices of exactly `B::AREA` length
320        let block = unsafe { BlockMut::new(self.current_coords(), chunk) };
321        self.block_col += 1;
322        if self.block_col == self.col_blocks {
323            self.block_row += 1;
324            self.block_col = 0;
325        }
326        Some(block)
327    }
328
329    #[inline]
330    fn size_hint(&self) -> (usize, Option<usize>) {
331        self.chunks.size_hint()
332    }
333
334    #[inline]
335    fn count(self) -> usize {
336        self.chunks.count()
337    }
338}
339
340impl<T, B: BlockDim> ExactSizeIterator for BlockIterMut<'_, T, B> {
341    #[inline]
342    fn len(&self) -> usize {
343        self.chunks.len()
344    }
345}
346
347impl<T, B: BlockDim> FusedIterator for BlockIterMut<'_, T, B> {}
348
349impl<'a, T, B: BlockDim> RowMajorIter<'a, T, B> {
350    pub(crate) fn new(grid: &'a BlockGrid<T, B>) -> Self {
351        Self {
352            row: 0,
353            col: 0,
354            grid,
355        }
356    }
357}
358
359impl<T, B: BlockDim> CoordsIterator for RowMajorIter<'_, T, B> {
360    #[inline]
361    fn current_coords(&self) -> Coords {
362        (self.row, self.col)
363    }
364}
365
366impl<'a, T, B: BlockDim> Iterator for RowMajorIter<'a, T, B> {
367    type Item = &'a T;
368
369    #[inline]
370    fn next(&mut self) -> Option<Self::Item> {
371        if self.row >= self.grid.rows() {
372            return None;
373        }
374        // SAFETY: Method logic ensures `(self.row, self.col)` is a valid index
375        let x = unsafe { self.grid.get_unchecked((self.row, self.col)) };
376        self.col += 1;
377        if self.col == self.grid.cols() {
378            self.row += 1;
379            self.col = 0;
380        }
381        Some(x)
382    }
383
384    #[inline]
385    fn size_hint(&self) -> (usize, Option<usize>) {
386        let idx = self.row * self.grid.cols() + self.col;
387        let k = self.grid.size() - idx;
388        (k, Some(k))
389    }
390
391    #[inline]
392    fn count(self) -> usize {
393        self.len()
394    }
395}
396
397impl<T, B: BlockDim> ExactSizeIterator for RowMajorIter<'_, T, B> {}
398
399impl<T, B: BlockDim> FusedIterator for RowMajorIter<'_, T, B> {}
400
401impl<'a, T, B: BlockDim> RowMajorIterMut<'a, T, B> {
402    pub(crate) fn new(grid: &'a mut BlockGrid<T, B>) -> Self {
403        Self {
404            row: 0,
405            col: 0,
406            grid: grid.into(),
407            _phantom: PhantomData,
408        }
409    }
410}
411
412impl<T, B: BlockDim> CoordsIterator for RowMajorIterMut<'_, T, B> {
413    #[inline]
414    fn current_coords(&self) -> Coords {
415        (self.row, self.col)
416    }
417}
418
419impl<'a, T, B: BlockDim> Iterator for RowMajorIterMut<'a, T, B> {
420    type Item = &'a mut T;
421
422    #[inline]
423    fn next(&mut self) -> Option<Self::Item> {
424        // SAFETY: `self.grid` is a valid pointer
425        let (rows, cols) = unsafe {
426            let grid = self.grid.as_ref();
427            (grid.rows(), grid.cols())
428        };
429        if self.row >= rows {
430            return None;
431        }
432        // SAFETY: `self.grid` is a valid mutable pointer and method logic ensures
433        //         `(self.row, self.col)` is a valid index
434        let x = unsafe { (&mut *self.grid.as_ptr()).get_unchecked_mut((self.row, self.col)) };
435        self.col += 1;
436        if self.col == cols {
437            self.row += 1;
438            self.col = 0;
439        }
440        Some(x)
441    }
442
443    #[inline]
444    fn size_hint(&self) -> (usize, Option<usize>) {
445        // SAFETY: `self.grid` is a valid pointer
446        let grid = unsafe { self.grid.as_ref() };
447        let idx = self.row * grid.cols() + self.col;
448        let k = grid.size() - idx;
449        (k, Some(k))
450    }
451
452    #[inline]
453    fn count(self) -> usize {
454        self.len()
455    }
456}
457
458impl<T, B: BlockDim> ExactSizeIterator for RowMajorIterMut<'_, T, B> {}
459
460impl<T, B: BlockDim> FusedIterator for RowMajorIterMut<'_, T, B> {}
461
462impl<I: CoordsIterator> Iterator for WithCoordsIter<I> {
463    type Item = (Coords, I::Item);
464
465    #[inline]
466    fn next(&mut self) -> Option<Self::Item> {
467        let c = self.iter.current_coords();
468        self.iter.next().map(|x| (c, x))
469    }
470
471    #[inline]
472    fn size_hint(&self) -> (usize, Option<usize>) {
473        self.iter.size_hint()
474    }
475
476    #[inline]
477    fn count(self) -> usize {
478        self.iter.count()
479    }
480
481    #[inline]
482    fn nth(&mut self, n: usize) -> Option<Self::Item> {
483        if n >= 1 {
484            self.iter.nth(n - 1)?;
485        }
486        self.next()
487    }
488}
489
490impl<I: CoordsIterator + ExactSizeIterator> ExactSizeIterator for WithCoordsIter<I> {
491    #[inline]
492    fn len(&self) -> usize {
493        self.iter.len()
494    }
495}
496
497impl<I: CoordsIterator + FusedIterator> FusedIterator for WithCoordsIter<I> {}
498
499/// Prevent users from implementing the `CoordsIterator` trait.
500mod private {
501    use super::*;
502    pub trait Sealed {}
503    impl<T, B: BlockDim> Sealed for EachIter<'_, T, B> {}
504    impl<T, B: BlockDim> Sealed for EachIterMut<'_, T, B> {}
505    impl<T, B: BlockDim> Sealed for BlockIter<'_, T, B> {}
506    impl<T, B: BlockDim> Sealed for BlockIterMut<'_, T, B> {}
507    impl<T, B: BlockDim> Sealed for RowMajorIter<'_, T, B> {}
508    impl<T, B: BlockDim> Sealed for RowMajorIterMut<'_, T, B> {}
509}
block_grid/iters.rs

block_grid/
iters.rs
