Change the return type of at_most_one() to AtMostOneResult

src/at_most_one.rs

@@ -0,0 +1,75 @@
+use crate::size_hint;
+use core::iter::ExactSizeIterator;
+use core::mem;
+
+/// The enum returned by `at_most_one`, depending on the number of remaining
+/// elements in the iterator.
+/// 
+/// See [`.at_most_one()`](crate::Itertools::at_most_one) for more detail.
+#[derive(PartialEq, Eq, Clone, Debug)]
+pub enum AtMostOneResult<I: Iterator> {
+    /// The iterator was empty and therefore had zero elements.
+    Zero,
+
+    /// The iterator had exactly one element.
+    One(I::Item),
+
+    /// The iterator had more than one element. 
+    /// [`MoreThanOne`](crate::MoreThanOne) is an iterator which yields the same elements as the original iterator. 
+    MoreThanOne(MoreThanOne<I>),
+}
+
+#[derive(PartialEq, Eq, Clone, Debug)]
+enum IterSource<T> {
+    FirstElement(T, T),
+    SecondElement(T),
+    InnerIter,
+}
+
+/// The iterator returned by [`.at_most_one()`](crate::Itertools::at_most_one), if the original iterator
+/// had at least two elements remaining. Yields the same elements as the original iterator.
+#[derive(PartialEq, Eq, Clone, Debug)]
+pub struct MoreThanOne<I: Iterator> {
+    next_source: IterSource<I::Item>,
+    inner: I,
+}
+
+impl<I: Iterator> MoreThanOne<I> {
+    pub(crate) fn new(first_two: [I::Item; 2], inner: I) -> Self {
+        let [first, second] = first_two;
+        let next_source = IterSource::FirstElement(first, second);
+
+        Self { next_source, inner }
+    }
+
+    fn additional_len(&self) -> usize {
+        match self.next_source {
+            IterSource::FirstElement(_, _) => 2,
+            IterSource::SecondElement(_) => 1,
+            IterSource::InnerIter => 0,
+        }
+    }
+}
+
+impl<I: Iterator> Iterator for MoreThanOne<I> {
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let source = mem::replace(&mut self.next_source, IterSource::InnerIter);
+
+        match source {
+            IterSource::FirstElement(first, second) => {
+                self.next_source = IterSource::SecondElement(second);
+                Some(first)
+            }
+            IterSource::SecondElement(second) => Some(second),
+            IterSource::InnerIter => self.inner.next(),
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        size_hint::add_scalar(self.inner.size_hint(), self.additional_len())
+    }
+}
+
+impl<I> ExactSizeIterator for MoreThanOne<I> where I: ExactSizeIterator {}

src/lib.rs

@@ -112,6 +112,7 @@ pub mod structs {
     pub use crate::adaptors::{MapResults, Step};
     #[cfg(feature = "use_alloc")]
     pub use crate::adaptors::MultiProduct;
+    pub use crate::at_most_one::{AtMostOneResult, MoreThanOne};
     #[cfg(feature = "use_alloc")]
     pub use crate::combinations::Combinations;
     #[cfg(feature = "use_alloc")]
@@ -183,6 +184,7 @@ pub use crate::with_position::Position;
 pub use crate::unziptuple::{multiunzip, MultiUnzip};
 pub use crate::ziptuple::multizip;
 mod adaptors;
+mod at_most_one;
 mod either_or_both;
 pub use crate::either_or_both::EitherOrBoth;
 #[doc(hidden)]
@@ -3690,39 +3692,50 @@ pub trait Itertools : Iterator {
         }
     }
 
-    /// If the iterator yields no elements, Ok(None) will be returned. If the iterator yields
-    /// exactly one element, that element will be returned, otherwise an error will be returned
-    /// containing an iterator that has the same output as the input iterator.
-    ///
-    /// This provides an additional layer of validation over just calling `Iterator::next()`.
-    /// If your assumption that there should be at most one element yielded is false this provides
-    /// the opportunity to detect and handle that, preventing errors at a distance.
+    /// If the iterator yields no elements, [`AtMostOneResult::Zero`] will be returned. 
+    /// If the iterator yields exactly one element, that element will be returned as [`AtMostOneResult::One`]. 
+    /// Otherwise [`AtMostOneResult::MoreThanOne`] will be returned, containing an iterator that has the same 
+    /// output as the input iterator.
     ///
