From 3b80809fb448ed5409f022cc40432be071f02a7b Mon Sep 17 00:00:00 2001 From: M Alghazwi Date: Fri, 18 Oct 2024 12:01:01 +0200 Subject: [PATCH] add sample_cells bench --- codex-plonky2-circuits/Cargo.toml | 4 + .../benches/sample_cells.rs | 129 ++++++++++++++++++ codex-plonky2-circuits/src/circuits/params.rs | 13 +- .../src/circuits/prove_single_cell.rs | 18 --- .../src/circuits/sample_cells.rs | 84 +++++++++++- 5 files changed, 225 insertions(+), 23 deletions(-) create mode 100644 codex-plonky2-circuits/benches/sample_cells.rs diff --git a/codex-plonky2-circuits/Cargo.toml b/codex-plonky2-circuits/Cargo.toml index 2502b1d..42ddfea 100644 --- a/codex-plonky2-circuits/Cargo.toml +++ b/codex-plonky2-circuits/Cargo.toml @@ -29,3 +29,7 @@ harness = false [[bench]] name = "prove_cells" harness = false + +[[bench]] +name = "sample_cells" +harness = false diff --git a/codex-plonky2-circuits/benches/sample_cells.rs b/codex-plonky2-circuits/benches/sample_cells.rs new file mode 100644 index 0000000..a24083d --- /dev/null +++ b/codex-plonky2-circuits/benches/sample_cells.rs @@ -0,0 +1,129 @@ +use criterion::{criterion_group, criterion_main, Criterion}; +use anyhow::Result; +use std::time::{Duration, Instant}; +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 plonky2::plonk::circuit_builder::CircuitBuilder; +use codex_plonky2_circuits::circuits::params::TESTING_SLOT_INDEX; +use codex_plonky2_circuits::circuits::sample_cells::DatasetTreeCircuit; + +macro_rules! pretty_print { + ($($arg:tt)*) => { + print!("\x1b[0;36mINFO ===========>\x1b[0m "); + println!($($arg)*); + } +} + +// Hash function used +type HF = PoseidonHash; + +fn prepare_data< + F: RichField + Extendable + Poseidon2, + C: GenericConfig, + const D: usize, + H: Hasher + AlgebraicHasher, +>() -> Result<( + DatasetTreeCircuit, + usize, + usize, +)> { + // Initialize the dataset tree with testing data + let mut dataset_t = DatasetTreeCircuit::::new_for_testing(); + + let slot_index = TESTING_SLOT_INDEX; + let entropy = 123; + + Ok((dataset_t, slot_index, entropy)) +} + +fn build_circuit< + F: RichField + Extendable + Poseidon2, + C: GenericConfig, + const D: usize, + H: Hasher + AlgebraicHasher, +>( + dataset_tree: &mut DatasetTreeCircuit, + slot_index: usize, + entropy: usize, + // proofs: &[MerkleProof], +) -> Result<(CircuitData, PartialWitness)> +{ + // Create the circuit + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + + let mut targets = dataset_tree.sample_slot_circuit(&mut builder); + + // Create a PartialWitness + let mut pw = PartialWitness::new(); + dataset_tree.sample_slot_assign_witness(&mut pw, &mut targets,slot_index,entropy); + + // Build the circuit + let data = builder.build::(); + + Ok((data, pw)) +} + +fn sampling_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("Sampling Benchmark"); + + // Circuit parameters + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + type H = PoseidonHash; + + // Prepare the data that will be used in all steps + let (mut dataset_tree, slot_index, entropy) = prepare_data::().unwrap(); + + // Benchmark the circuit building + group.bench_function("Single Cell Proof Build", |b| { + b.iter(|| { + build_circuit::(&mut dataset_tree, slot_index, entropy).unwrap(); + }) + }); + + // Build the circuit + let (data, pw) = build_circuit::(&mut dataset_tree, slot_index, entropy).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 = sampling_benchmark); +criterion_main!(benches); diff --git a/codex-plonky2-circuits/src/circuits/params.rs b/codex-plonky2-circuits/src/circuits/params.rs index 8a0b33e..305d87e 100644 --- a/codex-plonky2-circuits/src/circuits/params.rs +++ b/codex-plonky2-circuits/src/circuits/params.rs @@ -4,15 +4,20 @@ use plonky2::hash::poseidon::PoseidonHash; // constants and types used throughout the circuit -pub const N_FIELD_ELEMS_PER_CELL: usize = 4; +pub const N_FIELD_ELEMS_PER_CELL: usize = 256; pub const BOT_DEPTH: usize = 5; // block depth - depth of the block merkle tree -pub const MAX_DEPTH: usize = 8; // depth of big tree (slot tree depth + block tree depth) +pub const MAX_DEPTH: usize = 16; // depth of big tree (slot tree depth + block tree depth) pub const N_CELLS_IN_BLOCKS: usize = 1< and causing a lot of headache diff --git a/codex-plonky2-circuits/src/circuits/prove_single_cell.rs b/codex-plonky2-circuits/src/circuits/prove_single_cell.rs index 4b6ebc8..c73417d 100644 --- a/codex-plonky2-circuits/src/circuits/prove_single_cell.rs +++ b/codex-plonky2-circuits/src/circuits/prove_single_cell.rs @@ -91,23 +91,6 @@ impl< .collect::>(); // create slot tree let slot_tree = MerkleTree::::new(&block_roots, zero).unwrap(); - // let mt = - // MerkleTree::{ - // tree: slot_tree, - // block_trees, - // cell_data, - // cell_hash: leaves, - // } - - // create block circuits - // let block_circuits = block_trees.iter() - // .map(|b_tree| { - // // 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 - // MerkleTreeCircuit::{ tree:b_tree.clone(), _phantom:Default::default()}, - // }) - // .collect::>(); Self{ tree: MerkleTreeCircuit::{ tree:slot_tree, _phantom:Default::default()}, @@ -124,7 +107,6 @@ impl< const D: usize, H: Hasher + AlgebraicHasher, > SlotTreeCircuit { - /// same as default but with supplied cell data pub fn new(cell_data: Vec>) -> Self{ let leaves: Vec> = cell_data diff --git a/codex-plonky2-circuits/src/circuits/sample_cells.rs b/codex-plonky2-circuits/src/circuits/sample_cells.rs index 2cbbbbd..54405cb 100644 --- a/codex-plonky2-circuits/src/circuits/sample_cells.rs +++ b/codex-plonky2-circuits/src/circuits/sample_cells.rs @@ -28,7 +28,7 @@ use plonky2::plonk::config::PoseidonGoldilocksConfig; use plonky2::hash::hashing::PlonkyPermutation; use crate::circuits::prove_single_cell::{SingleCellTargets, SlotTreeCircuit}; -use crate::circuits::params::{MAX_DEPTH, BOT_DEPTH, N_FIELD_ELEMS_PER_CELL, N_CELLS_IN_BLOCKS, N_BLOCKS, N_CELLS, HF, DATASET_DEPTH, N_SAMPLES}; +use crate::circuits::params::{MAX_DEPTH, BOT_DEPTH, N_FIELD_ELEMS_PER_CELL, N_CELLS_IN_BLOCKS, N_BLOCKS, N_CELLS, HF, DATASET_DEPTH, 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}; @@ -93,6 +93,45 @@ impl< const D: usize, H: Hasher + AlgebraicHasher, > DatasetTreeCircuit { + /// 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 { + let mut slot_trees = vec![]; + let n_slots = 1<{ + tree: MerkleTreeCircuit { + tree: MerkleTree::::new(&[zero.clone()], zero.clone()).unwrap(), + _phantom: Default::default(), + }, + block_trees: vec![], + cell_data: vec![], + cell_hash: vec![], + }; + for i in 0..n_slots { + if(i == TESTING_SLOT_INDEX) { + slot_trees.push(SlotTreeCircuit::::default()); + }else{ + slot_trees.push(zero_slot.clone()); + } + + } + // get the roots or slot trees + let slot_roots = slot_trees.iter() + .map(|t| { + t.tree.tree.root().unwrap() + }) + .collect::>(); + let dataset_tree = MerkleTree::::new(&slot_roots, zero).unwrap(); + Self{ + tree: MerkleTreeCircuit::{ tree:dataset_tree, _phantom:Default::default()}, + slot_trees, + } + } + /// same as default but with supplied slot trees pub fn new(slot_trees: Vec>) -> Self{ // get the roots or slot trees @@ -297,4 +336,47 @@ mod tests { Ok(()) } + + #[test] + fn test_sample_cells_circuit_from_selected_slot() -> Result<()> { + + let mut dataset_t = DatasetTreeCircuit::::new_for_testing(); + + let slot_index = TESTING_SLOT_INDEX; + 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