Skip to content

Commit

Permalink
changed interface to mixed merkle tree
Browse files Browse the repository at this point in the history
<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/starkware-libs/stwo/456)
<!-- Reviewable:end -->
  • Loading branch information
ohad-starkware authored Mar 21, 2024
1 parent f482ba8 commit e827cc9
Show file tree
Hide file tree
Showing 3 changed files with 175 additions and 163 deletions.
16 changes: 8 additions & 8 deletions src/commitment_scheme/merkle_input.rs
Original file line number Diff line number Diff line change
Expand Up @@ -98,31 +98,31 @@ impl<'a, F: Field> MerkleTreeInput<'a, F> {
self.columns_to_inject.len()
}

// Returns the structure of the merkle tree. i.e. for each depth, the length of the columns
// Returns the column layout of the merkle tree. i.e. for each depth, the length of the columns
// assigned to it.
// TODO(Ohad): implement this logic for the verifier.
pub fn configuration(&self) -> MerkleTreeConfig {
pub fn column_layout(&self) -> MerkleTreeColumnLayout {
let column_sizes = self
.columns_to_inject
.iter()
.map(|col| col.len())
.collect::<Vec<usize>>();
MerkleTreeConfig {
MerkleTreeColumnLayout {
column_sizes,
injected_depths_map: self.injected_depths_map.clone(),
}
}
}

/// The structure of a mixed degree merkle tree.
/// 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.
pub struct MerkleTreeConfig {
pub struct MerkleTreeColumnLayout {
column_sizes: Vec<usize>,
injected_depths_map: Vec<Vec<usize>>,
}

impl MerkleTreeConfig {
impl MerkleTreeColumnLayout {
pub fn sort_queries_by_layer(&self, queries: &[Vec<usize>]) -> Vec<Vec<Vec<usize>>> {
let mut queries_to_layers = vec![vec![]; self.height()];
(1..=self.height()).for_each(|i| {
Expand Down Expand Up @@ -261,7 +261,7 @@ mod tests {
merkle_input.insert_column(2, &column_length_4);
merkle_input.insert_column(3, &column_length_16);

let merkle_config = merkle_input.configuration();
let merkle_config = merkle_input.column_layout();

assert_eq!(merkle_config.column_lengths_at_depth(3), vec![4, 16]);
assert_eq!(merkle_config.column_lengths_at_depth(2), vec![8, 4]);
Expand All @@ -280,7 +280,7 @@ mod tests {

let queries = vec![vec![0], vec![1], vec![2], vec![3], vec![4], vec![5]];

let merkle_config = merkle_input.configuration();
let merkle_config = merkle_input.column_layout();
let sorted_queries = merkle_config.sort_queries_by_layer(&queries);

assert_eq!(
Expand Down
60 changes: 20 additions & 40 deletions src/commitment_scheme/mixed_degree_decommitment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,14 @@ use itertools::Itertools;
use merging_iterator::MergeIter;

use super::hasher::Hasher;
use super::merkle_input::MerkleTreeConfig;
use super::merkle_input::MerkleTreeColumnLayout;
use super::mixed_degree_merkle_tree::queried_nodes_in_layer;
use crate::core::fields::{Field, IntoSlice};

/// A Merkle proof of queried indices.
/// Used for storing a all the paths from the query leaves to the root.
/// 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 structure.
/// of the tree, proving the queried values and the tree's column layout.
// TODO(Ohad): write printing functions.
pub struct MixedDecommitment<F: Field, H: Hasher> {
pub hashes: Vec<H::Hash>,
Expand All @@ -34,24 +34,24 @@ impl<F: Field, H: Hasher> MixedDecommitment<F, H> {
pub fn verify(
&self,
root: H::Hash,
structure: &MerkleTreeConfig,
column_layout: &MerkleTreeColumnLayout,
queries: &[Vec<usize>],
mut queried_values: impl Iterator<Item = F>,
) -> bool
where
F: IntoSlice<H::NativeType>,
{
let mut witness_hashes = self.hashes.iter();
let sorted_queries_by_layer = structure.sort_queries_by_layer(queries);
let sorted_queries_by_layer = 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..=structure.height()).rev() {
for i in (1..=column_layout.height()).rev() {
(next_layer_hashes, ancestor_indices) = Self::verify_single_layer(
i,
&sorted_queries_by_layer[i - 1],
structure,
column_layout,
ancestor_indices.iter().copied().peekable(),
queried_values.by_ref(),
&mut witness_elements,
Expand All @@ -67,7 +67,7 @@ impl<F: Field, H: Hasher> MixedDecommitment<F, H> {
fn verify_single_layer<'a>(
layer_depth: usize,
queries_to_layer: &[Vec<usize>],
structure: &MerkleTreeConfig,
column_layout: &MerkleTreeColumnLayout,
mut previous_layers_indices: Peekable<impl ExactSizeIterator<Item = usize> + Clone>,
mut queried_values: impl Iterator<Item = F>,
mut witness_elements: impl Iterator<Item = F>,
Expand All @@ -78,7 +78,7 @@ impl<F: Field, H: Hasher> MixedDecommitment<F, H> {
F: IntoSlice<H::NativeType>,
{
let directly_queried_node_indices =
queried_nodes_in_layer(queries_to_layer.iter(), structure, layer_depth);
queried_nodes_in_layer(queries_to_layer.iter(), column_layout, layer_depth);
let mut node_indices = MergeIter::new(
directly_queried_node_indices.iter().copied(),
previous_layers_indices.clone().map(|q| q / 2),
Expand All @@ -97,7 +97,7 @@ impl<F: Field, H: Hasher> MixedDecommitment<F, H> {
for &node_index in &node_indices {
// Push correct child hashes to the hasher.
match previous_layers_indices.next_if(|hash_index| *hash_index / 2 == node_index) {
None if layer_depth < structure.height() => {
None if layer_depth < column_layout.height() => {
hasher.update(witness_hashes_iter.next().unwrap().as_ref());
hasher.update(witness_hashes_iter.next().unwrap().as_ref());
}
Expand Down Expand Up @@ -129,7 +129,7 @@ impl<F: Field, H: Hasher> MixedDecommitment<F, H> {

// Chunk size - according to the column's length and the current depth, we calculate the
// number of elements from that column 'injected' to the current node.
for (chunk_size, column_queries) in structure
for (chunk_size, column_queries) in column_layout
.column_lengths_at_depth(layer_depth)
.iter()
.map(|&column_length| column_length >> (layer_depth - 1))
Expand All @@ -154,9 +154,7 @@ mod tests {
use crate::commitment_scheme::blake3_hash::Blake3Hasher;
use crate::commitment_scheme::hasher::Hasher;
use crate::commitment_scheme::merkle_input::MerkleTreeInput;
use crate::commitment_scheme::mixed_degree_merkle_tree::{
MixedDegreeMerkleTree, MixedDegreeMerkleTreeConfig,
};
use crate::commitment_scheme::mixed_degree_merkle_tree::MixedDegreeMerkleTree;
use crate::core::fields::m31::M31;

#[test]
Expand All @@ -168,17 +166,11 @@ 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.configuration();
let mut tree = MixedDegreeMerkleTree::<M31, Blake3Hasher>::new(
input,
MixedDegreeMerkleTreeConfig {
multi_layer_sizes: [3, 1].to_vec(),
},
);
let commitment = tree.commit();
let config = input.column_layout();
let (tree, commitment) = MixedDegreeMerkleTree::<M31, Blake3Hasher>::commit_default(&input);
let queries: Vec<Vec<usize>> = vec![vec![2], vec![0_usize], vec![4, 7]];

let decommitment = tree.decommit(&queries);
let decommitment = tree.decommit(&input, &queries);
assert!(decommitment.verify(
commitment,
&config,
Expand All @@ -197,18 +189,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.configuration();
let mut tree = MixedDegreeMerkleTree::<M31, Blake3Hasher>::new(
input,
MixedDegreeMerkleTreeConfig {
multi_layer_sizes: [3, 1].to_vec(),
},
);
tree.commit();
let config = input.column_layout();
let (tree, _) = MixedDegreeMerkleTree::<M31, Blake3Hasher>::commit_default(&input);
let false_commitment = Blake3Hasher::hash(b"false_commitment");

let queries: Vec<Vec<usize>> = vec![vec![2], vec![0_usize], vec![4, 7]];
let decommitment = tree.decommit(&queries);
let decommitment = tree.decommit(&input, &queries);

assert!(decommitment.verify(
false_commitment,
Expand All @@ -228,17 +214,11 @@ 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.configuration();
let mut tree = MixedDegreeMerkleTree::<M31, Blake3Hasher>::new(
input,
MixedDegreeMerkleTreeConfig {
multi_layer_sizes: vec![3, 1],
},
);
let commitment = tree.commit();
let config = input.column_layout();
let (tree, commitment) = MixedDegreeMerkleTree::<M31, Blake3Hasher>::commit_default(&input);

let queries: Vec<Vec<usize>> = vec![vec![2], vec![0_usize], vec![4, 7]];
let mut decommitment = tree.decommit(&queries);
let mut decommitment = tree.decommit(&input, &queries);
decommitment.hashes[0] = Blake3Hasher::hash(b"false_hash");

assert!(decommitment.verify(
Expand Down
Loading

0 comments on commit e827cc9

Please sign in to comment.