Add an iterator that ensures known size

Introduce the `KnownSize` iterator wrapper, which allows providing the
size at construction time. This provides an `ExactSizeIterator`
implemenation so we can check a generator's value count during testing.
This commit is contained in:
Trevor Gross
2024-12-30 06:12:16 +00:00
parent 9a7a519370
commit b78f7b7b48

View File

@@ -5,6 +5,43 @@ pub mod domain_logspace;
pub mod edge_cases;
pub mod random;
/// A wrapper to turn any iterator into an `ExactSizeIterator`. Asserts the final result to ensure
/// the provided size was correct.
#[derive(Debug)]
pub struct KnownSize<I> {
total: u64,
current: u64,
iter: I,
}
impl<I> KnownSize<I> {
pub fn new(iter: I, total: u64) -> Self {
Self { total, current: 0, iter }
}
}
impl<I: Iterator> Iterator for KnownSize<I> {
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> {
let next = self.iter.next();
if next.is_some() {
self.current += 1;
return next;
}
assert_eq!(self.current, self.total, "total items did not match expected");
None
}
fn size_hint(&self) -> (usize, Option<usize>) {
let remaining = usize::try_from(self.total - self.current).unwrap();
(remaining, Some(remaining))
}
}
impl<I: Iterator> ExactSizeIterator for KnownSize<I> {}
/// Helper type to turn any reusable input into a generator.
#[derive(Clone, Debug, Default)]
pub struct CachedInput {