feat: add checked add/mul/sub

starknet-io · Dec 25, 2024 · df5532b · df5532b
1 parent 7ef3023
commit df5532b
Showing 1 changed file with 49 additions and 1 deletion.
diff --git a/crates/starknet-types-core/src/felt/num_traits_impl.rs b/crates/starknet-types-core/src/felt/num_traits_impl.rs
@@ -1,6 +1,8 @@
 use super::Felt;
 use num_bigint::{ToBigInt, ToBigUint};
-use num_traits::{FromPrimitive, Inv, One, Pow, ToPrimitive, Zero};
+use num_traits::{
+    CheckedAdd, CheckedMul, CheckedSub, FromPrimitive, Inv, One, Pow, ToPrimitive, Zero,
+};
 
 impl ToBigInt for Felt {
     /// Converts the value of `self` to a [`BigInt`].
@@ -121,6 +123,39 @@ impl Pow<u128> for Felt {
     }
 }
 
+impl CheckedAdd for Felt {
+    fn checked_add(&self, v: &Self) -> Option<Self> {
+        let res = self + v;
+        if res < *self {
+            None
+        } else {
+            Some(res)
+        }
+    }
+}
+
+impl CheckedMul for Felt {
+    fn checked_mul(&self, v: &Self) -> Option<Self> {
+        let res = self * v;
+        if res < *self {
+            None
+        } else {
+            Some(res)
+        }
+    }
+}
+
+impl CheckedSub for Felt {
+    fn checked_sub(&self, v: &Self) -> Option<Self> {
+        let res = self - v;
+        if res > *self {
+            None
+        } else {
+            Some(res)
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -139,4 +174,17 @@ mod tests {
     fn default_is_zero() {
         assert!(Felt::default().is_zero());
     }
+
+    #[test]
+    fn checked_ops() {
+        assert_eq!(Felt::ONE.checked_add(&Felt::TWO), Some(Felt::THREE));
+        assert!(Felt::MAX.checked_add(&Felt::ONE).is_none());
+        assert_eq!(
+            Felt::TWO.checked_mul(&Felt::THREE),
+            Some(Felt::from_hex_unchecked("0x6"))
+        );
+        assert!(Felt::MAX.checked_mul(&Felt::TWO).is_none());
+        assert_eq!(Felt::THREE.checked_sub(&Felt::TWO), Some(Felt::ONE));
+        assert!(Felt::ONE.checked_sub(&Felt::TWO).is_none());
+    }
 }