Merge pull request #47 from TAdev0/generic_bitshift

[dev] make utils::{shr, shl} generic over input types

src/utils.cairo

@@ -1,34 +1,109 @@
-use core::traits::Into;
-use core::traits::TryInto;
+use core::num::traits::{Zero, One, BitSize};
+use core::starknet::secp256_trait::Secp256PointTrait;
 
-// Bitwise shift left for u256
-pub fn shl(value: u256, shift: u32) -> u256 {
-    value * fast_pow(2.into(), shift.into())
+pub fn shl<
+    T,
+    U,
+    +Zero<T>,
+    +Zero<U>,
+    +One<T>,
+    +One<U>,
+    +Add<T>,
+    +Add<U>,
+    +Sub<U>,
+    +Mul<T>,
+    +Div<U>,
+    +Rem<U>,
+    +Copy<T>,
+    +Copy<U>,
+    +Drop<T>,
+    +Drop<U>,
+    +PartialOrd<U>,
+    +PartialEq<U>,
+    +BitSize<T>,
+    +Into<usize, U>
+>(
+    self: T, shift: U,
+) -> T {
+    if shift > BitSize::<T>::bits().into() - One::one() {
+        return Zero::zero();
+    }
+    let two = One::one() + One::one();
+    self * fast_pow(two, shift)
 }
 
-// Bitwise shift right for u256
-pub fn shr(value: u256, shift: u32) -> u256 {
-    value / fast_pow(2.into(), shift.into())
+pub fn shr<
+    T,
+    U,
+    +Zero<T>,
+    +Zero<U>,
+    +One<T>,
+    +One<U>,
+    +Add<T>,
+    +Add<U>,
+    +Sub<U>,
+    +Div<T>,
+    +Mul<T>,
+    +Div<U>,
+    +Rem<U>,
+    +Copy<T>,
+    +Copy<U>,
+    +Drop<T>,
+    +Drop<U>,
+    +PartialOrd<U>,
+    +PartialEq<U>,
+    +BitSize<T>,
+    +Into<usize, U>
+>(
+    self: T, shift: U
+) -> T {
+    if shift > BitSize::<T>::bits().try_into().unwrap() - One::one() {
+        return Zero::zero();
+    }
+
+    let two = One::one() + One::one();
+    self / fast_pow(two, shift)
 }
 
+
 // Fast exponentiation using the square-and-multiply algorithm
 // Reference:
 // https://github.com/keep-starknet-strange/alexandria/blob/bcdca70afdf59c9976148e95cebad5cf63d75a7f/packages/math/src/fast_power.cairo#L12
-pub fn fast_pow(base: u256, exp: u32) -> u256 {
-    if exp == 0 {
-        return 1_u256;
+pub fn fast_pow<
+    T,
+    U,
+    +Zero<T>,
+    +Zero<U>,
+    +One<T>,
+    +One<U>,
+    +Add<U>,
+    +Mul<T>,
+    +Rem<U>,
+    +Div<U>,
+    +Copy<T>,
+    +Copy<U>,
+    +Drop<T>,
+    +Drop<U>,
+    +PartialEq<U>,
+>(
+    base: T, exp: U
+) -> T {
+    if exp == Zero::zero() {
+        return One::one();
     }
 
-    let mut res: u256 = 1_u256;
-    let mut base: u256 = base;
-    let mut exp: u32 = exp;
+    let mut res: T = One::one();
+    let mut base: T = base;
+    let mut exp: U = exp;
+
+    let two: U = One::one() + One::one();
 
     loop {
-        if exp % 2 == 1 {
+        if exp % two == One::one() {
             res = res * base;
         }
-        exp = exp / 2;
-        if exp == 0 {
+        exp = exp / two;
+        if exp == Zero::zero() {
             break res;
         }
         base = base * base;

src/validation.cairo

@@ -190,7 +190,7 @@ pub fn target_to_bits(target: u256) -> Result<u32, felt252> {
     let size_u256: u256 = size.into();
 
     // Combine size and mantissa
-    let result: u32 = (shl(size_u256, 24) + mantissa.into()).try_into().unwrap();
+    let result: u32 = (shl(size_u256, 24_u32) + mantissa.into()).try_into().unwrap();
 
     Result::Ok(result)
 }
@@ -219,16 +219,18 @@ fn validate_coinbase(block: @Block, total_fees: u256) -> Result<(), ByteArray> {
 
 // Return BTC reward in SATS
 fn compute_block_reward(block_height: u32) -> u64 {
-    shr(5000000000, block_height / 210_000).try_into().unwrap()
+    shr(5000000000_u256, (block_height / 210000_u32)).try_into().unwrap()
 }
+
+
 #[cfg(test)]
 mod tests {
+    use raito::state::{Header, Transaction, TxIn, TxOut};
     use super::{
-        validate_target, validate_timestamp, validate_proof_of_work, compute_block_reward,
-        compute_total_work, compute_work_from_target, shr, shl, REWARD_INITIAL, POW_SATS_AMOUNT
+        Block, ChainState, UtreexoState, REWARD_INITIAL, POW_SATS_AMOUNT, compute_block_reward,
+        compute_work_from_target, compute_total_work, validate_proof_of_work, validate_target,
+        validate_timestamp
     };
-    use super::{Block, ChainState, UtreexoState};
-    use super::super::state::{Header, Transaction, TxIn, TxOut};
 
     #[test]
     fn test_validate_target() {
@@ -298,27 +300,31 @@ mod tests {
         let work = compute_work_from_target(target);
         assert(expected_work == work, 'Failed to compute target');
     }
+
     #[test]
     fn test_compute_work_from_target2() {
         let expected_work = 0x26d946e509ac00026d;
         let target: u256 = 0x00000000000000000696f4000000000000000000000000000000000000000000;
         let work = compute_work_from_target(target);
         assert(expected_work == work, 'Failed to compute target');
     }
+
     #[test]
     fn test_compute_work_from_target3() {
         let expected_work = 0xe10005c64415f04ef3e387b97db388404db9fdfaab2b1918f6783471d;
         let target: u256 = 0x12345600;
         let work = compute_work_from_target(target);
         assert(expected_work == work, 'Failed to compute target');
     }
+
     #[test]
     fn test_compute_work_from_target4() {
         let expected_work = 0x1c040c95a099201bcaf85db4e7f2e21e18707c8d55a887643b95afb2f;
         let target: u256 = 0x92340000;
         let work = compute_work_from_target(target);
         assert(expected_work == work, 'Failed to compute target');
     }
+
     #[test]
     fn test_compute_work_from_target5() {
         let expected_work = 0x21809b468faa88dbe34f;