Rollup merge of #144275 - Qelxiros:saturating-arithmetic, r=tgross35

implement Sum and Product for Saturating(u*)

ACP: rust-lang/libs-team#604

`@rustbot` label +needs-fcp

2 files changed, 93 insertions, 1 deletions

diff --git a/library/core/src/iter/traits/accum.rs b/library/core/src/iter/traits/accum.rs
index 12e2b8b393a..73122369b41 100644
--- a/library/core/src/iter/traits/accum.rs
+++ b/library/core/src/iter/traits/accum.rs
@@ -1,5 +1,5 @@
 use crate::iter;
-use crate::num::Wrapping;
+use crate::num::{Saturating, Wrapping};
 
 /// Trait to represent types that can be created by summing up an iterator.
 ///
@@ -98,6 +98,61 @@ macro_rules! integer_sum_product {
     );
 }
 
+macro_rules! saturating_integer_sum_product {
+    (@impls $zero:expr, $one:expr, $doc:expr, #[$attr:meta], $($a:ty)*) => ($(
+        #[$attr]
+        #[doc = $doc]
+        impl Sum for $a {
+            fn sum<I: Iterator<Item=Self>>(iter: I) -> Self {
+                iter.fold(
+                    $zero,
+                    |a, b| a + b,
+                )
+            }
+        }
+
+        #[$attr]
+        #[doc = $doc]
+        impl Product for $a {
+            fn product<I: Iterator<Item=Self>>(iter: I) -> Self {
+                iter.fold(
+                    $one,
+                    |a, b| a * b,
+                )
+            }
+        }
+
+        #[$attr]
+        #[doc = $doc]
+        impl<'a> Sum<&'a $a> for $a {
+            fn sum<I: Iterator<Item=&'a Self>>(iter: I) -> Self {
+                iter.fold(
+                    $zero,
+                    |a, b| a + b,
+                )
+            }
+        }
+
+        #[$attr]
+        #[doc = $doc]
+        impl<'a> Product<&'a $a> for $a {
+            fn product<I: Iterator<Item=&'a Self>>(iter: I) -> Self {
+                iter.fold(
+                    $one,
+                    |a, b| a * b,
+                )
+            }
+        }
+    )*);
+    ($($a:ty)*) => (
+        saturating_integer_sum_product!(@impls Saturating(0), Saturating(1),
+                "The short-circuiting behavior of this implementation is unspecified. If you care about \
+                short-circuiting, use [`Iterator::fold`] directly.",
+                #[stable(feature = "saturating_iter_arith", since = "CURRENT_RUSTC_VERSION")],
+                $(Saturating<$a>)*);
+    );
+}
+
 macro_rules! float_sum_product {
     ($($a:ident)*) => ($(
         #[stable(feature = "iter_arith_traits", since = "1.12.0")]
@@ -147,6 +202,7 @@ macro_rules! float_sum_product {
 }
 
 integer_sum_product! { i8 i16 i32 i64 i128 isize u8 u16 u32 u64 u128 usize }
+saturating_integer_sum_product! { u8 u16 u32 u64 u128 usize }
 float_sum_product! { f32 f64 }
 
 #[stable(feature = "iter_arith_traits_result", since = "1.16.0")]
diff --git a/library/coretests/tests/iter/traits/accum.rs b/library/coretests/tests/iter/traits/accum.rs
index f3eeb31fe58..95f299d2680 100644
--- a/library/coretests/tests/iter/traits/accum.rs
+++ b/library/coretests/tests/iter/traits/accum.rs
@@ -1,4 +1,5 @@
 use core::iter::*;
+use std::num::Saturating;
 
 #[test]
 fn test_iterator_sum() {
@@ -64,3 +65,38 @@ fn test_iterator_product_option() {
     let v: &[Option<i32>] = &[Some(1), None, Some(3), Some(4)];
     assert_eq!(v.iter().cloned().product::<Option<i32>>(), None);
 }
+
+#[test]
+fn test_saturating_sum_product() {
+    let v = (1u32..=10).map(|i| Saturating(i));
+    assert_eq!(v.sum::<Saturating<u32>>(), Saturating(55));
+    let v = (1u32..=10).map(|i| Saturating(i));
+    assert_eq!(v.product::<Saturating<u32>>(), Saturating(3628800));
+    let v = [Saturating(usize::MAX), Saturating(2)];
+    assert_eq!(v.iter().copied().sum::<Saturating<usize>>(), Saturating(usize::MAX));
+    assert_eq!(v.iter().copied().product::<Saturating<usize>>(), Saturating(usize::MAX));
+
+    let mut cnt = 0;
+    let v = 250..=u8::MAX;
+    assert_eq!(
+        v.map(|i| {
+            cnt += 1;
+            Saturating(i)
+        })
+        .sum::<Saturating<u8>>(),
+        Saturating(u8::MAX)
+    );
+    assert_eq!(cnt, 6); // no short-circuiting
+
+    let mut cnt = 0;
+    let v = (250..=u8::MAX).chain(0..5);
+    assert_eq!(
+        v.map(|i| {
+            cnt += 1;
+            Saturating(i)
+        })
+        .product::<Saturating<u8>>(),
+        Saturating(0)
+    );
+    assert_eq!(cnt, 11); // no short-circuiting
+}