diff --git a/library/compiler-builtins/libm/crates/libm-test/src/gen.rs b/library/compiler-builtins/libm/crates/libm-test/src/gen.rs index 2d15915d91ae..2305d2a23773 100644 --- a/library/compiler-builtins/libm/crates/libm-test/src/gen.rs +++ b/library/compiler-builtins/libm/crates/libm-test/src/gen.rs @@ -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 { + total: u64, + current: u64, + iter: I, +} + +impl KnownSize { + pub fn new(iter: I, total: u64) -> Self { + Self { total, current: 0, iter } + } +} + +impl Iterator for KnownSize { + type Item = I::Item; + + fn next(&mut self) -> Option { + 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) { + let remaining = usize::try_from(self.total - self.current).unwrap(); + (remaining, Some(remaining)) + } +} + +impl ExactSizeIterator for KnownSize {} + /// Helper type to turn any reusable input into a generator. #[derive(Clone, Debug, Default)] pub struct CachedInput {