add proving cells circuit and clean up
This commit is contained in:
parent
79904f9251
commit
19789fd112
|
@ -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
|
||||
|
|
|
@ -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<F, H>(N: usize) -> Result<(
|
||||
SlotTree<F, H>,
|
||||
Vec<usize>,
|
||||
Vec<MerkleProof<F, H>>,
|
||||
)>
|
||||
where
|
||||
F: RichField + Extendable<2> + Poseidon2,
|
||||
H: Hasher<F> + AlgebraicHasher<F> + Hasher<F>,
|
||||
{
|
||||
// Initialize the slot tree with default data
|
||||
let slot_tree = SlotTree::<F, H>::default();
|
||||
|
||||
// Select N leaf indices to prove
|
||||
let leaf_indices: Vec<usize> = (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<F, C, const D: usize, H>(
|
||||
slot_tree: &SlotTree<F, H>,
|
||||
leaf_indices: &[usize],
|
||||
proofs: &[MerkleProof<F, H>],
|
||||
) -> Result<(CircuitData<F, C, D>, PartialWitness<F>)>
|
||||
where
|
||||
F: RichField + Extendable<D> + Poseidon2,
|
||||
C: GenericConfig<D, F = F>,
|
||||
H: Hasher<F> + AlgebraicHasher<F> + Hasher<F>,
|
||||
{
|
||||
// Create the circuit
|
||||
let config = CircuitConfig::standard_recursion_config();
|
||||
let mut builder = CircuitBuilder::<F, D>::new(config);
|
||||
|
||||
// Create a PartialWitness
|
||||
let mut pw = PartialWitness::new();
|
||||
|
||||
// Initialize the circuit instance
|
||||
let mut circuit_instance = MerkleTreeCircuit::<F, C, D, H> {
|
||||
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::<C>();
|
||||
|
||||
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 = <C as GenericConfig<D>>::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::<F, H>(N).unwrap();
|
||||
|
||||
// Benchmark the circuit building
|
||||
group.bench_function("Single Cell Proof Build", |b| {
|
||||
b.iter(|| {
|
||||
build_circuit::<F, C, D, H>(&slot_tree, &leaf_indices, &proofs).unwrap();
|
||||
})
|
||||
});
|
||||
|
||||
// Build the circuit
|
||||
let (data, pw) = build_circuit::<F, C, D, H>(&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);
|
|
@ -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<F, H>(N: usize) -> Result<(
|
||||
MerkleTree<F, H>,
|
||||
Vec<HashOut<F>>,
|
||||
Vec<usize>,
|
||||
Vec<MerkleProof<F, H>>,
|
||||
HashOut<F>,
|
||||
)>
|
||||
where
|
||||
F: RichField + Extendable<2> + Poseidon2,
|
||||
H: Hasher<F> + AlgebraicHasher<F> + Hasher<F>,
|
||||
{
|
||||
// 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::<Vec<_>>();
|
||||
|
||||
// Hash the data to obtain leaf hashes
|
||||
let leaves: Vec<HashOut<F>> = data
|
||||
.iter()
|
||||
.map(|&element| {
|
||||
PoseidonHash::hash_no_pad(&[element])
|
||||
})
|
||||
.collect();
|
||||
|
||||
let zero_hash = HashOut {
|
||||
elements: [F::ZERO; 4],
|
||||
};
|
||||
let tree = MerkleTree::<F, H>::new(&leaves, zero_hash)?;
|
||||
|
||||
// Select N leaf indices to prove
|
||||
let leaf_indices: Vec<usize> = (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::<Result<Vec<_>, _>>()?;
|
||||
|
||||
// Expected Merkle root
|
||||
let expected_root = tree.root()?;
|
||||
|
||||
Ok((tree, leaves, leaf_indices, proofs, expected_root))
|
||||
}
|
||||
|
||||
fn build_circuit<F, C, const D: usize, H>(
|
||||
tree: &MerkleTree<F, H>,
|
||||
leaf_indices: &[usize],
|
||||
) -> Result<(CircuitData<F, C, D>, PartialWitness<F>)>
|
||||
where
|
||||
F: RichField + Extendable<D> + Poseidon2,
|
||||
C: GenericConfig<D, F = F>,
|
||||
H: Hasher<F> + AlgebraicHasher<F> + Hasher<F>,
|
||||
{
|
||||
// Create the circuit
|
||||
let config = CircuitConfig::standard_recursion_config();
|
||||
let mut builder = CircuitBuilder::<F, D>::new(config);
|
||||
|
||||
// Create a PartialWitness
|
||||
let mut pw = PartialWitness::new();
|
||||
|
||||
// Initialize the circuit instance
|
||||
let mut circuit_instance = MerkleTreeCircuit::<F, C, D, H> {
|
||||
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::<C>();
|
||||
|
||||
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 = <C as GenericConfig<D>>::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::<F, H>(N).unwrap();
|
||||
|
||||
// Benchmark the circuit building
|
||||
group.bench_function("Merkle Proof Build", |b| {
|
||||
b.iter(|| {
|
||||
build_circuit::<F, C, D, H>(&tree, &leaf_indices).unwrap();
|
||||
})
|
||||
});
|
||||
|
||||
// Build the circuit once to get the data for the proving and verifying steps
|
||||
let (data, pw) = build_circuit::<F, C, D, H>(&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);
|
|
@ -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::<F, D>) -> MerkleTreeTargets<F, C, D, H>{
|
||||
|
||||
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};
|
||||
|
||||
|
|
|
@ -1,2 +1,3 @@
|
|||
pub mod capped_tree_circuit;
|
||||
pub mod safe_tree_circuit;
|
||||
pub mod prove_single_cell;
|
|
@ -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<<BOT_DEPTH; //2^BOT_DEPTH
|
||||
const N_BLOCKS: usize = 1<<(MAX_DEPTH - BOT_DEPTH); // 2^(MAX_DEPTH - BOT_DEPTH)
|
||||
const N_CELLS: usize = N_CELLS_IN_BLOCKS * N_BLOCKS;
|
||||
// hash function used. this is hackish way of doing it because
|
||||
// H::Hash is not consistent with HashOut<F> and causing a lot of headache
|
||||
// will look into this later.
|
||||
type HF = PoseidonHash;
|
||||
|
||||
// ------ Slot Tree --------
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct SlotTree<F: RichField, H: Hasher<F>> {
|
||||
pub tree: MerkleTree<F,H>, // slot tree
|
||||
pub block_trees: Vec<MerkleTree<F,H>>, // vec of block trees
|
||||
pub cell_data: Vec<Vec<F>>, // cell data as field elements
|
||||
pub cell_hash: Vec<HashOut<F>>, // hash of above
|
||||
}
|
||||
|
||||
impl<F: RichField, H: Hasher<F>> Default for SlotTree<F,H>{
|
||||
/// 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::<Vec<_>>()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
// hash it
|
||||
let leaves: Vec<HashOut<F>> = 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::<Vec<_>>();
|
||||
// get the roots or block trees
|
||||
let block_roots = block_trees.iter()
|
||||
.map(|t| {
|
||||
t.root().unwrap()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
// create slot tree
|
||||
let slot_tree = MerkleTree::<F, H>::new(&block_roots, zero).unwrap();
|
||||
Self{
|
||||
tree: slot_tree,
|
||||
block_trees,
|
||||
cell_data,
|
||||
cell_hash: leaves,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<F: RichField, H: Hasher<F>> SlotTree<F, H> {
|
||||
|
||||
/// same as default but with supplied cell data
|
||||
pub fn new(cell_data: Vec<Vec<F>>) -> Self{
|
||||
let leaves: Vec<HashOut<F>> = 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::<Vec<_>>();
|
||||
let block_roots = block_trees.iter()
|
||||
.map(|t| {
|
||||
t.root().unwrap()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let slot_tree = MerkleTree::<F, H>::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<F, H> {
|
||||
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::<F, H> {
|
||||
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<F, H>, root: HashOut<F>) -> Result<bool>{
|
||||
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::<F,H>::reconstruct_root2(leaf_hash,block_path_bits.clone(),block_last_bits.clone(),block_path);
|
||||
let reconstructed_root = MerkleProof::<F,H>::reconstruct_root2(block_res.unwrap(),slot_path_bits,slot_last_bits,slot_path);
|
||||
|
||||
Ok(reconstructed_root.unwrap() == root)
|
||||
}
|
||||
|
||||
fn get_block_tree(leaves: &Vec<HashOut<F>>) -> MerkleTree<F, H> {
|
||||
let zero = HashOut {
|
||||
elements: [F::ZERO; 4],
|
||||
};
|
||||
// Build the Merkle tree
|
||||
let block_tree = MerkleTree::<F, H>::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<bool> {
|
||||
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<D> + Poseidon2,
|
||||
C: GenericConfig<D, F = F>,
|
||||
const D: usize,
|
||||
H: Hasher<F> + AlgebraicHasher<F>,
|
||||
> {
|
||||
pub expected_slot_root_target: HashOutTarget,
|
||||
pub proof_target: MerkleProofTarget,
|
||||
pub leaf_target: Vec<Target>,
|
||||
pub path_bits: Vec<BoolTarget>,
|
||||
pub last_bits: Vec<BoolTarget>,
|
||||
_phantom: PhantomData<(C,H)>,
|
||||
}
|
||||
|
||||
//------- circuit impl --------
|
||||
|
||||
impl<
|
||||
F: RichField + Extendable<D> + Poseidon2,
|
||||
C: GenericConfig<D, F=F>,
|
||||
const D: usize,
|
||||
H: Hasher<F> + AlgebraicHasher<F> + Hasher<F>,
|
||||
> MerkleTreeCircuit<F, C, D, H> {
|
||||
|
||||
pub fn prove_single_cell2(
|
||||
&mut self,
|
||||
builder: &mut CircuitBuilder::<F, D>
|
||||
) -> SingleCellTargets<F, C, D, H> {
|
||||
|
||||
// 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::<Vec<_>>();
|
||||
|
||||
let mut perm_inputs:Vec<Target>= Vec::new();
|
||||
perm_inputs.extend_from_slice(&leaf);
|
||||
let leaf_hash = builder.hash_n_to_hash_no_pad::<H>(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::<Vec<_>>();
|
||||
let mut slot_path_bits = (0..(depth - BOT_DEPTH)).map(|_| builder.add_virtual_bool_target_safe()).collect::<Vec<_>>();
|
||||
|
||||
// last bits (binary decomposition of last_index = nleaves - 1)
|
||||
let block_last_bits = (0..BOT_DEPTH).map(|_| builder.add_virtual_bool_target_safe()).collect::<Vec<_>>();
|
||||
let slot_last_bits = (0..(depth-BOT_DEPTH)).map(|_| builder.add_virtual_bool_target_safe()).collect::<Vec<_>>();
|
||||
|
||||
// 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<F>,
|
||||
targets: &mut SingleCellTargets<F, C, D, H>,
|
||||
leaf_index: usize,
|
||||
leaf: &Vec<F>,
|
||||
proof: MerkleProof<F,H>,
|
||||
)-> 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<F, D >, leaf: &mut Vec<Target>){
|
||||
builder.hash_n_to_hash_no_pad::<H>(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::<F,H>::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 = <C as GenericConfig<D>>::F;
|
||||
type H = PoseidonHash;
|
||||
|
||||
let slot_t = SlotTree::<F,H>::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::<F, D>::new(config);
|
||||
let mut circuit_instance = MerkleTreeCircuit::<F, C, D, H> {
|
||||
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::<C>();
|
||||
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(())
|
||||
}
|
||||
}
|
||||
|
|
@ -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<D> + Poseidon2,
|
||||
C: GenericConfig<D, F = F>,
|
||||
const D: usize,
|
||||
H: Hasher<F> + AlgebraicHasher<F> + KeyedHasher<F>,
|
||||
H: Hasher<F> + AlgebraicHasher<F>,
|
||||
> {
|
||||
pub leaf: HashOutTarget,
|
||||
pub path_bits: Vec<BoolTarget>,
|
||||
pub last_bits: Vec<BoolTarget>,
|
||||
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<D> + Poseidon2,
|
||||
C: GenericConfig<D, F = F>,
|
||||
const D: usize,
|
||||
H: Hasher<F> + AlgebraicHasher<F> + KeyedHasher<F>,
|
||||
H: Hasher<F> + AlgebraicHasher<F>,
|
||||
> {
|
||||
pub tree: MerkleTree<F, H>,
|
||||
pub _phantom: PhantomData<C>,
|
||||
|
@ -59,9 +58,10 @@ impl<
|
|||
F: RichField + Extendable<D> + Poseidon2,
|
||||
C: GenericConfig<D, F=F>,
|
||||
const D: usize,
|
||||
H: Hasher<F> + AlgebraicHasher<F> + KeyedHasher<F>,
|
||||
H: Hasher<F> + AlgebraicHasher<F>,
|
||||
> MerkleTreeCircuit<F, C, D, H> {
|
||||
|
||||
/// defines the computations inside the circuit and returns the targets used
|
||||
pub fn build_circuit(
|
||||
&mut self,
|
||||
builder: &mut CircuitBuilder::<F, D>
|
||||
|
@ -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<F, D>,
|
||||
targets: &mut MerkleTreeTargets<F, C, D, H>,
|
||||
) {
|
||||
) -> 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<D> + Poseidon2,
|
||||
C: GenericConfig<D, F = F>,
|
||||
const D: usize,
|
||||
H: Hasher<F> + AlgebraicHasher<F> + KeyedHasher<F>,
|
||||
H: Hasher<F> + AlgebraicHasher<F>,
|
||||
> MerkleTreeCircuit<F, C, D, H> {
|
||||
/// 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<bool> {
|
||||
pub(crate) fn usize_to_bits_le_padded(&self, index: usize, bit_length: usize) -> Vec<bool> {
|
||||
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::<F, H>::new(&leaves, zero_hash, H::compress)?;
|
||||
let tree = MerkleTree::<F, H>::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::<F, H>::new(&leaves, zero_hash, H::compress)?;
|
||||
let tree = MerkleTree::<F, H>::new(&leaves, zero_hash)?;
|
||||
|
||||
let expected_root = tree.root()?;
|
||||
|
||||
|
|
|
@ -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<F: RichField>: Hasher<F> {
|
||||
fn compress(x: Self::Hash, y: Self::Hash, key: u64) -> Self::Hash;
|
||||
}
|
||||
|
||||
impl KeyedHasher<GoldilocksField> 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<F> 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<F: RichField, H: KeyedHasher<F>> {
|
||||
pub layers: Vec<Vec<H::Hash>>,
|
||||
pub compress: fn(H::Hash, H::Hash, u64) -> H::Hash,
|
||||
pub zero: H::Hash,
|
||||
pub struct MerkleTree<F: RichField, H: Hasher<F>> {
|
||||
pub layers: Vec<Vec<HashOut<F>>>,
|
||||
pub zero: HashOut<F>,
|
||||
phantom_data: PhantomData<H>
|
||||
}
|
||||
|
||||
impl<F: RichField, H: KeyedHasher<F>> MerkleTree<F, H> {
|
||||
impl<F: RichField, H: Hasher<F>> MerkleTree<F, H> {
|
||||
/// 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<F>],
|
||||
zero: HashOut<F>,
|
||||
) -> Result<Self> {
|
||||
let layers = merkle_tree_worker::<F,H>(leaves, zero, compress, true)?;
|
||||
let layers = merkle_tree_worker::<F,H>(leaves, zero, true)?;
|
||||
Ok(Self {
|
||||
layers,
|
||||
compress,
|
||||
// compress,
|
||||
zero,
|
||||
phantom_data: Default::default(),
|
||||
})
|
||||
}
|
||||
|
||||
|
@ -66,7 +57,7 @@ impl<F: RichField, H: KeyedHasher<F>> MerkleTree<F, H> {
|
|||
}
|
||||
|
||||
/// Returns the root hash of the Merkle tree.
|
||||
pub fn root(&self) -> Result<H::Hash> {
|
||||
pub fn root(&self) -> Result<HashOut<F>> {
|
||||
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<F: RichField, H: KeyedHasher<F>> MerkleTree<F, H> {
|
|||
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<F: RichField>(x: HashOut<F>, y: HashOut<F>, key: u64) -> HashOut<F> {
|
||||
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<F: RichField, H: KeyedHasher<F>>(
|
||||
xs: &[H::Hash],
|
||||
zero: H::Hash,
|
||||
compress: fn(H::Hash, H::Hash, u64) -> H::Hash,
|
||||
fn merkle_tree_worker<F: RichField, H: Hasher<F>>(
|
||||
xs: &[HashOut<F>],
|
||||
zero: HashOut<F>,
|
||||
is_bottom_layer: bool,
|
||||
) -> Result<Vec<Vec<H::Hash>>> {
|
||||
) -> Result<Vec<Vec<HashOut<F>>>> {
|
||||
let m = xs.len();
|
||||
if !is_bottom_layer && m == 1 {
|
||||
return Ok(vec![xs.to_vec()]);
|
||||
|
@ -125,7 +125,7 @@ fn merkle_tree_worker<F: RichField, H: KeyedHasher<F>>(
|
|||
|
||||
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::<F>(xs[2 * i], xs[2 * i + 1], key);
|
||||
ys.push(h);
|
||||
}
|
||||
|
||||
|
@ -135,12 +135,12 @@ fn merkle_tree_worker<F: RichField, H: KeyedHasher<F>>(
|
|||
} else {
|
||||
KEY_ODD
|
||||
};
|
||||
let h = compress(xs[n], zero, key);
|
||||
let h = key_compress::<F>(xs[n], zero, key);
|
||||
ys.push(h);
|
||||
}
|
||||
|
||||
let mut layers = vec![xs.to_vec()];
|
||||
let mut upper_layers = merkle_tree_worker::<F,H>(&ys, zero, compress, false)?;
|
||||
let mut upper_layers = merkle_tree_worker::<F,H>(&ys, zero, false)?;
|
||||
layers.append(&mut upper_layers);
|
||||
|
||||
Ok(layers)
|
||||
|
@ -148,12 +148,12 @@ fn merkle_tree_worker<F: RichField, H: KeyedHasher<F>>(
|
|||
|
||||
/// Merkle proof struct, containing the index, path, and other necessary data.
|
||||
#[derive(Clone)]
|
||||
pub struct MerkleProof<F: RichField, H: KeyedHasher<F>> {
|
||||
pub struct MerkleProof<F: RichField, H: Hasher<F>> {
|
||||
pub index: usize, // Index of the leaf
|
||||
pub path: Vec<H::Hash>, // Sibling hashes from the leaf to the root
|
||||
pub path: Vec<HashOut<F>>, // 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<F>,
|
||||
pub(crate) phantom_data: PhantomData<H>
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
|
@ -162,9 +162,9 @@ pub struct MerkleProofTarget {
|
|||
pub path: Vec<HashOutTarget>,
|
||||
}
|
||||
|
||||
impl<F: RichField, H: KeyedHasher<F>> MerkleProof<F, H> {
|
||||
impl<F: RichField, H: Hasher<F>> MerkleProof<F, H> {
|
||||
/// Reconstructs the root hash from the proof and the given leaf.
|
||||
pub fn reconstruct_root(&self, leaf: H::Hash) -> Result<H::Hash> {
|
||||
pub fn reconstruct_root(&self, leaf: HashOut<F>) -> Result<HashOut<F>> {
|
||||
let mut m = self.nleaves;
|
||||
let mut j = self.index;
|
||||
let mut h = leaf;
|
||||
|
@ -174,14 +174,14 @@ impl<F: RichField, H: KeyedHasher<F>> MerkleProof<F, H> {
|
|||
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::<F>(*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::<F>(h, *p, bottom_flag + 2);
|
||||
} else {
|
||||
// Even node
|
||||
h = (self.compress)(h, *p, bottom_flag);
|
||||
h = key_compress::<F>(h, *p, bottom_flag);
|
||||
}
|
||||
}
|
||||
bottom_flag = KEY_NONE;
|
||||
|
@ -192,30 +192,75 @@ impl<F: RichField, H: KeyedHasher<F>> MerkleProof<F, H> {
|
|||
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<F>, path_bits: Vec<bool>, last_bits:Vec<bool>, path: Vec<HashOut<F>>) -> Result<HashOut<F>> {
|
||||
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::<F>(*p, h, key);
|
||||
} else {
|
||||
h = key_compress::<F>(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<bool> {
|
||||
pub fn verify(&self, leaf: HashOut<F>, root: HashOut<F>) -> Result<bool> {
|
||||
let reconstructed_root = self.reconstruct_root(leaf)?;
|
||||
Ok(reconstructed_root == root)
|
||||
}
|
||||
}
|
||||
|
||||
///helper function to compute is_last
|
||||
fn compute_is_last(path_bits: Vec<bool>, last_bits: Vec<bool>) -> Vec<bool> {
|
||||
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<GoldilocksField>,
|
||||
y: HashOut<GoldilocksField>,
|
||||
x: HashOut<F>,
|
||||
y: HashOut<F>,
|
||||
key: u64,
|
||||
) -> HashOut<GoldilocksField> {
|
||||
let key_field = GoldilocksField::from_canonical_u64(key);
|
||||
) -> HashOut<F> {
|
||||
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<GoldilocksField>,
|
||||
) -> Result<MerkleTree<GoldilocksField, PoseidonHash>> {
|
||||
let compress_fn = PoseidonHash::compress;
|
||||
|
||||
data: &[F],
|
||||
zero: HashOut<F>,
|
||||
) -> Result<MerkleTree<F, H>> {
|
||||
// Hash the data to obtain leaf hashes
|
||||
let leaves: Vec<HashOut<GoldilocksField>> = 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::<GoldilocksField, PoseidonHash>::new(&leaves, zero, compress_fn)
|
||||
MerkleTree::<F, H>::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::<Vec<_>>();
|
||||
|
||||
// Hash the data to obtain leaf hashes
|
||||
let leaves: Vec<HashOut<GoldilocksField>> = data
|
||||
let leaves: Vec<HashOut<F>> = 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::<GoldilocksField, PoseidonHash>::new(&leaves, zero, compress_fn)?;
|
||||
let tree = MerkleTree::<F, H>::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::<Vec<_>>();
|
||||
|
||||
// Hash the data to get leaf hashes
|
||||
let leaf_hashes: Vec<HashOut<GoldilocksField>> = data
|
||||
let leaf_hashes: Vec<HashOut<F>> = 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::<Vec<_>>();
|
||||
|
||||
// Hash the data to get leaf hashes
|
||||
let leaf_hashes: Vec<HashOut<GoldilocksField>> = data
|
||||
let leaf_hashes: Vec<HashOut<F>> = 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::<Vec<_>>();
|
||||
|
||||
// Hash the data to get leaf hashes
|
||||
let leaf_hashes: Vec<HashOut<GoldilocksField>> = data
|
||||
let leaf_hashes: Vec<HashOut<F>> = 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::<Vec<_>>();
|
||||
|
||||
// Hash the data to get leaf hashes
|
||||
let leaf_hashes: Vec<HashOut<GoldilocksField>> = data
|
||||
let leaf_hashes: Vec<HashOut<F>> = 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::<GoldilocksField, PoseidonHash>::new(&leaf_hashes, zero, compress_fn)?;
|
||||
let tree = MerkleTree::<F, H>::new(&leaf_hashes, zero)?;
|
||||
|
||||
// Get the root
|
||||
let expected_root = tree.root()?;
|
||||
|
|
Loading…
Reference in New Issue