advent_of_code/common/
chunk_iterator.rs1pub 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}