Commit multiple columns using a single Merkle Tree

Cargo.toml

@@ -8,13 +8,14 @@ rust-version = "1.66"
 
 [dependencies]
 rand = "0.8.5"
-lambdaworks-math = { git = "https://github.com/lambdaclass/lambdaworks", rev = "ca0c68a" }
-lambdaworks-crypto = { git = "https://github.com/lambdaclass/lambdaworks", rev = "ca0c68a" }
-lambdaworks-fft = { git = "https://github.com/lambdaclass/lambdaworks", rev = "ca0c68a" }
+lambdaworks-math = { git = "https://github.com/lambdaclass/lambdaworks", rev = "133ffb3" }
+lambdaworks-crypto = { git = "https://github.com/lambdaclass/lambdaworks", rev = "133ffb3" }
+lambdaworks-fft = { git = "https://github.com/lambdaclass/lambdaworks", rev = "133ffb3" }
 thiserror = "1.0.38"
 log = "0.4.17"
 bincode = { version = "2.0.0-rc.2", tag = "v2.0.0-rc.2", git = "https://github.com/bincode-org/bincode.git" }
 cairo-vm = { git = "https://github.com/lambdaclass/cairo-rs/" }
+sha3 = "0.10.6"
 
 [dev-dependencies]
 hex = "0.4.3"

README.md

@@ -14,6 +14,5 @@ To be added:
 - Grinding
 - Skipping FRI layers
 - Optimizing verifier operations
-- Batch proofs in Merkle trees
 - Range check and Pedersen built-ins
 - Different layouts

src/fri/fri_commitment.rs

@@ -6,12 +6,15 @@ use lambdaworks_math::{
     traits::ByteConversion,
 };
 
-use super::HASHER;
 pub use super::{FriMerkleTree, Polynomial};
 use lambdaworks_fft::polynomial::FFTPoly;
 
 #[derive(Clone)]
