From 494cb8a248dc302cde5f78afd392b1b04514a4da Mon Sep 17 00:00:00 2001 From: M Alghazwi Date: Fri, 18 Oct 2024 10:09:41 +0200 Subject: [PATCH] refactor dataset sampling --- codex-plonky2-circuits/src/circuits/params.rs | 2 +- .../src/circuits/safe_tree_circuit.rs | 45 ----------- .../src/circuits/sample_cells.rs | 74 ++++++++++++++----- codex-plonky2-circuits/src/circuits/utils.rs | 39 +++++++++- 4 files changed, 94 insertions(+), 66 deletions(-) diff --git a/codex-plonky2-circuits/src/circuits/params.rs b/codex-plonky2-circuits/src/circuits/params.rs index 1bfb1c4..8a0b33e 100644 --- a/codex-plonky2-circuits/src/circuits/params.rs +++ b/codex-plonky2-circuits/src/circuits/params.rs @@ -6,7 +6,7 @@ use plonky2::hash::poseidon::PoseidonHash; // 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) +pub const MAX_DEPTH: usize = 8; // depth of big tree (slot tree depth + block tree depth) pub const N_CELLS_IN_BLOCKS: usize = 1<, - pw: PartialWitness - ) -> Result>{ - let proof = data.prove(pw); - return proof - } - - /// verify given verifier data, public input, and proof - pub fn verify( - &mut self, - verifier_data: &VerifierCircuitData, - public_inputs: Vec, - proof: Proof - )-> Result<()> { - verifier_data.verify(ProofWithPublicInputs { - proof, - public_inputs, - }) - } - /// assign the witness values in the circuit targets /// this takes leaf_index and fills all required circuit targets(inputs) pub fn assign_witness( @@ -234,28 +211,6 @@ impl< } -// --------- helper functions --------- -// impl< -// F: RichField + Extendable + Poseidon2, -// C: GenericConfig, -// const D: usize, -// H: Hasher + AlgebraicHasher, -// > MerkleTreeCircuit { -// /// Converts an index to a vector of bits (LSB first) with padding. -// pub(crate) fn usize_to_bits_le_padded(&self, index: usize, bit_length: usize) -> Vec { -// let mut bits = Vec::with_capacity(bit_length); -// for i in 0..bit_length { -// bits.push(((index >> i) & 1) == 1); -// } -// // If index requires fewer bits, pad with `false` -// while bits.len() < bit_length { -// bits.push(false); -// } -// bits -// } -// } - - // 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)] diff --git a/codex-plonky2-circuits/src/circuits/sample_cells.rs b/codex-plonky2-circuits/src/circuits/sample_cells.rs index c968f32..2cbbbbd 100644 --- a/codex-plonky2-circuits/src/circuits/sample_cells.rs +++ b/codex-plonky2-circuits/src/circuits/sample_cells.rs @@ -176,17 +176,8 @@ pub struct DatasetTargets< const D: usize, H: Hasher + AlgebraicHasher, > { - // pub expected_dataset_root_target: HashOutTarget, - // pub slot_index: Target, - // pub entropy_target: Target, - // pub slot_root: HashOutTarget, pub slot_proofs: Vec>, - - // pub proof_target: MerkleProofTarget, - // pub leaf_target: Vec, - // pub path_bits: Vec, - // pub last_bits: Vec, _phantom: PhantomData<(C,H)>, } @@ -204,8 +195,6 @@ impl< builder: &mut CircuitBuilder::, )-> DatasetTargets{ - // Create virtual targets - // let slot_root = builder.add_virtual_hash(); let mut slot_proofs =vec![]; for i in 0..N_SAMPLES{ let proof_i = SlotTreeCircuit::::prove_single_cell(builder); @@ -213,26 +202,30 @@ impl< } DatasetTargets::{ - // expected_dataset_root_target: HashOutTarget {}, - // slot_index: Default::default(), - // entropy_target: Default::default(), - // slot_root: HashOutTarget {}, slot_proofs, _phantom: Default::default(), } } - // assign the witnesses to the target - // takes the dataset tree, slot index, and entropy + // assign the witnesses to the targets + // takes pw, the dataset targets, slot index, and entropy pub fn sample_slot_assign_witness( &mut self, pw: &mut PartialWitness, - targets: DatasetTargets, - dataset_tree: DatasetTreeCircuit, + targets: &mut DatasetTargets, slot_index:usize, entropy:usize, ){ + let slot = &self.slot_trees[slot_index]; + let slot_root = slot.tree.tree.root().unwrap(); + for i in 0..N_SAMPLES { + let cell_index_bits = calculate_cell_index_bits(entropy, slot_root, i); + let cell_index = bits_le_padded_to_usize(&cell_index_bits); + 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()); + } } @@ -261,4 +254,47 @@ mod tests { let res = dataset_t.verify_sampling(proof).unwrap(); assert_eq!(res, true); } + + #[test] + fn test_sample_cells_circuit() -> Result<()> { + + let mut dataset_t = DatasetTreeCircuit::::default(); + + let slot_index = 2; + let entropy = 123; + + // sanity check + let proof = dataset_t.sample_slot(slot_index,entropy); + let slot_root = dataset_t.slot_trees[slot_index].tree.tree.root().unwrap(); + let res = dataset_t.verify_sampling(proof).unwrap(); + assert_eq!(res, true); + + // create the circuit + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + + let mut targets = dataset_t.sample_slot_circuit(&mut builder); + + // create a PartialWitness and assign + let mut pw = PartialWitness::new(); + dataset_t.sample_slot_assign_witness(&mut pw, &mut targets,slot_index,entropy); + + // build the circuit + let data = builder.build::(); + println!("circuit size = {:?}", data.common.degree_bits()); + + // Prove the circuit with the assigned witness + let start_time = Instant::now(); + let proof_with_pis = data.prove(pw)?; + println!("prove_time = {:?}", start_time.elapsed()); + + // verify the proof + let verifier_data = data.verifier_data(); + assert!( + verifier_data.verify(proof_with_pis).is_ok(), + "Merkle proof verification failed" + ); + + Ok(()) + } } \ No newline at end of file diff --git a/codex-plonky2-circuits/src/circuits/utils.rs b/codex-plonky2-circuits/src/circuits/utils.rs index e493b4a..8cd5835 100644 --- a/codex-plonky2-circuits/src/circuits/utils.rs +++ b/codex-plonky2-circuits/src/circuits/utils.rs @@ -1,6 +1,12 @@ use plonky2::hash::hash_types::{HashOut, RichField}; -use plonky2::plonk::config::{GenericHashOut, Hasher}; +use plonky2::iop::witness::PartialWitness; +use plonky2::plonk::circuit_data::{CircuitData, VerifierCircuitData}; +use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, GenericHashOut, Hasher}; +use plonky2::plonk::proof::{Proof, ProofWithPublicInputs}; +use plonky2_field::extension::Extendable; +use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; use crate::circuits::params::{HF, MAX_DEPTH}; +use anyhow::Result; // --------- helper functions --------- @@ -45,4 +51,35 @@ pub(crate) fn bits_le_padded_to_usize(bits: &[bool]) -> usize { acc } }) +} + +/// prove given the circuit data and partial witness +pub fn prove< + F: RichField + Extendable + Poseidon2, + C: GenericConfig, + const D: usize, + H: Hasher + AlgebraicHasher, +>( + data: CircuitData, + pw: PartialWitness +) -> Result>{ + let proof = data.prove(pw); + return proof +} + +/// verify given verifier data, public input, and proof +pub fn verify< + F: RichField + Extendable + Poseidon2, + C: GenericConfig, + const D: usize, + H: Hasher + AlgebraicHasher, +>( + verifier_data: &VerifierCircuitData, + public_inputs: Vec, + proof: Proof +)-> Result<()> { + verifier_data.verify(ProofWithPublicInputs { + proof, + public_inputs, + }) } \ No newline at end of file