From 19789fd112dc649f87e985e55f56f383ee99a84e Mon Sep 17 00:00:00 2001 From: M Alghazwi Date: Tue, 15 Oct 2024 14:03:46 +0200 Subject: [PATCH] add proving cells circuit and clean up --- codex-plonky2-circuits/Cargo.toml | 8 + codex-plonky2-circuits/benches/prove_cells.rs | 156 ++++++ .../benches/safe_circuit.rs | 164 +++++++ .../src/circuits/capped_tree_circuit.rs | 16 +- codex-plonky2-circuits/src/circuits/mod.rs | 3 +- .../src/circuits/prove_single_cell.rs | 443 ++++++++++++++++++ .../src/circuits/safe_tree_circuit.rs | 58 +-- .../src/merkle_tree/merkle_safe.rs | 213 +++++---- 8 files changed, 928 insertions(+), 133 deletions(-) create mode 100644 codex-plonky2-circuits/benches/prove_cells.rs create mode 100644 codex-plonky2-circuits/benches/safe_circuit.rs create mode 100644 codex-plonky2-circuits/src/circuits/prove_single_cell.rs diff --git a/codex-plonky2-circuits/Cargo.toml b/codex-plonky2-circuits/Cargo.toml index d68fee0..2502b1d 100644 --- a/codex-plonky2-circuits/Cargo.toml +++ b/codex-plonky2-circuits/Cargo.toml @@ -21,3 +21,11 @@ rand = "0.8.5" [dev-dependencies] criterion = { version = "0.5.1", default-features = false } tynm = { version = "0.1.6", default-features = false } + +[[bench]] +name = "safe_circuit" +harness = false + +[[bench]] +name = "prove_cells" +harness = false diff --git a/codex-plonky2-circuits/benches/prove_cells.rs b/codex-plonky2-circuits/benches/prove_cells.rs new file mode 100644 index 0000000..e52f6b3 --- /dev/null +++ b/codex-plonky2-circuits/benches/prove_cells.rs @@ -0,0 +1,156 @@ +use criterion::{criterion_group, criterion_main, Criterion}; +use anyhow::Result; +use std::time::{Duration, Instant}; + +use codex_plonky2_circuits::{ + merkle_tree::merkle_safe::MerkleProof, + circuits::safe_tree_circuit::MerkleTreeCircuit, +}; +use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData}; +use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, Hasher, PoseidonGoldilocksConfig}; +use plonky2::iop::witness::PartialWitness; +use plonky2::hash::poseidon::PoseidonHash; +use plonky2::field::extension::Extendable; +use plonky2::hash::hash_types::RichField; +use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; +use std::marker::PhantomData; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use codex_plonky2_circuits::circuits::prove_single_cell::SlotTree; + +macro_rules! pretty_print { + ($($arg:tt)*) => { + print!("\x1b[0;36mINFO ===========>\x1b[0m "); + println!($($arg)*); + } +} + +// Hash function used +type HF = PoseidonHash; + +fn prepare_data(N: usize) -> Result<( + SlotTree, + Vec, + Vec>, +)> +where + F: RichField + Extendable<2> + Poseidon2, + H: Hasher + AlgebraicHasher + Hasher, +{ + // Initialize the slot tree with default data + let slot_tree = SlotTree::::default(); + + // Select N leaf indices to prove + let leaf_indices: Vec = (0..N).collect(); + + // Get the Merkle proofs for the selected leaves + let proofs: Vec<_> = leaf_indices + .iter() + .map(|&leaf_index| slot_tree.get_proof(leaf_index)) + .collect(); + + Ok((slot_tree, leaf_indices, proofs)) +} + +fn build_circuit( + slot_tree: &SlotTree, + leaf_indices: &[usize], + proofs: &[MerkleProof], +) -> Result<(CircuitData, PartialWitness)> +where + F: RichField + Extendable + Poseidon2, + C: GenericConfig, + H: Hasher + AlgebraicHasher + Hasher, +{ + // Create the circuit + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + + // Create a PartialWitness + let mut pw = PartialWitness::new(); + + // Initialize the circuit instance + let mut circuit_instance = MerkleTreeCircuit:: { + tree: slot_tree.tree.clone(), + _phantom: PhantomData, + }; + + // For each proof, create targets, add constraints, and assign witnesses + for (i, &leaf_index) in leaf_indices.iter().enumerate() { + // Build the circuit for each proof + let mut targets = circuit_instance.prove_single_cell2(&mut builder); + + // Assign witnesses for each proof + circuit_instance.single_cell_assign_witness( + &mut pw, + &mut targets, + leaf_index, + &slot_tree.cell_data[leaf_index], + proofs[i].clone(), + )?; + } + + // Build the circuit + let data = builder.build::(); + + Ok((data, pw)) +} + +fn single_cell_proof_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("Single Cell Proof Benchmark"); + + // Circuit parameters + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + type H = PoseidonHash; + + // Prepare the data that will be used in all steps + let N = 5; // Number of leaves to prove + let (slot_tree, leaf_indices, proofs) = prepare_data::(N).unwrap(); + + // Benchmark the circuit building + group.bench_function("Single Cell Proof Build", |b| { + b.iter(|| { + build_circuit::(&slot_tree, &leaf_indices, &proofs).unwrap(); + }) + }); + + // Build the circuit + let (data, pw) = build_circuit::(&slot_tree, &leaf_indices, &proofs).unwrap(); + + pretty_print!( + "Circuit size: 2^{} gates", + data.common.degree_bits() + ); + + let start_time = Instant::now(); + let proof_with_pis = data.prove(pw.clone()).unwrap(); + println!("prove_time = {:?}", start_time.elapsed()); + + // Benchmark the proving time + group.bench_function("Single Cell Proof Prove", |b| { + b.iter(|| { + let _proof_with_pis = data.prove(pw.clone()).unwrap(); + }) + }); + + // Generate the proof + let proof_with_pis = data.prove(pw.clone()).unwrap(); + let verifier_data = data.verifier_data(); + + pretty_print!("Proof size: {} bytes", proof_with_pis.to_bytes().len()); + + // Benchmark the verification time + group.bench_function("Single Cell Proof Verify", |b| { + b.iter(|| { + verifier_data.verify(proof_with_pis.clone()).unwrap(); + }) + }); + + group.finish(); +} + +criterion_group!(name = benches; + config = Criterion::default().sample_size(10); + targets = single_cell_proof_benchmark); +criterion_main!(benches); diff --git a/codex-plonky2-circuits/benches/safe_circuit.rs b/codex-plonky2-circuits/benches/safe_circuit.rs new file mode 100644 index 0000000..1fdff3f --- /dev/null +++ b/codex-plonky2-circuits/benches/safe_circuit.rs @@ -0,0 +1,164 @@ +use criterion::{criterion_group, criterion_main, Criterion}; +use anyhow::Result; + +use codex_plonky2_circuits::{merkle_tree::merkle_safe::MerkleTree, circuits::safe_tree_circuit::MerkleTreeCircuit}; +use plonky2::field::types::Field; +use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData}; +use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, Hasher, PoseidonGoldilocksConfig}; +use plonky2::iop::witness::PartialWitness; +use plonky2::hash::hash_types::HashOut; +use plonky2::hash::poseidon::PoseidonHash; +use plonky2::field::extension::Extendable; +use plonky2::hash::hash_types::RichField; +use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; +use std::marker::PhantomData; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use codex_plonky2_circuits::merkle_tree::merkle_safe::MerkleProof; + +macro_rules! pretty_print { + ($($arg:tt)*) => { + print!("\x1b[0;36mINFO ===========>\x1b[0m "); + println!($($arg)*); + } +} + +fn prepare_data(N: usize) -> Result<( + MerkleTree, + Vec>, + Vec, + Vec>, + HashOut, +)> + where + F: RichField + Extendable<2> + Poseidon2, + H: Hasher + AlgebraicHasher + Hasher, +{ + // Total number of leaves in the Merkle tree + let nleaves = 1u64 << 16; + + // Generate leaf data + let data = (0..nleaves) + .map(|i| F::from_canonical_u64(i as u64)) + .collect::>(); + + // Hash the data to obtain leaf hashes + let leaves: Vec> = data + .iter() + .map(|&element| { + PoseidonHash::hash_no_pad(&[element]) + }) + .collect(); + + let zero_hash = HashOut { + elements: [F::ZERO; 4], + }; + let tree = MerkleTree::::new(&leaves, zero_hash)?; + + // Select N leaf indices to prove + let leaf_indices: Vec = (0..N).collect(); + + // Get the Merkle proofs for the selected leaves + let proofs: Vec<_> = leaf_indices + .iter() + .map(|&leaf_index| tree.get_proof(leaf_index)) + .collect::, _>>()?; + + // Expected Merkle root + let expected_root = tree.root()?; + + Ok((tree, leaves, leaf_indices, proofs, expected_root)) +} + +fn build_circuit( + tree: &MerkleTree, + leaf_indices: &[usize], +) -> Result<(CircuitData, PartialWitness)> + where + F: RichField + Extendable + Poseidon2, + C: GenericConfig, + H: Hasher + AlgebraicHasher + Hasher, +{ + // Create the circuit + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + + // Create a PartialWitness + let mut pw = PartialWitness::new(); + + // Initialize the circuit instance + let mut circuit_instance = MerkleTreeCircuit:: { + tree: tree.clone(), + _phantom: PhantomData, + }; + + // For each proof, create targets, add constraints, and assign witnesses + for &leaf_index in leaf_indices.iter() { + // Build the circuit for each proof + let mut targets = circuit_instance.build_circuit(&mut builder); + + // Assign witnesses for each proof + circuit_instance.assign_witness(&mut pw, &mut targets, leaf_index)?; + } + + // Build the circuit + let data = builder.build::(); + + Ok((data, pw)) +} + +fn merkle_proof_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("Merkle Proof Benchmark"); + + // Circuit parameters + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + type H = PoseidonHash; + + // Prepare the data that will be used in all steps + let N = 5; // Number of leaves to prove + let (tree, _leaves, leaf_indices, _proofs, _expected_root) = prepare_data::(N).unwrap(); + + // Benchmark the circuit building + group.bench_function("Merkle Proof Build", |b| { + b.iter(|| { + build_circuit::(&tree, &leaf_indices).unwrap(); + }) + }); + + // Build the circuit once to get the data for the proving and verifying steps + let (data, pw) = build_circuit::(&tree, &leaf_indices).unwrap(); + + pretty_print!( + "circuit size: 2^{} gates", + data.common.degree_bits() + ); + + // Benchmark the proving time + group.bench_function("Merkle Proof Prove", |b| { + b.iter(|| { + let _proof_with_pis = data.prove(pw.clone()).unwrap(); + }) + }); + + // Generate the proof once for verification + let proof_with_pis = data.prove(pw.clone()).unwrap(); + let verifier_data = data.verifier_data(); + + pretty_print!("proof size: {}", proof_with_pis.to_bytes().len()); + + // Benchmark the verification time + group.bench_function("Merkle Proof Verify", |b| { + b.iter(|| { + verifier_data.verify(proof_with_pis.clone()).unwrap(); + }) + }); + + group.finish(); +} + +// criterion_group!(benches, merkle_proof_benchmark); +criterion_group!(name = benches; + config = Criterion::default().sample_size(10); + targets = merkle_proof_benchmark); +criterion_main!(benches); diff --git a/codex-plonky2-circuits/src/circuits/capped_tree_circuit.rs b/codex-plonky2-circuits/src/circuits/capped_tree_circuit.rs index 4f3bce2..39773eb 100644 --- a/codex-plonky2-circuits/src/circuits/capped_tree_circuit.rs +++ b/codex-plonky2-circuits/src/circuits/capped_tree_circuit.rs @@ -1,7 +1,10 @@ +// circuit for regular merkle tree implementation (non-safe version) +// the circuit uses caps in similar way as in Plonky2 Merkle tree implementation +// NOTE: this might be deleted at later time, since we don't use it for codex + use anyhow::Result; use plonky2::field::extension::Extendable; use plonky2::hash::hash_types::RichField; -use plonky2::hash::hashing::hash_n_to_m_no_pad; use plonky2::iop::target::{BoolTarget, Target}; use plonky2::iop::witness::{PartialWitness, WitnessWrite, Witness}; use plonky2::plonk::circuit_builder::CircuitBuilder; @@ -15,16 +18,14 @@ use crate::merkle_tree::capped_tree::MerkleTree; use plonky2::hash::poseidon::PoseidonHash; use plonky2::hash::hash_types::{HashOutTarget, MerkleCapTarget, NUM_HASH_OUT_ELTS}; -use crate::merkle_tree::capped_tree::{MerkleProof, MerkleProofTarget}; -use plonky2_poseidon2::poseidon2_hash::poseidon2::{Poseidon2, Poseidon2Hash}; +use crate::merkle_tree::capped_tree::MerkleProofTarget; +use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; -use plonky2::field::goldilocks_field::GoldilocksField; use plonky2::plonk::config::PoseidonGoldilocksConfig; use plonky2::plonk::proof::Proof; use plonky2::hash::hashing::PlonkyPermutation; use plonky2::plonk::circuit_data::VerifierCircuitTarget; -use crate::merkle_tree::capped_tree::MerkleCap; // size of leaf data (in number of field elements) pub const LEAF_LEN: usize = 4; @@ -66,7 +67,6 @@ impl< } // build the circuit and returns the circuit data - // note, this fn generate circuit data with pub fn build_circuit(&mut self, builder: &mut CircuitBuilder::) -> MerkleTreeTargets{ let proof_t = MerkleProofTarget { @@ -89,8 +89,6 @@ impl< ); MerkleTreeTargets{ - // depth: 0, - // cap_height: 0, proof_target: proof_t, cap_target: cap_t, leaf: leaf_t.to_vec(), @@ -360,7 +358,7 @@ pub mod tests { use super::*; use plonky2::field::types::Field; use crate::merkle_tree::capped_tree::MerkleTree; - use plonky2::iop::witness::{PartialWitness, WitnessWrite}; + use plonky2::iop::witness::PartialWitness; use plonky2::plonk::circuit_data::CircuitConfig; use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig}; diff --git a/codex-plonky2-circuits/src/circuits/mod.rs b/codex-plonky2-circuits/src/circuits/mod.rs index cf32426..f4e17f6 100644 --- a/codex-plonky2-circuits/src/circuits/mod.rs +++ b/codex-plonky2-circuits/src/circuits/mod.rs @@ -1,2 +1,3 @@ pub mod capped_tree_circuit; -pub mod safe_tree_circuit; \ No newline at end of file +pub mod safe_tree_circuit; +pub mod prove_single_cell; \ No newline at end of file diff --git a/codex-plonky2-circuits/src/circuits/prove_single_cell.rs b/codex-plonky2-circuits/src/circuits/prove_single_cell.rs new file mode 100644 index 0000000..cdc53b9 --- /dev/null +++ b/codex-plonky2-circuits/src/circuits/prove_single_cell.rs @@ -0,0 +1,443 @@ +// prove single cell +// consistent with: +// https://github.com/codex-storage/codex-storage-proofs-circuits/blob/master/circuit/codex/single_cell.circom +// circuit consists of: +// - reconstruct the block merkle root +// - use merkle root as leaf and reconstruct slot root +// - check equality with given slot root + +use anyhow::Result; +use plonky2::field::extension::Extendable; +use plonky2::hash::hash_types::{HashOut, RichField}; +use plonky2::iop::target::{BoolTarget, Target}; +use plonky2::iop::witness::{PartialWitness, WitnessWrite, Witness}; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use plonky2::plonk::circuit_data::CircuitConfig; +use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, Hasher, GenericHashOut}; +use std::marker::PhantomData; +use itertools::Itertools; + +use crate::merkle_tree::merkle_safe::MerkleTree; +use plonky2::hash::poseidon::PoseidonHash; + +use plonky2::hash::hash_types::{HashOutTarget, NUM_HASH_OUT_ELTS}; +use crate::merkle_tree::merkle_safe::{MerkleProof, MerkleProofTarget}; +use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; + +use plonky2::field::goldilocks_field::GoldilocksField; +use plonky2::plonk::config::PoseidonGoldilocksConfig; + +use plonky2::hash::hashing::PlonkyPermutation; +use crate::circuits::safe_tree_circuit::{MerkleTreeCircuit, MerkleTreeTargets}; + +// constants and types used throughout the circuit +pub const N_FIELD_ELEMS_PER_CELL: usize = 4; +pub const BOT_DEPTH: usize = 5; // block depth - depth of the block merkle tree +pub const MAX_DEPTH: usize = 16; // depth of big tree (slot tree depth + block tree depth) +const N_CELLS_IN_BLOCKS: usize = 1< and causing a lot of headache +// will look into this later. +type HF = PoseidonHash; + +// ------ Slot Tree -------- + +#[derive(Clone)] +pub struct SlotTree> { + pub tree: MerkleTree, // slot tree + pub block_trees: Vec>, // vec of block trees + pub cell_data: Vec>, // cell data as field elements + pub cell_hash: Vec>, // hash of above +} + +impl> Default for SlotTree{ + /// slot tree with fake data, for testing only + fn default() -> Self { + // generate fake cell data + let mut cell_data = (0..N_CELLS) + .map(|i|{ + (0..N_FIELD_ELEMS_PER_CELL) + .map(|j| F::from_canonical_u64((j+i) as u64)) + .collect::>() + }) + .collect::>(); + // hash it + let leaves: Vec> = cell_data + .iter() + .map(|element| { + HF::hash_no_pad(&element) + }) + .collect(); + // zero hash + let zero = HashOut { + elements: [F::ZERO; 4], + }; + // create block tree + let block_trees = (0..N_BLOCKS) + .map(|i| { + let start = i * N_CELLS_IN_BLOCKS; + let end = (i + 1) * N_CELLS_IN_BLOCKS; + Self::get_block_tree(&leaves[start..end].to_vec()) // use helper function + }) + .collect::>(); + // get the roots or block trees + let block_roots = block_trees.iter() + .map(|t| { + t.root().unwrap() + }) + .collect::>(); + // create slot tree + let slot_tree = MerkleTree::::new(&block_roots, zero).unwrap(); + Self{ + tree: slot_tree, + block_trees, + cell_data, + cell_hash: leaves, + } + } +} + +impl> SlotTree { + + /// same as default but with supplied cell data + pub fn new(cell_data: Vec>) -> Self{ + let leaves: Vec> = cell_data + .iter() + .map(|element| { + HF::hash_no_pad(element) + }) + .collect(); + let zero = HashOut { + elements: [F::ZERO; 4], + }; + let block_trees = (0..N_BLOCKS as usize) + .map(|i| { + let start = i * N_CELLS_IN_BLOCKS; + let end = (i + 1) * N_CELLS_IN_BLOCKS; + Self::get_block_tree(&leaves[start..end].to_vec()) + }) + .collect::>(); + let block_roots = block_trees.iter() + .map(|t| { + t.root().unwrap() + }) + .collect::>(); + let slot_tree = MerkleTree::::new(&block_roots, zero).unwrap(); + Self{ + tree: slot_tree, + block_trees, + cell_data, + cell_hash: leaves, + } + } + + /// generates a proof for given leaf index + /// the path in the proof is a combined block and slot path to make up the full path + pub fn get_proof(&self, index: usize) -> MerkleProof { + let block_index = index/ N_CELLS_IN_BLOCKS; + let leaf_index = index % N_CELLS_IN_BLOCKS; + let block_proof = self.block_trees[block_index].get_proof(leaf_index).unwrap(); + let slot_proof = self.tree.get_proof(block_index).unwrap(); + + // Combine the paths from the block and slot proofs + let mut combined_path = block_proof.path.clone(); + combined_path.extend(slot_proof.path.clone()); + + MerkleProof:: { + index: index, + path: combined_path, + nleaves: self.cell_hash.len(), + zero: block_proof.zero.clone(), + phantom_data: Default::default(), + } + + } + + /// verify the given proof for slot tree, checks equality with given root + pub fn verify_cell_proof(&self, proof: MerkleProof, root: HashOut) -> Result{ + let mut block_path_bits = self.usize_to_bits_le_padded(proof.index, MAX_DEPTH); + let last_index = N_CELLS - 1; + let mut block_last_bits = self.usize_to_bits_le_padded(last_index, MAX_DEPTH); + + let split_point = BOT_DEPTH; + + let slot_last_bits = block_last_bits.split_off(split_point); + let slot_path_bits = block_path_bits.split_off(split_point); + + let leaf_hash = self.cell_hash[proof.index]; + + let mut block_path = proof.path; + let slot_path = block_path.split_off(split_point); + + let block_res = MerkleProof::::reconstruct_root2(leaf_hash,block_path_bits.clone(),block_last_bits.clone(),block_path); + let reconstructed_root = MerkleProof::::reconstruct_root2(block_res.unwrap(),slot_path_bits,slot_last_bits,slot_path); + + Ok(reconstructed_root.unwrap() == root) + } + + fn get_block_tree(leaves: &Vec>) -> MerkleTree { + let zero = HashOut { + elements: [F::ZERO; 4], + }; + // Build the Merkle tree + let block_tree = MerkleTree::::new(leaves, zero).unwrap(); + return block_tree; + } + + /// Converts an index to a vector of bits (LSB first) with padding. + pub(crate) fn usize_to_bits_le_padded(&self, index: usize, bit_length: usize) -> Vec { + let mut bits = Vec::with_capacity(bit_length); + for i in 0..bit_length { + bits.push(((index >> i) & 1) == 1); + } + // If index requires fewer bits, pad with `false` + while bits.len() < bit_length { + bits.push(false); + } + bits + } +} + +//------- single cell struct ------ +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SingleCellTargets< + F: RichField + Extendable + Poseidon2, + C: GenericConfig, + const D: usize, + H: Hasher + AlgebraicHasher, +> { + pub expected_slot_root_target: HashOutTarget, + pub proof_target: MerkleProofTarget, + pub leaf_target: Vec, + pub path_bits: Vec, + pub last_bits: Vec, + _phantom: PhantomData<(C,H)>, +} + +//------- circuit impl -------- + +impl< + F: RichField + Extendable + Poseidon2, + C: GenericConfig, + const D: usize, + H: Hasher + AlgebraicHasher + Hasher, +> MerkleTreeCircuit { + + pub fn prove_single_cell2( + &mut self, + builder: &mut CircuitBuilder:: + ) -> SingleCellTargets { + + // Retrieve tree depth + let depth = MAX_DEPTH; + + // Create virtual targets + let mut leaf = (0..N_FIELD_ELEMS_PER_CELL).map(|_| builder.add_virtual_target()).collect::>(); + + let mut perm_inputs:Vec= Vec::new(); + perm_inputs.extend_from_slice(&leaf); + let leaf_hash = builder.hash_n_to_hash_no_pad::(perm_inputs); + + // path bits (binary decomposition of leaf_index) + let mut block_path_bits = (0..BOT_DEPTH).map(|_| builder.add_virtual_bool_target_safe()).collect::>(); + let mut slot_path_bits = (0..(depth - BOT_DEPTH)).map(|_| builder.add_virtual_bool_target_safe()).collect::>(); + + // last bits (binary decomposition of last_index = nleaves - 1) + let block_last_bits = (0..BOT_DEPTH).map(|_| builder.add_virtual_bool_target_safe()).collect::>(); + let slot_last_bits = (0..(depth-BOT_DEPTH)).map(|_| builder.add_virtual_bool_target_safe()).collect::>(); + + // Merkle path (sibling hashes from leaf to root) + let mut block_merkle_path = MerkleProofTarget { + path: (0..BOT_DEPTH).map(|_| builder.add_virtual_hash()).collect(), + }; + let mut slot_merkle_path = MerkleProofTarget { + path: (0..(depth - BOT_DEPTH)).map(|_| builder.add_virtual_hash()).collect(), + }; + + // expected Merkle root + let slot_expected_root = builder.add_virtual_hash(); + + let mut block_targets = MerkleTreeTargets { + leaf: leaf_hash, + path_bits:block_path_bits, + last_bits: block_last_bits, + merkle_path: block_merkle_path, + _phantom: PhantomData, + }; + + // reconstruct block root + let block_root = self.reconstruct_merkle_root_circuit(builder, &mut block_targets); + + // create MerkleTreeTargets struct + let mut slot_targets = MerkleTreeTargets { + leaf: block_root, + path_bits:slot_path_bits, + last_bits:slot_last_bits, + merkle_path:slot_merkle_path, + _phantom: PhantomData, + }; + + // reconstruct slot root with block root as leaf + let slot_root = self.reconstruct_merkle_root_circuit(builder, &mut slot_targets); + + // check equality with expected root + for i in 0..NUM_HASH_OUT_ELTS { + builder.connect(slot_expected_root.elements[i], slot_root.elements[i]); + } + + let mut proof_target = MerkleProofTarget{ + path: block_targets.merkle_path.path, + }; + proof_target.path.extend_from_slice(&slot_targets.merkle_path.path); + + let mut path_bits = block_targets.path_bits; + path_bits.extend_from_slice(&slot_targets.path_bits); + + let mut last_bits = block_targets.last_bits; + last_bits.extend_from_slice(&slot_targets.last_bits); + + let mut cell_targets = SingleCellTargets { + expected_slot_root_target: slot_expected_root, + proof_target, + leaf_target: leaf, + path_bits, + last_bits, + _phantom: Default::default(), + }; + + // Return MerkleTreeTargets + cell_targets + } + + /// assign the witness values in the circuit targets + /// this takes leaf_index, leaf, and proof (generated from slot_tree) + /// and fills all required circuit targets(circuit inputs) + pub fn single_cell_assign_witness( + &mut self, + pw: &mut PartialWitness, + targets: &mut SingleCellTargets, + leaf_index: usize, + leaf: &Vec, + proof: MerkleProof, + )-> Result<()> { + + // Assign the leaf to the leaf target + for i in 0..targets.leaf_target.len(){ + pw.set_target(targets.leaf_target[i], leaf[i]); + } + + // Convert `leaf_index` to binary bits and assign as path_bits + let path_bits = self.usize_to_bits_le_padded(leaf_index, MAX_DEPTH); + for (i, bit) in path_bits.iter().enumerate() { + pw.set_bool_target(targets.path_bits[i], *bit); + } + + // get `last_index` (nleaves - 1) in binary bits and assign + let last_index = N_CELLS - 1; + let last_bits = self.usize_to_bits_le_padded(last_index, MAX_DEPTH); + for (i, bit) in last_bits.iter().enumerate() { + pw.set_bool_target(targets.last_bits[i], *bit); + } + + // assign the Merkle path (sibling hashes) to the targets + for (i, sibling_hash) in proof.path.iter().enumerate() { + // This is a bit hacky because it should be HashOutTarget, but it is H:Hash + // pw.set_hash_target(targets.merkle_path.path[i],sibling_hash); + // TODO: fix this HashOutTarget later + let sibling_hash_out = sibling_hash.to_vec(); + for j in 0..sibling_hash_out.len() { + pw.set_target(targets.proof_target.path[i].elements[j], sibling_hash_out[j]); + } + } + + // assign the expected Merkle root to the target + let expected_root = self.tree.root()?; + // TODO: fix this HashOutTarget later same issue as above + let expected_root_hash_out = expected_root.to_vec(); + for j in 0..expected_root_hash_out.len() { + pw.set_target(targets.expected_slot_root_target.elements[j], expected_root_hash_out[j]); + } + + Ok(()) + } + + fn hash_leaf(builder: &mut CircuitBuilder, leaf: &mut Vec){ + builder.hash_n_to_hash_no_pad::(leaf.to_owned()); + } +} + +#[cfg(test)] +mod tests { + use std::time::Instant; + use super::*; + use plonky2::plonk::circuit_data::CircuitConfig; + use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig}; + use plonky2::iop::witness::PartialWitness; + + //types for tests + type F = GoldilocksField; + type H = PoseidonHash; + + #[test] + fn test_prove_single_cell(){ + let slot_t = SlotTree::::default(); + let index = 8; + let proof = slot_t.get_proof(index); + let res = slot_t.verify_cell_proof(proof,slot_t.tree.root().unwrap()).unwrap(); + assert_eq!(res, true); + } + + #[test] + fn test_cell_build_circuit() -> Result<()> { + // circuit params + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + type H = PoseidonHash; + + let slot_t = SlotTree::::default(); + + // select leaf index to prove + let leaf_index: usize = 8; + + let proof = slot_t.get_proof(leaf_index); + // get the expected Merkle root + let expected_root = slot_t.tree.root().unwrap(); + let res = slot_t.verify_cell_proof(proof.clone(),expected_root).unwrap(); + assert_eq!(res, true); + + // create the circuit + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + let mut circuit_instance = MerkleTreeCircuit:: { + tree: slot_t.tree.clone(), + _phantom: PhantomData, + }; + let mut targets = circuit_instance.prove_single_cell2(&mut builder); + + // create a PartialWitness and assign + let mut pw = PartialWitness::new(); + circuit_instance.single_cell_assign_witness(&mut pw, &mut targets, leaf_index, &slot_t.cell_data[leaf_index], proof)?; + + // build the circuit + let data = builder.build::(); + println!("circuit size = {:?}", data.common.degree_bits()); + + // Prove the circuit with the assigned witness + let start_time = Instant::now(); + let proof_with_pis = data.prove(pw)?; + println!("prove_time = {:?}", start_time.elapsed()); + + // verify the proof + let verifier_data = data.verifier_data(); + assert!( + verifier_data.verify(proof_with_pis).is_ok(), + "Merkle proof verification failed" + ); + + Ok(()) + } +} + diff --git a/codex-plonky2-circuits/src/circuits/safe_tree_circuit.rs b/codex-plonky2-circuits/src/circuits/safe_tree_circuit.rs index 6f024b5..04c8f75 100644 --- a/codex-plonky2-circuits/src/circuits/safe_tree_circuit.rs +++ b/codex-plonky2-circuits/src/circuits/safe_tree_circuit.rs @@ -6,21 +6,21 @@ use anyhow::Result; use plonky2::field::extension::Extendable; use plonky2::field::goldilocks_field::GoldilocksField; use plonky2::field::types::Field; -use plonky2::hash::hash_types::{HashOut, HashOutTarget, MerkleCapTarget, RichField, NUM_HASH_OUT_ELTS}; -use plonky2::hash::hashing::{hash_n_to_m_no_pad, PlonkyPermutation}; +use plonky2::hash::hash_types::{HashOut, HashOutTarget, RichField, NUM_HASH_OUT_ELTS}; +use plonky2::hash::hashing::PlonkyPermutation; use plonky2::hash::poseidon::PoseidonHash; use plonky2::iop::target::{BoolTarget, Target}; use plonky2::iop::witness::{PartialWitness, Witness, WitnessWrite}; use plonky2::plonk::circuit_builder::CircuitBuilder; -use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, VerifierCircuitData, VerifierCircuitTarget}; +use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, VerifierCircuitData}; use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, GenericHashOut, Hasher, PoseidonGoldilocksConfig}; use plonky2::plonk::proof::{Proof, ProofWithPublicInputs}; use std::marker::PhantomData; use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; use serde::Serialize; -use crate::merkle_tree::merkle_safe::{MerkleTree, MerkleProof, MerkleProofTarget, KeyedHasher}; -use crate::merkle_tree::merkle_safe::{KEY_NONE,KEY_BOTTOM_LAYER,KEY_ODD,KEY_ODD_AND_BOTTOM_LAYER}; +use crate::merkle_tree::merkle_safe::{MerkleTree, MerkleProofTarget}; +use crate::merkle_tree::merkle_safe::{KEY_NONE,KEY_BOTTOM_LAYER}; /// Merkle tree targets representing the input to the circuit @@ -32,14 +32,13 @@ pub struct MerkleTreeTargets< F: RichField + Extendable + Poseidon2, C: GenericConfig, const D: usize, - H: Hasher + AlgebraicHasher + KeyedHasher, + H: Hasher + AlgebraicHasher, > { pub leaf: HashOutTarget, pub path_bits: Vec, pub last_bits: Vec, pub merkle_path: MerkleProofTarget, - pub expected_root: HashOutTarget, - _phantom: PhantomData<(C, H)>, + pub _phantom: PhantomData<(C, H)>, } /// Merkle tree circuit contains the tree and functions for @@ -49,7 +48,7 @@ pub struct MerkleTreeCircuit< F: RichField + Extendable + Poseidon2, C: GenericConfig, const D: usize, - H: Hasher + AlgebraicHasher + KeyedHasher, + H: Hasher + AlgebraicHasher, > { pub tree: MerkleTree, pub _phantom: PhantomData, @@ -59,9 +58,10 @@ impl< F: RichField + Extendable + Poseidon2, C: GenericConfig, const D: usize, - H: Hasher + AlgebraicHasher + KeyedHasher, + H: Hasher + AlgebraicHasher, > MerkleTreeCircuit { + /// defines the computations inside the circuit and returns the targets used pub fn build_circuit( &mut self, builder: &mut CircuitBuilder:: @@ -83,21 +83,17 @@ impl< path: (0..depth).map(|_| builder.add_virtual_hash()).collect(), }; - // expected Merkle root - let expected_root = builder.add_virtual_hash(); - // create MerkleTreeTargets struct let mut targets = MerkleTreeTargets { leaf, path_bits, last_bits, merkle_path, - expected_root, _phantom: PhantomData, }; // Add Merkle proof verification constraints to the circuit - self.verify_merkle_proof_circuit2(builder, &mut targets); + self.reconstruct_merkle_root_circuit(builder, &mut targets); // Return MerkleTreeTargets targets @@ -171,24 +167,16 @@ impl< } } - // assign the expected Merkle root to the target - let expected_root = self.tree.root()?; - // TODO: fix this HashOutTarget later same issue as above - let expected_root_hash_out = expected_root.to_vec(); - for j in 0..expected_root_hash_out.len() { - pw.set_target(targets.expected_root.elements[j], expected_root_hash_out[j]); - } - Ok(()) } - /// Verifies a Merkle proof within the circuit. /// takes the params from the targets struct - pub fn verify_merkle_proof_circuit2( + /// outputs the reconstructed merkle root + pub fn reconstruct_merkle_root_circuit( &self, builder: &mut CircuitBuilder, targets: &mut MerkleTreeTargets, - ) { + ) -> HashOutTarget { let max_depth = targets.path_bits.len(); let mut state: HashOutTarget = targets.leaf; let zero = builder.zero(); @@ -239,11 +227,7 @@ impl< i += 1; } - // check equality with expected root - for i in 0..NUM_HASH_OUT_ELTS { - builder.connect(targets.expected_root.elements[i], state.elements[i]); - } - + return state; } @@ -254,10 +238,10 @@ impl< F: RichField + Extendable + Poseidon2, C: GenericConfig, const D: usize, - H: Hasher + AlgebraicHasher + KeyedHasher, + H: Hasher + AlgebraicHasher, > MerkleTreeCircuit { /// Converts an index to a vector of bits (LSB first) with padding. - fn usize_to_bits_le_padded(&self, index: usize, bit_length: usize) -> Vec { + pub(crate) fn usize_to_bits_le_padded(&self, index: usize, bit_length: usize) -> Vec { let mut bits = Vec::with_capacity(bit_length); for i in 0..bit_length { bits.push(((index >> i) & 1) == 1); @@ -270,6 +254,9 @@ impl< } } + +// NOTE: for now these tests don't check the reconstructed root is equal to expected_root +// will be fixed later, but for that test check the prove_single_cell tests #[cfg(test)] mod tests { use super::*; @@ -277,7 +264,6 @@ mod tests { use plonky2::plonk::circuit_data::CircuitConfig; use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig}; use plonky2::iop::witness::PartialWitness; - use rand::Rng; #[test] fn test_build_circuit() -> Result<()> { @@ -305,7 +291,7 @@ mod tests { let zero_hash = HashOut { elements: [GoldilocksField::ZERO; 4], }; - let tree = MerkleTree::::new(&leaves, zero_hash, H::compress)?; + let tree = MerkleTree::::new(&leaves, zero_hash)?; // select leaf index to prove let leaf_index: usize = 8; @@ -372,7 +358,7 @@ mod tests { let zero_hash = HashOut { elements: [GoldilocksField::ZERO; 4], }; - let tree = MerkleTree::::new(&leaves, zero_hash, H::compress)?; + let tree = MerkleTree::::new(&leaves, zero_hash)?; let expected_root = tree.root()?; diff --git a/codex-plonky2-circuits/src/merkle_tree/merkle_safe.rs b/codex-plonky2-circuits/src/merkle_tree/merkle_safe.rs index 2f52caa..9c1af56 100644 --- a/codex-plonky2-circuits/src/merkle_tree/merkle_safe.rs +++ b/codex-plonky2-circuits/src/merkle_tree/merkle_safe.rs @@ -2,6 +2,7 @@ // consistent with the one in codex: // https://github.com/codex-storage/nim-codex/blob/master/codex/merkletree/merkletree.nim +use std::marker::PhantomData; use anyhow::{ensure, Result}; use plonky2::field::goldilocks_field::GoldilocksField; use plonky2::hash::hash_types::{HashOut, HashOutTarget, RichField}; @@ -10,48 +11,38 @@ use plonky2::plonk::config::Hasher; use std::ops::Shr; use plonky2_field::types::Field; + // Constants for the keys used in compression pub const KEY_NONE: u64 = 0x0; pub const KEY_BOTTOM_LAYER: u64 = 0x1; pub const KEY_ODD: u64 = 0x2; pub const KEY_ODD_AND_BOTTOM_LAYER: u64 = 0x3; -/// Trait for a hash function that supports keyed compression. -pub trait KeyedHasher: Hasher { - fn compress(x: Self::Hash, y: Self::Hash, key: u64) -> Self::Hash; -} - -impl KeyedHasher for PoseidonHash { - fn compress(x: Self::Hash, y: Self::Hash, key: u64) -> Self::Hash { - let key_field = GoldilocksField::from_canonical_u64(key); - let mut inputs = Vec::new(); - inputs.extend_from_slice(&x.elements); - inputs.extend_from_slice(&y.elements); - inputs.push(key_field); - PoseidonHash::hash_no_pad(&inputs) // TODO: double-check this function - } -} +// hash function used. this is hackish way of doing it because +// H::Hash is not consistent with HashOut and causing a lot of headache +// will look into this later. +type HF = PoseidonHash; /// Merkle tree struct, containing the layers, compression function, and zero hash. #[derive(Clone)] -pub struct MerkleTree> { - pub layers: Vec>, - pub compress: fn(H::Hash, H::Hash, u64) -> H::Hash, - pub zero: H::Hash, +pub struct MerkleTree> { + pub layers: Vec>>, + pub zero: HashOut, + phantom_data: PhantomData } -impl> MerkleTree { +impl> MerkleTree { /// Constructs a new Merkle tree from the given leaves. pub fn new( - leaves: &[H::Hash], - zero: H::Hash, - compress: fn(H::Hash, H::Hash, u64) -> H::Hash, + leaves: &[HashOut], + zero: HashOut, ) -> Result { - let layers = merkle_tree_worker::(leaves, zero, compress, true)?; + let layers = merkle_tree_worker::(leaves, zero, true)?; Ok(Self { layers, - compress, + // compress, zero, + phantom_data: Default::default(), }) } @@ -66,7 +57,7 @@ impl> MerkleTree { } /// Returns the root hash of the Merkle tree. - pub fn root(&self) -> Result { + pub fn root(&self) -> Result> { let last_layer = self.layers.last().ok_or_else(|| anyhow::anyhow!("Empty tree"))?; ensure!(last_layer.len() == 1, "Invalid Merkle tree"); Ok(last_layer[0]) @@ -99,19 +90,28 @@ impl> MerkleTree { index, path, nleaves, - compress: self.compress, zero: self.zero, + phantom_data: Default::default(), }) } } +/// compress input (x and y) with key using the define HF hash function +fn key_compress(x: HashOut, y: HashOut, key: u64) -> HashOut { + let key_field = F::from_canonical_u64(key); + let mut inputs = Vec::new(); + inputs.extend_from_slice(&x.elements); + inputs.extend_from_slice(&y.elements); + inputs.push(key_field); + HF::hash_no_pad(&inputs) // TODO: double-check this function +} + /// Build the Merkle tree layers. -fn merkle_tree_worker>( - xs: &[H::Hash], - zero: H::Hash, - compress: fn(H::Hash, H::Hash, u64) -> H::Hash, +fn merkle_tree_worker>( + xs: &[HashOut], + zero: HashOut, is_bottom_layer: bool, -) -> Result>> { +) -> Result>>> { let m = xs.len(); if !is_bottom_layer && m == 1 { return Ok(vec![xs.to_vec()]); @@ -125,7 +125,7 @@ fn merkle_tree_worker>( for i in 0..halfn { let key = if is_bottom_layer { KEY_BOTTOM_LAYER } else { KEY_NONE }; - let h = compress(xs[2 * i], xs[2 * i + 1], key); + let h = key_compress::(xs[2 * i], xs[2 * i + 1], key); ys.push(h); } @@ -135,12 +135,12 @@ fn merkle_tree_worker>( } else { KEY_ODD }; - let h = compress(xs[n], zero, key); + let h = key_compress::(xs[n], zero, key); ys.push(h); } let mut layers = vec![xs.to_vec()]; - let mut upper_layers = merkle_tree_worker::(&ys, zero, compress, false)?; + let mut upper_layers = merkle_tree_worker::(&ys, zero, false)?; layers.append(&mut upper_layers); Ok(layers) @@ -148,12 +148,12 @@ fn merkle_tree_worker>( /// Merkle proof struct, containing the index, path, and other necessary data. #[derive(Clone)] -pub struct MerkleProof> { +pub struct MerkleProof> { pub index: usize, // Index of the leaf - pub path: Vec, // Sibling hashes from the leaf to the root + pub path: Vec>, // Sibling hashes from the leaf to the root pub nleaves: usize, // Total number of leaves - pub compress: fn(H::Hash, H::Hash, u64) -> H::Hash, // compression function - TODO: make it generic instead - pub zero: H::Hash, + pub zero: HashOut, + pub(crate) phantom_data: PhantomData } #[derive(Clone, Debug, Eq, PartialEq)] @@ -162,9 +162,9 @@ pub struct MerkleProofTarget { pub path: Vec, } -impl> MerkleProof { +impl> MerkleProof { /// Reconstructs the root hash from the proof and the given leaf. - pub fn reconstruct_root(&self, leaf: H::Hash) -> Result { + pub fn reconstruct_root(&self, leaf: HashOut) -> Result> { let mut m = self.nleaves; let mut j = self.index; let mut h = leaf; @@ -174,14 +174,14 @@ impl> MerkleProof { let odd_index = (j & 1) != 0; if odd_index { // The index of the child is odd - h = (self.compress)(*p, h, bottom_flag); + h = key_compress::(*p, h, bottom_flag); } else { if j == m - 1 { // Single child -> so odd node - h = (self.compress)(h, *p, bottom_flag + 2); + h = key_compress::(h, *p, bottom_flag + 2); } else { // Even node - h = (self.compress)(h, *p, bottom_flag); + h = key_compress::(h, *p, bottom_flag); } } bottom_flag = KEY_NONE; @@ -192,30 +192,75 @@ impl> MerkleProof { Ok(h) } + /// reconstruct the root using path_bits and last_bits in similar way as the circuit + /// this is used for testing - sanity check + pub fn reconstruct_root2(leaf: HashOut, path_bits: Vec, last_bits:Vec, path: Vec>) -> Result> { + let is_last = compute_is_last(path_bits.clone(),last_bits); + + let mut h = leaf; + let mut i = 0; + + for p in &path { + let bottom = if(i==0){ + KEY_BOTTOM_LAYER + }else{ + KEY_NONE + }; + + let odd = (is_last[i] as usize) * (1-(path_bits[i] as usize)); + + let key = bottom + (2 * (odd as u64)); + let odd_index = path_bits[i]; + if odd_index { + h = key_compress::(*p, h, key); + } else { + h = key_compress::(h, *p, key); + } + i += 1; + } + + Ok(h) + } + /// Verifies the proof against a given root and leaf. - pub fn verify(&self, leaf: H::Hash, root: H::Hash) -> Result { + pub fn verify(&self, leaf: HashOut, root: HashOut) -> Result { let reconstructed_root = self.reconstruct_root(leaf)?; Ok(reconstructed_root == root) } } +///helper function to compute is_last +fn compute_is_last(path_bits: Vec, last_bits: Vec) -> Vec { + let max_depth = path_bits.len(); + + // Initialize isLast vector + let mut is_last = vec![false; max_depth + 1]; + is_last[max_depth] = true; // Set isLast[max_depth] to 1 (true) + + // Iterate over eq and isLast in reverse order + for i in (0..max_depth).rev() { + let eq_out = path_bits[i] == last_bits[i]; // eq[i].out + is_last[i] = is_last[i + 1] && eq_out; // isLast[i] = isLast[i+1] * eq[i].out + } + + is_last +} + #[cfg(test)] mod tests { use super::*; use plonky2::field::types::Field; - // Constants for the keys used in compression - // const KEY_NONE: u64 = 0x0; - // const KEY_BOTTOM_LAYER: u64 = 0x1; - // const KEY_ODD: u64 = 0x2; - // const KEY_ODD_AND_BOTTOM_LAYER: u64 = 0x3; + // types used in all tests + type F = GoldilocksField; + type H = PoseidonHash; fn compress( - x: HashOut, - y: HashOut, + x: HashOut, + y: HashOut, key: u64, - ) -> HashOut { - let key_field = GoldilocksField::from_canonical_u64(key); + ) -> HashOut { + let key_field = F::from_canonical_u64(key); let mut inputs = Vec::new(); inputs.extend_from_slice(&x.elements); inputs.extend_from_slice(&y.elements); @@ -224,46 +269,42 @@ mod tests { } fn make_tree( - data: &[GoldilocksField], - zero: HashOut, - ) -> Result> { - let compress_fn = PoseidonHash::compress; - + data: &[F], + zero: HashOut, + ) -> Result> { // Hash the data to obtain leaf hashes let leaves: Vec> = data .iter() .map(|&element| { // Hash each field element to get the leaf hash - PoseidonHash::hash_no_pad(&[element]) + H::hash_no_pad(&[element]) }) .collect(); - MerkleTree::::new(&leaves, zero, compress_fn) + MerkleTree::::new(&leaves, zero) } #[test] fn single_proof_test() -> Result<()> { let data = (1u64..=8) - .map(|i| GoldilocksField::from_canonical_u64(i)) + .map(|i| F::from_canonical_u64(i)) .collect::>(); // Hash the data to obtain leaf hashes - let leaves: Vec> = data + let leaves: Vec> = data .iter() .map(|&element| { // Hash each field element to get the leaf hash - PoseidonHash::hash_no_pad(&[element]) + H::hash_no_pad(&[element]) }) .collect(); let zero = HashOut { - elements: [GoldilocksField::ZERO; 4], + elements: [F::ZERO; 4], }; - let compress_fn = PoseidonHash::compress; - // Build the Merkle tree - let tree = MerkleTree::::new(&leaves, zero, compress_fn)?; + let tree = MerkleTree::::new(&leaves, zero)?; // Get the root let root = tree.root()?; @@ -282,18 +323,18 @@ mod tests { fn test_correctness_even_bottom_layer() -> Result<()> { // Data for the test (field elements) let data = (1u64..=8) - .map(|i| GoldilocksField::from_canonical_u64(i)) + .map(|i| F::from_canonical_u64(i)) .collect::>(); // Hash the data to get leaf hashes - let leaf_hashes: Vec> = data + let leaf_hashes: Vec> = data .iter() - .map(|&element| PoseidonHash::hash_no_pad(&[element])) + .map(|&element| H::hash_no_pad(&[element])) .collect(); // zero hash let zero = HashOut { - elements: [GoldilocksField::ZERO; 4], + elements: [F::ZERO; 4], }; let expected_root = @@ -343,18 +384,18 @@ mod tests { fn test_correctness_odd_bottom_layer() -> Result<()> { // Data for the test (field elements) let data = (1u64..=7) - .map(|i| GoldilocksField::from_canonical_u64(i)) + .map(|i| F::from_canonical_u64(i)) .collect::>(); // Hash the data to get leaf hashes - let leaf_hashes: Vec> = data + let leaf_hashes: Vec> = data .iter() - .map(|&element| PoseidonHash::hash_no_pad(&[element])) + .map(|&element| H::hash_no_pad(&[element])) .collect(); // zero hash let zero = HashOut { - elements: [GoldilocksField::ZERO; 4], + elements: [F::ZERO; 4], }; let expected_root = @@ -404,18 +445,18 @@ mod tests { fn test_correctness_even_bottom_odd_upper_layers() -> Result<()> { // Data for the test (field elements) let data = (1u64..=10) - .map(|i| GoldilocksField::from_canonical_u64(i)) + .map(|i| F::from_canonical_u64(i)) .collect::>(); // Hash the data to get leaf hashes - let leaf_hashes: Vec> = data + let leaf_hashes: Vec> = data .iter() - .map(|&element| PoseidonHash::hash_no_pad(&[element])) + .map(|&element| H::hash_no_pad(&[element])) .collect(); // zero hash let zero = HashOut { - elements: [GoldilocksField::ZERO; 4], + elements: [F::ZERO; 4], }; let expected_root = compress( @@ -480,24 +521,22 @@ mod tests { fn test_proofs() -> Result<()> { // Data for the test (field elements) let data = (1u64..=10) - .map(|i| GoldilocksField::from_canonical_u64(i)) + .map(|i| F::from_canonical_u64(i)) .collect::>(); // Hash the data to get leaf hashes - let leaf_hashes: Vec> = data + let leaf_hashes: Vec> = data .iter() - .map(|&element| PoseidonHash::hash_no_pad(&[element])) + .map(|&element| H::hash_no_pad(&[element])) .collect(); // zero hash let zero = HashOut { - elements: [GoldilocksField::ZERO; 4], + elements: [F::ZERO; 4], }; - let compress_fn = PoseidonHash::compress; - // Build the tree - let tree = MerkleTree::::new(&leaf_hashes, zero, compress_fn)?; + let tree = MerkleTree::::new(&leaf_hashes, zero)?; // Get the root let expected_root = tree.root()?;