Add PoW check to channel (#172)

stwo_cairo_verifier/src/channel.cairo

@@ -1,14 +1,20 @@
 use core::array::SpanTrait;
+use core::num::traits::{WrappingMul, WrappingSub};
 use core::poseidon::{poseidon_hash_span, hades_permutation};
 use core::traits::DivRem;
 use stwo_cairo_verifier::fields::qm31::QM31Trait;
 use stwo_cairo_verifier::utils::pack4;
 
 use stwo_cairo_verifier::{BaseField, SecureField};
 
-const M31_SHIFT: felt252 = 0x80000000; // 2**31.
-const M31_SHIFT_NZ_U256: NonZero<u256> = 0x80000000; // 2**31.
+/// Equals `2^31`.
+const M31_SHIFT: felt252 = 0x80000000;
+
+/// Equals `2^31`.
+const M31_SHIFT_NZ_U256: NonZero<u256> = 0x80000000;
+
 pub const EXTENSION_FELTS_PER_HASH: usize = 2;
+
 pub const FELTS_PER_HASH: usize = 8;
 
 #[derive(Default, Drop)]
@@ -137,6 +143,22 @@ pub impl ChannelImpl of ChannelTrait {
             };
         bytes
     }
+
+    fn check_proof_of_work(self: @Channel, n_bits: u32) -> bool {
+        let u256 { low, .. } = (*self.digest).into();
+        low & gen_bit_mask(n_bits) == 0
+    }
+}
+
+/// Generates a bit mask with the least significant `n_bits` set to 1.
+fn gen_bit_mask(n_bits: u32) -> u128 {
+    assert!(n_bits <= 128);
+    let mut mask = 1;
+    for _ in 0..n_bits {
+        mask = mask.wrapping_mul(2);
+    };
+    mask = mask.wrapping_sub(1);
+    mask
 }
 
 #[inline]
@@ -155,7 +177,7 @@ fn extract_m31<const N: usize>(ref num: u256) -> BaseField {
 #[cfg(test)]
 mod tests {
     use stwo_cairo_verifier::fields::qm31::qm31;
-    use super::ChannelTrait;
+    use super::{Channel, ChannelTrait, gen_bit_mask};
 
     #[test]
     fn test_initialize_channel() {
@@ -332,4 +354,37 @@ mod tests {
         let second_result = channel.draw_random_bytes();
         assert_ne!(first_result, second_result);
     }
+
+    #[test]
+    fn test_gen_bit_mask_with_0_bits() {
+        assert_eq!(gen_bit_mask(0), 0);
+    }
+
+    #[test]
+    fn test_gen_bit_mask_with_8_bits() {
+        assert_eq!(gen_bit_mask(8), 0b11111111);
+    }
+
+    #[test]
+    fn test_gen_bit_mask_with_128_bits() {
+        assert_eq!(gen_bit_mask(128), 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF);
+    }
+
+    #[test]
+    fn test_check_proof_of_work() {
+        let channel = Channel { digest: 0b1000, channel_time: Default::default(), };
+
+        let res = channel.check_proof_of_work(3);
+
+        assert!(res);
+    }
+
+    #[test]
+    fn test_check_proof_of_work_with_invalid_n_bits() {
+        let channel = Channel { digest: 0b1000, channel_time: Default::default(), };
+
+        let res = channel.check_proof_of_work(4);
+
+        assert!(!res);
+    }
 }