-pub struct FriLayer<F: IsField> {
+pub struct FriLayer<F>
+where
+    F: IsField,
+    FieldElement<F>: ByteConversion,
+{
     pub poly: Polynomial<FieldElement<F>>,
     pub evaluation: Vec<FieldElement<F>>,
     pub merkle_tree: FriMerkleTree<F>,
@@ -33,7 +36,7 @@ where
             .evaluate_offset_fft(1, Some(domain_size), coset_offset)
             .unwrap(); // TODO: return error
 
-        let merkle_tree = FriMerkleTree::build(&evaluation, Box::new(HASHER));
+        let merkle_tree = FriMerkleTree::build(&evaluation);
 
         Self {
             poly,

src/fri/fri_decommit.rs

@@ -3,10 +3,12 @@ use lambdaworks_crypto::merkle_tree::proof::Proof;
 use lambdaworks_math::field::element::FieldElement;
 use lambdaworks_math::field::traits::IsField;
 
+use super::Commitment;
+
 #[derive(Debug, Clone)]
 pub struct FriDecommitment<F: IsField> {
-    pub layers_auth_paths_sym: Vec<Proof<F>>,
+    pub layers_auth_paths_sym: Vec<Proof<Commitment>>,
     pub layers_evaluations_sym: Vec<FieldElement<F>>,
     pub first_layer_evaluation: FieldElement<F>,
-    pub first_layer_auth_path: Proof<F>,
+    pub first_layer_auth_path: Proof<Commitment>,
 }

src/fri/mod.rs

@@ -4,9 +4,9 @@ mod fri_functions;
 use crate::air::traits::AIR;
 use crate::fri::fri_commitment::FriLayer;
 use crate::{transcript_to_field, transcript_to_usize};
-use lambdaworks_crypto::hash::sha3::Sha3Hasher;
 
 pub use lambdaworks_crypto::fiat_shamir::transcript::Transcript;
+use lambdaworks_crypto::merkle_tree::merkle::FieldElementBackend;
 pub use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
 use lambdaworks_math::field::traits::{IsFFTField, IsField};
 use lambdaworks_math::traits::ByteConversion;
@@ -18,8 +18,9 @@ pub use lambdaworks_math::{
 use self::fri_decommit::FriDecommitment;
 use self::fri_functions::fold_polynomial;
 
-pub type FriMerkleTree<F> = MerkleTree<F>;
-pub(crate) const HASHER: Sha3Hasher = Sha3Hasher::new();
+pub type Commitment = [u8; 32];
+pub type FriMerkleBackend<F> = FieldElementBackend<F>;
+pub type FriMerkleTree<F> = MerkleTree<FieldElementBackend<F>>;
 
 pub fn fri_commit_phase<F: IsField + IsFFTField, T: Transcript>(
     number_layers: usize,
@@ -38,7 +39,7 @@ where
     fri_layer_list.push(current_layer.clone());
 
     // >>>> Send commitment: [p₀]
-    transcript.append(&current_layer.merkle_tree.root.to_bytes_be());
+    transcript.append(&current_layer.merkle_tree.root);
 
     let mut coset_offset = coset_offset.clone();
 
@@ -51,7 +52,7 @@ where
         // Compute layer polynomial and domain
         let next_poly = fold_polynomial(&current_layer.poly, &zeta);
         current_layer = FriLayer::new(next_poly, &coset_offset, domain_size);
-        let new_data = &current_layer.merkle_tree.root.to_bytes_be();
+        let new_data = &current_layer.merkle_tree.root;
         fri_layer_list.push(current_layer.clone()); // TODO: remove this clone
 
         // >>>> Send commitment: [pₖ]

src/lib.rs

@@ -6,17 +6,22 @@ pub mod proof;
 pub mod prover;
 pub mod verifier;
 
+use std::marker::PhantomData;
+
 use air::traits::AIR;
-use lambdaworks_crypto::fiat_shamir::transcript::Transcript;
+use lambdaworks_crypto::{
+    fiat_shamir::transcript::Transcript, merkle_tree::traits::IsMerkleTreeBackend,
+};
 use lambdaworks_fft::roots_of_unity::get_powers_of_primitive_root_coset;
 use lambdaworks_math::{
     field::{
         element::FieldElement,
         fields::fft_friendly::stark_252_prime_field::Stark252PrimeField,
-        traits::{IsFFTField, IsPrimeField},
+        traits::{IsFFTField, IsField, IsPrimeField},
     },
     traits::ByteConversion,
 };
+use sha3::{Digest, Sha3_256};
 
 pub struct ProofConfig {
     pub count_queries: usize,
@@ -145,6 +150,50 @@ impl<F: IsFFTField> Domain<F> {
     }
 }
 
+/// A Merkle tree backend for vectors of field elements.
+/// This is used by the Stark prover to commit to
+/// multiple trace columns using a single Merkle tree.
+#[derive(Clone)]
+pub struct BatchStarkProverBackend<F> {
+    phantom: PhantomData<F>,
+}
+
+impl<F> Default for BatchStarkProverBackend<F> {
+    fn default() -> Self {
+        Self {
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<F> IsMerkleTreeBackend for BatchStarkProverBackend<F>
+where
+    F: IsField,
+    FieldElement<F>: ByteConversion,
+{
+    type Node = [u8; 32];
+    type Data = Vec<FieldElement<F>>;
+
+    fn hash_data(&self, input: &Vec<FieldElement<F>>) -> [u8; 32] {
+        let mut hasher = Sha3_256::new();
+        for element in input.iter() {
+            hasher.update(element.to_bytes_be());
+        }
+        let mut result_hash = [0_u8; 32];
+        result_hash.copy_from_slice(&hasher.finalize());
+        result_hash
+    }
+
+    fn hash_new_parent(&self, left: &[u8; 32], right: &[u8; 32]) -> [u8; 32] {
+        let mut hasher = Sha3_256::new();
+        hasher.update(left);
+        hasher.update(right);
+        let mut result_hash = [0_u8; 32];
+        result_hash.copy_from_slice(&hasher.finalize());
+        result_hash
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use lambdaworks_math::{

src/proof.rs

@@ -1,35 +1,35 @@
 use lambdaworks_crypto::merkle_tree::proof::Proof;
 use lambdaworks_math::field::{element::FieldElement, traits::IsFFTField};
 
-use crate::{air::frame::Frame, fri::fri_decommit::FriDecommitment};
+use crate::{
+    air::frame::Frame,
+    fri::{fri_decommit::FriDecommitment, Commitment},
+};
 
 #[derive(Debug, Clone)]
 pub struct DeepPolynomialOpenings<F: IsFFTField> {
-    pub lde_composition_poly_even_proof: Proof<F>,
+    pub lde_composition_poly_proof: Proof<Commitment>,
     pub lde_composition_poly_even_evaluation: FieldElement<F>,
-    pub lde_composition_poly_odd_proof: Proof<F>,
     pub lde_composition_poly_odd_evaluation: FieldElement<F>,
-    pub lde_trace_merkle_proofs: Vec<Proof<F>>,
+    pub lde_trace_merkle_proofs: Vec<Proof<Commitment>>,
     pub lde_trace_evaluations: Vec<FieldElement<F>>,
 }
 
 #[derive(Debug)]
 pub struct StarkProof<F: IsFFTField> {
     // Commitments of the trace columns
     // [tⱼ]
-    pub lde_trace_merkle_roots: Vec<FieldElement<F>>,
+    pub lde_trace_merkle_roots: Vec<Commitment>,
     // tⱼ(zgᵏ)
     pub trace_ood_frame_evaluations: Frame<F>,
-    // [H₁]
-    pub composition_poly_even_root: FieldElement<F>,
+    // [H₁] and [H₂]
+    pub composition_poly_root: Commitment,
     // H₁(z²)
     pub composition_poly_even_ood_evaluation: FieldElement<F>,
-    // [H₂]
-    pub composition_poly_odd_root: FieldElement<F>,
     // H₂(z²)
     pub composition_poly_odd_ood_evaluation: FieldElement<F>,
     // [pₖ]
-    pub fri_layers_merkle_roots: Vec<FieldElement<F>>,
+    pub fri_layers_merkle_roots: Vec<Commitment>,
     // pₙ
     pub fri_last_value: FieldElement<F>,
     // Open(p₀(D₀), 𝜐ₛ), Opwn(pₖ(Dₖ), −𝜐ₛ^(2ᵏ))