codex-storage
/
proof-aggregation
mirror of https://github.com/codex-storage/proof-aggregation.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

remove phantom date and refactor

This commit is contained in:
M Alghazwi 2024-11-03 11:45:20 +01:00
parent cba274b64b
commit c725043c4d
4 changed files with 98 additions and 155 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

88
codex-plonky2-circuits/src/circuits/prove_single_cell.rs
View File

@ -36,21 +36,17 @@ use crate::circuits::params::{MAX_DEPTH, BOT_DEPTH, N_FIELD_ELEMS_PER_CELL, N_CE
#[derive(Clone)]
pub struct SlotTreeCircuit<
F: RichField + Extendable<D> + Poseidon2,
C: GenericConfig<D, F = F>,
const D: usize,
H: Hasher<F> + AlgebraicHasher<F>,
> {
pub tree: MerkleTreeCircuit<F,C,D,H>, // slot tree
pub block_trees: Vec<MerkleTreeCircuit<F,C,D,H>>, // vec of block trees
pub tree: MerkleTreeCircuit<F,D>, // slot tree
pub block_trees: Vec<MerkleTreeCircuit<F,D>>, // vec of block trees
pub cell_data: Vec<Vec<F>>, // cell data as field elements
}
impl<
F: RichField + Extendable<D> + Poseidon2,
C: GenericConfig<D, F = F>,
const D: usize,
H: Hasher<F> + AlgebraicHasher<F>,
> Default for SlotTreeCircuit<F,C,D,H>{
> Default for SlotTreeCircuit<F,D>{
/// slot tree with fake data, for testing only
fn default() -> Self {
// generate fake cell data
@ -78,7 +74,7 @@ impl<
let start = i * N_CELLS_IN_BLOCKS;
let end = (i + 1) * N_CELLS_IN_BLOCKS;
let b_tree = Self::get_block_tree(&leaves[start..end].to_vec()); // use helper function
MerkleTreeCircuit::<F,C,D,H>{ tree:b_tree, _phantom:Default::default()}
MerkleTreeCircuit::<F,D>{ tree:b_tree}
})
.collect::<Vec<_>>();
// get the roots or block trees
@ -88,10 +84,10 @@ impl<
})
.collect::<Vec<_>>();
// create slot tree
let slot_tree = MerkleTree::<F, H>::new(&block_roots, zero).unwrap();
let slot_tree = MerkleTree::<F>::new(&block_roots, zero).unwrap();
Self{
tree: MerkleTreeCircuit::<F,C,D,H>{ tree:slot_tree, _phantom:Default::default()},
tree: MerkleTreeCircuit::<F,D>{ tree:slot_tree},
block_trees,
cell_data,
}
@ -100,10 +96,8 @@ impl<
impl<
F: RichField + Extendable<D> + Poseidon2,
C: GenericConfig<D, F = F>,
const D: usize,
H: Hasher<F> + AlgebraicHasher<F>,
> SlotTreeCircuit<F,C,D, H> {
> SlotTreeCircuit<F,D> {
/// Slot tree with fake data, for testing only
pub fn new_for_testing() -> Self {
@ -133,9 +127,9 @@ impl<
let b_tree = Self::get_block_tree(&leaves_block);
// Create a block tree circuit
let block_tree_circuit = MerkleTreeCircuit::<F, C, D, H> {
let block_tree_circuit = MerkleTreeCircuit::<F, D> {
tree: b_tree,
_phantom: Default::default(),
// _phantom: Default::default(),
};
// Now replicate this block tree for all N_BLOCKS blocks
@ -148,16 +142,15 @@ impl<
.collect::<Vec<_>>();
// Create the slot tree from block roots
let slot_tree = MerkleTree::<F, H>::new(&block_roots, zero).unwrap();
let slot_tree = MerkleTree::<F>::new(&block_roots, zero).unwrap();
// Create the full cell data and cell hash by repeating the block data
let cell_data = vec![cell_data_block.clone(); N_BLOCKS].concat();
// Return the constructed Self
Self {
tree: MerkleTreeCircuit::<F, C, D, H> {
tree: MerkleTreeCircuit::<F, D> {
tree: slot_tree,
_phantom: Default::default(),
},
block_trees,
cell_data,
@ -181,7 +174,7 @@ impl<
let start = i * N_CELLS_IN_BLOCKS;
let end = (i + 1) * N_CELLS_IN_BLOCKS;
let b_tree = Self::get_block_tree(&leaves[start..end].to_vec());
MerkleTreeCircuit::<F,C,D,H>{ tree:b_tree, _phantom:Default::default()}
MerkleTreeCircuit::<F,D>{ tree:b_tree}
})
.collect::<Vec<_>>();
let block_roots = block_trees.iter()
@ -189,9 +182,9 @@ impl<
t.tree.root().unwrap()
})
.collect::<Vec<_>>();
let slot_tree = MerkleTree::<F, H>::new(&block_roots, zero).unwrap();
let slot_tree = MerkleTree::<F>::new(&block_roots, zero).unwrap();
Self{
tree: MerkleTreeCircuit::<F,C,D,H>{ tree:slot_tree, _phantom:Default::default()},
tree: MerkleTreeCircuit::<F,D>{ tree:slot_tree},
block_trees,
cell_data,
}
@ -199,7 +192,7 @@ impl<
/// 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> {
pub fn get_proof(&self, index: usize) -> MerkleProof<F> {
let block_index = index/ N_CELLS_IN_BLOCKS;
let leaf_index = index % N_CELLS_IN_BLOCKS;
let block_proof = self.block_trees[block_index].tree.get_proof(leaf_index).unwrap();
@ -209,18 +202,17 @@ impl<
let mut combined_path = block_proof.path.clone();
combined_path.extend(slot_proof.path.clone());
MerkleProof::<F, H> {
MerkleProof::<F> {
index: index,
path: combined_path,
nleaves: self.cell_data.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>{
pub fn verify_cell_proof(&self, proof: MerkleProof<F>, root: HashOut<F>) -> Result<bool>{
let mut block_path_bits = usize_to_bits_le_padded(proof.index, MAX_DEPTH);
let last_index = N_CELLS - 1;
let mut block_last_bits = usize_to_bits_le_padded(last_index, MAX_DEPTH);
@ -235,51 +227,42 @@ impl<
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);
let block_res = MerkleProof::<F>::reconstruct_root2(leaf_hash,block_path_bits.clone(),block_last_bits.clone(),block_path);
let reconstructed_root = MerkleProof::<F>::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> {
fn get_block_tree(leaves: &Vec<HashOut<F>>) -> MerkleTree<F> {
let zero = HashOut {
elements: [F::ZERO; 4],
};
// Build the Merkle tree
let block_tree = MerkleTree::<F, H>::new(leaves, zero).unwrap();
return block_tree;
let block_tree = MerkleTree::<F>::new(leaves, zero).unwrap();
block_tree
}
}
//------- 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 struct SingleCellTargets{
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>,
> SlotTreeCircuit<F, C, D, H> {
> SlotTreeCircuit<F, D> {
pub fn prove_single_cell(
// &mut self,
builder: &mut CircuitBuilder::<F, D>
) -> SingleCellTargets<F, C, D, H> {
) -> SingleCellTargets {
// Retrieve tree depth
let depth = MAX_DEPTH;
@ -289,7 +272,7 @@ impl<
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);
let leaf_hash = builder.hash_n_to_hash_no_pad::<HF>(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<_>>();
@ -315,11 +298,10 @@ impl<
path_bits:block_path_bits,
last_bits: block_last_bits,
merkle_path: block_merkle_path,
_phantom: PhantomData,
};
// reconstruct block root
let block_root = MerkleTreeCircuit::<F,C,D,H>::reconstruct_merkle_root_circuit(builder, &mut block_targets);
let block_root = MerkleTreeCircuit::<F,D>::reconstruct_merkle_root_circuit(builder, &mut block_targets);
// create MerkleTreeTargets struct
let mut slot_targets = MerkleTreeTargets {
@ -327,11 +309,10 @@ impl<
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 = MerkleTreeCircuit::<F,C,D,H>::reconstruct_merkle_root_circuit(builder, &mut slot_targets);
let slot_root = MerkleTreeCircuit::<F,D>::reconstruct_merkle_root_circuit(builder, &mut slot_targets);
// check equality with expected root
for i in 0..NUM_HASH_OUT_ELTS {
@ -355,7 +336,6 @@ impl<
leaf_target: leaf,
path_bits,
last_bits,
_phantom: Default::default(),
};
// Return MerkleTreeTargets
@ -368,10 +348,10 @@ impl<
pub fn single_cell_assign_witness(
&self,
pw: &mut PartialWitness<F>,
targets: &mut SingleCellTargets<F, C, D, H>,
targets: &mut SingleCellTargets,
leaf_index: usize,
leaf: &Vec<F>,
proof: MerkleProof<F,H>,
proof: MerkleProof<F>,
)-> Result<()> {
// Assign the leaf to the leaf target
@ -415,7 +395,7 @@ impl<
}
fn hash_leaf(builder: &mut CircuitBuilder<F, D >, leaf: &mut Vec<Target>){
builder.hash_n_to_hash_no_pad::<H>(leaf.to_owned());
builder.hash_n_to_hash_no_pad::<HF>(leaf.to_owned());
}
}
@ -435,7 +415,7 @@ mod tests {
#[test]
fn test_prove_single_cell(){
let slot_t = SlotTreeCircuit::<F,C,D,H>::default();
let slot_t = SlotTreeCircuit::<F,D>::default();
let index = 8;
let proof = slot_t.get_proof(index);
let res = slot_t.verify_cell_proof(proof,slot_t.tree.tree.root().unwrap()).unwrap();
@ -445,7 +425,7 @@ mod tests {
#[test]
fn test_cell_build_circuit() -> Result<()> {
let slot_t = SlotTreeCircuit::<F,C,D,H>::default();
let slot_t = SlotTreeCircuit::<F,D>::default();
// select leaf index to prove
let leaf_index: usize = 8;
@ -460,7 +440,7 @@ mod tests {
let config = CircuitConfig::standard_recursion_config();
let mut builder = CircuitBuilder::<F, D>::new(config);
let mut targets = SlotTreeCircuit::<F,C,D,H>::prove_single_cell(&mut builder);
let mut targets = SlotTreeCircuit::<F,D>::prove_single_cell(&mut builder);
// create a PartialWitness and assign
let mut pw = PartialWitness::new();

39
codex-plonky2-circuits/src/circuits/safe_tree_circuit.rs
View File

@ -18,6 +18,7 @@ use plonky2::plonk::proof::{Proof, ProofWithPublicInputs};
use std::marker::PhantomData;
use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
use serde::Serialize;
use crate::circuits::params::HF;
use crate::circuits::utils::usize_to_bits_le_padded;
use crate::merkle_tree::merkle_safe::{MerkleTree, MerkleProofTarget};
@ -29,17 +30,11 @@ use crate::merkle_tree::merkle_safe::{KEY_NONE,KEY_BOTTOM_LAYER};
// uses the Plonk's permutation argument to check that two elements are equal.
// TODO: double check the need for mask
#[derive(Clone)]
pub struct MerkleTreeTargets<
F: RichField + Extendable<D> + Poseidon2,
C: GenericConfig<D, F = F>,
const D: usize,
H: Hasher<F> + AlgebraicHasher<F>,
> {
pub struct MerkleTreeTargets{
pub leaf: HashOutTarget,
pub path_bits: Vec<BoolTarget>,
pub last_bits: Vec<BoolTarget>,
pub merkle_path: MerkleProofTarget,
pub _phantom: PhantomData<(C, H)>,
}
/// Merkle tree circuit contains the tree and functions for
@ -47,26 +42,21 @@ pub struct MerkleTreeTargets<
#[derive(Clone)]
pub struct MerkleTreeCircuit<
F: RichField + Extendable<D> + Poseidon2,
C: GenericConfig<D, F = F>,
const D: usize,
H: Hasher<F> + AlgebraicHasher<F>,
> {
pub tree: MerkleTree<F, H>,
pub _phantom: PhantomData<C>,
pub tree: MerkleTree<F>,
}
impl<
F: RichField + Extendable<D> + Poseidon2,
C: GenericConfig<D, F=F>,
const D: usize,
H: Hasher<F> + AlgebraicHasher<F>,
> MerkleTreeCircuit<F, C, D, H> {
> MerkleTreeCircuit<F, D> {
/// defines the computations inside the circuit and returns the targets used
pub fn build_circuit(
&mut self,
builder: &mut CircuitBuilder::<F, D>
) -> (MerkleTreeTargets<F, C, { D }, H>, HashOutTarget) {
) -> (MerkleTreeTargets, HashOutTarget) {
// Retrieve tree depth
let depth = self.tree.depth();
@ -90,7 +80,6 @@ impl<
path_bits,
last_bits,
merkle_path,
_phantom: PhantomData,
};
// Add Merkle proof verification constraints to the circuit
@ -105,7 +94,7 @@ impl<
pub fn assign_witness(
&mut self,
pw: &mut PartialWitness<F>,
targets: &mut MerkleTreeTargets<F, C, D, H>,
targets: &mut MerkleTreeTargets,
leaf_index: usize,
)-> Result<()> {
// Get the total number of leaves and tree depth
@ -151,9 +140,8 @@ impl<
/// takes the params from the targets struct
/// outputs the reconstructed merkle root
pub fn reconstruct_merkle_root_circuit(
// &self,
builder: &mut CircuitBuilder<F, D>,
targets: &mut MerkleTreeTargets<F, C, D, H>,
targets: &mut MerkleTreeTargets,
) -> HashOutTarget {
let max_depth = targets.path_bits.len();
let mut state: HashOutTarget = targets.leaf;
@ -200,7 +188,7 @@ impl<
perm_inputs.extend_from_slice(&left);
perm_inputs.extend_from_slice(&right);
perm_inputs.push(key);
state = builder.hash_n_to_hash_no_pad::<H>(perm_inputs);
state = builder.hash_n_to_hash_no_pad::<HF>(perm_inputs);
i += 1;
}
@ -245,7 +233,7 @@ mod tests {
let zero_hash = HashOut {
elements: [GoldilocksField::ZERO; 4],
};
let tree = MerkleTree::<F, H>::new(&leaves, zero_hash)?;
let tree = MerkleTree::<F>::new(&leaves, zero_hash)?;
// select leaf index to prove
let leaf_index: usize = 8;
@ -262,9 +250,9 @@ mod tests {
// 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> {
let mut circuit_instance = MerkleTreeCircuit::<F, D> {
tree: tree.clone(),
_phantom: PhantomData,
// _phantom: PhantomData,
};
let (mut targets, expected_root_target) = circuit_instance.build_circuit(&mut builder);
@ -312,15 +300,14 @@ mod tests {
let zero_hash = HashOut {
elements: [GoldilocksField::ZERO; 4],
};
let tree = MerkleTree::<F, H>::new(&leaves, zero_hash)?;
let tree = MerkleTree::<F>::new(&leaves, zero_hash)?;
let expected_root = tree.root()?;
let config = CircuitConfig::standard_recursion_config();
let mut builder = CircuitBuilder::<F, D>::new(config);
let mut circuit_instance = MerkleTreeCircuit::<F, C, D, H> {
let mut circuit_instance = MerkleTreeCircuit::<F, D> {
tree: tree.clone(),
_phantom: PhantomData,
};
let (mut targets, expected_root_target) = circuit_instance.build_circuit(&mut builder);

89
codex-plonky2-circuits/src/circuits/sample_cells.rs
View File

@ -26,7 +26,7 @@ use plonky2::plonk::config::PoseidonGoldilocksConfig;
use plonky2::hash::hashing::PlonkyPermutation;
use crate::circuits::prove_single_cell::{SingleCellTargets, SlotTreeCircuit};
use crate::circuits::params::{BOT_DEPTH, DATASET_DEPTH, MAX_DEPTH, N_FIELD_ELEMS_PER_CELL, N_SAMPLES, TESTING_SLOT_INDEX};
use crate::circuits::params::{BOT_DEPTH, DATASET_DEPTH, HF, MAX_DEPTH, N_FIELD_ELEMS_PER_CELL, N_SAMPLES, TESTING_SLOT_INDEX};
use crate::circuits::safe_tree_circuit::{MerkleTreeCircuit, MerkleTreeTargets};
use crate::circuits::utils::{bits_le_padded_to_usize, calculate_cell_index_bits};
@ -36,35 +36,31 @@ use crate::circuits::utils::{bits_le_padded_to_usize, calculate_cell_index_bits}
#[derive(Clone)]
pub struct DatasetTreeCircuit<
F: RichField + Extendable<D> + Poseidon2,
C: GenericConfig<D, F = F>,
const D: usize,
H: Hasher<F> + AlgebraicHasher<F>,
> {
pub tree: MerkleTreeCircuit<F, C, D, H>, // dataset tree
pub slot_trees: Vec<SlotTreeCircuit<F,C,D,H>>, // vec of slot trees
pub tree: MerkleTreeCircuit<F, D>, // dataset tree
pub slot_trees: Vec<SlotTreeCircuit<F,D>>, // vec of slot trees
}
/// Dataset Merkle proof struct, containing the dataset proof and N_SAMPLES proofs.
#[derive(Clone)]
pub struct DatasetMerkleProof<F: RichField, H: Hasher<F>> {
pub struct DatasetMerkleProof<F: RichField> {
pub slot_index: usize,
pub entropy: usize,
pub dataset_proof: MerkleProof<F,H>, // proof for dataset level tree
pub slot_proofs: Vec<MerkleProof<F,H>>, // proofs for sampled slot, contains N_SAMPLES proofs
pub dataset_proof: MerkleProof<F>, // proof for dataset level tree
pub slot_proofs: Vec<MerkleProof<F>>, // proofs for sampled slot, contains N_SAMPLES proofs
}
impl<
F: RichField + Extendable<D> + Poseidon2,
C: GenericConfig<D, F = F>,
const D: usize,
H: Hasher<F> + AlgebraicHasher<F>,
> Default for DatasetTreeCircuit<F,C,D,H> {
> Default for DatasetTreeCircuit<F,D> {
/// dataset tree with fake data, for testing only
fn default() -> Self {
let mut slot_trees = vec![];
let n_slots = 1<<DATASET_DEPTH;
for i in 0..n_slots {
slot_trees.push(SlotTreeCircuit::<F,C,D,H>::default());
slot_trees.push(SlotTreeCircuit::<F,D>::default());
}
// get the roots or slot trees
let slot_roots = slot_trees.iter()
@ -76,9 +72,9 @@ impl<
let zero = HashOut {
elements: [F::ZERO; 4],
};
let dataset_tree = MerkleTree::<F, H>::new(&slot_roots, zero).unwrap();
let dataset_tree = MerkleTree::<F>::new(&slot_roots, zero).unwrap();
Self{
tree: MerkleTreeCircuit::<F,C,D,H>{ tree:dataset_tree, _phantom:Default::default()},
tree: MerkleTreeCircuit::<F,D>{ tree:dataset_tree},
slot_trees,
}
}
@ -87,10 +83,8 @@ impl<
impl<
F: RichField + Extendable<D> + Poseidon2,
C: GenericConfig<D, F = F>,
const D: usize,
H: Hasher<F> + AlgebraicHasher<F>,
> DatasetTreeCircuit<F,C,D,H> {
> DatasetTreeCircuit<F,D> {
/// dataset tree with fake data, for testing only
/// create data for only the TESTING_SLOT_INDEX in params file
pub fn new_for_testing() -> Self {
@ -100,17 +94,16 @@ impl<
let zero = HashOut {
elements: [F::ZERO; 4],
};
let zero_slot = SlotTreeCircuit::<F,C,D,H>{
let zero_slot = SlotTreeCircuit::<F,D>{
tree: MerkleTreeCircuit {
tree: MerkleTree::<F,H>::new(&[zero.clone()], zero.clone()).unwrap(),
_phantom: Default::default(),
tree: MerkleTree::<F>::new(&[zero.clone()], zero.clone()).unwrap(),
},
block_trees: vec![],
cell_data: vec![],
};
for i in 0..n_slots {
if(i == TESTING_SLOT_INDEX) {
slot_trees.push(SlotTreeCircuit::<F, C, D, H>::new_for_testing());
slot_trees.push(SlotTreeCircuit::<F, D>::new_for_testing());
}else{
slot_trees.push(zero_slot.clone());
}
@ -122,15 +115,15 @@ impl<
t.tree.tree.root().unwrap()
})
.collect::<Vec<_>>();
let dataset_tree = MerkleTree::<F, H>::new(&slot_roots, zero).unwrap();
let dataset_tree = MerkleTree::<F>::new(&slot_roots, zero).unwrap();
Self{
tree: MerkleTreeCircuit::<F,C,D,H>{ tree:dataset_tree, _phantom:Default::default()},
tree: MerkleTreeCircuit::<F,D>{ tree:dataset_tree},
slot_trees,
}
}
/// same as default but with supplied slot trees
pub fn new(slot_trees: Vec<SlotTreeCircuit<F,C,D,H>>) -> Self{
pub fn new(slot_trees: Vec<SlotTreeCircuit<F,D>>) -> Self{
// get the roots or slot trees
let slot_roots = slot_trees.iter()
.map(|t| {
@ -141,23 +134,23 @@ impl<
let zero = HashOut {
elements: [F::ZERO; 4],
};
let dataset_tree = MerkleTree::<F, H>::new(&slot_roots, zero).unwrap();
let dataset_tree = MerkleTree::<F>::new(&slot_roots, zero).unwrap();
Self{
tree: MerkleTreeCircuit::<F,C,D,H>{ tree:dataset_tree, _phantom:Default::default()},
tree: MerkleTreeCircuit::<F,D>{ tree:dataset_tree},
slot_trees,
}
}
/// generates a dataset level proof for given slot index
/// just a regular merkle tree proof
pub fn get_proof(&self, index: usize) -> MerkleProof<F, H> {
pub fn get_proof(&self, index: usize) -> MerkleProof<F> {
let dataset_proof = self.tree.tree.get_proof(index).unwrap();
dataset_proof
}
/// generates a proof for given slot index
/// also takes entropy so it can use it sample the slot
pub fn sample_slot(&self, index: usize, entropy: usize) -> DatasetMerkleProof<F, H> {
pub fn sample_slot(&self, index: usize, entropy: usize) -> DatasetMerkleProof<F> {
let dataset_proof = self.tree.tree.get_proof(index).unwrap();
let slot = &self.slot_trees[index];
let slot_root = slot.tree.tree.root().unwrap();
@ -178,7 +171,7 @@ impl<
}
// verify the sampling - non-circuit version
pub fn verify_sampling(&self, proof: DatasetMerkleProof<F,H>) -> Result<bool>{
pub fn verify_sampling(&self, proof: DatasetMerkleProof<F>) -> Result<bool>{
let slot = &self.slot_trees[proof.slot_index];
let slot_root = slot.tree.tree.root().unwrap();
// check dataset level proof
@ -204,12 +197,7 @@ impl<
}
#[derive(Clone)]
pub struct DatasetTargets<
F: RichField + Extendable<D> + Poseidon2,
C: GenericConfig<D, F = F>,
const D: usize,
H: Hasher<F> + AlgebraicHasher<F>,
> {
pub struct DatasetTargets{
pub dataset_proof: MerkleProofTarget, // proof that slot_root in dataset tree
pub dataset_root: HashOutTarget,
@ -219,23 +207,20 @@ pub struct DatasetTargets<
pub slot_root: HashOutTarget,
pub slot_proofs: Vec<MerkleProofTarget>,
_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>,
> DatasetTreeCircuit<F, C, D, H> {
> DatasetTreeCircuit<F, D> {
// in-circuit sampling
// TODO: make it more modular
pub fn sample_slot_circuit(
&mut self,
builder: &mut CircuitBuilder::<F, D>,
)-> DatasetTargets<F,C,D,H>{
)-> DatasetTargets {
// constants
let zero = builder.zero();
@ -271,12 +256,11 @@ impl<
path_bits: d_path_bits,
last_bits: d_last_bits,
merkle_path: d_merkle_path,
_phantom: PhantomData,
};
// dataset reconstructed root
let d_reconstructed_root =
MerkleTreeCircuit::<F,C,D,H>::reconstruct_merkle_root_circuit(builder, &mut d_targets);
MerkleTreeCircuit::<F,D>::reconstruct_merkle_root_circuit(builder, &mut d_targets);
// expected Merkle root
let d_expected_root = builder.add_virtual_hash();
@ -309,7 +293,7 @@ impl<
let mut perm_inputs:Vec<Target>= Vec::new();
perm_inputs.extend_from_slice(&data_i);
let data_i_hash = builder.hash_n_to_hash_no_pad::<H>(perm_inputs);
let data_i_hash = builder.hash_n_to_hash_no_pad::<HF>(perm_inputs);
// counter constant
let ctr_target = builder.constant(F::from_canonical_u64((i+1) as u64));
let mut ctr = builder.add_virtual_hash();
@ -337,22 +321,20 @@ impl<
path_bits:b_path_bits,
last_bits: b_last_bits.clone(),
merkle_path: b_merkle_path,
_phantom: PhantomData,
};
// reconstruct block root
let b_root = MerkleTreeCircuit::<F,C,D,H>::reconstruct_merkle_root_circuit(builder, &mut block_targets);
let b_root = MerkleTreeCircuit::<F,D>::reconstruct_merkle_root_circuit(builder, &mut block_targets);
let mut slot_targets = MerkleTreeTargets {
leaf: b_root,
path_bits:s_path_bits,
last_bits:s_last_bits.clone(),
merkle_path:s_merkle_path,
_phantom: PhantomData,
};
// reconstruct slot root with block root as leaf
let slot_reconstructed_root = MerkleTreeCircuit::<F,C,D,H>::reconstruct_merkle_root_circuit(builder, &mut slot_targets);
let slot_reconstructed_root = MerkleTreeCircuit::<F,D>::reconstruct_merkle_root_circuit(builder, &mut slot_targets);
// check equality with expected root
for i in 0..NUM_HASH_OUT_ELTS {
@ -370,7 +352,7 @@ impl<
}
DatasetTargets::<F,C,D,H>{
DatasetTargets{
dataset_proof: d_targets.merkle_path,
dataset_root: d_expected_root,
cell_data: data_targets,
@ -378,7 +360,6 @@ impl<
slot_index,
slot_root: d_targets.leaf,
slot_proofs: slot_sample_proofs,
_phantom: Default::default(),
}
}
@ -387,7 +368,7 @@ impl<
hash_inputs.extend_from_slice(&entropy.elements);
hash_inputs.extend_from_slice(&slot_root.elements);
hash_inputs.extend_from_slice(&ctr.elements);
let hash_out = builder.hash_n_to_hash_no_pad::<H>(hash_inputs);
let hash_out = builder.hash_n_to_hash_no_pad::<HF>(hash_inputs);
let cell_index_bits = builder.low_bits(hash_out.elements[0], MAX_DEPTH, 64);
cell_index_bits
@ -396,7 +377,7 @@ impl<
pub fn sample_slot_assign_witness(
&mut self,
pw: &mut PartialWitness<F>,
targets: &mut DatasetTargets<F,C,D,H>,
targets: &mut DatasetTargets,
slot_index:usize,
entropy:usize,
){
@ -475,7 +456,7 @@ mod tests {
#[test]
fn test_sample_cells() {
let dataset_t = DatasetTreeCircuit::<F,C,D,H>::default();
let dataset_t = DatasetTreeCircuit::<F,D>::default();
let slot_index = 2;
let entropy = 123;
let proof = dataset_t.sample_slot(slot_index,entropy);
@ -488,7 +469,7 @@ mod tests {
#[test]
fn test_sample_cells_circuit() -> Result<()> {
let mut dataset_t = DatasetTreeCircuit::<F,C,D,H>::default();
let mut dataset_t = DatasetTreeCircuit::<F,D>::default();
let slot_index = 2;
let entropy = 123;
@ -532,7 +513,7 @@ mod tests {
#[test]
fn test_sample_cells_circuit_from_selected_slot() -> Result<()> {
let mut dataset_t = DatasetTreeCircuit::<F,C,D,H>::new_for_testing();
let mut dataset_t = DatasetTreeCircuit::<F,D>::new_for_testing();
let slot_index = TESTING_SLOT_INDEX;
let entropy = 123;

29
codex-plonky2-circuits/src/merkle_tree/merkle_safe.rs
View File

@ -25,24 +25,21 @@ type HF = PoseidonHash;
/// Merkle tree struct, containing the layers, compression function, and zero hash.
#[derive(Clone)]
pub struct MerkleTree<F: RichField, H: Hasher<F>> {
pub struct MerkleTree<F: RichField> {
pub layers: Vec<Vec<HashOut<F>>>,
pub zero: HashOut<F>,
phantom_data: PhantomData<H>
}
impl<F: RichField, H: Hasher<F>> MerkleTree<F, H> {
impl<F: RichField> MerkleTree<F> {
/// Constructs a new Merkle tree from the given leaves.
pub fn new(
leaves: &[HashOut<F>],
zero: HashOut<F>,
) -> Result<Self> {
let layers = merkle_tree_worker::<F,H>(leaves, zero, true)?;
let layers = merkle_tree_worker::<F>(leaves, zero, true)?;
Ok(Self {
layers,
// compress,
zero,
phantom_data: Default::default(),
})
}
@ -64,7 +61,7 @@ impl<F: RichField, H: Hasher<F>> MerkleTree<F, H> {
}
/// Generates a Merkle proof for a given leaf index.
pub fn get_proof(&self, index: usize) -> Result<MerkleProof<F, H>> {
pub fn get_proof(&self, index: usize) -> Result<MerkleProof<F>> {
let depth = self.depth();
let nleaves = self.leaves_count();
@ -91,7 +88,6 @@ impl<F: RichField, H: Hasher<F>> MerkleTree<F, H> {
path,
nleaves,
zero: self.zero,
phantom_data: Default::default(),
})
}
}
@ -107,7 +103,7 @@ fn key_compress<F: RichField>(x: HashOut<F>, y: HashOut<F>, key: u64) -> HashOut
}
/// Build the Merkle tree layers.
fn merkle_tree_worker<F: RichField, H: Hasher<F>>(
fn merkle_tree_worker<F: RichField>(
xs: &[HashOut<F>],
zero: HashOut<F>,
is_bottom_layer: bool,
@ -140,7 +136,7 @@ fn merkle_tree_worker<F: RichField, H: Hasher<F>>(
}
let mut layers = vec![xs.to_vec()];
let mut upper_layers = merkle_tree_worker::<F,H>(&ys, zero, false)?;
let mut upper_layers = merkle_tree_worker::<F>(&ys, zero, false)?;
layers.append(&mut upper_layers);
Ok(layers)
@ -148,12 +144,11 @@ fn merkle_tree_worker<F: RichField, H: Hasher<F>>(
/// Merkle proof struct, containing the index, path, and other necessary data.
#[derive(Clone)]
pub struct MerkleProof<F: RichField, H: Hasher<F>> {
pub struct MerkleProof<F: RichField> {
pub index: usize, // Index of the leaf
pub path: Vec<HashOut<F>>, // Sibling hashes from the leaf to the root
pub nleaves: usize, // Total number of leaves
pub zero: HashOut<F>,
pub(crate) phantom_data: PhantomData<H>
}
#[derive(Clone, Debug, Eq, PartialEq)]
@ -162,7 +157,7 @@ pub struct MerkleProofTarget {
pub path: Vec<HashOutTarget>,
}
impl<F: RichField, H: Hasher<F>> MerkleProof<F, H> {
impl<F: RichField> MerkleProof<F> {
/// Reconstructs the root hash from the proof and the given leaf.
pub fn reconstruct_root(&self, leaf: HashOut<F>) -> Result<HashOut<F>> {
let mut m = self.nleaves;
@ -271,7 +266,7 @@ mod tests {
fn make_tree(
data: &[F],
zero: HashOut<F>,
) -> Result<MerkleTree<F, H>> {
) -> Result<MerkleTree<F>> {
// Hash the data to obtain leaf hashes
let leaves: Vec<HashOut<GoldilocksField>> = data
.iter()
@ -281,7 +276,7 @@ mod tests {
})
.collect();
MerkleTree::<F, H>::new(&leaves, zero)
MerkleTree::<F>::new(&leaves, zero)
}
#[test]
@ -304,7 +299,7 @@ mod tests {
};
// Build the Merkle tree
let tree = MerkleTree::<F, H>::new(&leaves, zero)?;
let tree = MerkleTree::<F>::new(&leaves, zero)?;
// Get the root
let root = tree.root()?;
@ -536,7 +531,7 @@ mod tests {
};
// Build the tree
let tree = MerkleTree::<F, H>::new(&leaf_hashes, zero)?;
let tree = MerkleTree::<F>::new(&leaf_hashes, zero)?;
// Get the root
let expected_root = tree.root()?;