+    /// This function is especially useful in `match` statements, to clearly assert the number of elements
+    /// in an iterator or take different paths depending on this number. 
+    /// For example, using [`.find()`](std::iter::Iterator::find) will stop at the first matched element, 
+    /// but you may additionally want to validate that no more than a single element satisfies the predicate.
+    /// 
     /// # Examples
     /// ```
-    /// use itertools::Itertools;
-    ///
-    /// assert_eq!((0..10).filter(|&x| x == 2).at_most_one().unwrap(), Some(2));
-    /// assert!((0..10).filter(|&x| x > 1 && x < 4).at_most_one().unwrap_err().eq(2..4));
-    /// assert!((0..10).filter(|&x| x > 1 && x < 5).at_most_one().unwrap_err().eq(2..5));
-    /// assert_eq!((0..10).filter(|&_| false).at_most_one().unwrap(), None);
+    /// use itertools::{Itertools, AtMostOneResult};
+    /// use core::iter::{empty, once};
+    /// 
+    /// let numbers = [-9, -1, -7, 4, -38, -21].iter().copied();
+    /// let positive_number: Option<i32> = match numbers.filter(|&num| num >= 0).at_most_one() {
+    ///     AtMostOneResult::Zero => None,
+    ///     AtMostOneResult::One(n) => Some(n),
+    ///     AtMostOneResult::MoreThanOne(_) => panic!("expected no more than one positive number"),
+    /// };
+    /// assert_eq!(positive_number, Some(4));
+    /// 
+    /// let zero_elements = empty::<i32>().at_most_one();
+    /// assert!(matches!(zero_elements, AtMostOneResult::Zero));
+    /// 
+    /// let one_element = once(5).at_most_one();
+    /// assert!(matches!(one_element, AtMostOneResult::One(5)));
+    /// 
+    /// let many_elements = (1..=10).at_most_one();
+    /// let AtMostOneResult::MoreThanOne(more_than_one) = many_elements else { panic!() };
+    /// assert!(more_than_one.eq(1..=10));
     /// ```
-    fn at_most_one(mut self) -> Result<Option<Self::Item>, ExactlyOneError<Self>>
+    fn at_most_one(mut self) -> AtMostOneResult<Self>
     where
         Self: Sized,
     {
-        match self.next() {
-            Some(first) => {
-                match self.next() {
-                    Some(second) => {
-                        Err(ExactlyOneError::new(Some(Either::Left([first, second])), self))
-                    }
-                    None => {
-                        Ok(Some(first))
-                    }
-                }
+        match (self.next(), self.next()) {
+            (None, _) => AtMostOneResult::Zero,
+            (Some(first), None) => AtMostOneResult::One(first),
+            (Some(first), Some(second)) => {
+                let more_than_one = MoreThanOne::new([first, second], self); 
+                AtMostOneResult::MoreThanOne(more_than_one)
             }
-            None => Ok(None),
         }
     }
 

tests/quick.rs

@@ -1300,12 +1300,32 @@ quickcheck! {
 
 quickcheck! {
     fn at_most_one_i32(a: Vec<i32>) -> TestResult {
-        let ret = a.iter().cloned().at_most_one();
-        match a.len() {
-            0 => TestResult::from_bool(ret.unwrap() == None),
-            1 => TestResult::from_bool(ret.unwrap() == Some(a[0])),
-            _ => TestResult::from_bool(ret.unwrap_err().eq(a.iter().cloned())),
-        }
+        use itertools::AtMostOneResult;
+
+        let iter = a.iter().copied();
+
+        let (first, second, tail) = {
+            let mut iter = iter.clone();
+            (iter.next(), iter.next(), iter)
+        };
+
+        let amo_result = iter.at_most_one();
+        let expected = match a.len() {
+            0 => matches!(amo_result, AtMostOneResult::Zero),
+            1 => matches!(amo_result, AtMostOneResult::One(n) if first == Some(n)),
+            _ => {
+                let AtMostOneResult::MoreThanOne(mut more_than_one) = amo_result else {
+                    return TestResult::failed();
+                };
+
+                let actual_first = more_than_one.next();
+                let actual_second = more_than_one.next();
+
+                first == actual_first && second == actual_second && more_than_one.eq(tail)
+            },
+        };
+
+        TestResult::from_bool(expected)
     }
 }