Implement `iter::Sum` and `iter::Product` for `Option`

This is similar to the existing implementation for `Result`. It will
take each item into the accumulator unless a `None` is returned.
This commit is contained in:
John Downey 2019-03-06 12:17:26 -06:00
parent f22dca0a1b
commit a35cdd4e41
2 changed files with 129 additions and 0 deletions

View File

@ -223,3 +223,116 @@ impl<T, U, E> Product<Result<U, E>> for Result<T, E>
ResultShunt::process(iter, |i| i.product())
}
}
/// An iterator adapter that produces output as long as the underlying
/// iterator produces `Option::Some` values.
struct OptionShunt<I> {
iter: I,
exited_early: bool,
}
impl<I, T> OptionShunt<I>
where
I: Iterator<Item = Option<T>>,
{
/// Process the given iterator as if it yielded a `T` instead of a
/// `Option<T>`. Any `None` value will stop the inner iterator and
/// the overall result will be a `None`.
pub fn process<F, U>(iter: I, mut f: F) -> Option<U>
where
F: FnMut(&mut Self) -> U,
{
let mut shunt = OptionShunt::new(iter);
let value = f(shunt.by_ref());
shunt.reconstruct(value)
}
fn new(iter: I) -> Self {
OptionShunt {
iter,
exited_early: false,
}
}
/// Consume the adapter and rebuild a `Option` value.
fn reconstruct<U>(self, val: U) -> Option<U> {
if self.exited_early {
None
} else {
Some(val)
}
}
}
impl<I, T> Iterator for OptionShunt<I>
where
I: Iterator<Item = Option<T>>,
{
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
match self.iter.next() {
Some(Some(v)) => Some(v),
Some(None) => {
self.exited_early = true;
None
}
None => None,
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
if self.exited_early {
(0, Some(0))
} else {
let (_, upper) = self.iter.size_hint();
(0, upper)
}
}
}
#[stable(feature = "iter_arith_traits_option", since = "1.34.0")]
impl<T, U> Sum<Option<U>> for Option<T>
where
T: Sum<U>,
{
/// Takes each element in the `Iterator`: if it is a `None`, no further
/// elements are taken, and the `None` is returned. Should no `None` occur,
/// the sum of all elements is returned.
///
/// # Examples
///
/// This sums up every integer in a vector, rejecting the sum if a negative
/// element is encountered:
///
/// ```
/// let v = vec![1, 2];
/// let res: Option<i32> = v.iter().map(|&x: &i32|
/// if x < 0 { None }
/// else { Some(x) }
/// ).sum();
/// assert_eq!(res, Some(3));
/// ```
fn sum<I>(iter: I) -> Option<T>
where
I: Iterator<Item = Option<U>>,
{
OptionShunt::process(iter, |i| i.sum())
}
}
#[stable(feature = "iter_arith_traits_option", since = "1.34.0")]
impl<T, U> Product<Option<U>> for Option<T>
where
T: Product<U>,
{
/// Takes each element in the `Iterator`: if it is a `None`, no further
/// elements are taken, and the `None` is returned. Should no `None` occur,
/// the product of all elements is returned.
fn product<I>(iter: I) -> Option<T>
where
I: Iterator<Item = Option<U>>,
{
OptionShunt::process(iter, |i| i.product())
}
}

View File

@ -1066,6 +1066,14 @@ fn test_iterator_sum_result() {
assert_eq!(v.iter().cloned().sum::<Result<i32, _>>(), Err(()));
}
#[test]
fn test_iterator_sum_option() {
let v: &[Option<i32>] = &[Some(1), Some(2), Some(3), Some(4)];
assert_eq!(v.iter().cloned().sum::<Option<i32>>(), Some(10));
let v: &[Option<i32>] = &[Some(1), None, Some(3), Some(4)];
assert_eq!(v.iter().cloned().sum::<Option<i32>>(), None);
}
#[test]
fn test_iterator_product() {
let v: &[i32] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
@ -1082,6 +1090,14 @@ fn test_iterator_product_result() {
assert_eq!(v.iter().cloned().product::<Result<i32, _>>(), Err(()));
}
#[test]
fn test_iterator_product_option() {
let v: &[Option<i32>] = &[Some(1), Some(2), Some(3), Some(4)];
assert_eq!(v.iter().cloned().product::<Option<i32>>(), Some(24));
let v: &[Option<i32>] = &[Some(1), None, Some(3), Some(4)];
assert_eq!(v.iter().cloned().product::<Option<i32>>(), None);
}
#[test]
fn test_iterator_max() {
let v: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];