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.

2 files changed, 129 insertions, 0 deletions

diff --git a/src/libcore/iter/traits/accum.rs b/src/libcore/iter/traits/accum.rs
index dfe1d2a1006..b67714ca779 100644
--- a/src/libcore/iter/traits/accum.rs
+++ b/src/libcore/iter/traits/accum.rs
@@ -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())
+    }
+}
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index d880abb181c..814f4ec4ca5 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -1067,6 +1067,14 @@ fn test_iterator_sum_result() {
 }
 
 #[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];
     assert_eq!(v[..4].iter().cloned().product::<i32>(), 0);
@@ -1083,6 +1091,14 @@ fn test_iterator_product_result() {
 }
 
 #[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];
     assert_eq!(v[..4].iter().cloned().max(), Some(3));