From 9dac4d69e0f142e9a2dc4b61ea49365a8cae3f4b Mon Sep 17 00:00:00 2001
From: yancy <github@yancy.lol>
Date: Fri, 23 May 2025 14:58:03 -0500
Subject: [PATCH] Implement CheckedSum for Weight iterator

Expose `checked_sum` method to sum an iterator of Weights checked.
---
 units/src/lib.rs    |  3 ++-
 units/src/weight.rs | 21 +++++++++++++++++++++
 units/tests/api.rs  |  1 +
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/units/src/lib.rs b/units/src/lib.rs
index e6b36a7dd..643120223 100644
--- a/units/src/lib.rs
+++ b/units/src/lib.rs
@@ -78,11 +78,12 @@ pub trait CheckedSum<R>: sealed::Sealed<R> {
 }
 
 mod sealed {
-    use super::{Amount, SignedAmount};
+    use super::{Amount, SignedAmount, Weight};
 
     /// Used to seal the `CheckedSum` trait
     pub trait Sealed<A> {}
 
     impl<T> Sealed<Amount> for T where T: Iterator<Item = Amount> {}
     impl<T> Sealed<SignedAmount> for T where T: Iterator<Item = SignedAmount> {}
+    impl<T> Sealed<Weight> for T where T: Iterator<Item = Weight> {}
 }
diff --git a/units/src/weight.rs b/units/src/weight.rs
index d801dd254..6649b4d64 100644
--- a/units/src/weight.rs
+++ b/units/src/weight.rs
@@ -10,6 +10,8 @@ use arbitrary::{Arbitrary, Unstructured};
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 
+use crate::CheckedSum;
+
 /// The factor that non-witness serialization data is multiplied by during weight calculation.
 pub const WITNESS_SCALE_FACTOR: usize = 4;
 
@@ -241,6 +243,13 @@ impl ops::RemAssign<u64> for Weight {
     fn rem_assign(&mut self, rhs: u64) { *self = Weight::from_wu(self.to_wu() % rhs); }
 }
 
+impl<T> CheckedSum<Weight> for T
+where
+    T: Iterator<Item = Weight>,
+{
+    fn checked_sum(mut self) -> Option<Weight> { self.try_fold(Weight::ZERO, Weight::checked_add) }
+}
+
 impl core::iter::Sum for Weight {
     fn sum<I>(iter: I) -> Self
     where
@@ -524,4 +533,16 @@ mod tests {
         weight %= 3;
         assert_eq!(weight, Weight::from_wu(1));
     }
+
+    #[test]
+    #[cfg(feature = "alloc")]
+    fn checked_sum_weights() {
+        assert_eq!([].into_iter().checked_sum(), Some(Weight::ZERO));
+
+        let sum = alloc::vec![0, 1, 2].iter().map(|&w| Weight::from_wu(w)).checked_sum().unwrap();
+        assert_eq!(sum, Weight::from_wu(3));
+
+        let sum = alloc::vec![1, u64::MAX].iter().map(|&w| Weight::from_wu(w)).checked_sum();
+        assert!(sum.is_none());
+    }
 }
diff --git a/units/tests/api.rs b/units/tests/api.rs
index cf5f5def4..c44025295 100644
--- a/units/tests/api.rs
+++ b/units/tests/api.rs
@@ -273,6 +273,7 @@ fn dyn_compatible() {
     // If this builds then traits are dyn compatible.
     struct Traits {
         a: Box<dyn CheckedSum<Amount>>,
+        b: Box<dyn CheckedSum<Weight>>,
         // These traits are explicitly not dyn compatible.
         // b: Box<dyn amount::serde::SerdeAmount>,
         // c: Box<dyn amount::serde::SerdeAmountForOpt>,