feat: saturating traits (#5621)

starkware-libs · May 23, 2024 · 75d3994 · 75d3994
1 parent 99b739a
commit 75d3994
Show file tree

Hide file tree

Showing 5 changed files with 173 additions and 1 deletion.
diff --git a/corelib/src/integer.cairo b/corelib/src/integer.cairo
@@ -3135,6 +3135,42 @@ impl U64CheckedMul = core::num::traits::ops::checked::overflow_based::TCheckedMu
 impl U128CheckedMul = core::num::traits::ops::checked::overflow_based::TCheckedMul<u128>;
 impl U256CheckedMul = core::num::traits::ops::checked::overflow_based::TCheckedMul<u256>;
 
+
+// SaturatingAdd implementations
+impl U8SaturatingAdd = core::num::traits::ops::saturating::overflow_based::TSaturatingAdd<u8>;
+impl U16SaturatingAdd = core::num::traits::ops::saturating::overflow_based::TSaturatingAdd<u16>;
+impl U32SaturatingAdd = core::num::traits::ops::saturating::overflow_based::TSaturatingAdd<u32>;
+impl U64SaturatingAdd = core::num::traits::ops::saturating::overflow_based::TSaturatingAdd<u64>;
+impl U128SaturatingAdd = core::num::traits::ops::saturating::overflow_based::TSaturatingAdd<u128>;
+impl U256SaturatingAdd = core::num::traits::ops::saturating::overflow_based::TSaturatingAdd<u256>;
+impl I8SaturatingAdd = core::num::traits::ops::saturating::overflow_based::TSaturatingAdd<i8>;
+impl I16SaturatingAdd = core::num::traits::ops::saturating::overflow_based::TSaturatingAdd<i16>;
+impl I32SaturatingAdd = core::num::traits::ops::saturating::overflow_based::TSaturatingAdd<i32>;
+impl I64SaturatingAdd = core::num::traits::ops::saturating::overflow_based::TSaturatingAdd<i64>;
+impl I128SaturatingAdd = core::num::traits::ops::saturating::overflow_based::TSaturatingAdd<i128>;
+
+// SaturatingSub implementations
+impl U8SaturatingSub = core::num::traits::ops::saturating::overflow_based::TSaturatingSub<u8>;
+impl U16SaturatingSub = core::num::traits::ops::saturating::overflow_based::TSaturatingSub<u16>;
+impl U32SaturatingSub = core::num::traits::ops::saturating::overflow_based::TSaturatingSub<u32>;
+impl U64SaturatingSub = core::num::traits::ops::saturating::overflow_based::TSaturatingSub<u64>;
+impl U128SaturatingSub = core::num::traits::ops::saturating::overflow_based::TSaturatingSub<u128>;
+impl U256SaturatingSub = core::num::traits::ops::saturating::overflow_based::TSaturatingSub<u256>;
+impl I8SaturatingSub = core::num::traits::ops::saturating::overflow_based::TSaturatingSub<i8>;
+impl I16SaturatingSub = core::num::traits::ops::saturating::overflow_based::TSaturatingSub<i16>;
+impl I32SaturatingSub = core::num::traits::ops::saturating::overflow_based::TSaturatingSub<i32>;
+impl I64SaturatingSub = core::num::traits::ops::saturating::overflow_based::TSaturatingSub<i64>;
+impl I128SaturatingSub = core::num::traits::ops::saturating::overflow_based::TSaturatingSub<i128>;
+
+// SaturatingMul implementations
+impl U8SaturatingMul = core::num::traits::ops::saturating::overflow_based::TSaturatingMul<u8>;
+impl U16SaturatingMul = core::num::traits::ops::saturating::overflow_based::TSaturatingMul<u16>;
+impl U32SaturatingMul = core::num::traits::ops::saturating::overflow_based::TSaturatingMul<u32>;
+impl U64SaturatingMul = core::num::traits::ops::saturating::overflow_based::TSaturatingMul<u64>;
+impl U128SaturatingMul = core::num::traits::ops::saturating::overflow_based::TSaturatingMul<u128>;
+impl U256SaturatingMul = core::num::traits::ops::saturating::overflow_based::TSaturatingMul<u256>;
+
+
 /// Internal trait for easier finding of absolute values.
 pub(crate) trait AbsAndSign<Signed, Unsigned> {
     /// Returns the absolute value of the `Signed` value as `Unsigned` and the original sign.

diff --git a/corelib/src/num/traits.cairo b/corelib/src/num/traits.cairo
@@ -11,3 +11,4 @@ pub mod ops;
 pub use ops::overflowing::{OverflowingAdd, OverflowingSub, OverflowingMul};
 pub use ops::wrapping::{WrappingAdd, WrappingSub, WrappingMul};
 pub use ops::checked::{CheckedAdd, CheckedSub, CheckedMul};
+pub use ops::saturating::{SaturatingAdd, SaturatingSub, SaturatingMul};
diff --git a/corelib/src/num/traits/ops.cairo b/corelib/src/num/traits/ops.cairo
@@ -1,3 +1,4 @@
 pub mod overflowing;
 pub mod wrapping;
 pub mod checked;
+pub mod saturating;
diff --git a/corelib/src/num/traits/ops/saturating.cairo b/corelib/src/num/traits/ops/saturating.cairo
@@ -0,0 +1,58 @@
+/// Performs addition that saturates at the numeric bounds instead of overflowing.
+pub trait SaturatingAdd<T> {
+    /// Saturating addition. Computes `self + other`, saturating at the relevant high or low
+    /// boundary of the type.
+    fn saturating_add(self: T, other: T) -> T;
+}
+
+/// Performs subtraction that saturates at the numeric bounds instead of overflowing.
+pub trait SaturatingSub<T> {
+    /// Saturating subtraction. Computes `self - other`, saturating at the relevant high or low
+    /// boundary of the type.
+    fn saturating_sub(self: T, other: T) -> T;
+}
+
+/// Performs multiplication that saturates at the numeric bounds instead of overflowing.
+pub trait SaturatingMul<T> {
+    /// Saturating multiplication. Computes `self * other`, saturating at the relevant high or low
+    /// boundary of the type.
+    fn saturating_mul(self: T, other: T) -> T;
+}
+
+pub(crate) mod overflow_based {
+    pub(crate) impl TSaturatingAdd<
+        T, +Drop<T>, +core::num::traits::OverflowingAdd<T>, +core::integer::BoundedInt<T>
+    > of core::num::traits::SaturatingAdd<T> {
+        fn saturating_add(self: T, other: T) -> T {
+            let (result, overflow) = self.overflowing_add(other);
+            match overflow {
+                true => core::integer::BoundedInt::max(),
+                false => result,
+            }
+        }
+    }
+
+    pub(crate) impl TSaturatingSub<
+        T, +Drop<T>, +core::num::traits::OverflowingSub<T>, +core::integer::BoundedInt<T>
+    > of core::num::traits::SaturatingSub<T> {
+        fn saturating_sub(self: T, other: T) -> T {
+            let (result, overflow) = self.overflowing_sub(other);
+            match overflow {
+                true => core::integer::BoundedInt::min(),
+                false => result,
+            }
+        }
+    }
+
+    pub(crate) impl TSaturatingMul<
+        T, +Drop<T>, +core::num::traits::OverflowingMul<T>, +core::integer::BoundedInt<T>
+    > of core::num::traits::SaturatingMul<T> {
+        fn saturating_mul(self: T, other: T) -> T {
+            let (result, overflow) = self.overflowing_mul(other);
+            match overflow {
+                true => core::integer::BoundedInt::max(),
+                false => result,
+            }
+        }
+    }
+}
diff --git a/corelib/src/test/num_test.cairo b/corelib/src/test/num_test.cairo
@@ -1,7 +1,7 @@
 use core::num::traits::BitSize;
 use core::num::traits::{
     OverflowingAdd, OverflowingSub, OverflowingMul, WrappingAdd, WrappingSub, WrappingMul,
-    CheckedAdd, CheckedSub, CheckedMul
+    CheckedAdd, CheckedSub, CheckedMul, SaturatingAdd, SaturatingSub, SaturatingMul
 };
 use core::integer::BoundedInt;
 
@@ -341,3 +341,79 @@ fn test_checked_mul_unsigned_integers() {
     assert_eq!(2_u256.checked_mul(3), Option::Some(6));
     assert!(BoundedInt::<u256>::max().checked_mul(2).is_none());
 }
+
+#[test]
+fn test_saturating_add_unsigned_integers() {
+    assert_eq!(1_u8.saturating_add(2), 3);
+    assert_eq!(BoundedInt::<u8>::max().saturating_add(1), BoundedInt::<u8>::max());
+    assert_eq!(1_u16.saturating_add(2), 3);
+    assert_eq!(BoundedInt::<u16>::max().saturating_add(1), BoundedInt::<u16>::max());
+    assert_eq!(1_u32.saturating_add(2), 3);
+    assert_eq!(BoundedInt::<u32>::max().saturating_add(1), BoundedInt::<u32>::max());
+    assert_eq!(1_u64.saturating_add(2), 3);
+    assert_eq!(BoundedInt::<u64>::max().saturating_add(1), BoundedInt::<u64>::max());
+    assert_eq!(1_u128.saturating_add(2), 3);
+    assert_eq!(BoundedInt::<u128>::max().saturating_add(1), BoundedInt::<u128>::max());
+    assert_eq!(1_u256.saturating_add(2), 3);
+    assert_eq!(BoundedInt::<u256>::max().saturating_add(1), BoundedInt::<u256>::max());
+}
+
+#[test]
+fn test_saturating_add_signed_integers() {
+    assert_eq!(1_i8.saturating_add(2), 3);
+    assert_eq!(BoundedInt::<i8>::max().saturating_add(1), BoundedInt::<i8>::max());
+    assert_eq!(1_i16.saturating_add(2), 3);
+    assert_eq!(BoundedInt::<i16>::max().saturating_add(1), BoundedInt::<i16>::max());
+    assert_eq!(1_i32.saturating_add(2), 3);
+    assert_eq!(BoundedInt::<i32>::max().saturating_add(1), BoundedInt::<i32>::max());
+    assert_eq!(1_i64.saturating_add(2), 3);
+    assert_eq!(BoundedInt::<i64>::max().saturating_add(1), BoundedInt::<i64>::max());
+    assert_eq!(1_i128.saturating_add(2), 3);
+    assert_eq!(BoundedInt::<i128>::max().saturating_add(1), BoundedInt::<i128>::max());
+}
+
+#[test]
+fn test_saturating_sub_unsigned_integers() {
+    assert_eq!(3_u8.saturating_sub(2), 1);
+    assert_eq!(0_u8.saturating_sub(1), 0);
+    assert_eq!(3_u16.saturating_sub(2), 1);
+    assert_eq!(0_u16.saturating_sub(1), 0);
+    assert_eq!(3_u32.saturating_sub(2), 1);
+    assert_eq!(0_u32.saturating_sub(1), 0);
+    assert_eq!(3_u64.saturating_sub(2), 1);
+    assert_eq!(0_u64.saturating_sub(1), 0);
+    assert_eq!(3_u128.saturating_sub(2), 1);
+    assert_eq!(0_u128.saturating_sub(1), 0);
+    assert_eq!(3_u256.saturating_sub(2), 1);
+    assert_eq!(0_u256.saturating_sub(1), 0);
+}
+
+#[test]
+fn test_saturating_sub_signed_integers() {
+    assert_eq!(3_i8.saturating_sub(2), 1);
+    assert_eq!(BoundedInt::<i8>::min().saturating_sub(1), BoundedInt::<i8>::min());
+    assert_eq!(3_i16.saturating_sub(2), 1);
+    assert_eq!(BoundedInt::<i16>::min().saturating_sub(1), BoundedInt::<i16>::min());
+    assert_eq!(3_i32.saturating_sub(2), 1);
+    assert_eq!(BoundedInt::<i32>::min().saturating_sub(1), BoundedInt::<i32>::min());
+    assert_eq!(3_i64.saturating_sub(2), 1);
+    assert_eq!(BoundedInt::<i64>::min().saturating_sub(1), BoundedInt::<i64>::min());
+    assert_eq!(3_i128.saturating_sub(2), 1);
+    assert_eq!(BoundedInt::<i128>::min().saturating_sub(1), BoundedInt::<i128>::min());
+}
+
+#[test]
+fn test_saturating_mul_unsigned_integers() {
+    assert_eq!(2_u8.saturating_mul(3), 6);
+    assert_eq!(BoundedInt::<u8>::max().saturating_mul(2), BoundedInt::<u8>::max());
+    assert_eq!(2_u16.saturating_mul(3), 6);
+    assert_eq!(BoundedInt::<u16>::max().saturating_mul(2), BoundedInt::<u16>::max());
+    assert_eq!(2_u32.saturating_mul(3), 6);
+    assert_eq!(BoundedInt::<u32>::max().saturating_mul(2), BoundedInt::<u32>::max());
+    assert_eq!(2_u64.saturating_mul(3), 6);
+    assert_eq!(BoundedInt::<u64>::max().saturating_mul(2), BoundedInt::<u64>::max());
+    assert_eq!(2_u128.saturating_mul(3), 6);
+    assert_eq!(BoundedInt::<u128>::max().saturating_mul(2), BoundedInt::<u128>::max());
+    assert_eq!(2_u256.saturating_mul(3), 6);
+    assert_eq!(BoundedInt::<u256>::max().saturating_mul(2), BoundedInt::<u256>::max());
+}