feat: improve integration of PRNG for zk

crypto/src/merkle/mod.rs

@@ -7,11 +7,11 @@ use alloc::{
     collections::{BTreeMap, BTreeSet},
     vec::Vec,
 };
-use core::slice;
+use core::{marker::PhantomData, slice};
 
 use rand::{
     distributions::{Distribution, Standard},
-    thread_rng, Rng, RngCore, SeedableRng,
+    Rng, RngCore, SeedableRng,
 };
 
 use crate::{
@@ -481,20 +481,21 @@ impl<H: Hasher> VectorCommitment<H> for MerkleTree<H> {
 // SALTED MERKLE TREE
 // ================================================================================================
 
-pub struct SaltedMerkleTree<H: Hasher> {
+pub struct SaltedMerkleTree<H: Hasher, P: RngCore + SeedableRng> {
     leaves: Vec<H::Digest>,
     tree: MerkleTree<H>,
     salts: Vec<H::Digest>,
+    _prng: PhantomData<P>,
 }
 
-impl<H: Hasher> SaltedMerkleTree<H>
+impl<H: Hasher, P: RngCore + SeedableRng> SaltedMerkleTree<H, P>
 where
     Standard: Distribution<<H as Hasher>::Digest>,
 {
     // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
 
-    pub fn new<R: RngCore>(leaves: Vec<H::Digest>, prng: &mut R) -> Result<Self, MerkleTreeError> {
+    pub fn new(leaves: Vec<H::Digest>, prng: &mut P) -> Result<Self, MerkleTreeError> {
         if leaves.len() < 2 {
             return Err(MerkleTreeError::TooFewLeaves(2, leaves.len()));
         }
@@ -513,23 +514,27 @@ where
 
         let tree = MerkleTree::new(salted_leaves)?;
 
-        Ok(Self { tree, leaves, salts })
+        Ok(Self { tree, leaves, salts, _prng: PhantomData })
     }
 
     /// Returns the root of the tree.
     pub fn root(&self) -> &H::Digest {
         self.tree.root()
     }
 
+    /// Returns the depth of the tree.
     pub fn depth(&self) -> usize {
         self.tree.depth()
     }
 
+    /// Returns a Merkle proof to a leaf at the specified `index`.
     pub fn prove(&self, index: usize) -> Result<SaltedMerkleTreeOpening<H>, MerkleTreeError> {
         let (_, proof) = self.tree.prove(index)?;
         Ok((self.leaves[index], (self.salts[index], proof)))
     }
 
+    /// Computes Merkle proofs for the provided indexes, compresses the proofs into a single batch
+    /// and returns the batch proof alongside the leaves at the provided indexes.
     pub fn prove_batch(
         &self,
         indexes: &[usize],
@@ -540,6 +545,7 @@ where
         Ok((leaves_at_indices, (salts_at_indices, proof)))
     }
 
+    /// Checks whether the `proof` for the given `leaf` at the specified `index` is valid.
     pub fn verify(
         root: H::Digest,
         index: usize,
@@ -552,15 +558,6 @@ where
     }
 
     /// Checks whether the batch proof contains Merkle paths for the of the specified `indexes`.
-    ///
-    /// # Errors
-    /// Returns an error if:
-    /// * No indexes were provided (i.e., `indexes` is an empty slice).
-    /// * Number of provided indexes is greater than 255.
-    /// * Any of the specified `indexes` is greater than or equal to the number of leaves in the
-    ///   tree from which the batch proof was generated.
-    /// * List of indexes contains duplicates.
-    /// * Any of the paths in the batch proof does not resolve to the specified `root`.
     pub fn verify_batch(
         root: &H::Digest,
         indexes: &[usize],
@@ -586,7 +583,7 @@ impl Distribution<ByteDigest<32>> for Standard {
     }
 }
 
-impl<H: Hasher> VectorCommitment<H> for SaltedMerkleTree<H>
+impl<H: Hasher, P: RngCore + SeedableRng> VectorCommitment<H> for SaltedMerkleTree<H, P>
 where
     Standard: Distribution<<H as Hasher>::Digest>,
 {
@@ -599,16 +596,16 @@ where
     type Error = MerkleTreeError;
 
     fn new(items: Vec<H::Digest>) -> Result<Self, Self::Error> {
-        let mut _prng = thread_rng();
-        let seed = [0_u8; 32];
-        let mut prng = rand_chacha::ChaCha20Rng::from_seed(seed);
+        // TODO: make random
+        let seed = P::Seed::default();
+        let mut prng = P::from_seed(seed);
         SaltedMerkleTree::new(items, &mut prng)
     }
 
     fn with_options(items: Vec<H::Digest>, _options: Self::Options) -> Result<Self, Self::Error> {
-        let mut _prng = thread_rng();
-        let seed = [0_u8; 32];
-        let mut prng = rand_chacha::ChaCha20Rng::from_seed(seed);
+        // TODO: make random
+        let seed = P::Seed::default();
+        let mut prng = P::from_seed(seed);
         Self::new(items, &mut prng)
     }
 
@@ -645,7 +642,7 @@ where
         item: H::Digest,
         proof: &Self::Proof,
     ) -> Result<(), Self::Error> {
-        SaltedMerkleTree::<H>::verify(commitment, index, item, proof.0, &proof.1)
+        SaltedMerkleTree::<H, P>::verify(commitment, index, item, proof.0, &proof.1)
     }
 
     fn verify_many(
@@ -654,6 +651,6 @@ where
         items: &[H::Digest],
         proof: &Self::MultiProof,
     ) -> Result<(), Self::Error> {
-        SaltedMerkleTree::<H>::verify_batch(&commitment, indexes, items, &proof.0, &proof.1)
+        SaltedMerkleTree::<H, P>::verify_batch(&commitment, indexes, items, &proof.0, &proof.1)
     }
 }

crypto/src/merkle/tests.rs

@@ -5,6 +5,7 @@
 
 use math::fields::f128::BaseElement;
 use proptest::prelude::*;
+use rand_chacha::ChaCha20Rng;
 
 use super::*;
 
@@ -258,22 +259,50 @@ fn from_proofs() {
 fn verify_salted() {
     // depth 4
     let leaves = Digest256::bytes_as_digests(&LEAVES4).to_vec();
-    let mut prng = thread_rng();
-    let tree: SaltedMerkleTree<Blake3_256> = SaltedMerkleTree::new(leaves, &mut prng).unwrap();
+    let mut prng = ChaCha20Rng::from_entropy();
+    let tree: SaltedMerkleTree<Blake3_256, _> = SaltedMerkleTree::new(leaves, &mut prng).unwrap();
     let (leaf, (salt, proof)) = tree.prove(1).unwrap();
-    assert!(SaltedMerkleTree::<Blake3_256>::verify(*tree.root(), 1, leaf, salt, &proof).is_ok());
+    assert!(SaltedMerkleTree::<Blake3_256, ChaCha20Rng>::verify(
+        *tree.root(),
+        1,
+        leaf,
+        salt,
+        &proof
+    )
+    .is_ok());
 
     let (leaf, (salt, proof)) = tree.prove(2).unwrap();
-    assert!(SaltedMerkleTree::<Blake3_256>::verify(*tree.root(), 2, leaf, salt, &proof).is_ok());
+    assert!(SaltedMerkleTree::<Blake3_256, ChaCha20Rng>::verify(
+        *tree.root(),
+        2,
+        leaf,
+        salt,
+        &proof
+    )
+    .is_ok());
 
     // depth 5
     let leaf = Digest256::bytes_as_digests(&LEAVES8).to_vec();
-    let tree: SaltedMerkleTree<Blake3_256> = SaltedMerkleTree::new(leaf, &mut prng).unwrap();
+    let tree: SaltedMerkleTree<Blake3_256, _> = SaltedMerkleTree::new(leaf, &mut prng).unwrap();
     let (leaf, (salt, proof)) = tree.prove(1).unwrap();
-    assert!(SaltedMerkleTree::<Blake3_256>::verify(*tree.root(), 1, leaf, salt, &proof).is_ok());
+    assert!(SaltedMerkleTree::<Blake3_256, ChaCha20Rng>::verify(
+        *tree.root(),
+        1,
+        leaf,
+        salt,
+        &proof
+    )
+    .is_ok());
 
     let (leaf, (salt, proof)) = tree.prove(6).unwrap();
-    assert!(SaltedMerkleTree::<Blake3_256>::verify(*tree.root(), 6, leaf, salt, &proof).is_ok());
+    assert!(SaltedMerkleTree::<Blake3_256, ChaCha20Rng>::verify(
+        *tree.root(),
+        6,
+        leaf,
+        salt,
+        &proof
+    )
+    .is_ok());
 }
 
 proptest! {

examples/src/fibonacci/fib2/mod.rs

@@ -108,7 +108,7 @@ where
                 });
 
         // generate the proof
-        prover.prove(trace).unwrap()
+        prover.prove(trace, None).unwrap()
     }
 
     fn verify(&self, proof: Proof) -> Result<(), VerifierError> {

examples/src/fibonacci/fib2/prover.rs

@@ -4,11 +4,10 @@
 // LICENSE file in the root directory of this source tree.
 
 use air::ZkParameters;
-use rand_chacha::{rand_core::SeedableRng, ChaCha20Rng};
 use winterfell::{
     crypto::MerkleTree, matrix::ColMatrix, AuxRandElements, ConstraintCompositionCoefficients,
-    DefaultConstraintEvaluator, DefaultTraceLde, PartitionOptions, StarkDomain, Trace, TraceInfo,
-    TracePolyTable, TraceTable,
+    DefaultConstraintEvaluator, DefaultTraceLde, MockPrng, PartitionOptions, StarkDomain, Trace,
+    TraceInfo, TracePolyTable, TraceTable,
 };
 
 use super::{
@@ -64,6 +63,7 @@ where
         DefaultTraceLde<E, Self::HashFn, Self::VC>;
     type ConstraintEvaluator<'a, E: FieldElement<BaseField = Self::BaseField>> =
         DefaultConstraintEvaluator<'a, Self::Air, E>;
+    type ZkPrng = MockPrng;
 
     fn get_pub_inputs(&self, trace: &Self::Trace) -> BaseElement {
         let last_step = trace.length() - 1;
@@ -81,16 +81,9 @@ where
         domain: &StarkDomain<Self::BaseField>,
         partition_option: PartitionOptions,
         zk_parameters: Option<ZkParameters>,
+        prng: &mut Option<Self::ZkPrng>,
     ) -> (Self::TraceLde<E>, TracePolyTable<E>) {
-        let mut prng = ChaCha20Rng::from_entropy();
-        DefaultTraceLde::new(
-            trace_info,
-            main_trace,
-            domain,
-            partition_option,
-            zk_parameters,
-            &mut prng,
-        )
+        DefaultTraceLde::new(trace_info, main_trace, domain, partition_option, zk_parameters, prng)
     }
 
     fn new_evaluator<'a, E: FieldElement<BaseField = Self::BaseField>>(

examples/src/fibonacci/fib8/mod.rs

@@ -108,7 +108,7 @@ where
                 });
 
         // generate the proof
-        prover.prove(trace).unwrap()
+        prover.prove(trace, None).unwrap()
     }
 
     fn verify(&self, proof: Proof) -> Result<(), VerifierError> {

examples/src/fibonacci/fib8/prover.rs

@@ -4,11 +4,10 @@
 // LICENSE file in the root directory of this source tree.
 
 use air::ZkParameters;
-use rand_chacha::{rand_core::SeedableRng, ChaCha20Rng};
 use winterfell::{
     crypto::MerkleTree, matrix::ColMatrix, AuxRandElements, ConstraintCompositionCoefficients,
-    DefaultConstraintEvaluator, DefaultTraceLde, PartitionOptions, StarkDomain, Trace, TraceInfo,
-    TracePolyTable, TraceTable,
+    DefaultConstraintEvaluator, DefaultTraceLde, MockPrng, PartitionOptions, StarkDomain, Trace,
+    TraceInfo, TracePolyTable, TraceTable,
 };
 
 use super::{
@@ -79,6 +78,7 @@ where
         DefaultTraceLde<E, Self::HashFn, Self::VC>;
     type ConstraintEvaluator<'a, E: FieldElement<BaseField = Self::BaseField>> =
         DefaultConstraintEvaluator<'a, Self::Air, E>;
+    type ZkPrng = MockPrng;
 
     fn get_pub_inputs(&self, trace: &Self::Trace) -> BaseElement {
         let last_step = trace.length() - 1;
@@ -96,16 +96,9 @@ where
         domain: &StarkDomain<Self::BaseField>,
         partition_option: PartitionOptions,
         zk_parameters: Option<ZkParameters>,
+        prng: &mut Option<Self::ZkPrng>,
     ) -> (Self::TraceLde<E>, TracePolyTable<E>) {
-        let mut prng = ChaCha20Rng::from_entropy();
-        DefaultTraceLde::new(
-            trace_info,
-            main_trace,
-            domain,
-            partition_option,
-            zk_parameters,
-            &mut prng,
-        )
+        DefaultTraceLde::new(trace_info, main_trace, domain, partition_option, zk_parameters, prng)
     }
 
     fn new_evaluator<'a, E: FieldElement<BaseField = Self::BaseField>>(

examples/src/fibonacci/fib_small/mod.rs

@@ -119,7 +119,7 @@ where
                 });
 
         // generate the proof
-        prover.prove(trace).unwrap()
+        prover.prove(trace, None).unwrap()
     }
 
     fn verify(&self, proof: Proof) -> Result<(), VerifierError> {

examples/src/fibonacci/fib_small/prover.rs

@@ -3,11 +3,10 @@
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
 use air::ZkParameters;
-use rand_chacha::{rand_core::SeedableRng, ChaCha20Rng};
 use winterfell::{
     crypto::MerkleTree, matrix::ColMatrix, AuxRandElements, ConstraintCompositionCoefficients,
-    DefaultConstraintEvaluator, DefaultTraceLde, PartitionOptions, StarkDomain, Trace, TraceInfo,
-    TracePolyTable, TraceTable,
+    DefaultConstraintEvaluator, DefaultTraceLde, MockPrng, PartitionOptions, StarkDomain, Trace,
+    TraceInfo, TracePolyTable, TraceTable,
 };
 
 use super::{
@@ -69,6 +68,7 @@ where
         DefaultTraceLde<E, Self::HashFn, Self::VC>;
     type ConstraintEvaluator<'a, E: FieldElement<BaseField = Self::BaseField>> =
         DefaultConstraintEvaluator<'a, Self::Air, E>;
+    type ZkPrng = MockPrng;
 
     fn get_pub_inputs(&self, trace: &Self::Trace) -> BaseElement {
         let last_step = trace.length() - 1;
@@ -86,16 +86,9 @@ where
         domain: &StarkDomain<Self::BaseField>,
         partition_option: PartitionOptions,
         zk_parameters: Option<ZkParameters>,
+        prng: &mut Option<Self::ZkPrng>,
     ) -> (Self::TraceLde<E>, TracePolyTable<E>) {
-        let mut prng = ChaCha20Rng::from_entropy();
-        DefaultTraceLde::new(
-            trace_info,
-            main_trace,
-            domain,
-            partition_option,
-            zk_parameters,
-            &mut prng,
-        )
+        DefaultTraceLde::new(trace_info, main_trace, domain, partition_option, zk_parameters, prng)
     }
 
     fn new_evaluator<'a, E: FieldElement<BaseField = Self::BaseField>>(

examples/src/fibonacci/mulfib2/mod.rs

@@ -108,7 +108,7 @@ where
                 });
 
         // generate the proof
-        prover.prove(trace).unwrap()
+        prover.prove(trace, None).unwrap()
     }
 
     fn verify(&self, proof: Proof) -> Result<(), VerifierError> {