Merge pull request #27 from succinctlabs/kevjue/u32_related_hash_func…

…tions

feat: add u32 related hash functions

challenger/src/serializing_challenger.rs

@@ -21,9 +21,9 @@ use crate::{
 /// -  Samples a field element in a prime field of size `p` by sampling uniformly an element in the
 ///    range (0..1 << log_2(p)). This avoids modulo bias.
 #[derive(Clone, Debug)]
-pub struct SerializingChallenger32<F, Inner> {
+pub struct SerializingChallenger32<F, T, Inner> {
     inner: Inner,
-    _marker: PhantomData<F>,
+    _marker: PhantomData<(F, T)>,
 }
 
 /// Given a challenger that can observe and sample bytes, produces a challenger that is able to
@@ -41,7 +41,7 @@ pub struct SerializingChallenger64<F, Inner> {
     _marker: PhantomData<F>,
 }
 
-impl<F: PrimeField32, Inner: CanObserve<u8>> SerializingChallenger32<F, Inner> {
+impl<F: PrimeField32, T, Inner: CanObserve<T>> SerializingChallenger32<F, T, Inner> {
     pub fn new(inner: Inner) -> Self {
         Self {
             inner,
@@ -50,34 +50,45 @@ impl<F: PrimeField32, Inner: CanObserve<u8>> SerializingChallenger32<F, Inner> {
     }
 }
 
-impl<F, H> SerializingChallenger32<F, HashChallenger<u8, H, 32>>
+impl<F, T, H, const N: usize> SerializingChallenger32<F, T, HashChallenger<T, H, N>>
 where
     F: PrimeField32,
-    H: CryptographicHasher<u8, [u8; 32]>,
+    H: CryptographicHasher<T, [T; N]>,
+    T: Clone,
 {
-    pub fn from_hasher(initial_state: Vec<u8>, hasher: H) -> Self {
+    pub fn from_hasher(initial_state: Vec<T>, hasher: H) -> Self {
         Self::new(HashChallenger::new(initial_state, hasher))
     }
 }
 
-impl<F: PrimeField32, Inner: CanObserve<u8>> CanObserve<F> for SerializingChallenger32<F, Inner> {
+impl<F: PrimeField32, Inner: CanObserve<u8>> CanObserve<F>
+    for SerializingChallenger32<F, u8, Inner>
+{
     fn observe(&mut self, value: F) {
         self.inner
             .observe_slice(&value.as_canonical_u32().to_le_bytes());
     }
 }
 
-impl<F: PrimeField32, const N: usize, Inner: CanObserve<u8>> CanObserve<Hash<F, u8, N>>
-    for SerializingChallenger32<F, Inner>
+impl<F: PrimeField32, Inner: CanObserve<u32>> CanObserve<F>
+    for SerializingChallenger32<F, u32, Inner>
 {
-    fn observe(&mut self, values: Hash<F, u8, N>) {
+    fn observe(&mut self, value: F) {
+        self.inner.observe(value.as_canonical_u32());
+    }
+}
+
+impl<F: PrimeField32, T, const N: usize, Inner: CanObserve<T>> CanObserve<Hash<F, T, N>>
+    for SerializingChallenger32<F, T, Inner>
+{
+    fn observe(&mut self, values: Hash<F, T, N>) {
         for value in values {
             self.inner.observe(value);
         }
     }
 }
 
-impl<F, EF, Inner> CanSample<EF> for SerializingChallenger32<F, Inner>
+impl<F, EF, Inner> CanSample<EF> for SerializingChallenger32<F, u8, Inner>
 where
     F: PrimeField32,
     EF: ExtensionField<F>,
@@ -99,7 +110,29 @@ where
     }
 }
 
-impl<F, Inner> CanSampleBits<usize> for SerializingChallenger32<F, Inner>
+impl<F, EF, Inner> CanSample<EF> for SerializingChallenger32<F, u32, Inner>
+where
+    F: PrimeField32,
+    EF: ExtensionField<F>,
+    Inner: CanSample<u32>,
+{
+    fn sample(&mut self) -> EF {
+        let modulus = F::ORDER_U64 as u32;
+        let log_size = log2_ceil_u64(F::ORDER_U64);
+        let pow_of_two_bound = (1 << log_size) - 1;
+        // Perform rejection sampling over the uniform range (0..log2_ceil(p))
+        let sample_base = |inner: &mut Inner| loop {
+            let value = inner.sample();
+            let value = value & pow_of_two_bound;
+            if value < modulus {
+                return F::from_canonical_u32(value);
+            }
+        };
+        EF::from_base_fn(|_| sample_base(&mut self.inner))
+    }
+}
+
+impl<F, Inner> CanSampleBits<usize> for SerializingChallenger32<F, u8, Inner>
 where
     F: PrimeField32,
     Inner: CanSample<u8>,
@@ -113,7 +146,21 @@ where
     }
 }
 
-impl<F, Inner> GrindingChallenger for SerializingChallenger32<F, Inner>
+impl<F, Inner> CanSampleBits<usize> for SerializingChallenger32<F, u32, Inner>
+where
+    F: PrimeField32,
+    Inner: CanSample<u32>,
+{
+    fn sample_bits(&mut self, bits: usize) -> usize {
+        debug_assert!(bits < (usize::BITS as usize));
+        // Limiting the number of bits to the field size
+        debug_assert!((1 << bits) <= F::ORDER_U64 as usize);
+        let rand_usize = self.inner.sample() as usize;
+        rand_usize & ((1 << bits) - 1)
+    }
+}
+
+impl<F, Inner> GrindingChallenger for SerializingChallenger32<F, u8, Inner>
 where
     F: PrimeField32,
     Inner: CanSample<u8> + CanObserve<u8> + Clone + Send + Sync,
@@ -132,13 +179,39 @@ where
     }
 }
 
-impl<F, Inner> FieldChallenger<F> for SerializingChallenger32<F, Inner>
+impl<F, Inner> GrindingChallenger for SerializingChallenger32<F, u32, Inner>
+where
+    F: PrimeField32,
+    Inner: CanSample<u32> + CanObserve<u32> + Clone + Send + Sync,
+{
+    type Witness = F;
+
+    #[instrument(name = "grind for proof-of-work witness", skip_all)]
+    fn grind(&mut self, bits: usize) -> Self::Witness {
+        let witness = (0..F::ORDER_U64)
+            .into_par_iter()
+            .map(|i| F::from_canonical_u64(i))
+            .find_any(|witness| self.clone().check_witness(bits, *witness))
+            .expect("failed to find witness");
+        assert!(self.check_witness(bits, witness));
+        witness
+    }
+}
+
+impl<F, Inner> FieldChallenger<F> for SerializingChallenger32<F, u8, Inner>
 where
     F: PrimeField32,
     Inner: CanSample<u8> + CanObserve<u8> + Clone + Send + Sync,
 {
 }
 
+impl<F, Inner> FieldChallenger<F> for SerializingChallenger32<F, u32, Inner>
+where
+    F: PrimeField32,
+    Inner: CanSample<u32> + CanObserve<u32> + Clone + Send + Sync,
+{
+}
+
 impl<F: PrimeField64, Inner: CanObserve<u8>> SerializingChallenger64<F, Inner> {
     pub fn new(inner: Inner) -> Self {
         Self {

keccak-air/examples/prove_baby_bear_keccak.rs

@@ -48,7 +48,7 @@ fn main() -> Result<(), VerificationError> {
     type Dft = Radix2DitParallel;
     let dft = Dft {};
 
-    type Challenger = SerializingChallenger32<Val, HashChallenger<u8, ByteHash, 32>>;
+    type Challenger = SerializingChallenger32<Val, u8, HashChallenger<u8, ByteHash, 32>>;
 
     let fri_config = FriConfig {
         log_blowup: 1,