feat: add u32 related hash functions

blake3/Cargo.toml

@@ -7,6 +7,7 @@ license = "MIT OR Apache-2.0"
 [dependencies]
 p3-symmetric = { path = "../symmetric" }
 blake3 = "1.5"
+blake3_zkvm = { git = "https://github.com/sp1-patches/BLAKE3.git", branch = "patch-blake3_zkvm/v.1.0.0" }
 
 [features]
 neon = ["blake3/neon"]

blake3/src/lib.rs

@@ -35,3 +35,35 @@ impl CryptographicHasher<u8, [u8; 32]> for Blake3 {
         hasher.finalize().into()
     }
 }
+
+impl CryptographicHasher<u32, [u32; 8]> for Blake3 {
+    fn hash_iter<I>(&self, input: I) -> [u32; 8]
+    where
+        I: IntoIterator<Item = u32>,
+    {
+        let mut input = input.into_iter().collect::<Vec<_>>();
+        if input.len() <= blake3_zkvm::BLOCK_LEN {
+            let size = input.len();
+            input.resize(blake3_zkvm::BLOCK_LEN, 0u32);
+            blake3_zkvm::hash_single_block(input.as_slice().try_into().unwrap(), size)
+        } else {
+            let ret = self.hash_iter_slices([input.as_slice()]);
+            ret
+        }
+    }
+
+    fn hash_iter_slices<'a, I>(&self, input: I) -> [u32; 8]
+    where
+        I: IntoIterator<Item = &'a [u32]>,
+    {
+        let mut zkvm_hasher = blake3_zkvm::Hasher::new();
+
+        for chunk in input.into_iter() {
+            zkvm_hasher.update(chunk);
+        }
+        let mut out: [u32; 8] = [0u32; 8];
+        zkvm_hasher.finalize(&mut out);
+
+        out
+    }
+}

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,56 +50,77 @@ 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, 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, u8, N>) {
+    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>,
     Inner: CanSample<u8>,
 {
     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 log_size = log2_ceil_u64(F::ORDER_U64).saturating_sub(1);
+        let bound = (1 << log_size) - 1;
+        let sample_base = |inner: &mut Inner| {
             let value = u32::from_le_bytes(inner.sample_array::<4>());
-            let value = value & pow_of_two_bound;
-            if value < modulus {
-                return F::from_canonical_u32(value);
-            }
+            F::from_canonical_u32(value & bound)
+        };
+        EF::from_base_fn(|_| sample_base(&mut self.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 log_size = log2_ceil_u64(F::ORDER_U64).saturating_sub(1);
+        let bound = (1 << log_size) - 1;
+        let sample_base = |inner: &mut Inner| {
+            let value = inner.sample();
+            F::from_canonical_u32(value & bound)
         };
         EF::from_base_fn(|_| sample_base(&mut self.inner))
     }
 }
 
-impl<F, Inner> CanSampleBits<usize> for SerializingChallenger32<F, Inner>
+
+impl<F, Inner> CanSampleBits<usize> for SerializingChallenger32<F, u8, Inner>
 where
     F: PrimeField32,
     Inner: CanSample<u8>,
@@ -113,15 +134,31 @@ 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,
 {
     type Witness = F;
 
     #[instrument(name = "grind for proof-of-work witness", skip_all)]
-    fn grind(&mut self, bits: usize) -> Self::Witness {
+    fn grind(&mut self, bits: usize) -> Self::Witness
+    {
         let witness = (0..F::ORDER_U64)
             .into_par_iter()
             .map(|i| F::from_canonical_u64(i))
@@ -132,13 +169,40 @@ 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 {
@@ -182,16 +246,11 @@ where
     Inner: CanSample<u8>,
 {
     fn sample(&mut self) -> EF {
-        let modulus = F::ORDER_U64;
-        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 log_size = log2_ceil_u64(F::ORDER_U64).saturating_sub(1);
+        let bound = (1 << log_size) - 1;
+        let sample_base = |inner: &mut Inner| {
             let value = u64::from_le_bytes(inner.sample_array::<8>());
-            let value = value & pow_of_two_bound;
-            if value < modulus {
-                return F::from_canonical_u64(value);
-            }
+            F::from_canonical_u64(value & bound)
         };
         EF::from_base_fn(|_| sample_base(&mut self.inner))
     }
@@ -235,4 +294,4 @@ where
     F: PrimeField64,
     Inner: CanSample<u8> + CanObserve<u8> + Clone + Send + Sync,
 {
-}
+}

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,

symmetric/src/serializing_hasher.rs

@@ -96,4 +96,4 @@ where
                 .map(|x| PW::from_fn(|i| x.as_slice()[i].as_canonical_u64())),
         )
     }
-}
+}