dev: optimized bitshifts by using a lookup table for powers of two

kkrt-labs · Sep 30, 2024 · 7c374e9 · 7c374e9
1 parent 9b490d5
commit 7c374e9
Show file tree

Hide file tree

Showing 2 changed files with 269 additions and 3 deletions.
diff --git a/crates/utils/src/constants.cairo b/crates/utils/src/constants.cairo
@@ -210,6 +210,267 @@ pub const POW_2: [
     0x80000000000000000000000000000000
 ];
 
+pub const POW_2_256: [
+    u256
+    ; 256] = [
+    0x1,
+    0x2,
+    0x4,
+    0x8,
+    0x10,
+    0x20,
+    0x40,
+    0x80,
+    0x100,
+    0x200,
+    0x400,
+    0x800,
+    0x1000,
+    0x2000,
+    0x4000,
+    0x8000,
+    0x10000,
+    0x20000,
+    0x40000,
+    0x80000,
+    0x100000,
+    0x200000,
+    0x400000,
+    0x800000,
+    0x1000000,
+    0x2000000,
+    0x4000000,
+    0x8000000,
+    0x10000000,
+    0x20000000,
+    0x40000000,
+    0x80000000,
+    0x100000000,
+    0x200000000,
+    0x400000000,
+    0x800000000,
+    0x1000000000,
+    0x2000000000,
+    0x4000000000,
+    0x8000000000,
+    0x10000000000,
+    0x20000000000,
+    0x40000000000,
+    0x80000000000,
+    0x100000000000,
+    0x200000000000,
+    0x400000000000,
+    0x800000000000,
+    0x1000000000000,
+    0x2000000000000,
+    0x4000000000000,
+    0x8000000000000,
+    0x10000000000000,
+    0x20000000000000,
+    0x40000000000000,
+    0x80000000000000,
+    0x100000000000000,
+    0x200000000000000,
+    0x400000000000000,
+    0x800000000000000,
+    0x1000000000000000,
+    0x2000000000000000,
+    0x4000000000000000,
+    0x8000000000000000,
+    0x10000000000000000,
+    0x20000000000000000,
+    0x40000000000000000,
+    0x80000000000000000,
+    0x100000000000000000,
+    0x200000000000000000,
+    0x400000000000000000,
+    0x800000000000000000,
+    0x1000000000000000000,
+    0x2000000000000000000,
+    0x4000000000000000000,
+    0x8000000000000000000,
+    0x10000000000000000000,
+    0x20000000000000000000,
+    0x40000000000000000000,
+    0x80000000000000000000,
+    0x100000000000000000000,
+    0x200000000000000000000,
+    0x400000000000000000000,
+    0x800000000000000000000,
+    0x1000000000000000000000,
+    0x2000000000000000000000,
+    0x4000000000000000000000,
+    0x8000000000000000000000,
+    0x10000000000000000000000,
+    0x20000000000000000000000,
+    0x40000000000000000000000,
+    0x80000000000000000000000,
+    0x100000000000000000000000,
+    0x200000000000000000000000,
+    0x400000000000000000000000,
+    0x800000000000000000000000,
+    0x1000000000000000000000000,
+    0x2000000000000000000000000,
+    0x4000000000000000000000000,
+    0x8000000000000000000000000,
+    0x10000000000000000000000000,
+    0x20000000000000000000000000,
+    0x40000000000000000000000000,
+    0x80000000000000000000000000,
+    0x100000000000000000000000000,
+    0x200000000000000000000000000,
+    0x400000000000000000000000000,
+    0x800000000000000000000000000,
+    0x1000000000000000000000000000,
+    0x2000000000000000000000000000,
+    0x4000000000000000000000000000,
+    0x8000000000000000000000000000,
+    0x10000000000000000000000000000,
+    0x20000000000000000000000000000,
+    0x40000000000000000000000000000,
+    0x80000000000000000000000000000,
+    0x100000000000000000000000000000,
+    0x200000000000000000000000000000,
+    0x400000000000000000000000000000,
+    0x800000000000000000000000000000,
+    0x1000000000000000000000000000000,
+    0x2000000000000000000000000000000,
+    0x4000000000000000000000000000000,
+    0x8000000000000000000000000000000,
+    0x10000000000000000000000000000000,
+    0x20000000000000000000000000000000,
+    0x40000000000000000000000000000000,
+    0x80000000000000000000000000000000,
+    0x100000000000000000000000000000000,
+    0x200000000000000000000000000000000,
+    0x400000000000000000000000000000000,
+    0x800000000000000000000000000000000,
+    0x1000000000000000000000000000000000,
+    0x2000000000000000000000000000000000,
+    0x4000000000000000000000000000000000,
+    0x8000000000000000000000000000000000,
+    0x10000000000000000000000000000000000,
+    0x20000000000000000000000000000000000,
+    0x40000000000000000000000000000000000,
+    0x80000000000000000000000000000000000,
+    0x100000000000000000000000000000000000,
+    0x200000000000000000000000000000000000,
+    0x400000000000000000000000000000000000,
+    0x800000000000000000000000000000000000,
+    0x1000000000000000000000000000000000000,
+    0x2000000000000000000000000000000000000,
+    0x4000000000000000000000000000000000000,
+    0x8000000000000000000000000000000000000,
+    0x10000000000000000000000000000000000000,
+    0x20000000000000000000000000000000000000,
+    0x40000000000000000000000000000000000000,
+    0x80000000000000000000000000000000000000,
+    0x100000000000000000000000000000000000000,
+    0x200000000000000000000000000000000000000,
+    0x400000000000000000000000000000000000000,
+    0x800000000000000000000000000000000000000,
+    0x1000000000000000000000000000000000000000,
+    0x2000000000000000000000000000000000000000,
+    0x4000000000000000000000000000000000000000,
+    0x8000000000000000000000000000000000000000,
+    0x10000000000000000000000000000000000000000,
+    0x20000000000000000000000000000000000000000,
+    0x40000000000000000000000000000000000000000,
+    0x80000000000000000000000000000000000000000,
+    0x100000000000000000000000000000000000000000,
+    0x200000000000000000000000000000000000000000,
+    0x400000000000000000000000000000000000000000,
+    0x800000000000000000000000000000000000000000,
+    0x1000000000000000000000000000000000000000000,
+    0x2000000000000000000000000000000000000000000,
+    0x4000000000000000000000000000000000000000000,
+    0x8000000000000000000000000000000000000000000,
+    0x10000000000000000000000000000000000000000000,
+    0x20000000000000000000000000000000000000000000,
+    0x40000000000000000000000000000000000000000000,
+    0x80000000000000000000000000000000000000000000,
+    0x100000000000000000000000000000000000000000000,
+    0x200000000000000000000000000000000000000000000,
+    0x400000000000000000000000000000000000000000000,
+    0x800000000000000000000000000000000000000000000,
+    0x1000000000000000000000000000000000000000000000,
+    0x2000000000000000000000000000000000000000000000,
+    0x4000000000000000000000000000000000000000000000,
+    0x8000000000000000000000000000000000000000000000,
+    0x10000000000000000000000000000000000000000000000,
+    0x20000000000000000000000000000000000000000000000,
+    0x40000000000000000000000000000000000000000000000,
+    0x80000000000000000000000000000000000000000000000,
+    0x100000000000000000000000000000000000000000000000,
+    0x200000000000000000000000000000000000000000000000,
+    0x400000000000000000000000000000000000000000000000,
+    0x800000000000000000000000000000000000000000000000,
+    0x1000000000000000000000000000000000000000000000000,
+    0x2000000000000000000000000000000000000000000000000,
+    0x4000000000000000000000000000000000000000000000000,
+    0x8000000000000000000000000000000000000000000000000,
+    0x10000000000000000000000000000000000000000000000000,
+    0x20000000000000000000000000000000000000000000000000,
+    0x40000000000000000000000000000000000000000000000000,
+    0x80000000000000000000000000000000000000000000000000,
+    0x100000000000000000000000000000000000000000000000000,
+    0x200000000000000000000000000000000000000000000000000,
+    0x400000000000000000000000000000000000000000000000000,
+    0x800000000000000000000000000000000000000000000000000,
+    0x1000000000000000000000000000000000000000000000000000,
+    0x2000000000000000000000000000000000000000000000000000,
+    0x4000000000000000000000000000000000000000000000000000,
+    0x8000000000000000000000000000000000000000000000000000,
+    0x10000000000000000000000000000000000000000000000000000,
+    0x20000000000000000000000000000000000000000000000000000,
+    0x40000000000000000000000000000000000000000000000000000,
+    0x80000000000000000000000000000000000000000000000000000,
+    0x100000000000000000000000000000000000000000000000000000,
+    0x200000000000000000000000000000000000000000000000000000,
+    0x400000000000000000000000000000000000000000000000000000,
+    0x800000000000000000000000000000000000000000000000000000,
+    0x1000000000000000000000000000000000000000000000000000000,
+    0x2000000000000000000000000000000000000000000000000000000,
+    0x4000000000000000000000000000000000000000000000000000000,
+    0x8000000000000000000000000000000000000000000000000000000,
+    0x10000000000000000000000000000000000000000000000000000000,
+    0x20000000000000000000000000000000000000000000000000000000,
+    0x40000000000000000000000000000000000000000000000000000000,
+    0x80000000000000000000000000000000000000000000000000000000,
+    0x100000000000000000000000000000000000000000000000000000000,
+    0x200000000000000000000000000000000000000000000000000000000,
+    0x400000000000000000000000000000000000000000000000000000000,
+    0x800000000000000000000000000000000000000000000000000000000,
+    0x1000000000000000000000000000000000000000000000000000000000,
+    0x2000000000000000000000000000000000000000000000000000000000,
+    0x4000000000000000000000000000000000000000000000000000000000,
+    0x8000000000000000000000000000000000000000000000000000000000,
+    0x10000000000000000000000000000000000000000000000000000000000,
+    0x20000000000000000000000000000000000000000000000000000000000,
+    0x40000000000000000000000000000000000000000000000000000000000,
+    0x80000000000000000000000000000000000000000000000000000000000,
+    0x100000000000000000000000000000000000000000000000000000000000,
+    0x200000000000000000000000000000000000000000000000000000000000,
+    0x400000000000000000000000000000000000000000000000000000000000,
+    0x800000000000000000000000000000000000000000000000000000000000,
+    0x1000000000000000000000000000000000000000000000000000000000000,
+    0x2000000000000000000000000000000000000000000000000000000000000,
+    0x4000000000000000000000000000000000000000000000000000000000000,
+    0x8000000000000000000000000000000000000000000000000000000000000,
+    0x10000000000000000000000000000000000000000000000000000000000000,
+    0x20000000000000000000000000000000000000000000000000000000000000,
+    0x40000000000000000000000000000000000000000000000000000000000000,
+    0x80000000000000000000000000000000000000000000000000000000000000,
+    0x100000000000000000000000000000000000000000000000000000000000000,
+    0x200000000000000000000000000000000000000000000000000000000000000,
+    0x400000000000000000000000000000000000000000000000000000000000000,
+    0x800000000000000000000000000000000000000000000000000000000000000,
+    0x1000000000000000000000000000000000000000000000000000000000000000,
+    0x2000000000000000000000000000000000000000000000000000000000000000,
+    0x4000000000000000000000000000000000000000000000000000000000000000,
+    0x8000000000000000000000000000000000000000000000000000000000000000,
+];
+
 pub const POW_2_0: u128 = 0x1;
 pub const POW_2_8: u128 = 0x100;
 pub const POW_2_16: u128 = 0x10000;

