advent_of_code/common/
chunk_iterator.rs

1pub struct ChunksIterator<I, T, const SIZE: usize>
2where
3    I: Iterator<Item = T>,
4    T: Copy + Default,
5{
6    iterator: I,
7    items: [T; SIZE],
8    current_idx: usize,
9}
10
11impl<I, T, const SIZE: usize> Iterator for ChunksIterator<I, T, SIZE>
12where
13    I: Iterator<Item = T>,
14    T: Copy + Default,
15{
16    type Item = [T; SIZE];
17
18    fn next(&mut self) -> Option<Self::Item> {
19        while self.current_idx < SIZE {
20            self.items[self.current_idx] = self.iterator.next()?;
21            self.current_idx += 1;
22        }
23        self.current_idx = 0;
24        Some(self.items)
25    }
26}
27
28impl<I, T, const SIZE: usize> ExactSizeIterator for ChunksIterator<I, T, SIZE>
29where
30    I: ExactSizeIterator<Item = T>,
31    T: Copy + Default,
32{
33    fn len(&self) -> usize {
34        (self.current_idx + self.iterator.len()) / SIZE
35    }
36}
37
38pub trait ChunkIteratorExt: Iterator {
39    fn chunks_exact<const SIZE: usize>(self) -> ChunksIterator<Self, Self::Item, SIZE>
40    where
41        Self::Item: Copy + Default,
42        Self: Sized,
43    {
44        ChunksIterator {
45            iterator: self,
46            items: [Self::Item::default(); SIZE],
47            current_idx: 0,
48        }
49    }
50}
51
52impl<I: Iterator> ChunkIteratorExt for I {}
53
54#[test]
55fn test() {
56    let a = [1, 2, 3, 4, 5, 6];
57    let mut it = a.into_iter().chunks_exact::<3>();
58    assert_eq!(2, it.len());
59    assert_eq!(it.next(), Some([1, 2, 3]));
60    assert_eq!(1, it.len());
61    assert_eq!(it.next(), Some([4, 5, 6]));
62    assert_eq!(0, it.len());
63    assert_eq!(it.next(), None);
64    assert_eq!(0, it.len());
65
66    let a = [1, 2, 3, 4, 5, 6, 7];
67    let mut it = a.into_iter().chunks_exact::<3>();
68    assert_eq!(2, it.len());
69    assert_eq!(it.next(), Some([1, 2, 3]));
70    assert_eq!(1, it.len());
71    assert_eq!(it.next(), Some([4, 5, 6]));
72    assert_eq!(0, it.len());
73    assert_eq!(it.next(), None);
74    assert_eq!(0, it.len());
75}