From 2cc205fd216ed821db9f717cb111b75cb622bad2 Mon Sep 17 00:00:00 2001
From: Ori Ziv <oriz@starkware.co>
Date: Tue, 19 Nov 2024 10:58:16 +0200
Subject: [PATCH] Added `Pow` trait and implementations for base numeric types
 and `usize`.

---
 corelib/src/num/traits.cairo         |  1 +
 corelib/src/num/traits/ops.cairo     |  1 +
 corelib/src/num/traits/ops/pow.cairo | 43 ++++++++++++++++++++++++++++
 corelib/src/test/num_test.cairo      | 29 ++++++++++++++++++-
 4 files changed, 73 insertions(+), 1 deletion(-)
 create mode 100644 corelib/src/num/traits/ops/pow.cairo

diff --git a/corelib/src/num/traits.cairo b/corelib/src/num/traits.cairo
index 8a202b41a9f..4e797bb423e 100644
--- a/corelib/src/num/traits.cairo
+++ b/corelib/src/num/traits.cairo
@@ -15,6 +15,7 @@ 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::pow::Pow;
 pub use ops::saturating::{SaturatingAdd, SaturatingSub, SaturatingMul};
 pub use ops::widemul::WideMul;
 pub use ops::widesquare::WideSquare;
diff --git a/corelib/src/num/traits/ops.cairo b/corelib/src/num/traits/ops.cairo
index eda726a270c..9c9016b0eb2 100644
--- a/corelib/src/num/traits/ops.cairo
+++ b/corelib/src/num/traits/ops.cairo
@@ -2,6 +2,7 @@ pub mod overflowing;
 pub mod wrapping;
 pub mod checked;
 pub mod saturating;
+pub mod pow;
 pub(crate) mod sqrt;
 pub(crate) mod widemul;
 pub(crate) mod widesquare;
diff --git a/corelib/src/num/traits/ops/pow.cairo b/corelib/src/num/traits/ops/pow.cairo
new file mode 100644
index 00000000000..668ddb43a55
--- /dev/null
+++ b/corelib/src/num/traits/ops/pow.cairo
@@ -0,0 +1,43 @@
+/// A trait for calculating a base to the power of an exponent.
+pub trait Pow<Base, Exp> {
+    type Output;
+
+    /// Returns `self` to the power `exp`.
+    fn pow(self: Base, exp: Exp) -> Self::Output;
+}
+
+mod mul_based {
+    /// Square and multiply implementation for `Pow`.
+    pub impl PowByMul<
+        Base, +Mul<Base>, +Copy<Base>, +Drop<Base>, +core::num::traits::One<Base>,
+    > of super::Pow<Base, usize> {
+        type Output = Base;
+
+        fn pow(mut self: Base, mut exp: usize) -> Base {
+            let (tail_exp, head_exp) = DivRem::div_rem(exp, 2);
+            let tail_result = if tail_exp == 0 {
+                core::num::traits::One::one()
+            } else {
+                Self::pow(self * self, tail_exp)
+            };
+            if head_exp == 0 {
+                tail_result
+            } else {
+                tail_result * self
+            }
+        }
+    }
+}
+
+impl PowFelt252 = mul_based::PowByMul<felt252>;
+impl PowI8 = mul_based::PowByMul<i8>;
+impl PowU8 = mul_based::PowByMul<u8>;
+impl PowI16 = mul_based::PowByMul<i16>;
+impl PowU16 = mul_based::PowByMul<u16>;
+impl PowI32 = mul_based::PowByMul<i32>;
+impl PowU32 = mul_based::PowByMul<u32>;
+impl PowI64 = mul_based::PowByMul<i64>;
+impl PowU64 = mul_based::PowByMul<u64>;
+impl PowI128 = mul_based::PowByMul<i128>;
+impl PowU128 = mul_based::PowByMul<u128>;
+impl PowU256 = mul_based::PowByMul<u256>;
diff --git a/corelib/src/test/num_test.cairo b/corelib/src/test/num_test.cairo
index e786a68555e..4d4eb3548f6 100644
--- a/corelib/src/test/num_test.cairo
+++ b/corelib/src/test/num_test.cairo
@@ -1,7 +1,7 @@
 use crate::num::traits::BitSize;
 use crate::num::traits::{
     OverflowingAdd, OverflowingSub, OverflowingMul, WrappingAdd, WrappingSub, WrappingMul,
-    CheckedAdd, CheckedSub, CheckedMul, SaturatingAdd, SaturatingSub, SaturatingMul,
+    CheckedAdd, CheckedSub, CheckedMul, SaturatingAdd, SaturatingSub, SaturatingMul, Pow,
 };
 use crate::num::traits::Bounded;
 
@@ -441,3 +441,30 @@ fn test_saturating_mul_unsigned_integers() {
     assert_eq!(2_u256.saturating_mul(3), 6);
     assert_eq!(Bounded::<u256>::MAX.saturating_mul(2), Bounded::<u256>::MAX);
 }
+
+#[test]
+fn test_pow() {
+    assert_eq!((-2_i8).pow(0), 1);
+    assert_eq!((-2_i8).pow(1), -2);
+    assert_eq!((-2_i8).pow(2), 4);
+    assert_eq!((-2_i8).pow(3), -8);
+    assert_eq!((-2_i8).pow(4), 16);
+    assert_eq!((-2_i8).pow(5), -32);
+    assert_eq!((-2_i8).pow(6), 64);
+
+    assert_eq!(0.pow(0), 1);
+    assert_eq!(0.pow(1), 0);
+    assert_eq!(0.pow(2), 0);
+
+    assert_eq!(2.pow(0), 0b1);
+    assert_eq!(2.pow(1), 0b10);
+    assert_eq!(2.pow(2), 0b100);
+    assert_eq!(2.pow(3), 0b1000);
+    assert_eq!(2.pow(4), 0b10000);
+    assert_eq!(2.pow(5), 0b100000);
+    assert_eq!(2.pow(6), 0b1000000);
+    assert_eq!(2.pow(7), 0b10000000);
+    assert_eq!(2.pow(8), 0b100000000);
+    assert_eq!(2.pow(9), 0b1000000000);
+    assert_eq!(2.pow(10), 0b10000000000);
+}