Generic Merkle Channel in Prover (#466)

Author: Gali-StarkWare

crates/cairo_air/src/privacy_test.rs

@@ -19,6 +19,7 @@ use num_traits::Zero;
 use stwo::core::fields::qm31::QM31;
 use stwo::core::fri::FriConfig;
 use stwo::core::pcs::PcsConfig;
+use stwo::core::vcs_lifted::blake2_merkle::Blake2sM31MerkleHasher;
 
 use crate::privacy::{privacy_cairo_verifier_config, privacy_components};
 use crate::test::verify_cairo_with_component_set;
@@ -28,7 +29,7 @@ use crate::verify::build_cairo_verifier_circuit;
 /// Verifies with a circuit a proof of execution of another circuit.
 fn verify_circuit_proof(
     preprocessed_circuit: &PreprocessedCircuit,
-    circuit_proof: CircuitProof,
+    circuit_proof: CircuitProof<Blake2sM31MerkleHasher>,
     preprocessed_root: HashValue<QM31>,
 ) -> Context<QM31> {
     let preprocessed_column_ids = preprocessed_circuit.preprocessed_trace.ids();

crates/circuit_prover/src/prover.rs

@@ -18,8 +18,7 @@ use circuits_stark_verifier::proof_from_stark_proof::{
 use itertools::chain;
 use num_traits::Zero;
 use stwo::core::air::Component;
-use stwo::core::channel::Blake2sM31Channel;
-use stwo::core::channel::Channel;
+use stwo::core::channel::{Channel, MerkleChannel};
 use stwo::core::fields::qm31::QM31;
 use stwo::core::pcs::PcsConfig;
 use stwo::core::poly::circle::CanonicCoset;
@@ -28,6 +27,7 @@ use stwo::core::proof_of_work::GrindOps;
 use stwo::core::utils::MaybeOwned;
 use stwo::core::vcs_lifted::blake2_merkle::Blake2sM31MerkleChannel;
 use stwo::core::vcs_lifted::blake2_merkle::Blake2sM31MerkleHasher;
+use stwo::core::vcs_lifted::merkle_hasher::MerkleHasherLifted;
 use stwo::prover::CommitmentSchemeProver;
 use stwo::prover::CommitmentTreeProver;
 use stwo::prover::ComponentProver;
@@ -39,13 +39,13 @@ use stwo::prover::{ProvingError, prove_ex};
 
 const COMPOSITION_POLYNOMIAL_LOG_DEGREE_BOUND: u32 = 1;
 
-pub struct CircuitProof {
+pub struct CircuitProof<H: MerkleHasherLifted> {
     pub pcs_config: PcsConfig,
     pub claim: CircuitClaim,
     pub interaction_pow_nonce: u64,
     pub interaction_claim: CircuitInteractionClaim,
     pub components: Vec<Box<dyn Component>>,
-    pub stark_proof: Result<ExtendedStarkProof<Blake2sM31MerkleHasher>, ProvingError>,
+    pub stark_proof: Result<ExtendedStarkProof<H>, ProvingError>,
     pub channel_salt: u32,
 }
 
@@ -82,7 +82,25 @@ pub fn prove_circuit_assignment(
     preprocessed_circuit: &PreprocessedCircuit,
     base_column_pool: &BaseColumnPool<SimdBackend>,
     pcs_config: PcsConfig,
-) -> CircuitProof {
+) -> CircuitProof<Blake2sM31MerkleHasher> {
+    prove_circuit_assignment_with_channel::<Blake2sM31MerkleChannel>(
+        values,
+        preprocessed_circuit,
+        base_column_pool,
+        pcs_config,
+    )
+}
+
+pub fn prove_circuit_assignment_with_channel<MC>(
+    values: &[QM31],
+    preprocessed_circuit: &PreprocessedCircuit,
+    base_column_pool: &BaseColumnPool<SimdBackend>,
+    pcs_config: PcsConfig,
+) -> CircuitProof<MC::H>
+where
+    MC: MerkleChannel,
+    SimdBackend: stwo::prover::backend::BackendForChannel<MC>,
+{
     let trace_log_size = preprocessed_circuit.params.trace_log_size;
     let lifting_log_size = trace_log_size + pcs_config.fri_config.log_blowup_factor;
     let pcs_config = PcsConfig { lifting_log_size: Some(lifting_log_size), ..pcs_config };
@@ -106,7 +124,7 @@ pub fn prove_circuit_assignment(
     let preprocessed_trace_polys = SimdBackend::interpolate_columns(preprocessed_trace, &twiddles);
 
     let store_polynomials_coefficients = true;
-    let preprocessed_tree = CommitmentTreeProver::<SimdBackend, Blake2sM31MerkleChannel>::new(
+    let preprocessed_tree = CommitmentTreeProver::<SimdBackend, MC>::new(
         preprocessed_trace_polys,
         pcs_config.fri_config.log_blowup_factor,
         &twiddles,
@@ -115,7 +133,7 @@ pub fn prove_circuit_assignment(
         base_column_pool,
     );
 
-    prove_circuit_with_precompute(
+    prove_circuit_with_precompute::<MC>(
         base_column_pool,
         &twiddles,
         preprocessed_circuit,
@@ -125,14 +143,18 @@ pub fn prove_circuit_assignment(
     )
 }
 
-pub fn prove_circuit_with_precompute<'a>(
+pub fn prove_circuit_with_precompute<'a, MC>(
     base_column_pool: &BaseColumnPool<SimdBackend>,
     twiddles: &TwiddleTree<SimdBackend>,
     preprocessed_circuit: &PreprocessedCircuit,
-    preprocessed_tree: MaybeOwned<'a, CommitmentTreeProver<SimdBackend, Blake2sM31MerkleChannel>>,
+    preprocessed_tree: MaybeOwned<'a, CommitmentTreeProver<SimdBackend, MC>>,
     values: &[QM31],
     pcs_config: PcsConfig,
-) -> CircuitProof {
+) -> CircuitProof<MC::H>
+where
+    MC: MerkleChannel,
+    SimdBackend: stwo::prover::backend::BackendForChannel<MC>,
+{
     let PreprocessedCircuit { preprocessed_trace, params } = preprocessed_circuit;
     let CircuitParams { first_permutation_row, n_blake_gates, output_addresses, .. } = params;
     let trace_generator = TraceGenerator {
@@ -142,18 +164,17 @@ pub fn prove_circuit_with_precompute<'a>(
     };
 
     // Setup protocol.
-    let channel = &mut Blake2sM31Channel::default();
+    let channel = &mut MC::C::default();
 
     // Mix channel salt. Note that we first reduce it modulo `M31::P`, then cast it as QM31.
     let channel_salt = 0_u32;
     channel.mix_felts(&[channel_salt.into()]);
     pcs_config.mix_into(channel);
-    let mut commitment_scheme =
-        CommitmentSchemeProver::<SimdBackend, Blake2sM31MerkleChannel>::with_memory_pool(
-            pcs_config,
-            twiddles,
-            base_column_pool,
-        );
+    let mut commitment_scheme = CommitmentSchemeProver::<SimdBackend, MC>::with_memory_pool(
+        pcs_config,
+        twiddles,
+        base_column_pool,
+    );
 
     commitment_scheme.set_store_polynomials_coefficients();
 
@@ -225,7 +246,7 @@ pub fn prove_circuit_with_precompute<'a>(
 }
 
 pub fn prepare_circuit_proof_for_circuit_verifier(
-    circuit_proof: CircuitProof,
+    circuit_proof: CircuitProof<Blake2sM31MerkleHasher>,
     proof_config: &ProofConfig,
 ) -> (Proof<QM31>, CircuitPublicData) {
     let CircuitProof {

crates/circuit_prover/src/prover_test.rs

@@ -21,7 +21,7 @@ use stwo::core::channel::Channel;
 use stwo::core::fields::qm31::QM31;
 use stwo::core::pcs::{CommitmentSchemeVerifier, PcsConfig, TreeVec};
 use stwo::core::vcs_lifted::blake2_merkle::Blake2sM31MerkleChannel;
-
+use stwo::core::vcs_lifted::blake2_merkle::Blake2sM31MerkleHasher;
 // Not a power of 2 so that we can test component padding.
 const N: usize = 1030;
 
@@ -105,7 +105,10 @@ pub fn build_m31_to_u32_context() -> Context<QM31> {
 
 /// Verifies a [`CircuitProof`] using the stwo verifier. Asserts that the proof is valid
 /// and that the logup sum is zero.
-fn stwo_verify(circuit_proof: CircuitProof, preprocessed_circuit: &PreprocessedCircuit) {
+fn stwo_verify(
+    circuit_proof: CircuitProof<Blake2sM31MerkleHasher>,
+    preprocessed_circuit: &PreprocessedCircuit,
+) {
     let CircuitProof {
         components,
         claim,
@@ -231,7 +234,7 @@ fn test_prove_and_stark_verify_m31_to_u32_context() {
 /// Verifies a [`CircuitProof`] using the circuit verifier. Requires the expected
 /// `preprocessed_root` of the preprocessed trace.
 fn circuit_verify(
-    circuit_proof: CircuitProof,
+    circuit_proof: CircuitProof<Blake2sM31MerkleHasher>,
     preprocessed_circuit: &PreprocessedCircuit,
     preprocessed_root: [u32; 8],
 ) {

crates/circuit_prover/src/witness/trace.rs

@@ -25,9 +25,10 @@ use circuit_common::preprocessed::PreProcessedTrace;
 use itertools::Itertools;
 use num_traits::Zero;
 use rayon::scope;
+use stwo::core::channel::MerkleChannel;
 use stwo::core::fields::qm31::QM31;
-use stwo::core::vcs_lifted::blake2_merkle::Blake2sM31MerkleChannel;
 use stwo::prover::TreeBuilder;
+use stwo::prover::backend::BackendForChannel;
 use stwo::prover::backend::simd::SimdBackend;
 use stwo::prover::poly::circle::PolyOps;
 use stwo::prover::poly::twiddles::TwiddleTree;
@@ -36,14 +37,17 @@ pub struct TraceGenerator {
     pub qm31_ops_trace_generator: Qm31OpsTraceGenerator,
 }
 
-pub fn write_trace(
+pub fn write_trace<MC: MerkleChannel>(
     context_values: &[QM31],
     preprocessed_trace: Arc<PreProcessedTrace>,
     output_addresses: &[usize],
-    tree_builder: &mut TreeBuilder<'_, '_, SimdBackend, Blake2sM31MerkleChannel>,
+    tree_builder: &mut TreeBuilder<'_, '_, SimdBackend, MC>,
     trace_generator: &TraceGenerator,
     twiddles: &TwiddleTree<SimdBackend>,
-) -> (CircuitClaim, CircuitInteractionClaimGenerator) {
+) -> (CircuitClaim, CircuitInteractionClaimGenerator)
+where
+    SimdBackend: BackendForChannel<MC>,
+{
     let preprocessed_trace_ref = preprocessed_trace.as_ref();
 
     // Parent scope: eq/qm31_ops traces run as spawns alongside everything else.
@@ -445,13 +449,16 @@ pub struct CircuitInteractionClaimGenerator {
     pub range_check_16: range_check_16::InteractionClaimGenerator,
 }
 
-pub fn write_interaction_trace(
+pub fn write_interaction_trace<MC: MerkleChannel>(
     circuit_claim: &CircuitClaim,
     circuit_interaction_claim_generator: CircuitInteractionClaimGenerator,
-    tree_builder: &mut TreeBuilder<'_, '_, SimdBackend, Blake2sM31MerkleChannel>,
+    tree_builder: &mut TreeBuilder<'_, '_, SimdBackend, MC>,
     interaction_elements: &CircuitInteractionElements,
     twiddles: &TwiddleTree<SimdBackend>,
-) -> CircuitInteractionClaim {
+) -> CircuitInteractionClaim
+where
+    SimdBackend: BackendForChannel<MC>,
+{
     let CircuitClaim { log_sizes, output_values: _ } = circuit_claim;
     let mut component_log_size_iter = log_sizes.iter();