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.

1 files changed, 37 insertions, 0 deletions

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 2d15915d91a..2305d2a2377 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<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 {