refactor dataset sampling

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<<BOT_DEPTH; //2^BOT_DEPTH
 pub const N_BLOCKS: usize = 1<<(MAX_DEPTH - BOT_DEPTH); // 2^(MAX_DEPTH - BOT_DEPTH)
 pub const N_CELLS: usize = N_CELLS_IN_BLOCKS * N_BLOCKS;

codex-plonky2-circuits/src/circuits/safe_tree_circuit.rs

@@ -100,29 +100,6 @@ impl<
         targets
     }
 
-    /// prove given the circuit data and partial witness
-    pub fn prove(
-        &mut self,
-        data: CircuitData<F, C, D>,
-        pw: PartialWitness<F>
-    ) -> Result<ProofWithPublicInputs<F, C, D>>{
-        let proof = data.prove(pw);
-        return proof
-    }
-
-    /// verify given verifier data, public input, and proof
-    pub fn verify(
-        &mut self,
-        verifier_data: &VerifierCircuitData<F, C, D>,
-        public_inputs: Vec<F>,
-        proof: Proof<F, C, D>
-    )-> 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<D> + Poseidon2,
-//     C: GenericConfig<D, F = F>,
-//     const D: usize,
-//     H: Hasher<F> + AlgebraicHasher<F>,
-// > MerkleTreeCircuit<F, C, D, H> {
-//     /// 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
-//     }
-// }
-
-
 // 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)]

codex-plonky2-circuits/src/circuits/sample_cells.rs

@@ -176,17 +176,8 @@ pub struct DatasetTargets<
     const D: usize,
     H: Hasher<F> + AlgebraicHasher<F>,
 > {
-    // pub expected_dataset_root_target: HashOutTarget,
-    // pub slot_index: Target,
-    // pub entropy_target: Target,
-    // pub slot_root: HashOutTarget,
     pub slot_proofs: Vec<SingleCellTargets<F, C, D, H>>,
 
-
-    // pub proof_target: MerkleProofTarget,
-    // pub leaf_target: Vec<Target>,
-    // pub path_bits: Vec<BoolTarget>,
-    // pub last_bits: Vec<BoolTarget>,
     _phantom: PhantomData<(C,H)>,
 }
 
@@ -204,8 +195,6 @@ impl<
         builder: &mut CircuitBuilder::<F, D>,
     )-> DatasetTargets<F,C,D,H>{
 
-        // 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::<F,C,D,H>::prove_single_cell(builder);
@@ -213,26 +202,30 @@ impl<
         }
 
         DatasetTargets::<F,C,D,H>{
-            // 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<F>,
-        targets: DatasetTargets<F,C,D,H>,
-        dataset_tree: DatasetTreeCircuit<F,C,D,H>,
+        targets: &mut DatasetTargets<F,C,D,H>,
         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::<F,C,D,H>::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::<F, D>::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::<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(())
+    }
 }

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 ---------
 
@@ -46,3 +52,34 @@ pub(crate) fn bits_le_padded_to_usize(bits: &[bool]) -> usize {
         }
     })
 }
+
+/// prove given the circuit data and partial witness
+pub fn prove<
+    F: RichField + Extendable<D> + Poseidon2,
+    C: GenericConfig<D, F = F>,
+    const D: usize,
+    H: Hasher<F> + AlgebraicHasher<F>,
+>(
+    data: CircuitData<F, C, D>,
+    pw: PartialWitness<F>
+) -> Result<ProofWithPublicInputs<F, C, D>>{
+    let proof = data.prove(pw);
+    return proof
+}
+
+/// verify given verifier data, public input, and proof
+pub fn verify<
+    F: RichField + Extendable<D> + Poseidon2,
+    C: GenericConfig<D, F = F>,
+    const D: usize,
+    H: Hasher<F> + AlgebraicHasher<F>,
+>(
+    verifier_data: &VerifierCircuitData<F, C, D>,
+    public_inputs: Vec<F>,
+    proof: Proof<F, C, D>
+)-> Result<()> {
+    verifier_data.verify(ProofWithPublicInputs {
+        proof,
+        public_inputs,
+    })
+}