diff --git a/src/libcore/iter/traits/accum.rs b/src/libcore/iter/traits/accum.rs index 7815fe9c59d08..4eac5cbc8e662 100644 --- a/src/libcore/iter/traits/accum.rs +++ b/src/libcore/iter/traits/accum.rs @@ -223,3 +223,113 @@ impl Product> for Result ResultShunt::process(iter, |i| i.product()) } } + +/// An iterator adapter that produces output as long as the underlying +/// iterator produces `Option::Some` values. +struct OptionShunt { + iter: I, + exited_early: bool, +} + +impl OptionShunt +where + I: Iterator>, +{ + /// Process the given iterator as if it yielded a `T` instead of a + /// `Option`. Any `None` value will stop the inner iterator and + /// the overall result will be a `None`. + pub fn process(iter: I, mut f: F) -> Option + 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(self, val: U) -> Option { + if self.exited_early { + None + } else { + Some(val) + } + } +} + +impl Iterator for OptionShunt +where + I: Iterator>, +{ + type Item = T; + + fn next(&mut self) -> Option { + match self.iter.next() { + Some(Some(v)) => Some(v), + Some(None) => { + self.exited_early = true; + None + } + None => None, + } + } + + fn size_hint(&self) -> (usize, Option) { + if self.exited_early { + (0, Some(0)) + } else { + let (_, upper) = self.iter.size_hint(); + (0, upper) + } + } +} + +#[stable(feature = "iter_arith_traits_option", since = "1.37.0")] +impl Sum> for Option +where + T: Sum, +{ + /// 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 the position of the character 'a' in a vector of strings, + /// if a word did not have the character 'a' the operation returns `None`: + /// + /// ``` + /// let words = vec!["have", "a", "great", "day"]; + /// let total: Option = words.iter().map(|w| w.find('a')).sum(); + /// assert_eq!(total, Some(5)); + /// ``` + fn sum(iter: I) -> Option + where + I: Iterator>, + { + OptionShunt::process(iter, |i| i.sum()) + } +} + +#[stable(feature = "iter_arith_traits_option", since = "1.37.0")] +impl Product> for Option +where + T: Product, +{ + /// 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(iter: I) -> Option + where + I: Iterator>, + { + OptionShunt::process(iter, |i| i.product()) + } +} diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs index 7dfb1adad9eed..bedb9e756129c 100644 --- a/src/libcore/tests/iter.rs +++ b/src/libcore/tests/iter.rs @@ -1084,6 +1084,14 @@ fn test_iterator_sum_result() { assert_eq!(v.iter().cloned().sum::>(), Err(())); } +#[test] +fn test_iterator_sum_option() { + let v: &[Option] = &[Some(1), Some(2), Some(3), Some(4)]; + assert_eq!(v.iter().cloned().sum::>(), Some(10)); + let v: &[Option] = &[Some(1), None, Some(3), Some(4)]; + assert_eq!(v.iter().cloned().sum::>(), None); +} + #[test] fn test_iterator_product() { let v: &[i32] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; @@ -1126,6 +1134,14 @@ impl Ord for Mod3 { } } +#[test] +fn test_iterator_product_option() { + let v: &[Option] = &[Some(1), Some(2), Some(3), Some(4)]; + assert_eq!(v.iter().cloned().product::>(), Some(24)); + let v: &[Option] = &[Some(1), None, Some(3), Some(4)]; + assert_eq!(v.iter().cloned().product::>(), None); +} + #[test] fn test_iterator_max() { let v: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];