Refactor generalize channel and merkleHasher (#759)

Co-authored-by: Sergio Chouhy <[email protected]>

crates/prover/benches/pcs.rs

@@ -29,7 +29,12 @@ fn benched_fn<B: Backend + MerkleOps<Blake2sMerkleHasher>>(
         .map(|eval| eval.interpolate_with_twiddles(twiddles))
         .collect();
 
-    CommitmentTreeProver::<B>::new(polys, LOG_BLOWUP_FACTOR, channel, twiddles);
+    CommitmentTreeProver::<B, Blake2sMerkleHasher>::new(
+        polys,
+        LOG_BLOWUP_FACTOR,
+        channel,
+        twiddles,
+    );
 }
 
 fn bench_pcs<B: Backend + MerkleOps<Blake2sMerkleHasher>>(c: &mut Criterion, id: &str) {

crates/prover/src/core/air/air_ext.rs

@@ -8,8 +8,7 @@ use crate::core::fields::qm31::SecureField;
 use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
 use crate::core::pcs::{CommitmentTreeProver, TreeVec};
 use crate::core::poly::circle::SecureCirclePoly;
-use crate::core::vcs::blake2_merkle::Blake2sMerkleHasher;
-use crate::core::vcs::ops::MerkleOps;
+use crate::core::vcs::ops::{MerkleHasher, MerkleOps};
 use crate::core::{ColumnVec, InteractionElements, LookupValues};
 
 pub trait AirExt: Air {
@@ -79,9 +78,9 @@ pub trait AirExt: Air {
         air_sizes
     }
 
-    fn component_traces<'a, B: Backend + MerkleOps<Blake2sMerkleHasher>>(
+    fn component_traces<'a, B: Backend + MerkleOps<H>, H: MerkleHasher>(
         &'a self,
-        trees: &'a [CommitmentTreeProver<B>],
+        trees: &'a [CommitmentTreeProver<B, H>],
     ) -> Vec<ComponentTrace<'_, B>> {
         let mut poly_iters = trees
             .iter()

crates/prover/src/core/channel/blake2s.rs

@@ -1,12 +1,11 @@
 use std::iter;
 
-use super::{Channel, ChannelTime};
+use super::{Channel, ChannelTime, Serializable};
 use crate::core::fields::m31::{BaseField, N_BYTES_FELT, P};
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
 use crate::core::fields::IntoSlice;
 use crate::core::vcs::blake2_hash::{Blake2sHash, Blake2sHasher};
-use crate::core::vcs::hasher::Hasher;
 
 pub const BLAKE_BYTES_PER_HASH: usize = 32;
 pub const FELTS_PER_HASH: usize = 8;
@@ -118,6 +117,12 @@ impl Channel for Blake2sChannel {
     }
 }
 
+impl Serializable for Blake2sHash {
+    fn to_bytes(self) -> Vec<u8> {
+        self.as_ref().to_vec()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::BTreeSet;

crates/prover/src/core/channel/mod.rs

@@ -1,9 +1,11 @@
 use super::fields::qm31::SecureField;
 
-mod blake2s;
 #[cfg(not(target_arch = "wasm32"))]
 mod poseidon252;
+#[cfg(not(target_arch = "wasm32"))]
+pub use poseidon252::Poseidon252Channel;
 
+mod blake2s;
 pub use blake2s::Blake2sChannel;
 
 pub const EXTENSION_FELTS_PER_HASH: usize = 2;
@@ -26,7 +28,7 @@ impl ChannelTime {
 }
 
 pub trait Channel {
-    type Digest;
+    type Digest: Serializable + Copy;
 
     const BYTES_PER_HASH: usize;
 
@@ -45,3 +47,7 @@ pub trait Channel {
     /// Returns a vector of random bytes of length `BYTES_PER_HASH`.
     fn draw_random_bytes(&mut self) -> Vec<u8>;
 }
+
+pub trait Serializable {
+    fn to_bytes(self) -> Vec<u8>;
+}

crates/prover/src/core/channel/poseidon252.rs

@@ -3,7 +3,7 @@ use std::iter;
 use starknet_crypto::{poseidon_hash, poseidon_hash_many};
 use starknet_ff::FieldElement as FieldElement252;
 
-use super::{Channel, ChannelTime};
+use super::{Channel, ChannelTime, Serializable};
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
@@ -122,6 +122,12 @@ impl Channel for Poseidon252Channel {
     }
 }
 
+impl Serializable for FieldElement252 {
+    fn to_bytes(self) -> Vec<u8> {
+        self.as_ref().to_bytes_be().to_vec()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::BTreeSet;

crates/prover/src/core/lookups/sumcheck.rs

@@ -213,7 +213,6 @@ mod tests {
     use crate::core::lookups::mle::Mle;
     use crate::core::lookups::sumcheck::{partially_verify, prove_batch};
     use crate::core::vcs::blake2_hash::Blake2sHasher;
-    use crate::core::vcs::hasher::Hasher;
 
     #[test]
     fn sumcheck_works() {

crates/prover/src/core/pcs/prover.rs

@@ -4,7 +4,6 @@ use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 use tracing::{span, Level};
 
-use super::super::channel::Blake2sChannel;
 use super::super::circle::CirclePoint;
 use super::super::fields::m31::BaseField;
 use super::super::fields::qm31::SecureField;
@@ -23,22 +22,17 @@ use crate::core::backend::Backend;
 use crate::core::channel::Channel;
 use crate::core::poly::circle::{CircleEvaluation, CirclePoly};
 use crate::core::poly::twiddles::TwiddleTree;
-use crate::core::vcs::blake2_hash::Blake2sHash;
-use crate::core::vcs::blake2_merkle::Blake2sMerkleHasher;
-use crate::core::vcs::ops::MerkleOps;
+use crate::core::vcs::ops::{MerkleHasher, MerkleOps};
 use crate::core::vcs::prover::{MerkleDecommitment, MerkleProver};
 
-type MerkleHasher = Blake2sMerkleHasher;
-type ProofChannel = Blake2sChannel;
-
 /// The prover side of a FRI polynomial commitment scheme. See [super].
-pub struct CommitmentSchemeProver<'a, B: Backend + MerkleOps<MerkleHasher>> {
-    pub trees: TreeVec<CommitmentTreeProver<B>>,
+pub struct CommitmentSchemeProver<'a, B: Backend + MerkleOps<H>, H: MerkleHasher> {
+    pub trees: TreeVec<CommitmentTreeProver<B, H>>,
     pub log_blowup_factor: u32,
     twiddles: &'a TwiddleTree<B>,
 }
 
-impl<'a, B: Backend + MerkleOps<MerkleHasher>> CommitmentSchemeProver<'a, B> {
+impl<'a, B: Backend + MerkleOps<H>, H: MerkleHasher> CommitmentSchemeProver<'a, B, H> {
     pub fn new(log_blowup_factor: u32, twiddles: &'a TwiddleTree<B>) -> Self {
         CommitmentSchemeProver {
             trees: TreeVec::default(),
@@ -47,22 +41,26 @@ impl<'a, B: Backend + MerkleOps<MerkleHasher>> CommitmentSchemeProver<'a, B> {
         }
     }
 
-    fn commit(&mut self, polynomials: ColumnVec<CirclePoly<B>>, channel: &mut ProofChannel) {
+    fn commit(
+        &mut self,
+        polynomials: ColumnVec<CirclePoly<B>>,
+        channel: &mut impl Channel<Digest = H::Hash>,
+    ) {
         let _span = span!(Level::INFO, "Commitment").entered();
         let tree =
             CommitmentTreeProver::new(polynomials, self.log_blowup_factor, channel, self.twiddles);
         self.trees.push(tree);
     }
 
-    pub fn tree_builder(&mut self) -> TreeBuilder<'_, 'a, B> {
+    pub fn tree_builder(&mut self) -> TreeBuilder<'_, 'a, B, H> {
         TreeBuilder {
             tree_index: self.trees.len(),
             commitment_scheme: self,
             polys: Vec::default(),
         }
     }
 
-    pub fn roots(&self) -> TreeVec<Blake2sHash> {
+    pub fn roots(&self) -> TreeVec<<H as MerkleHasher>::Hash> {
         self.trees.as_ref().map(|tree| tree.commitment.root())
     }
 
@@ -81,8 +79,8 @@ impl<'a, B: Backend + MerkleOps<MerkleHasher>> CommitmentSchemeProver<'a, B> {
     pub fn prove_values(
         &self,
         sampled_points: TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>>,
-        channel: &mut ProofChannel,
-    ) -> CommitmentSchemeProof {
+        channel: &mut impl Channel<Digest = H::Hash>,
+    ) -> CommitmentSchemeProof<H> {
         // Evaluate polynomials on open points.
         let span = span!(Level::INFO, "Evaluate columns out of domain").entered();
         let samples = self
@@ -109,8 +107,7 @@ impl<'a, B: Backend + MerkleOps<MerkleHasher>> CommitmentSchemeProver<'a, B> {
 
         // Run FRI commitment phase on the oods quotients.
         let fri_config = FriConfig::new(LOG_LAST_LAYER_DEGREE_BOUND, LOG_BLOWUP_FACTOR, N_QUERIES);
-        let fri_prover =
-            FriProver::<B, MerkleHasher>::commit(channel, fri_config, &quotients, self.twiddles);
+        let fri_prover = FriProver::commit(channel, fri_config, &quotients, self.twiddles);
 
         // Proof of work.
         let proof_of_work = ProofOfWork::new(PROOF_OF_WORK_BITS).prove(channel);
@@ -141,20 +138,20 @@ impl<'a, B: Backend + MerkleOps<MerkleHasher>> CommitmentSchemeProver<'a, B> {
 }
 
 #[derive(Debug, Serialize, Deserialize)]
-pub struct CommitmentSchemeProof {
+pub struct CommitmentSchemeProof<H: MerkleHasher> {
     pub sampled_values: TreeVec<ColumnVec<Vec<SecureField>>>,
-    pub decommitments: TreeVec<MerkleDecommitment<MerkleHasher>>,
+    pub decommitments: TreeVec<MerkleDecommitment<H>>,
     pub queried_values: TreeVec<ColumnVec<Vec<BaseField>>>,
     pub proof_of_work: ProofOfWorkProof,
-    pub fri_proof: FriProof<MerkleHasher>,
+    pub fri_proof: FriProof<H>,
 }
 
-pub struct TreeBuilder<'a, 'b, B: Backend + MerkleOps<MerkleHasher>> {
+pub struct TreeBuilder<'a, 'b, B: Backend + MerkleOps<H>, H: MerkleHasher> {
     tree_index: usize,
-    commitment_scheme: &'a mut CommitmentSchemeProver<'b, B>,
+    commitment_scheme: &'a mut CommitmentSchemeProver<'b, B, H>,
     polys: ColumnVec<CirclePoly<B>>,
 }
-impl<'a, 'b, B: Backend + MerkleOps<MerkleHasher>> TreeBuilder<'a, 'b, B> {
+impl<'a, 'b, B: Backend + MerkleOps<H>, H: MerkleHasher> TreeBuilder<'a, 'b, B, H> {
     pub fn extend_evals(
         &mut self,
         columns: ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>,
@@ -184,25 +181,25 @@ impl<'a, 'b, B: Backend + MerkleOps<MerkleHasher>> TreeBuilder<'a, 'b, B> {
         }
     }
 
-    pub fn commit(self, channel: &mut ProofChannel) {
+    pub fn commit(self, channel: &mut impl Channel<Digest = H::Hash>) {
         let _span = span!(Level::INFO, "Commitment").entered();
         self.commitment_scheme.commit(self.polys, channel);
     }
 }
 
 /// Prover data for a single commitment tree in a commitment scheme. The commitment scheme allows to
 /// commit on a set of polynomials at a time. This corresponds to such a set.
-pub struct CommitmentTreeProver<B: Backend + MerkleOps<MerkleHasher>> {
+pub struct CommitmentTreeProver<B: Backend + MerkleOps<H>, H: MerkleHasher> {
     pub polynomials: ColumnVec<CirclePoly<B>>,
     pub evaluations: ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>,
-    pub commitment: MerkleProver<B, MerkleHasher>,
+    pub commitment: MerkleProver<B, H>,
 }
 
-impl<B: Backend + MerkleOps<MerkleHasher>> CommitmentTreeProver<B> {
+impl<B: Backend + MerkleOps<H>, H: MerkleHasher> CommitmentTreeProver<B, H> {
     pub fn new(
         polynomials: ColumnVec<CirclePoly<B>>,
         log_blowup_factor: u32,
-        channel: &mut ProofChannel,
+        channel: &mut impl Channel<Digest = H::Hash>,
         twiddles: &TwiddleTree<B>,
     ) -> Self {
         let span = span!(Level::INFO, "Extension").entered();
@@ -236,7 +233,7 @@ impl<B: Backend + MerkleOps<MerkleHasher>> CommitmentTreeProver<B> {
     fn decommit(
         &self,
         queries: BTreeMap<u32, Vec<usize>>,
-    ) -> (ColumnVec<Vec<BaseField>>, MerkleDecommitment<MerkleHasher>) {
+    ) -> (ColumnVec<Vec<BaseField>>, MerkleDecommitment<H>) {
         let eval_vec = self
             .evaluations
             .iter()

crates/prover/src/core/pcs/verifier.rs

@@ -2,7 +2,6 @@ use std::iter::zip;
 
 use itertools::Itertools;
 
-use super::super::channel::Blake2sChannel;
 use super::super::circle::CirclePoint;
 use super::super::fields::qm31::SecureField;
 use super::super::fri::{CirclePolyDegreeBound, FriConfig, FriVerifier};
@@ -15,20 +14,17 @@ use super::utils::TreeVec;
 use super::CommitmentSchemeProof;
 use crate::core::channel::Channel;
 use crate::core::prover::VerificationError;
-use crate::core::vcs::blake2_hash::Blake2sHash;
-use crate::core::vcs::blake2_merkle::Blake2sMerkleHasher;
+use crate::core::vcs::ops::MerkleHasher;
 use crate::core::vcs::verifier::MerkleVerifier;
 use crate::core::ColumnVec;
 
-type ProofChannel = Blake2sChannel;
-
 /// The verifier side of a FRI polynomial commitment scheme. See [super].
 #[derive(Default)]
-pub struct CommitmentSchemeVerifier {
-    pub trees: TreeVec<MerkleVerifier<Blake2sMerkleHasher>>,
+pub struct CommitmentSchemeVerifier<H: MerkleHasher> {
+    pub trees: TreeVec<MerkleVerifier<H>>,
 }
 
-impl CommitmentSchemeVerifier {
+impl<H: MerkleHasher> CommitmentSchemeVerifier<H> {
     pub fn new() -> Self {
         Self::default()
     }
@@ -41,12 +37,11 @@ impl CommitmentSchemeVerifier {
     }
 
     /// Reads a commitment from the prover.
-    pub fn commit(
-        &mut self,
-        commitment: Blake2sHash,
-        log_sizes: &[u32],
-        channel: &mut ProofChannel,
-    ) {
+    pub fn commit<C>(&mut self, commitment: C::Digest, log_sizes: &[u32], channel: &mut C)
+    where
+        C: Channel,
+        H: MerkleHasher<Hash = C::Digest>,
+    {
         channel.mix_digest(commitment);
         let extended_log_sizes = log_sizes
             .iter()
@@ -56,12 +51,16 @@ impl CommitmentSchemeVerifier {
         self.trees.push(verifier);
     }
 
-    pub fn verify_values(
+    pub fn verify_values<C>(
         &self,
         sampled_points: TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>>,
-        proof: CommitmentSchemeProof,
-        channel: &mut ProofChannel,
-    ) -> Result<(), VerificationError> {
+        proof: CommitmentSchemeProof<H>,
+        channel: &mut C,
+    ) -> Result<(), VerificationError>
+    where
+        C: Channel,
+        H: MerkleHasher<Hash = C::Digest>,
+    {
         channel.mix_felts(&proof.sampled_values.clone().flatten_cols());
         let random_coeff = channel.draw_felt();
 

crates/prover/src/core/proof_of_work.rs

@@ -2,10 +2,10 @@ use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tracing::{span, Level};
 
-use super::channel::Blake2sChannel;
+use super::channel::Serializable;
+use super::vcs::blake2_hash::Blake2sHasher;
 use crate::core::channel::Channel;
-use crate::core::vcs::blake2_hash::{Blake2sHash, Blake2sHasher};
-use crate::core::vcs::hasher::Hasher;
+use crate::core::vcs::blake2_hash::Blake2sHash;
 
 // TODO(ShaharS): generalize to more channels and create a from function in the hash traits.
 pub struct ProofOfWork {
@@ -23,20 +23,20 @@ impl ProofOfWork {
         Self { n_bits }
     }
 
-    pub fn prove(&self, channel: &mut Blake2sChannel) -> ProofOfWorkProof {
+    pub fn prove(&self, channel: &mut impl Channel) -> ProofOfWorkProof {
         let _span = span!(Level::INFO, "Proof of work").entered();
-        let seed = channel.get_digest().as_ref().to_vec();
+        let seed = channel.get_digest().to_bytes();
         let proof = self.grind(seed);
         channel.mix_nonce(proof.nonce);
         proof
     }
 
     pub fn verify(
         &self,
-        channel: &mut Blake2sChannel,
+        channel: &mut impl Channel,
         proof: &ProofOfWorkProof,
     ) -> Result<(), ProofOfWorkVerificationError> {
-        let seed = channel.get_digest().as_ref().to_vec();
+        let seed = channel.get_digest().to_bytes();
         let verified = check_leading_zeros(
             self.hash_with_nonce(&seed, proof.nonce).as_ref(),
             self.n_bits,