Change RNG semantics to help Poseidon. (risc0#421)

risc0/zkp/src/core/blake2b.rs

@@ -117,14 +117,8 @@ impl<T: Blake2b> ConfigRng<BabyBear> for Blake2bRng<T> {
         self.current = T::blake2b(concat);
     }
 
-    fn random_u32(&mut self) -> u32 {
-        let next = T::blake2b(self.current);
-        self.current = next;
-
-        ((next[0] as u32) << 24)
-            + ((next[1] as u32) << 16)
-            + ((next[2] as u32) << 8)
-            + ((next[3] as u32) << 0)
+    fn random_bits(&mut self, bits: usize) -> u32 {
+        ((1 << bits) - 1) & self.next_u32()
     }
 
     fn random_elem(&mut self) -> BabyBearElem {
@@ -138,7 +132,12 @@ impl<T: Blake2b> ConfigRng<BabyBear> for Blake2bRng<T> {
 
 impl<T: Blake2b> RngCore for Blake2bRng<T> {
     fn next_u32(&mut self) -> u32 {
-        self.random_u32()
+        let next = T::blake2b(self.current);
+        self.current = next;
+        ((next[0] as u32) << 24)
+            + ((next[1] as u32) << 16)
+            + ((next[2] as u32) << 8)
+            + ((next[3] as u32) << 0)
     }
 
     fn next_u64(&mut self) -> u64 {

risc0/zkp/src/core/config.rs

@@ -62,8 +62,9 @@ pub trait ConfigRng<F: Field> {
     /// Mix in randomness from a Fiat-Shamir commitment.
     fn mix(&mut self, val: &Digest);
 
-    /// Get a cryptographically uniform u32
-    fn random_u32(&mut self) -> u32;
+    /// Get a cryptographically uniform set of bits, as the low order bits of a
+    /// u32
+    fn random_bits(&mut self, bits: usize) -> u32;
 
     /// Get a cryptographically uniform field element
     fn random_elem(&mut self) -> F::Elem;

risc0/zkp/src/core/poseidon/rng.rs

@@ -47,14 +47,15 @@ impl ConfigRng<BabyBear> for PoseidonRng {
         // Mix
         poseidon_mix(&mut self.cells);
     }
-    fn random_u32(&mut self) -> u32 {
-        let mut val: u64 = 0;
-        for _ in 0..6 {
-            val *= 15 * 2 ^ 27 + 1; // TODO: Export P from BabyBear?
-            val += self.random_elem().as_u32_montgomery() as u64;
-            val &= 0xffffffff;
+    fn random_bits(&mut self, bits: usize) -> u32 {
+        let mut val = self.random_elem().as_u32();
+        for _ in 0..3 {
+            let new_val = self.random_elem().as_u32();
+            if val == 0 {
+                val = new_val;
+            }
         }
-        val as u32
+        ((1 << bits) - 1) & val
     }
     fn random_elem(&mut self) -> Elem {
         if self.pool_used == CELLS_RATE {

risc0/zkp/src/core/sha_rng.rs

@@ -90,8 +90,8 @@ impl<S: Sha256, F: Field> ConfigRng<F> for ShaRng<S> {
     fn mix(&mut self, val: &Digest) {
         self.inner_mix(val);
     }
-    fn random_u32(&mut self) -> u32 {
-        self.next_u32()
+    fn random_bits(&mut self, bits: usize) -> u32 {
+        ((1 << bits) - 1) & self.next_u32()
     }
     fn random_elem(&mut self) -> F::Elem {
         F::Elem::random(self)

risc0/zkp/src/prove/fri.rs

@@ -116,8 +116,7 @@ pub fn fri_prove<H: Hal, F>(
     debug!("Doing Queries");
     for _ in 0..QUERIES {
         // Get a 'random' index.
-        let rng = iop.random_u32() as usize;
-        let mut pos = rng % orig_domain;
+        let mut pos = iop.random_bits(log2_ceil(orig_domain)) as usize;
         // Do the 'inner' proof for this index
         f(iop, pos);
         // Write the per-round proofs

risc0/zkp/src/prove/merkle.rs

@@ -131,7 +131,7 @@ mod tests {
         adapter::{MixState, PolyExt},
         core::{
             config::{ConfigRng, HashSuite, HashSuitePoseidon, HashSuiteSha256},
-            sha_cpu,
+            log2_ceil, sha_cpu,
         },
         hal::cpu::CpuHal,
         verify::{merkle::MerkleTreeVerifier, read_iop::ReadIOP, CpuVerifyHal, VerificationError},
@@ -196,7 +196,7 @@ mod tests {
         let mut iop = WriteIOP::<BabyBear, HS::Rng>::new();
         prover.commit(&mut iop);
         for _query in 0..queries {
-            let r_idx = (iop.rng.random_u32() as usize) % rows;
+            let r_idx = iop.rng.random_bits(log2_ceil(rows)) as usize;
             let col = prover.prove(&mut iop, r_idx);
             for c_idx in 0..cols {
                 assert_eq!(
@@ -215,7 +215,7 @@ mod tests {
         assert_eq!(verifier.root(), prover.root());
         let mut err = false;
         for query in 0..queries {
-            let r_idx = (r_iop.random_u32() as usize) % rows;
+            let r_idx = r_iop.random_bits(log2_ceil(rows)) as usize;
             if query == bad_query {
                 if rows == 1 {
                     assert!(false, "Cannot test for bad query if there is only one row");

risc0/zkp/src/prove/write_iop.rs

@@ -60,8 +60,8 @@ impl<F: Field, R: ConfigRng<F>> WriteIOP<F, R> {
     }
 
     /// Get a cryptographically uniform u32
-    pub fn random_u32(&mut self) -> u32 {
-        self.rng.random_u32()
+    pub fn random_bits(&mut self, bits: usize) -> u32 {
+        self.rng.random_bits(bits)
     }
 
     /// Get a cryptographically uniform field element

risc0/zkp/src/verify/fri.rs

@@ -113,8 +113,7 @@ where
     // Do queries
     let mut poly_buf: Vec<H::ExtElem> = Vec::with_capacity(degree);
     for _ in 0..QUERIES {
-        let rng = iop.random_u32();
-        let mut pos = rng as usize % orig_domain;
+        let mut pos = iop.random_bits(log2_ceil(orig_domain)) as usize;
         // Do the 'inner' verification for this index
         let mut goal = inner(iop, pos)?;
         // Verify the per-round proofs

risc0/zkp/src/verify/read_iop.rs

@@ -68,8 +68,8 @@ impl<'a, F: Field, R: ConfigRng<F>> ReadIOP<'a, F, R> {
     }
 
     /// Get a cryptographically uniform u32
-    pub fn random_u32(&mut self) -> u32 {
-        self.rng.random_u32()
+    pub fn random_bits(&mut self, bits: usize) -> u32 {
+        self.rng.random_bits(bits)
     }
 
     /// Get a cryptographically uniform field element

risc0/zkvm/src/prove/mod.rs

@@ -198,18 +198,26 @@ cfg_if::cfg_if! {
             let eval = CudaEvalCheck::new(hal.clone());
             (hal, eval)
         }
+
+        // TODO: default_poseidon_hal
     } else if #[cfg(feature = "metal")] {
-        use risc0_circuit_rv32im::metal::MetalEvalCheckSha256;
-        use risc0_zkp::hal::metal::MetalHalSha256;
+        use risc0_circuit_rv32im::metal::{MetalEvalCheck, MetalEvalCheckSha256};
+        use risc0_zkp::hal::metal::{MetalHalSha256, MetalHalPoseidon, MetalHashPoseidon};
 
         pub fn default_hal() -> (Rc<MetalHalSha256>, MetalEvalCheckSha256) {
             let hal = Rc::new(MetalHalSha256::new());
             let eval = MetalEvalCheckSha256::new(hal.clone());
             (hal, eval)
         }
+
+        pub fn default_poseidon_hal() -> (Rc<MetalHalPoseidon>, MetalEvalCheck<MetalHashPoseidon>) {
+            let hal = Rc::new(MetalHalPoseidon::new());
+            let eval = MetalEvalCheck::<MetalHashPoseidon>::new(hal.clone());
+            (hal, eval)
+        }
     } else {
         use risc0_circuit_rv32im::{CircuitImpl, cpu::CpuEvalCheck};
-        use risc0_zkp::hal::cpu::BabyBearSha256CpuHal;
+        use risc0_zkp::hal::cpu::{BabyBearSha256CpuHal, BabyBearPoseidonCpuHal};
 
         /// Returns the default HAL for the RISC Zero circuit
         ///
@@ -230,6 +238,13 @@ cfg_if::cfg_if! {
             let eval = CpuEvalCheck::new(&CIRCUIT);
             (hal, eval)
         }
+
+        pub fn default_poseidon_hal() -> (Rc<BabyBearPoseidonCpuHal>, CpuEvalCheck<'static, CircuitImpl>) {
+            let hal = Rc::new(BabyBearPoseidonCpuHal::new());
+            let eval = CpuEvalCheck::new(&CIRCUIT);
+            (hal, eval)
+        }
+
     }
 }