From 76ed312e2c0d3b0c3b49f344226139366911455b Mon Sep 17 00:00:00 2001 From: M Alghazwi Date: Tue, 22 Oct 2024 13:56:39 +0200 Subject: [PATCH] add entropy to cell sampling --- .../src/circuits/safe_tree_circuit.rs | 2 - .../src/circuits/sample_cells.rs | 193 ++++++++++++++++-- codex-plonky2-circuits/src/circuits/utils.rs | 17 +- 3 files changed, 180 insertions(+), 32 deletions(-) diff --git a/codex-plonky2-circuits/src/circuits/safe_tree_circuit.rs b/codex-plonky2-circuits/src/circuits/safe_tree_circuit.rs index c7d5af9..f86dcfa 100644 --- a/codex-plonky2-circuits/src/circuits/safe_tree_circuit.rs +++ b/codex-plonky2-circuits/src/circuits/safe_tree_circuit.rs @@ -207,8 +207,6 @@ impl< return state; } - - } // NOTE: for now these tests don't check the reconstructed root is equal to expected_root diff --git a/codex-plonky2-circuits/src/circuits/sample_cells.rs b/codex-plonky2-circuits/src/circuits/sample_cells.rs index aa565c7..76eb800 100644 --- a/codex-plonky2-circuits/src/circuits/sample_cells.rs +++ b/codex-plonky2-circuits/src/circuits/sample_cells.rs @@ -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::{ DATASET_DEPTH, N_SAMPLES, TESTING_SLOT_INDEX}; +use crate::circuits::params::{BOT_DEPTH, DATASET_DEPTH, 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}; @@ -203,8 +203,6 @@ impl< } } -//------- single cell struct ------ - #[derive(Clone)] pub struct DatasetTargets< F: RichField + Extendable + Poseidon2, @@ -212,9 +210,14 @@ pub struct DatasetTargets< const D: usize, H: Hasher + AlgebraicHasher, > { - pub dataset_proof: MerkleTreeTargets, + pub dataset_proof: MerkleProofTarget, // proof that slot_root in dataset tree pub dataset_root: HashOutTarget, - pub slot_proofs: Vec>, + + pub cell_data: Vec>, + pub entropy:Target, + pub slot_index: Target, + pub slot_root: HashOutTarget, + pub slot_proofs: Vec, _phantom: PhantomData<(C,H)>, } @@ -227,38 +230,159 @@ impl< H: Hasher + AlgebraicHasher + Hasher, > DatasetTreeCircuit { - // the in-circuit sampling of a slot in a dataset + // in-circuit sampling + // TODO: make it more modular pub fn sample_slot_circuit( &mut self, builder: &mut CircuitBuilder::, )-> DatasetTargets{ - let (dataset_proof, dataset_root_target) = self.tree.build_circuit(builder); + // constants + let one = builder.one(); + let two = builder.two(); + + // ***** prove slot root is in dataset tree ********* + + // Retrieve dataset tree depth + let d_depth = DATASET_DEPTH; + + // Create virtual target for slot root and index + let slot_root = builder.add_virtual_hash(); + let slot_index = builder.add_virtual_target(); + + // dataset path bits (binary decomposition of leaf_index) + let d_path_bits = builder.split_le(slot_index,d_depth); + + // dataset last bits (binary decomposition of last_index = nleaves - 1) + let depth_target = builder.constant(F::from_canonical_u64(d_depth as u64)); + let mut d_last_index = builder.exp(two,depth_target,d_depth); + d_last_index = builder.sub(d_last_index, one); + let d_last_bits = builder.split_le(d_last_index,d_depth); + + // dataset Merkle path (sibling hashes from leaf to root) + let d_merkle_path = MerkleProofTarget { + path: (0..d_depth).map(|_| builder.add_virtual_hash()).collect(), + }; + + // create MerkleTreeTargets struct + let mut d_targets = MerkleTreeTargets { + leaf: slot_root, + 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::::reconstruct_merkle_root_circuit(builder, &mut d_targets); // expected Merkle root - let dataset_expected_root = builder.add_virtual_hash(); + let d_expected_root = builder.add_virtual_hash(); // check equality with expected root for i in 0..NUM_HASH_OUT_ELTS { - builder.connect(dataset_expected_root.elements[i], dataset_root_target.elements[i]); + builder.connect(d_expected_root.elements[i], d_reconstructed_root.elements[i]); } - let mut slot_proofs =vec![]; + //*********** do the sampling ************ + + let mut data_targets =vec![]; + let mut slot_sample_proofs = vec![]; + let entropy_target = builder.add_virtual_target(); for i in 0..N_SAMPLES{ - let proof_i = SlotTreeCircuit::::prove_single_cell(builder); - slot_proofs.push(proof_i); + // cell data targets + let mut data_i = (0..N_FIELD_ELEMS_PER_CELL).map(|_| builder.add_virtual_target()).collect::>(); + + let mut perm_inputs:Vec= Vec::new(); + perm_inputs.extend_from_slice(&data_i); + let data_i_hash = builder.hash_n_to_hash_no_pad::(perm_inputs); + // counter constant + let ctr = builder.constant(F::from_canonical_u64(i as u64)); + // paths + let mut b_path_bits = Self::calculate_cell_index_bits(builder, &entropy_target, &d_targets.leaf, &ctr); + let mut s_path_bits = b_path_bits.split_off(BOT_DEPTH); + + //TODO: this can probably be optimized by supplying nCellsPerSlot as input to the circuit + let b_depth_target = builder.constant(F::from_canonical_u64(BOT_DEPTH as u64)); + let mut b_last_index = builder.exp(two,b_depth_target,BOT_DEPTH); + b_last_index = builder.sub(b_last_index, one); + let b_last_bits = builder.split_le(b_last_index,BOT_DEPTH); + + let s_depth_target = builder.constant(F::from_canonical_u64(MAX_DEPTH as u64)); + let mut s_last_index = builder.exp(two,s_depth_target,MAX_DEPTH); + s_last_index = builder.sub(s_last_index, one); + let s_last_bits = builder.split_le(s_last_index,MAX_DEPTH); + + let mut b_merkle_path = MerkleProofTarget { + path: (0..BOT_DEPTH).map(|_| builder.add_virtual_hash()).collect(), + }; + + let mut s_merkle_path = MerkleProofTarget { + path: (0..(MAX_DEPTH - BOT_DEPTH)).map(|_| builder.add_virtual_hash()).collect(), + }; + + let mut block_targets = MerkleTreeTargets { + leaf: data_i_hash, + path_bits:b_path_bits, + last_bits: b_last_bits, + merkle_path: b_merkle_path, + _phantom: PhantomData, + }; + + // reconstruct block root + let b_root = MerkleTreeCircuit::::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, + merkle_path:s_merkle_path, + _phantom: PhantomData, + }; + + // reconstruct slot root with block root as leaf + let slot_reconstructed_root = MerkleTreeCircuit::::reconstruct_merkle_root_circuit(builder, &mut slot_targets); + + // check equality with expected root + for i in 0..NUM_HASH_OUT_ELTS { + builder.connect( d_targets.leaf.elements[i], slot_reconstructed_root.elements[i]); + } + + // combine block and slot path to get the full path so we can assign it later. + let mut slot_sample_proof_target = MerkleProofTarget{ + path: block_targets.merkle_path.path, + }; + slot_sample_proof_target.path.extend_from_slice(&slot_targets.merkle_path.path); + + data_targets.push(data_i); + slot_sample_proofs.push(slot_sample_proof_target); + } DatasetTargets::{ - dataset_proof, - dataset_root: dataset_expected_root, - slot_proofs, + dataset_proof: d_targets.merkle_path, + dataset_root: d_expected_root, + cell_data: data_targets, + entropy: entropy_target, + slot_index, + slot_root: d_targets.leaf, + slot_proofs: slot_sample_proofs, _phantom: Default::default(), } } - // assign the witnesses to the targets - // takes pw, the dataset targets, slot index, and entropy + pub fn calculate_cell_index_bits(builder: &mut CircuitBuilder::, p0: &Target, p1: &HashOutTarget, p2: &Target) -> Vec { + let mut perm_inputs:Vec= Vec::new(); + perm_inputs.extend_from_slice(&p1.elements); + perm_inputs.push(*p0); + perm_inputs.push(*p2); + let data_i_hash = builder.hash_n_to_hash_no_pad::(perm_inputs); + let p_bits = builder.low_bits(data_i_hash.elements[NUM_HASH_OUT_ELTS-1], MAX_DEPTH, 64); + + p_bits + } + pub fn sample_slot_assign_witness( &mut self, pw: &mut PartialWitness, @@ -266,8 +390,19 @@ impl< slot_index:usize, entropy:usize, ){ - // assign witness for dataset level target (proving slot root is in dataset tree) - self.tree.assign_witness(pw,&mut targets.dataset_proof,slot_index); + // dataset proof + let d_proof = self.tree.tree.get_proof(slot_index).unwrap(); + + // assign dataset proof + for (i, sibling_hash) in d_proof.path.iter().enumerate() { + // 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.dataset_proof.path[i].elements[j], sibling_hash_out[j]); + } + } + // assign slot index + pw.set_target(targets.slot_index, F::from_canonical_u64(slot_index as u64)); // assign the expected Merkle root of dataset to the target let expected_root = self.tree.tree.root().unwrap(); @@ -279,14 +414,28 @@ impl< // the sampled slot let slot = &self.slot_trees[slot_index]; let slot_root = slot.tree.tree.root().unwrap(); + pw.set_hash_target(targets.slot_root, slot_root); + + // assign entropy + pw.set_target(targets.entropy, F::from_canonical_u64(entropy as u64)); // do the sample N times for i in 0..N_SAMPLES { - let cell_index_bits = calculate_cell_index_bits(entropy, slot_root, i); + let cell_index_bits = calculate_cell_index_bits(entropy,slot_root,i); let cell_index = bits_le_padded_to_usize(&cell_index_bits); + // assign cell data let leaf = &slot.cell_data[cell_index]; - let proof = slot.get_proof(cell_index); - slot.single_cell_assign_witness(pw, &mut targets.slot_proofs[i],cell_index,leaf, proof.clone()); + for j in 0..leaf.len(){ + pw.set_target(targets.cell_data[i][j], leaf[j]); + } + // assign proof for that cell + let cell_proof = slot.get_proof(cell_index); + for (k, sibling_hash) in cell_proof.path.iter().enumerate() { + let sibling_hash_out = sibling_hash.to_vec(); + for j in 0..sibling_hash_out.len() { + pw.set_target(targets.slot_proofs[i].path[k].elements[j], sibling_hash_out[j]); + } + } } } diff --git a/codex-plonky2-circuits/src/circuits/utils.rs b/codex-plonky2-circuits/src/circuits/utils.rs index 8cd5835..6436fc5 100644 --- a/codex-plonky2-circuits/src/circuits/utils.rs +++ b/codex-plonky2-circuits/src/circuits/utils.rs @@ -1,4 +1,4 @@ -use plonky2::hash::hash_types::{HashOut, RichField}; +use plonky2::hash::hash_types::{HashOut, NUM_HASH_OUT_ELTS, RichField}; use plonky2::iop::witness::PartialWitness; use plonky2::plonk::circuit_data::{CircuitData, VerifierCircuitData}; use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, GenericHashOut, Hasher}; @@ -22,18 +22,19 @@ pub(crate) fn usize_to_bits_le_padded(index: usize, bit_length: usize) -> Vec(p0: usize, p1: HashOut, p2: usize) -> Vec { - let p0_field = F::from_canonical_u64(p0 as u64); - let p2_field = F::from_canonical_u64(p2 as u64); +/// calculate the sampled cell index from entropy, slot root, and counter +pub(crate) fn calculate_cell_index_bits(entropy: usize, slot_root: HashOut, ctr: usize) -> Vec { + let p0_field = F::from_canonical_u64(entropy as u64); + let p2_field = F::from_canonical_u64(ctr as u64); let mut inputs = Vec::new(); - inputs.extend_from_slice(&p1.elements); + inputs.extend_from_slice(&slot_root.elements); inputs.push(p0_field); inputs.push(p2_field); let p_hash = HF::hash_no_pad(&inputs); - let p_bytes = p_hash.to_bytes(); + let p_bytes = p_hash.elements[NUM_HASH_OUT_ELTS - 1].to_canonical_u64(); - let p_bits = take_n_bits_from_bytes(&p_bytes, MAX_DEPTH); + // let p_bits = take_n_bits_from_bytes(&p_bytes, MAX_DEPTH); + let p_bits = usize_to_bits_le_padded(p_bytes as usize, MAX_DEPTH); p_bits } pub(crate) fn take_n_bits_from_bytes(bytes: &[u8], n: usize) -> Vec {