diff --git a/crates/utils/src/math.cairo b/crates/utils/src/math.cairo
@@ -2,6 +2,7 @@ use core::integer::{u512};
 use core::num::traits::{Zero, One, BitSize, OverflowingAdd, OverflowingMul, Bounded};
 use core::panic_with_felt252;
 use core::traits::{BitAnd};
+use utils::constants::POW_2_256;
 
 // === Exponentiation ===
 
@@ -236,6 +237,8 @@ impl BitshiftImpl<
     +PartialOrd<T>,
     +BitSize<T>,
     +TryInto<usize, T>,
+    +TryInto<T, usize>,
+    +TryInto<u256, T>,
 > of Bitshift<T> {
     fn shl(self: T, shift: usize) -> T {
         // if we shift by more than nb_bits of T, the result is 0
@@ -313,8 +316,10 @@ pub impl WrappingBitshiftImpl<
     +WrappingExponentiation<T>,
     +BitSize<T>,
     +Bounded<T>,
-    +Into<T, u256>,
     +TryInto<usize, T>,
+    +TryInto<T, usize>,
+    +TryInto<u256, T>,
+    +Into<T, u256>
 > of WrappingBitshift<T> {
     fn wrapping_shl(self: T, shift: usize) -> T {
         if shift <= BitSize::<u256>::bits() - One::<u32>::one() {
@@ -324,7 +329,7 @@ pub impl WrappingBitshiftImpl<
             return result;
         }
         let two = One::<T>::one() + One::<T>::one();
-        let (result, _) = self.overflowing_mul(two.wrapping_pow(shift));
+        let (result, _) = self.overflowing_mul(two.wrapping_pow(shift.try_into().unwrap()));
         result
     }
 
@@ -339,7 +344,7 @@ pub impl WrappingBitshiftImpl<
         if shift > BitSize::<T>::bits() - One::one() {
             return Zero::zero();
         }
-        self / two.pow(shift)
+        self / two.pow(shift.try_into().unwrap())
     }
 }