From 3524e6b1d20264f626298fdde787238dddfb14e7 Mon Sep 17 00:00:00 2001 From: Ohad <137686240+ohad-starkware@users.noreply.github.com> Date: Thu, 21 Mar 2024 14:12:10 +0200 Subject: [PATCH] integrated mixed degree merkle tree (#515) --- src/commitment_scheme/blake2_hash.rs | 2 +- src/commitment_scheme/blake3_hash.rs | 2 +- src/commitment_scheme/hasher.rs | 2 +- src/commitment_scheme/merkle_input.rs | 8 ++ .../mixed_degree_decommitment.rs | 24 ++---- .../mixed_degree_merkle_tree.rs | 3 +- src/core/commitment_scheme/mod.rs | 85 ++++++++++++++----- src/hash_functions/poseidon.rs | 1 + 8 files changed, 83 insertions(+), 44 deletions(-) diff --git a/src/commitment_scheme/blake2_hash.rs b/src/commitment_scheme/blake2_hash.rs index cc53828e6..8b29bb96a 100644 --- a/src/commitment_scheme/blake2_hash.rs +++ b/src/commitment_scheme/blake2_hash.rs @@ -64,7 +64,7 @@ impl super::hasher::Name for Blake2sHash { impl super::hasher::Hash for Blake2sHash {} // Wrapper for the blake2s Hashing functionalities. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Default)] pub struct Blake2sHasher { state: Blake2s256, } diff --git a/src/commitment_scheme/blake3_hash.rs b/src/commitment_scheme/blake3_hash.rs index 3c8f62bc5..e4dad5554 100644 --- a/src/commitment_scheme/blake3_hash.rs +++ b/src/commitment_scheme/blake3_hash.rs @@ -57,7 +57,7 @@ impl Name for Blake3Hash { impl super::hasher::Hash for Blake3Hash {} // Wrapper for the blake3 Hashing functionalities. -#[derive(Clone)] +#[derive(Clone, Default)] pub struct Blake3Hasher { state: blake3::Hasher, } diff --git a/src/commitment_scheme/hasher.rs b/src/commitment_scheme/hasher.rs index 08694b699..caf527a19 100644 --- a/src/commitment_scheme/hasher.rs +++ b/src/commitment_scheme/hasher.rs @@ -22,7 +22,7 @@ pub trait Name { /// assert_eq!(hash, Blake3Hasher::hash(&[1, 2, 3, 4, 5, 6])); /// ``` -pub trait Hasher: Sized { +pub trait Hasher: Sized + Default { type Hash: Hash; type NativeType: Sized + Eq; diff --git a/src/commitment_scheme/merkle_input.rs b/src/commitment_scheme/merkle_input.rs index 6053fce05..6c1180e61 100644 --- a/src/commitment_scheme/merkle_input.rs +++ b/src/commitment_scheme/merkle_input.rs @@ -117,6 +117,7 @@ impl<'a, F: Field> MerkleTreeInput<'a, F> { /// The column layout of a mixed degree merkle tree. /// The sizes of columns assigned to every layer, ordered as they were inserted & injected into hash /// blocks. +#[derive(Debug, Default)] pub struct MerkleTreeColumnLayout { column_sizes: Vec, injected_depths_map: Vec>, @@ -148,6 +149,13 @@ impl MerkleTreeColumnLayout { fn column_indices_at(&self, depth: usize) -> &[usize] { &self.injected_depths_map[depth - 1] } + + pub fn build_input<'a, F: Field>(&self, columns: &[&'a [F]]) -> MerkleTreeInput<'a, F> { + MerkleTreeInput { + columns_to_inject: columns.to_vec(), + injected_depths_map: self.injected_depths_map.clone(), + } + } } #[cfg(test)] diff --git a/src/commitment_scheme/mixed_degree_decommitment.rs b/src/commitment_scheme/mixed_degree_decommitment.rs index dde6a4930..cabd4613d 100644 --- a/src/commitment_scheme/mixed_degree_decommitment.rs +++ b/src/commitment_scheme/mixed_degree_decommitment.rs @@ -13,28 +13,24 @@ use crate::core::fields::{Field, IntoSlice}; /// A correctly generated decommitment should hold all the information needed to generate the root /// of the tree, proving the queried values and the tree's column layout. // TODO(Ohad): write printing functions. +#[derive(Debug, Default)] pub struct MixedDecommitment { pub hashes: Vec, pub witness_elements: Vec, - // TODO(Ohad): remove this in non-debug builds. + // TODO(Ohad): remove these in non-debug builds. pub queried_values: Vec, + pub column_layout: MerkleTreeColumnLayout, } -#[allow(clippy::new_without_default)] impl MixedDecommitment { pub fn new() -> Self { - Self { - hashes: vec![], - witness_elements: vec![], - queried_values: vec![], - } + Self::default() } pub fn verify( &self, root: H::Hash, - column_layout: &MerkleTreeColumnLayout, queries: &[Vec], mut queried_values: impl Iterator, ) -> bool @@ -42,16 +38,16 @@ impl MixedDecommitment { F: IntoSlice, { let mut witness_hashes = self.hashes.iter(); - let sorted_queries_by_layer = column_layout.sort_queries_by_layer(queries); + let sorted_queries_by_layer = self.column_layout.sort_queries_by_layer(queries); let mut next_layer_hashes = vec![]; let mut ancestor_indices = vec![]; let mut witness_elements = self.witness_elements.iter().copied(); - for i in (1..=column_layout.height()).rev() { + for i in (1..=self.column_layout.height()).rev() { (next_layer_hashes, ancestor_indices) = Self::verify_single_layer( i, &sorted_queries_by_layer[i - 1], - column_layout, + &self.column_layout, ancestor_indices.iter().copied().peekable(), queried_values.by_ref(), &mut witness_elements, @@ -166,14 +162,12 @@ mod tests { input.insert_column(TREE_HEIGHT, &column_length_8); input.insert_column(TREE_HEIGHT - 1, &column_length_4); input.insert_column(TREE_HEIGHT - 1, &column_length_8); - let config = input.column_layout(); let (tree, commitment) = MixedDegreeMerkleTree::::commit_default(&input); let queries: Vec> = vec![vec![2], vec![0_usize], vec![4, 7]]; let decommitment = tree.decommit(&input, &queries); assert!(decommitment.verify( commitment, - &config, &queries, decommitment.queried_values.iter().copied(), )); @@ -189,7 +183,6 @@ mod tests { input.insert_column(TREE_HEIGHT, &column_length_8); input.insert_column(TREE_HEIGHT - 1, &column_length_4); input.insert_column(TREE_HEIGHT - 1, &column_length_8); - let config = input.column_layout(); let (tree, _) = MixedDegreeMerkleTree::::commit_default(&input); let false_commitment = Blake3Hasher::hash(b"false_commitment"); @@ -198,7 +191,6 @@ mod tests { assert!(decommitment.verify( false_commitment, - &config, &queries, decommitment.queried_values.iter().copied(), )); @@ -214,7 +206,6 @@ mod tests { input.insert_column(TREE_HEIGHT, &column_length_8); input.insert_column(TREE_HEIGHT - 1, &column_length_4); input.insert_column(TREE_HEIGHT - 1, &column_length_8); - let config = input.column_layout(); let (tree, commitment) = MixedDegreeMerkleTree::::commit_default(&input); let queries: Vec> = vec![vec![2], vec![0_usize], vec![4, 7]]; @@ -223,7 +214,6 @@ mod tests { assert!(decommitment.verify( commitment, - &config, &queries, decommitment.queried_values.iter().copied(), )); diff --git a/src/commitment_scheme/mixed_degree_merkle_tree.rs b/src/commitment_scheme/mixed_degree_merkle_tree.rs index afbaf7204..47dc5a20f 100644 --- a/src/commitment_scheme/mixed_degree_merkle_tree.rs +++ b/src/commitment_scheme/mixed_degree_merkle_tree.rs @@ -179,6 +179,7 @@ where &mut decommitment, ); }); + decommitment.column_layout = input.column_layout(); decommitment } @@ -641,7 +642,6 @@ mod tests { input.insert_column(TREE_HEIGHT - 6, &column_2); input.insert_column(TREE_HEIGHT - 4, &column_1); input.insert_column(TREE_HEIGHT, &column_3); - let configuration = input.column_layout(); let (tree, commitment) = MixedDegreeMerkleTree::::commit_default(&input); let queries: Vec> = vec![ vec![2], @@ -656,7 +656,6 @@ mod tests { let test_decommitment = tree.decommit(&input, &queries); assert!(test_decommitment.verify( commitment, - &configuration, &queries, test_decommitment.queried_values.iter().copied() )); diff --git a/src/core/commitment_scheme/mod.rs b/src/core/commitment_scheme/mod.rs index 6334315a3..8b7a24569 100644 --- a/src/core/commitment_scheme/mod.rs +++ b/src/core/commitment_scheme/mod.rs @@ -5,9 +5,7 @@ //! the unique decoding regime. This is enough for a STARK proof though, where we only want to imply //! the existence of such polynomials, and are ok with having a small decoding list. //! Note: Opened points cannot come from the commitment domain. - pub mod utils; - use std::iter::zip; use std::ops::Deref; @@ -34,8 +32,9 @@ use super::prover::{ use super::queries::SparseSubCircleDomain; use super::ColumnVec; use crate::commitment_scheme::blake2_hash::{Blake2sHash, Blake2sHasher}; -use crate::commitment_scheme::merkle_decommitment::MerkleDecommitment; -use crate::commitment_scheme::merkle_tree::MerkleTree; +use crate::commitment_scheme::merkle_input::{MerkleTreeColumnLayout, MerkleTreeInput}; +use crate::commitment_scheme::mixed_degree_decommitment::MixedDecommitment; +use crate::commitment_scheme::mixed_degree_merkle_tree::MixedDegreeMerkleTree; use crate::core::channel::Channel; type MerkleHasher = Blake2sHasher; @@ -145,7 +144,7 @@ impl CommitmentSchemeProver { #[derive(Debug)] pub struct CommitmentSchemeProof { pub proved_values: TreeVec>>, - pub decommitments: TreeVec>, + pub decommitments: TreeVec>, pub queried_values: TreeVec>>, pub proof_of_work: ProofOfWorkProof, pub fri_proof: FriProof, @@ -156,8 +155,8 @@ pub struct CommitmentSchemeProof { pub struct CommitmentTreeProver { pub polynomials: ColumnVec, pub evaluations: ColumnVec>, - // TODO(AlonH): Change to mixed degree merkle and remove values clone. - commitment: MerkleTree, + pub commitment: MixedDegreeMerkleTree, + column_layout: MerkleTreeColumnLayout, } impl CommitmentTreeProver { @@ -174,18 +173,33 @@ impl CommitmentTreeProver { ) }) .collect_vec(); - let commitment = MerkleTree::::commit( - evaluations - .iter() - .map(|eval| eval.values.clone()) - .collect_vec(), + + let mut merkle_input = MerkleTreeInput::new(); + const LOG_N_BASEFIELD_ELEMENTS_IN_SACK: u32 = 4; + + // The desired depth for column of log_length n is such that Blake2s hashes are filled(64B). + // Explicitly: There are 2^(d-1) hash 'sacks' at depth d, hence, with elements of 4 bytes, + // 2^(d-1) = 2^n / 16, => d = n-3. + // Assuming rectangle trace, all columns go to the same depth. + // TOOD(AlonH): remove this assumption. + let inject_depth = std::cmp::max::( + evaluations[0].len().ilog2() as i32 - (LOG_N_BASEFIELD_ELEMENTS_IN_SACK as i32 - 1), + 1, ); - channel.mix_digest(commitment.root()); + for column in evaluations.iter().map(|eval| &eval.values) { + merkle_input.insert_column(inject_depth as usize, column); + } + let (tree, root) = + MixedDegreeMerkleTree::::commit_default(&merkle_input); + channel.mix_digest(root); + + let column_layout = merkle_input.column_layout(); CommitmentTreeProver { polynomials, evaluations, - commitment, + commitment: tree, + column_layout, } } @@ -196,20 +210,35 @@ impl CommitmentTreeProver { queries: Vec, ) -> ( ColumnVec>, - MerkleDecommitment, + MixedDecommitment, ) { let values = self .evaluations .iter() .map(|c| queries.iter().map(|p| c[*p]).collect()) .collect(); - let decommitment = self.commitment.generate_decommitment(queries); + // Assuming rectangle trace, queries should be similar for all columns. + // TOOD(AlonH): remove this assumption. + let queries = std::iter::repeat(queries.to_vec()) + .take(self.evaluations.len()) + .collect_vec(); + + // Rebuild the merkle input for now. + // TODO(Ohad): change after tree refactor. Consider removing the input struct and have the + // decommitment take queries and columns only. + let eval_vec = self + .evaluations + .iter() + .map(|eval| &eval.values[..]) + .collect_vec(); + let input = self.column_layout.build_input(&eval_vec); + let decommitment = self.commitment.decommit(&input, &queries); (values, decommitment) } } impl Deref for CommitmentTreeProver { - type Target = MerkleTree; + type Target = MixedDegreeMerkleTree; fn deref(&self) -> &Self::Target { &self.commitment @@ -277,10 +306,18 @@ impl CommitmentSchemeVerifier { .as_ref() .zip(&proof.decommitments) .map(|(tree, decommitment)| { - // TODO(spapini): Also very proved_values here. + // TODO(spapini): Also verify proved_values here. + // Assuming columns are of equal lengths, replicate queries for all columns. + // TOOD(AlonH): remove this assumption. tree.verify( decommitment, - &fri_query_domains[&(tree.log_sizes[0] + LOG_BLOWUP_FACTOR)].flatten(), + &std::iter::repeat( + fri_query_domains[&(tree.log_sizes[0] + LOG_BLOWUP_FACTOR)] + .flatten() + .clone(), + ) + .take(tree.log_sizes.len()) + .collect_vec(), ) }) .iter() @@ -370,9 +407,13 @@ impl CommitmentTreeVerifier { pub fn verify( &self, - decommitment: &MerkleDecommitment, - positions: &[usize], + decommitment: &MixedDecommitment, + queries: &[Vec], ) -> bool { - decommitment.verify(self.commitment, positions) + decommitment.verify( + self.commitment, + queries, + decommitment.queried_values.iter().copied(), + ) } } diff --git a/src/hash_functions/poseidon.rs b/src/hash_functions/poseidon.rs index 2bac2f32e..f0eb0a388 100644 --- a/src/hash_functions/poseidon.rs +++ b/src/hash_functions/poseidon.rs @@ -112,6 +112,7 @@ impl From for [BaseField; POSEIDON_CAPACITY] { impl hasher::Hash for PoseidonHash {} +#[derive(Default)] pub struct PoseidonHasher { params: PoseidonParams, state: [BaseField; POSEIDON_WIDTH],