70 lines
2 KiB
Rust
70 lines
2 KiB
Rust
pub trait AdventIterator: Iterator {
|
|
fn into_chunks<const N: usize>(self) -> impl Iterator<Item = Chunck<N, Self::Item>>
|
|
where
|
|
Self: Sized,
|
|
{
|
|
in_chunks::<N, Self::Item>(self)
|
|
}
|
|
}
|
|
|
|
impl<I> AdventIterator for I where I: Iterator {}
|
|
|
|
#[derive(Debug, PartialEq, Eq)]
|
|
pub enum Chunck<const N: usize, T> {
|
|
Next([T; N]),
|
|
Rest(Vec<T>),
|
|
}
|
|
|
|
pub fn in_chunks<const N: usize, T>(
|
|
mut iterator: impl Iterator<Item = T>,
|
|
) -> impl Iterator<Item = Chunck<N, T>> {
|
|
let mut buffer = Vec::new();
|
|
let mut done = false;
|
|
std::iter::from_fn(move || {
|
|
if done {
|
|
return None;
|
|
}
|
|
for _ in 0..N {
|
|
match iterator.next() {
|
|
Some(to_push) => buffer.push(to_push),
|
|
None => {
|
|
done = true;
|
|
return if buffer.is_empty() {
|
|
None
|
|
} else {
|
|
Some(Chunck::Rest(std::mem::take(&mut buffer)))
|
|
};
|
|
}
|
|
}
|
|
}
|
|
let array: [T; N] = std::mem::take(&mut buffer)
|
|
.try_into()
|
|
.unwrap_or_else(|_| panic!("Buffer must have the same size as the N ({})", N));
|
|
Some(Chunck::Next(array))
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod testing {
|
|
use crate::iterations::{in_chunks, Chunck};
|
|
|
|
#[test]
|
|
fn should_split_chunks() {
|
|
fn assert_case<const N: usize>(input: Vec<u32>, expected: Vec<Chunck<N, u32>>) {
|
|
let actual: Vec<Chunck<N, u32>> = in_chunks::<N, u32>(input.into_iter()).collect();
|
|
assert_eq!(expected, actual);
|
|
}
|
|
|
|
assert_case::<2>(vec![2], vec![Chunck::Rest(vec![2])]);
|
|
assert_case::<2>(vec![], vec![]);
|
|
assert_case::<2>(vec![2, 2], vec![Chunck::Next([2, 2])]);
|
|
assert_case::<2>(
|
|
vec![1, 2, 3],
|
|
vec![Chunck::Next([1, 2]), Chunck::Rest(vec![3])],
|
|
);
|
|
assert_case::<1>(
|
|
vec![1, 2, 3],
|
|
vec![Chunck::Next([1]), Chunck::Next([2]), Chunck::Next([3])],
|
|
);
|
|
}
|
|
}
|