From e53f6352081533c621415507447484cd8076d556 Mon Sep 17 00:00:00 2001
From: Kinto <kinto.behr@gmail.com>
Date: Sun, 5 May 2024 00:22:47 +1000
Subject: [PATCH] Implement two-element 'vector' for `flatten_ok` test

Previously the test used Option<u8> but the coverage was bad. We cannot use Vec<u8> because it is too slow.
---
 tests/specializations.rs | 64 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 62 insertions(+), 2 deletions(-)

diff --git a/tests/specializations.rs b/tests/specializations.rs
index cb2141053..712311472 100644
--- a/tests/specializations.rs
+++ b/tests/specializations.rs
@@ -1,7 +1,9 @@
 #![allow(unstable_name_collisions)]
 
 use itertools::Itertools;
+use quickcheck::Arbitrary;
 use quickcheck::{quickcheck, TestResult};
+use rand::Rng;
 use std::fmt::Debug;
 
 struct Unspecialized<I>(I);
@@ -452,8 +454,8 @@ quickcheck! {
         test_specializations(&v.into_iter().filter_map_ok(|i| if i < 20 { Some(i * 2) } else { None }));
     }
 
-    // `Option<u8>` because `Vec<u8>` would be very slow!! And we can't give `[u8; 3]`.
-    fn flatten_ok(v: Vec<Result<Option<u8>, char>>) -> () {
+    // `SmallIter2<u8>` because `Vec<u8>` is too slow and we get bad coverage from a singleton like Option<u8>
+    fn flatten_ok(v: Vec<Result<SmallIter2<u8>, char>>) -> () {
         let it = v.into_iter().flatten_ok();
         test_specializations(&it);
         test_double_ended_specializations(&it);
@@ -520,3 +522,61 @@ quickcheck! {
         }
     }
 }
+
+/// Like `VecIntoIter<T>` with maximum 2 elements.
+#[derive(Debug, Clone, Default)]
+enum SmallIter2<T> {
+    #[default]
+    Zero,
+    One(T),
+    Two(T, T),
+}
+
+impl<T: Arbitrary> Arbitrary for SmallIter2<T> {
+    fn arbitrary<G: quickcheck::Gen>(g: &mut G) -> Self {
+        match g.gen_range(0u8, 3) {
+            0 => Self::Zero,
+            1 => Self::One(T::arbitrary(g)),
+            2 => Self::Two(T::arbitrary(g), T::arbitrary(g)),
+            _ => unreachable!(),
+        }
+    }
+    // maybe implement shrink too, maybe not
+}
+
+impl<T> Iterator for SmallIter2<T> {
+    type Item = T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match std::mem::take(self) {
+            Self::Zero => None,
+            Self::One(val) => Some(val),
+            Self::Two(val, second) => {
+                *self = Self::One(second);
+                Some(val)
+            }
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = match self {
+            Self::Zero => 0,
+            Self::One(_) => 1,
+            Self::Two(_, _) => 2,
+        };
+        (len, Some(len))
+    }
+}
+
+impl<T> DoubleEndedIterator for SmallIter2<T> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        match std::mem::take(self) {
+            Self::Zero => None,
+            Self::One(val) => Some(val),
+            Self::Two(first, val) => {
+                *self = Self::One(first);
+                Some(val)
+            }
+        }
+    }
+}