add goldibear experiments

Cargo.toml

@@ -1,5 +1,6 @@
 [workspace]
 members = ["codex-plonky2-circuits","plonky2_poseidon2","proof-input","workflow"]
+exclude = ["goldibear_experiments"]
 resolver = "2"
 
 [workspace.dependencies]

goldibear_experiments/.gitignore

@@ -0,0 +1,13 @@
+#IDE Related
+.idea
+
+# Cargo build
+/target
+Cargo.lock
+
+# Profile-guided optimization
+/tmp
+pgo-data.profdata
+
+# MacOS nuisances
+.DS_Store

goldibear_experiments/README.md

@@ -0,0 +1,5 @@
+Plonky2 Goldibear Experiments
+================================
+
+Few experiments done to test the [Plonky2_Goldibear](https://github.com/telosnetwork/plonky2_goldibear/tree/main).
+The storage circuits are adapted for that version of plonky2 and tested.

goldibear_experiments/codex-plonky2-circuits/.gitignore

@@ -0,0 +1,13 @@
+#IDE Related
+.idea
+
+# Cargo build
+/target
+Cargo.lock
+
+# Profile-guided optimization
+/tmp
+pgo-data.profdata
+
+# MacOS nuisances
+.DS_Store

goldibear_experiments/codex-plonky2-circuits/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "codex-plonky2-circuits"
+description = "Codex storage proofs circuits for Plonky2"
+authors = ["Mohammed Alghazwi <m.ghazwi@gmail.com>"]
+readme = "README.md"
+version = "1.0.0"
+edition = "2021"
+
+[dependencies]
+anyhow     = { version = "1.0.89"}
+unroll = { version = "0.1.5"}
+serde      = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+plonky2    = { git = "https://github.com/telosnetwork/plonky2_goldibear.git"}
+plonky2_field = { git = "https://github.com/telosnetwork/plonky2_goldibear.git" }
+thiserror = "2.0.10"
+itertools = { version = "0.12.1"}
+plonky2_maybe_rayon = { git = "https://github.com/telosnetwork/plonky2_goldibear.git" }
+hashbrown = "0.14.5"
+
+[dev-dependencies]
+criterion = { version = "0.5.1", default-features = false }
+tynm = { version = "0.1.6", default-features = false }
+

goldibear_experiments/codex-plonky2-circuits/src/circuits/keyed_compress.rs

@@ -0,0 +1,34 @@
+use plonky2::hash::hash_types::{ HashOutTarget, RichField};
+use plonky2::hash::hashing::PlonkyPermutation;
+use plonky2::iop::target::Target;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::config::AlgebraicHasher;
+use plonky2_field::types::HasExtension;
+use crate::circuits::params::NUM_HASH_OUT_ELTS;
+
+/// Compression function which takes two 256 bit inputs (HashOutTarget) and key Target
+/// and returns a 256 bit output (HashOutTarget /  4 Targets).
+pub fn key_compress_circuit<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: AlgebraicHasher<F, NUM_HASH_OUT_ELTS>,
+>(
+    builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>,
+    x: HashOutTarget<NUM_HASH_OUT_ELTS>,
+    y: HashOutTarget<NUM_HASH_OUT_ELTS>,
+    key: Target,
+) -> HashOutTarget<NUM_HASH_OUT_ELTS> {
+    let zero = builder.zero();
+    let mut state = H::AlgebraicPermutation::new(core::iter::repeat(zero));
+
+    state.set_from_slice(&x.elements, 0);
+    state.set_from_slice(&y.elements, NUM_HASH_OUT_ELTS);
+    state.set_elt(key, NUM_HASH_OUT_ELTS*2);
+
+    state = builder.permute::<H>(state);
+
+    HashOutTarget {
+        elements: state.squeeze()[..NUM_HASH_OUT_ELTS].try_into().unwrap(),
+    }
+}
+

goldibear_experiments/codex-plonky2-circuits/src/circuits/merkle_circuit.rs

@@ -0,0 +1,166 @@
+// Plonky2 Circuit implementation of "safe" merkle tree
+// consistent with the one in codex:
+// https://github.com/codex-storage/codex-storage-proofs-circuits/blob/master/circuit/codex/merkle.circom
+
+use plonky2::{
+    hash::hash_types::{HashOutTarget, RichField, },
+    iop::target::BoolTarget,
+    plonk::circuit_builder::CircuitBuilder,
+};
+use std::marker::PhantomData;
+use plonky2::plonk::config::AlgebraicHasher;
+use plonky2_field::types::HasExtension;
+use crate::circuits::keyed_compress::key_compress_circuit;
+use crate::circuits::params::NUM_HASH_OUT_ELTS;
+use crate::circuits::utils::{add_assign_hash_out_target, mul_hash_out_target};
+use crate::Result;
+use crate::error::CircuitError;
+
+// Constants for the keys used in compression
+pub const KEY_NONE: u64 = 0x0;
+pub const KEY_BOTTOM_LAYER: u64 = 0x1;
+pub const KEY_ODD: u64 = 0x2;
+pub const KEY_ODD_AND_BOTTOM_LAYER: u64 = 0x3;
+
+/// Merkle tree targets representing the input to the circuit
+#[derive(Clone)]
+pub struct MerkleTreeTargets{
+    pub leaf: HashOutTarget<NUM_HASH_OUT_ELTS>,
+    pub path_bits: Vec<BoolTarget>,
+    pub last_bits: Vec<BoolTarget>,
+    pub mask_bits: Vec<BoolTarget>,
+    pub merkle_path: MerkleProofTarget,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct MerkleProofTarget {
+    /// The Merkle digest of each sibling subtree, staying from the bottommost layer.
+    pub path: Vec<HashOutTarget<NUM_HASH_OUT_ELTS>>,
+}
+
+/// Merkle tree circuit contains the functions for
+/// building, proving and verifying the circuit.
+#[derive(Clone)]
+pub struct MerkleTreeCircuit<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: AlgebraicHasher<F, NUM_HASH_OUT_ELTS>,
+> {
+    pub phantom_data: PhantomData<(F,H)>,
+}
+
+impl<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: AlgebraicHasher<F, NUM_HASH_OUT_ELTS>,
+> MerkleTreeCircuit<F, D, H> {
+
+
+    pub fn new() -> Self{
+        Self{
+            phantom_data: Default::default(),
+        }
+    }
+
+    /// Reconstructs the Merkle root from a leaf and Merkle path using a “mask” approach.
+    ///
+    /// # input
+    ///
+    /// * `builder`   - A circuit builder.
+    /// * `targets`   - The Merkle targets.
+    /// * `max_depth` - The maximum depth of the tree.
+    ///
+    /// # Returns
+    ///
+    /// A `HashOutTarget` representing the reconstructed Merkle root in-circuit.
+    ///
+    pub fn reconstruct_merkle_root_circuit_with_mask(
+        builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>,
+        targets: &mut MerkleTreeTargets,
+        max_depth: usize,
+    ) -> Result<HashOutTarget<NUM_HASH_OUT_ELTS>> {
+        let mut state: Vec<HashOutTarget<NUM_HASH_OUT_ELTS>> = Vec::with_capacity(max_depth+1);
+        state.push(targets.leaf);
+        let zero = builder.zero();
+        let one = builder.one();
+        let two = builder.two();
+
+        // --- Basic checks on input sizes.
+        let path_len = targets.path_bits.len();
+        let proof_len = targets.merkle_path.path.len();
+        let mask_len = targets.mask_bits.len();
+        let last_len = targets.last_bits.len();
+
+        if path_len != proof_len {
+            return Err(CircuitError::PathBitsLengthMismatch(path_len, proof_len));
+        }
+
+        if mask_len != path_len + 1 {
+            return Err(CircuitError::MaskBitsLengthMismatch(mask_len, path_len+1));
+        }
+
+        if last_len != path_len {
+            return Err(CircuitError::LastBitsLengthMismatch(last_len, path_len));
+        }
+
+        if path_len != max_depth {
+            return Err(CircuitError::PathBitsMaxDepthMismatch(path_len, max_depth));
+        }
+
+        // compute is_last
+        let mut is_last = vec![BoolTarget::new_unsafe(zero); max_depth + 1];
+        is_last[max_depth] = BoolTarget::new_unsafe(one); // set isLast[max_depth] to 1 (true)
+        for i in (0..max_depth).rev() {
+            let eq_out = builder.is_equal(targets.path_bits[i].target , targets.last_bits[i].target);
+            is_last[i] = builder.and( is_last[i + 1] , eq_out);
+        }
+
+        let mut i: usize = 0;
+        for (&bit, &sibling) in targets.path_bits.iter().zip(&targets.merkle_path.path) {
+
+            // logic: we add KEY_BOTTOM_LAYER if i == 0, otherwise KEY_NONE.
+            let bottom_key_val = if i == 0 {
+                KEY_BOTTOM_LAYER
+            } else {
+                KEY_NONE
+            };
+            let bottom = builder.constant(F::from_canonical_u64(bottom_key_val));
+
+            // compute: odd = isLast[i] * (1-pathBits[i]);
+            // compute: key = bottom + 2*odd
+            let mut odd = builder.sub(one, targets.path_bits[i].target);
+            odd = builder.mul(is_last[i].target, odd);
+            odd = builder.mul(two, odd);
+            let key = builder.add(bottom,odd);
+
+            // select left and right based on path_bit
+            let mut left = vec![];
+            let mut right = vec![];
+            for j in 0..NUM_HASH_OUT_ELTS {
+                left.push( builder.select(bit, sibling.elements[j], state[i].elements[j]));
+                right.push( builder.select(bit, state[i].elements[j], sibling.elements[j]));
+            }
+
+            // Compress them with a keyed-hash function
+            let combined_hash = key_compress_circuit::<F, D, H>
+                (builder,
+                HashOutTarget::from_vec(left),
+                HashOutTarget::from_vec(right),
+                key);
+            state.push(combined_hash);
+
+            i += 1;
+        }
+
+        // select the right layer using the mask bits
+        let mut reconstructed_root  = HashOutTarget::from_vec([builder.zero();4].to_vec());
+        for k in 0..max_depth {
+            let diff = builder.sub(targets.mask_bits[k].target, targets.mask_bits[k+1].target);
+            let mul_result = mul_hash_out_target(builder,&diff,&mut state[k+1]);
+            add_assign_hash_out_target(builder,&mut reconstructed_root, &mul_result);
+        }
+
+        Ok(reconstructed_root)
+
+    }
+}

goldibear_experiments/codex-plonky2-circuits/src/circuits/mod.rs

@@ -0,0 +1,6 @@
+pub mod merkle_circuit;
+pub mod sample_cells;
+pub mod utils;
+pub mod params;
+pub mod keyed_compress;
+pub mod sponge;

goldibear_experiments/codex-plonky2-circuits/src/circuits/params.rs

@@ -0,0 +1,64 @@
+// global params for the circuits
+
+use anyhow::{Context, Result};
+use std::env;
+use plonky2::hash::hash_types::GOLDILOCKS_NUM_HASH_OUT_ELTS;
+
+/// params used for the circuits
+/// should be defined prior to building the circuit
+#[derive(Clone, Debug)]
+pub struct CircuitParams{
+    pub max_depth: usize,
+    pub max_log2_n_slots: usize,
+    pub block_tree_depth: usize,
+    pub n_field_elems_per_cell: usize,
+    pub n_samples: usize,
+}
+
+impl CircuitParams {
+    /// Creates a new `CircuitParams` struct from environment.
+    ///
+    /// - `MAX_DEPTH`:The maximum slot depth
+    /// - `MAX_LOG2_N_SLOTS`:The maximum log2 number of slots
+    /// - `BLOCK_TREE_DEPTH`:The block tree depth
+    /// - `N_FIELD_ELEMS_PER_CELL`: The number of field elements per cell
+    /// - `N_SAMPLES`: number of samples
+    ///
+    /// Returns an error if any environment variable is missing or fails to parse.
+    pub fn from_env() -> Result<Self> {
+        let max_depth = env::var("MAX_DEPTH")
+            .context("MAX_DEPTH is not set")?
+            .parse::<usize>()
+            .context("MAX_DEPTH must be a valid usize")?;
+
+        let max_log2_n_slots = env::var("MAX_LOG2_N_SLOTS")
+            .context("MAX_LOG2_N_SLOTS is not set")?
+            .parse::<usize>()
+            .context("MAX_LOG2_N_SLOTS must be a valid usize")?;
+
+        let block_tree_depth = env::var("BLOCK_TREE_DEPTH")
+            .context("BLOCK_TREE_DEPTH is not set")?
+            .parse::<usize>()
+            .context("BLOCK_TREE_DEPTH must be a valid usize")?;
+
+        let n_field_elems_per_cell = env::var("N_FIELD_ELEMS_PER_CELL")
+            .context("N_FIELD_ELEMS_PER_CELL is not set")?
+            .parse::<usize>()
+            .context("N_FIELD_ELEMS_PER_CELL must be a valid usize")?;
+
+        let n_samples = env::var("N_SAMPLES")
+            .context("N_SAMPLES is not set")?
+            .parse::<usize>()
+            .context("N_SAMPLES must be a valid usize")?;
+
+        Ok(CircuitParams {
+            max_depth,
+            max_log2_n_slots,
+            block_tree_depth,
+            n_field_elems_per_cell,
+            n_samples,
+        })
+    }
+}
+
+pub const NUM_HASH_OUT_ELTS: usize = GOLDILOCKS_NUM_HASH_OUT_ELTS;

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

@@ -0,0 +1,404 @@
+// Sample cells
+// consistent with:
+// https://github.com/codex-storage/codex-storage-proofs-circuits/blob/master/circuit/codex/sample_cells.circom
+// circuit consists of:
+// - reconstruct the dataset merkle root using the slot root as leaf
+// - samples multiple cells by calling the sample_cells
+
+use std::marker::PhantomData;
+
+use plonky2::{
+    hash::{
+        hash_types::{HashOut, HashOutTarget, RichField},
+        hashing::PlonkyPermutation,
+    },
+    iop::{
+        target::{BoolTarget, Target},
+        witness::{PartialWitness, WitnessWrite},
+    },
+    plonk::circuit_builder::CircuitBuilder,
+};
+use plonky2::plonk::config::AlgebraicHasher;
+use plonky2_field::types::HasExtension;
+
+use crate::{
+    circuits::{
+    merkle_circuit::{MerkleProofTarget, MerkleTreeCircuit, MerkleTreeTargets},
+    params::{CircuitParams, NUM_HASH_OUT_ELTS},
+    sponge::{hash_n_no_padding, hash_n_with_padding},
+    utils::{assign_hash_out_targets, ceiling_log2},
+    },
+    Result,
+    error::CircuitError,
+};
+
+/// circuit for sampling a slot in a dataset merkle tree
+#[derive(Clone, Debug)]
+pub struct SampleCircuit<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: AlgebraicHasher<F, NUM_HASH_OUT_ELTS>,
+> {
+    params: CircuitParams,
+    phantom_data: PhantomData<(F,H)>,
+}
+
+impl<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: AlgebraicHasher<F, NUM_HASH_OUT_ELTS>,
+> SampleCircuit<F, D, H> {
+    pub fn new(params: CircuitParams) -> Self{
+        Self{
+            params,
+            phantom_data: Default::default(),
+        }
+    }
+}
+
+/// struct of input to the circuit as targets
+/// used to build the circuit and can be assigned after building
+#[derive(Clone, Debug)]
+pub struct SampleTargets {
+
+    pub entropy: HashOutTarget<NUM_HASH_OUT_ELTS>, // public input
+    pub dataset_root: HashOutTarget<NUM_HASH_OUT_ELTS>, // public input
+    pub slot_index: Target, // public input
+
+    pub slot_root: HashOutTarget<NUM_HASH_OUT_ELTS>,
+    pub n_cells_per_slot: Target,
+    pub n_slots_per_dataset: Target,
+
+    pub slot_proof: MerkleProofTarget,
+
+    pub cell_data: Vec<CellTarget>,
+    pub merkle_paths: Vec<MerkleProofTarget>,
+}
+
+/// circuit input as field elements
+#[derive(Clone, Debug, PartialEq)]
+pub struct SampleCircuitInput<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>{
+    pub entropy: HashOut<F, NUM_HASH_OUT_ELTS>, // public input
+    pub dataset_root: HashOut<F, NUM_HASH_OUT_ELTS>, // public input
+    pub slot_index: F, // public input
+
+    pub slot_root: HashOut<F, NUM_HASH_OUT_ELTS>,
+    pub n_cells_per_slot: F,
+    pub n_slots_per_dataset: F,
+
+    pub slot_proof: Vec<HashOut<F, NUM_HASH_OUT_ELTS>>,
+
+    pub cell_data: Vec<Cell<F,D>>,
+    pub merkle_paths: Vec<MerklePath<F,D>>,
+
+}
+
+/// merkle path from leaf to root as vec of HashOut (4 Goldilocks field elems)
+#[derive(Clone, Debug, PartialEq)]
+pub struct MerklePath<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+> {
+    pub path: Vec<HashOut<F, NUM_HASH_OUT_ELTS>>
+}
+
+/// a vec of cell targets
+#[derive(Clone, Debug, PartialEq)]
+pub struct CellTarget {
+    pub data: Vec<Target>
+}
+
+/// cell data as field elements
+#[derive(Clone, Debug, PartialEq)]
+pub struct Cell<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+> {
+    pub data: Vec<F>,
+}
+
+//------- circuit impl --------
+impl<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: AlgebraicHasher<F, NUM_HASH_OUT_ELTS>,
+> SampleCircuit<F, D, H> {
+
+    /// samples and registers the public input
+    pub fn sample_slot_circuit_with_public_input(
+        &self,
+        builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>,
+    ) -> Result<SampleTargets> {
+        let targets = self.sample_slot_circuit(builder)?;
+        let mut pub_targets = vec![];
+        pub_targets.push(targets.slot_index);
+        pub_targets.extend_from_slice(&targets.dataset_root.elements);
+        pub_targets.extend_from_slice(&targets.entropy.elements);
+        builder.register_public_inputs(&pub_targets);
+        Ok(targets)
+    }
+
+    /// in-circuit sampling
+    /// WARNING: no public input are registered when calling this function
+    pub fn sample_slot_circuit(
+        &self,
+        builder: &mut CircuitBuilder::<F, D, NUM_HASH_OUT_ELTS>,
+    ) -> Result<SampleTargets> {
+        // circuit params
+        let CircuitParams {
+            max_depth,
+            max_log2_n_slots,
+            block_tree_depth,
+            n_field_elems_per_cell,
+            n_samples,
+        } = self.params;
+
+        // constants
+        let zero = builder.zero();
+        let one = builder.one();
+
+        // ***** prove slot root is in dataset tree *********
+
+        // Create virtual target for slot root and index
+        let slot_root = builder.add_virtual_hash();
+        let slot_index = builder.add_virtual_target();// public input
+        // let slot_index = builder.add_virtual_public_input();// public input
+
+        // dataset path bits (binary decomposition of leaf_index)
+        let d_path_bits = builder.split_le(slot_index,max_log2_n_slots);
+
+        // create virtual target for n_slots_per_dataset
+        let n_slots_per_dataset = builder.add_virtual_target();
+
+        // dataset last bits and mask bits
+        let (d_last_bits, d_mask_bits) =
+            ceiling_log2(builder, n_slots_per_dataset, max_log2_n_slots);
+
+        // dataset Merkle path (sibling hashes from leaf to root)
+        let d_merkle_path = MerkleProofTarget {
+            path: (0..max_log2_n_slots).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,
+            mask_bits: d_mask_bits,
+            merkle_path: d_merkle_path,
+        };
+
+        // dataset reconstructed root
+        let d_reconstructed_root =
+            MerkleTreeCircuit::<F,D, H>::reconstruct_merkle_root_circuit_with_mask(builder, &mut d_targets, max_log2_n_slots)?;
+
+        // expected Merkle root
+        let d_expected_root = builder.add_virtual_hash(); // public input
+        // let d_expected_root = builder.add_virtual_hash_public_input(); // public input
+
+        // check equality with expected root
+        for i in 0..NUM_HASH_OUT_ELTS {
+            builder.connect(d_expected_root.elements[i], d_reconstructed_root.elements[i]);
+        }
+
+        //*********** do the sampling ************
+
+        let mut data_targets =vec![];
+        let mut slot_sample_proofs = vec![];
+        let entropy_target = builder.add_virtual_hash(); // public input
+        // let entropy_target = builder.add_virtual_hash_public_input(); // public input
+
+        // virtual target for n_cells_per_slot
+        let n_cells_per_slot = builder.add_virtual_target();
+
+        // calculate last index = n_cells_per_slot-1
+        let slot_last_index = builder.sub(n_cells_per_slot, one);
+
+        // create the mask bits
+        // TODO: re-use this for block and slot trees
+        let mask_bits = builder.split_le(slot_last_index,max_depth);
+
+        // last and mask bits for block tree
+        let mut b_last_bits = builder.split_le(slot_last_index,max_depth);
+        let mut b_mask_bits = builder.split_le(slot_last_index,max_depth);
+
+        // last and mask bits for the slot tree
+        let mut s_last_bits = b_last_bits.split_off(block_tree_depth);
+        let mut s_mask_bits = b_mask_bits.split_off(block_tree_depth);
+
+        // pad mask bits with 0
+        b_mask_bits.push(BoolTarget::new_unsafe(zero.clone()));
+        s_mask_bits.push(BoolTarget::new_unsafe(zero.clone()));
+
+        for i in 0..n_samples{
+            // cell data targets
+            let mut data_i = (0..n_field_elems_per_cell).map(|_| builder.add_virtual_target()).collect::<Vec<_>>();
+            // hash the cell data
+            let mut hash_inputs:Vec<Target>= Vec::new();
+            hash_inputs.extend_from_slice(&data_i);
+            // let data_i_hash = builder.hash_n_to_hash_no_pad::<HF>(hash_inputs);
+            let data_i_hash = hash_n_no_padding::<F,D,H>(builder, hash_inputs)?;
+            // make the counter into hash digest
+            let ctr_target = builder.constant(F::from_canonical_u64((i+1) as u64));
+            let mut ctr = builder.add_virtual_hash();
+            for i in 0..ctr.elements.len() {
+                if i==0 {
+                    ctr.elements[i] = ctr_target;
+                }else{
+                    ctr.elements[i] = zero.clone();
+                }
+            }
+            // paths for block and slot
+            let mut b_path_bits = self.calculate_cell_index_bits(builder, &entropy_target, &d_targets.leaf, &ctr, mask_bits.clone())?;
+            let mut s_path_bits = b_path_bits.split_off(block_tree_depth);
+
+            let mut b_merkle_path = MerkleProofTarget {
+                path: (0..block_tree_depth).map(|_| builder.add_virtual_hash()).collect(),
+            };
+
+            let mut s_merkle_path = MerkleProofTarget {
+                path: (0..(max_depth - block_tree_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.clone(),
+                mask_bits: b_mask_bits.clone(),
+                merkle_path: b_merkle_path,
+            };
+
+            // reconstruct block root
+            let b_root = MerkleTreeCircuit::<F,D,H>::reconstruct_merkle_root_circuit_with_mask(builder, &mut block_targets, block_tree_depth)?;
+
+            let mut slot_targets = MerkleTreeTargets {
+                leaf: b_root,
+                path_bits:s_path_bits,
+                last_bits:s_last_bits.clone(),
+                mask_bits:s_mask_bits.clone(),
+                merkle_path:s_merkle_path,
+            };
+
+            // reconstruct slot root with block root as leaf
+            let slot_reconstructed_root = MerkleTreeCircuit::<F,D,H>::reconstruct_merkle_root_circuit_with_mask(builder, &mut slot_targets, max_depth-block_tree_depth)?;
+
+            // 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);
+
+            let cell_i = CellTarget{
+                data: data_i
+            };
+            data_targets.push(cell_i);
+            slot_sample_proofs.push(slot_sample_proof_target);
+
+        }
+
+        let st = SampleTargets {
+            entropy: entropy_target,
+            dataset_root: d_expected_root,
+            slot_index,
+            slot_root: d_targets.leaf,
+            n_cells_per_slot,
+            n_slots_per_dataset,
+            slot_proof: d_targets.merkle_path,
+            cell_data: data_targets,
+            merkle_paths: slot_sample_proofs,
+        };
+
+        Ok(st)
+    }
+
+    /// calculate the cell index = H( entropy | slotRoot | counter ) `mod` nCells
+    pub fn calculate_cell_index_bits(&self, builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>, entropy: &HashOutTarget<NUM_HASH_OUT_ELTS>, slot_root: &HashOutTarget<NUM_HASH_OUT_ELTS>, ctr: &HashOutTarget<NUM_HASH_OUT_ELTS>, mask_bits: Vec<BoolTarget>) -> Result<Vec<BoolTarget>> {
+        let mut hash_inputs:Vec<Target>= Vec::new();
+        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 = hash_n_with_padding::<F,D,H>(builder, hash_inputs)?;
+        let cell_index_bits =  builder.low_bits(hash_out.elements[0], self.params.max_depth, true);
+
+        let mut masked_cell_index_bits = vec![];
+
+        // extract the lowest 32 bits using the bit mask
+        for i in 0..self.params.max_depth{
+            masked_cell_index_bits.push(BoolTarget::new_unsafe(builder.mul(mask_bits[i].target, cell_index_bits[i].target)));
+        }
+
+        Ok(masked_cell_index_bits)
+    }
+
+    /// helper method to assign the targets in the circuit to actual field elems
+    pub fn sample_slot_assign_witness(
+        &self,
+        pw: &mut PartialWitness<F>,
+        targets: &SampleTargets,
+        witnesses: &SampleCircuitInput<F, D>,
+    ) -> Result<()>{
+        // circuit params
+        let CircuitParams {
+            max_depth,
+            n_field_elems_per_cell,
+            n_samples,
+            ..
+        } = self.params;
+
+        // assign n_cells_per_slot
+        pw.set_target(targets.n_cells_per_slot, witnesses.n_cells_per_slot)
+            ;
+
+        // assign n_slots_per_dataset
+        pw.set_target(targets.n_slots_per_dataset, witnesses.n_slots_per_dataset)
+            ;
+
+        // assign dataset proof
+        for (i, sibling_hash) in witnesses.slot_proof.iter().enumerate() {
+            pw.set_hash_target(targets.slot_proof.path[i], *sibling_hash)
+                ;
+        }
+        // assign slot index
+        pw.set_target(targets.slot_index, witnesses.slot_index)
+            ;
+
+        // assign the expected Merkle root of dataset to the target
+        pw.set_hash_target(targets.dataset_root, witnesses.dataset_root)
+            ;
+
+        // assign the sampled slot
+        pw.set_hash_target(targets.slot_root, witnesses.slot_root)
+            ;
+
+        // assign entropy
+        assign_hash_out_targets(pw, &targets.entropy, &witnesses.entropy)?;
+
+        // do the sample N times
+        for i in 0..n_samples {
+            // assign cell data
+            let leaf = witnesses.cell_data[i].data.clone();
+            for j in 0..n_field_elems_per_cell{
+                pw.set_target(targets.cell_data[i].data[j], leaf[j])
+                    ;
+            }
+            // assign proof for that cell
+            let cell_proof = witnesses.merkle_paths[i].path.clone();
+            for k in 0..max_depth {
+                pw.set_hash_target(targets.merkle_paths[i].path[k], cell_proof[k])
+                    ;
+            }
+        }
+
+        Ok(())
+    }
+
+}

goldibear_experiments/codex-plonky2-circuits/src/circuits/sponge.rs

@@ -0,0 +1,185 @@
+use plonky2::hash::hash_types::{HashOutTarget, RichField};
+use plonky2::hash::hashing::PlonkyPermutation;
+use plonky2::iop::target::Target;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::config::AlgebraicHasher;
+use plonky2_field::types::HasExtension;
+use crate::error::CircuitError;
+use crate::Result;
+use crate::circuits::params::NUM_HASH_OUT_ELTS;
+
+/// hash n targets (field elements) into hash digest / HashOutTarget (4 Goldilocks field elements)
+/// this function uses the 10* padding
+pub fn hash_n_with_padding<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: AlgebraicHasher<F, NUM_HASH_OUT_ELTS>
+>(
+    builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>,
+    inputs: Vec<Target>,
+) -> Result<HashOutTarget<NUM_HASH_OUT_ELTS>> {
+    Ok(
+        HashOutTarget::from_vec(
+            hash_n_to_m_with_padding::<F,D,H>(builder, inputs, NUM_HASH_OUT_ELTS)?
+        )
+    )
+}
+
+pub fn hash_n_to_m_with_padding<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: AlgebraicHasher<F, NUM_HASH_OUT_ELTS>
+>(
+    builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>,
+    inputs: Vec<Target>,
+    num_outputs: usize,
+) -> Result<Vec<Target>> {
+    let rate = H::AlgebraicPermutation::RATE;
+    let width = H::AlgebraicPermutation::WIDTH; // rate + capacity
+    let zero = builder.zero();
+    let one = builder.one();
+    let mut state = H::AlgebraicPermutation::new(core::iter::repeat(zero).take(width));
+
+    // Set the domain separator at index 8
+    let dom_sep_value = rate as u64 + 256 * 12 + 65536 * 63;
+    let dom_sep = builder.constant(F::from_canonical_u64(dom_sep_value));
+    state.set_elt(dom_sep, 8);
+
+    let n = inputs.len();
+    let num_chunks = (n + rate) / rate; // 10* padding
+    let mut input_iter = inputs.iter();
+
+    // Process the first (num_chunks - 1) chunks
+    for _ in 0..(num_chunks - 1) {
+        let mut chunk = Vec::with_capacity(rate);
+        for _ in 0..rate {
+            if let Some(&input) = input_iter.next() {
+                chunk.push(input);
+            } else {
+                // should not happen here
+                return Err(CircuitError::InsufficientInputs(rate,chunk.len()));
+            }
+        }
+        // Add the chunk to the state
+        for j in 0..rate {
+            state.set_elt(builder.add(state.as_ref()[j], chunk[j]), j);
+        }
+        // Apply permutation
+        state = builder.permute::<H>(state);
+    }
+
+    // Process the last chunk with 10* padding
+    let rem = num_chunks * rate - n; // 0 < rem <= rate
+    let ofs = rate - rem; // Offset where padding starts
+
+    let mut last_chunk = Vec::with_capacity(rate);
+    for _ in 0..ofs {
+        if let Some(&input) = input_iter.next() {
+            last_chunk.push(input);
+        } else {
+            last_chunk.push(zero); // Pad zeros if no more inputs
+        }
+    }
+
+    // Add the '1' padding bit
+    last_chunk.push(one);
+
+    // Pad zeros to reach the full rate
+    while last_chunk.len() < rate {
+        last_chunk.push(zero);
+    }
+
+    // Add the last chunk to the state
+    for j in 0..rate {
+        state.set_elt(builder.add(state.as_ref()[j], last_chunk[j]), j);
+    }
+    // Apply permutation
+    state = builder.permute::<H>(state);
+
+    // Squeeze until we have the desired number of outputs
+    let mut outputs = Vec::with_capacity(num_outputs);
+    loop {
+        for &s in state.squeeze() {
+            outputs.push(s);
+            if outputs.len() == num_outputs {
+                return Ok(outputs);
+            }
+        }
+        state = builder.permute::<H>(state);
+    }
+}
+
+/// hash n targets (field elements) into hash digest / HashOutTarget (4 Goldilocks field elements)
+/// this function uses doesn't pad and expects input to be divisible by rate
+/// rate is fixed at 8 for now.
+pub fn hash_n_no_padding<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: AlgebraicHasher<F, NUM_HASH_OUT_ELTS>
+>(
+    builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>,
+    inputs: Vec<Target>,
+) -> Result<HashOutTarget<NUM_HASH_OUT_ELTS>> {
+    Ok(
+        HashOutTarget::from_vec(
+            hash_n_to_m_no_padding::<F, D, H>(builder, inputs, NUM_HASH_OUT_ELTS)?
+        )
+    )
+}
+
+pub fn hash_n_to_m_no_padding<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: AlgebraicHasher<F, NUM_HASH_OUT_ELTS>
+>(
+    builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>,
+    inputs: Vec<Target>,
+    num_outputs: usize,
+) -> Result<Vec<Target>> {
+    let rate = H::AlgebraicPermutation::RATE;
+    let width = H::AlgebraicPermutation::WIDTH; // rate + capacity
+    let zero = builder.zero();
+    let mut state = H::AlgebraicPermutation::new(core::iter::repeat(zero).take(width));
+
+    // Set the domain separator at index 8
+    let dom_sep_value = rate as u64 + 256 * 12 + 65536 * 8;
+    let dom_sep = builder.constant(F::from_canonical_u64(dom_sep_value));
+    state.set_elt(dom_sep, 8);
+
+    let n = inputs.len();
+    if n % rate != 0 {
+        return Err(CircuitError::SpongeInputLengthMismatch(n, rate));
+    }
+    let num_chunks = n / rate; // 10* padding
+    let mut input_iter = inputs.iter();
+
+    // Process all chunks
+    for _ in 0..num_chunks {
+        let mut chunk = Vec::with_capacity(rate);
+        for _ in 0..rate {
+            if let Some(&input) = input_iter.next() {
+                chunk.push(input);
+            } else {
+                // should not happen here
+                return Err(CircuitError::InsufficientInputs(rate,chunk.len()));
+            }
+        }
+        // Add the chunk to the state
+        for j in 0..rate {
+            state.set_elt(builder.add(state.as_ref()[j], chunk[j]), j);
+        }
+        // Apply permutation
+        state = builder.permute::<H>(state);
+    }
+    // Squeeze until we have the desired number of outputs
+    let mut outputs = Vec::with_capacity(num_outputs);
+    loop {
+        for &s in state.squeeze() {
+            outputs.push(s);
+            if outputs.len() == num_outputs {
+                return Ok(outputs);
+            }
+        }
+        state = builder.permute::<H>(state);
+    }
+}

goldibear_experiments/codex-plonky2-circuits/src/circuits/utils.rs

@@ -0,0 +1,144 @@
+use std::{fs, io};
+use std::path::Path;
+use itertools::Itertools;
+use plonky2::hash::hash_types::{HashOut, HashOutTarget, RichField};
+use plonky2::iop::witness::{PartialWitness, WitnessWrite};
+use plonky2::iop::target::{BoolTarget, Target};
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2_field::types::HasExtension;
+use crate::Result;
+use crate::error::CircuitError;
+use crate::circuits::params::NUM_HASH_OUT_ELTS;
+
+// --------- helper functions ---------
+
+/// computes the `last_index` (the binary decomposition of `inp-1`) and the `mask_bits`
+pub fn ceiling_log2<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(
+    builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>,
+    inp: Target,
+    n: usize,
+) -> (Vec<BoolTarget>, Vec<BoolTarget>){
+    let one = builder.one();
+    let zero = builder.zero();
+    let last_index = builder.sub(inp, one.clone());
+    let last_bits = builder.split_le(last_index,n);
+
+    let mut aux: Vec<BoolTarget> = vec![BoolTarget::new_unsafe(zero.clone()); n + 1];
+    aux[n] = BoolTarget::new_unsafe(one.clone());
+    let mut mask: Vec<BoolTarget> = vec![BoolTarget::new_unsafe(zero.clone()); n + 1];
+    for i in (0..n).rev(){
+        let diff = builder.sub(one.clone(), last_bits[i].target);
+        let aux_i = builder.mul( aux[i+1].target, diff);
+        aux[i] = BoolTarget::new_unsafe(aux_i);
+        mask[i] = BoolTarget::new_unsafe(builder.sub(one.clone(), aux[i].target));
+    }
+
+    (last_bits, mask)
+}
+
+/// assign a vec of bool values to a vec of BoolTargets
+pub fn assign_bool_targets<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(
+    pw: &mut PartialWitness<F>,
+    bool_targets: &Vec<BoolTarget>,
+    bools: Vec<bool>,
+) -> Result<()>{
+    if bools.len() > bool_targets.len() {
+        return Err(CircuitError::AssignmentLengthMismatch (
+                bool_targets.len(),
+                bools.len(),
+            )
+        );
+    }
+    for (i, bit) in bools.iter().enumerate() {
+        pw.set_bool_target(bool_targets[i], *bit);
+    }
+    Ok(())
+}
+
+/// assign a vec of field elems to hash out target elements
+/// TODO: change to HashOut
+pub fn assign_hash_out_targets<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(
+    pw: &mut PartialWitness<F>,
+    hash_out_elements_targets: &HashOutTarget<NUM_HASH_OUT_ELTS>,
+    hash_out_elements: &HashOut<F, NUM_HASH_OUT_ELTS>,
+) -> Result<()>{
+
+    // Assign each field element to its corresponding target
+    for (j, (&target, &element)) in hash_out_elements_targets.elements.iter().zip(hash_out_elements.elements.iter()).enumerate() {
+        pw.set_target(target, element)
+    }
+
+    Ok(())
+}
+
+/// helper fn to multiply a HashOutTarget by a Target
+pub fn mul_hash_out_target<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>, t: &Target, hash_target: &mut HashOutTarget<NUM_HASH_OUT_ELTS>) -> HashOutTarget<NUM_HASH_OUT_ELTS> {
+    let mut mul_elements = vec![];
+    for i in 0..NUM_HASH_OUT_ELTS {
+        mul_elements.push(builder.mul(hash_target.elements[i], *t));
+    }
+    HashOutTarget::from_vec(mul_elements)
+}
+
+/// helper fn to add AND assign a HashOutTarget (hot) to a mutable HashOutTarget (mut_hot)
+pub fn add_assign_hash_out_target<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>, mut_hot: &mut HashOutTarget<NUM_HASH_OUT_ELTS>, hot: &HashOutTarget<NUM_HASH_OUT_ELTS>) {
+    for i in 0..NUM_HASH_OUT_ELTS {
+        mut_hot.elements[i] = builder.add(mut_hot.elements[i], hot.elements[i]);
+    }
+}
+
+/// Reads the contents of the specified file and returns them as a vector of bytes using `std::fs::read`.
+pub fn read_bytes_from_file<P: AsRef<Path>>(path: P) -> io::Result<Vec<u8>> {
+    fs::read(path)
+}
+
+/// select hash helper method
+/// Computes `if b { h0 } else { h1 }`.
+pub fn select_hash<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(
+    builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>,
+    b: BoolTarget,
+    h0: HashOutTarget<NUM_HASH_OUT_ELTS>,
+    h1: HashOutTarget<NUM_HASH_OUT_ELTS>,
+) -> HashOutTarget<NUM_HASH_OUT_ELTS> {
+    HashOutTarget {
+        elements: core::array::from_fn(|i| builder.select(b, h0.elements[i], h1.elements[i])),
+    }
+}
+
+/// Converts a Vec<T> into a fixed-size array [T; N], returning an error if the lengths don't match.
+pub fn vec_to_array<const N: usize, T>(vec: Vec<T>) -> Result<[T; N]> {
+    vec.try_into().map_err(|v: Vec<T>| CircuitError::ArrayLengthMismatchError(format!(
+        "Expected exactly {} elements, got {}",
+        N,
+        v.len()
+    )))
+}
+
+/// Computes `if b { v0 } else { v1 }`.
+pub fn select_vec<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(builder: &mut CircuitBuilder<F, D, NUM_HASH_OUT_ELTS>, b: BoolTarget, v0: &[Target], v1: &[Target]) -> Vec<Target> {
+    v0.iter()
+        .zip_eq(v1)
+        .map(|(t0, t1)| builder.select(b, *t0, *t1))
+        .collect()
+}

goldibear_experiments/codex-plonky2-circuits/src/error.rs

@@ -0,0 +1,71 @@
+use thiserror::Error;
+
+/// Custom error types for the Circuits.
+#[derive(Error, Debug)]
+pub enum CircuitError {
+    #[error("Path bits length mismatch: expected {0}, found {1}")]
+    PathBitsLengthMismatch(usize, usize),
+
+    #[error("Mask bits length mismatch: expected {0}, found {1}")]
+    MaskBitsLengthMismatch(usize, usize),
+
+    #[error("Last bits length mismatch: expected {0}, found {1}")]
+    LastBitsLengthMismatch(usize, usize),
+
+    #[error("Path bits and max depth mismatch: path bits length {0}, max depth {1}")]
+    PathBitsMaxDepthMismatch(usize, usize),
+
+    #[error("Insufficient input elements for chunk; expected {0}, found {1}")]
+    InsufficientInputs (usize, usize),
+
+    #[error("Sponge: Input length ({0}) must be divisible by rate ({1}) for no padding")]
+    SpongeInputLengthMismatch(usize, usize),
+
+    #[error("Assignment length mismatch: expected at least {0}, found {1}")]
+    AssignmentLengthMismatch(usize, usize),
+
+    #[error("Failed to assign Target at index {0}: {1}")]
+    ArrayTargetAssignmentError(usize, String),
+
+    #[error("Failed to assign Target {0}: {1}")]
+    TargetAssignmentError(String, String),
+
+    #[error("Failed to assign BoolTarget at index {0}: {1}")]
+    ArrayBoolTargetAssignmentError(usize, String),
+
+    #[error("Failed to assign BoolTarget {0}: {1}")]
+    BoolTargetAssignmentError(String, String),
+
+    #[error("Failed to assign HashTarget {0}: {1}")]
+    HashTargetAssignmentError(String, String),
+
+    #[error("Failed to assign ProofTarget {0}: {1}")]
+    ProofTargetAssignmentError(String, String),
+
+    #[error("Failed to assign VerifierDataTarget {0}")]
+    VerifierDataTargetAssignmentError(String),
+
+    #[error("Array Length Mismatch Error {0}")]
+    ArrayLengthMismatchError(String),
+
+    #[error("Proof Verification Failed {0}")]
+    InvalidProofError(String),
+
+    #[error("Proof Generation Failed {0}")]
+    ProofGenerationError(String),
+
+    #[error("Error in Recursion Tree: {0}")]
+    RecursionTreeError(String),
+
+    #[error("Dummy Proof Generation Error: {0}")]
+    DummyProofGenerationError(String),
+
+    #[error("Conditional Verification Error: {0}")]
+    ConditionalVerificationError(String),
+
+    #[error("Recursive Proof VerifierData Check Failed: {0}")]
+    RecursiveProofVerifierDataCheckError(String),
+
+    #[error("Expected Option {0} to contain value")]
+    OptionError(String),
+}

goldibear_experiments/codex-plonky2-circuits/src/lib.rs

@@ -0,0 +1,4 @@
+pub mod circuits;
+pub mod error;
+
+pub type Result<T> = core::result::Result<T, error::CircuitError>;

goldibear_experiments/proof-input/.gitignore

@@ -0,0 +1,13 @@
+#IDE Related
+.idea
+
+# Cargo build
+/target
+Cargo.lock
+
+# Profile-guided optimization
+/tmp
+pgo-data.profdata
+
+# MacOS nuisances
+.DS_Store

goldibear_experiments/proof-input/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "proof-input"
+description = "proof input generation library"
+authors = ["Mohammed Alghazwi <m.ghazwi@gmail.com>"]
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+clap = { version = "4.0", features = ["derive"] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+anyhow = { version = "1.0.89"}
+plonky2  = { git = "https://github.com/telosnetwork/plonky2_goldibear.git"}
+plonky2_field = { git = "https://github.com/telosnetwork/plonky2_goldibear.git"}
+# --- local ---
+codex-plonky2-circuits = { path = "../codex-plonky2-circuits" }
+criterion = "0.5.1"
+
+[[bench]]
+name = "sample_cells"
+harness = false

goldibear_experiments/proof-input/benches/sample_cells.rs

@@ -0,0 +1,103 @@
+use anyhow::Result;
+use criterion::{criterion_group, criterion_main, Criterion};
+use plonky2::iop::witness::PartialWitness;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::CircuitConfig;
+use plonky2::plonk::config::GenericConfig;
+
+use codex_plonky2_circuits::circuits::sample_cells::SampleCircuit;
+use proof_input::gen_input::gen_testing_circuit_input;
+use proof_input::params::{D, C, F, HF, Params};
+use codex_plonky2_circuits::circuits::params::NUM_HASH_OUT_ELTS;
+
+/// Benchmark for building, proving, and verifying the Plonky2 circuit.
+fn bench_prove_verify(c: &mut Criterion) -> Result<()>{
+
+    let n_samples = 100;
+    // get default parameters
+    let params = Params::default();
+    let mut test_params = params.input_params;
+    test_params.n_samples = n_samples;
+
+    let mut circuit_params = params.circuit_params;
+    circuit_params.n_samples = n_samples;
+
+    // gen the circuit input
+    let circ_input = gen_testing_circuit_input::<F,D>(&test_params);
+
+    // Create the circuit configuration
+    let config = CircuitConfig::standard_recursion_config_gl();
+    let mut builder = CircuitBuilder::<F, D, NUM_HASH_OUT_ELTS>::new(config);
+
+    // Initialize the SampleCircuit with the parameters
+    let circ = SampleCircuit::<F,D,HF>::new(circuit_params.clone());
+    let targets = circ.sample_slot_circuit_with_public_input(&mut builder)?;
+
+    // Create a PartialWitness and assign the circuit input
+    let mut pw = PartialWitness::new();
+    circ.sample_slot_assign_witness(&mut pw, &targets, &circ_input.clone());
+
+    // Benchmark Group: Separate benchmarks for building, proving, and verifying
+    let mut group = c.benchmark_group("Sampling Circuit Benchmark");
+
+    // Benchmark the Circuit Building Phase
+    group.bench_function("Build Circuit", |b| {
+        b.iter(|| {
+            let config = CircuitConfig::standard_recursion_config_gl();
+            let mut local_builder = CircuitBuilder::<F, D, NUM_HASH_OUT_ELTS>::new(config);
+            let _targets = circ.sample_slot_circuit_with_public_input(&mut local_builder);
+            let _data = local_builder.build::<C>();
+        })
+    });
+
+    // Build the circuit once for proving and verifying benchmarks
+    let build_start = std::time::Instant::now();
+    let data = builder.build::<C>();
+    let build_duration = build_start.elapsed();
+    println!("Build time: {:?}", build_duration);
+    println!("Circuit size (degree bits): {:?}", data.common.degree_bits());
+
+    let num_constr: usize = data.common
+        .gates
+        .iter()
+        .map(|gate| gate.0.num_constraints())
+        .sum();
+
+    println!("Number of constraints: {}", num_constr);
+    println!("Number of gates used: {}", data.common.gates.len());
+
+    // Benchmark the Proving Phase
+    group.bench_function("Prove Circuit", |b| {
+        b.iter(|| {
+            let local_pw = pw.clone();
+            data.prove(local_pw).expect("Failed to prove circuit")
+        })
+    });
+
+    // Generate the proof once for verification benchmarking
+    let prove_start = std::time::Instant::now();
+    let proof_with_pis = data.prove(pw.clone()).expect("Failed to prove circuit");
+    let prove_duration = prove_start.elapsed();
+    println!("prove time: {:?}", prove_duration);
+    let verifier_data = data.verifier_data();
+
+    println!("Proof size: {} bytes", proof_with_pis.to_bytes().len());
+
+    // Benchmark the Verifying Phase
+    group.bench_function("Verify Proof", |b| {
+        b.iter(|| {
+            verifier_data.verify(proof_with_pis.clone()).expect("Failed to verify proof");
+        })
+    });
+
+    group.finish();
+    Ok(())
+}
+
+/// Criterion benchmark group
+criterion_group!{
+    name = prove_verify_benches;
+    config = Criterion::default().sample_size(10);
+    targets = bench_prove_verify
+}
+criterion_main!(prove_verify_benches);

goldibear_experiments/proof-input/src/data_structs.rs

@@ -0,0 +1,257 @@
+// Data structure used to generate the proof input
+
+use plonky2::hash::hash_types::{HashOut, RichField};
+use codex_plonky2_circuits::circuits::sample_cells::Cell;
+use plonky2_field::types::{HasExtension, Sample};
+use crate::merkle_tree::merkle_safe::{MerkleProof, MerkleTree};
+use crate::params::{InputParams, HF};
+use crate::sponge::hash_bytes_no_padding;
+use crate::utils::{bits_le_padded_to_usize, calculate_cell_index_bits, usize_to_bits_le};
+use codex_plonky2_circuits::circuits::params::NUM_HASH_OUT_ELTS;
+
+// ----------------- slot tree -----------------
+#[derive(Clone)]
+pub struct SlotTree<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+> {
+    pub tree: MerkleTree<F, D>,         // slot tree
+    pub block_trees: Vec<MerkleTree<F,D>>, // vec of block trees
+    pub cell_data: Vec<Cell<F, D>>,  // cell data as field elements
+    pub params: InputParams,              // parameters
+}
+
+impl<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+> SlotTree<F, D> {
+    /// Create a slot tree with fake data, for testing only
+    pub fn new_default(params: &InputParams) -> Self {
+        // generate fake cell data
+        let cell_data = (0..params.n_cells)
+            .map(|_| new_random_cell(params))
+            .collect::<Vec<_>>();
+        Self::new(cell_data, params.clone())
+    }
+
+    /// Create a new slot tree with the supplied cell data and parameters
+    pub fn new(cells: Vec<Cell<F, D>>, params: InputParams) -> Self {
+        let leaves: Vec<HashOut<F, NUM_HASH_OUT_ELTS>> = cells
+            .iter()
+            .map(|element| hash_bytes_no_padding::<F,D,HF>(&element.data))
+            .collect();
+        let zero = HashOut {
+            elements: [F::zero(); 4],
+        };
+        let n_blocks = params.n_blocks_test();
+        let n_cells_in_blocks = params.n_cells_in_blocks();
+
+        let block_trees = (0..n_blocks)
+            .map(|i| {
+                let start = i * n_cells_in_blocks;
+                let end = (i + 1) * n_cells_in_blocks;
+                Self::get_block_tree(&leaves[start..end].to_vec())
+            })
+            .collect::<Vec<_>>();
+        let block_roots = block_trees
+            .iter()
+            .map(|t| t.root().unwrap())
+            .collect::<Vec<_>>();
+        let slot_tree = MerkleTree::<F,D>::new(&block_roots, zero).unwrap();
+        Self {
+            tree: slot_tree,
+            block_trees,
+            cell_data: cells,
+            params,
+        }
+    }
+
+    /// Generates a proof for the 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,D> {
+        let block_index = index / self.params.n_cells_in_blocks();
+        let leaf_index = index % self.params.n_cells_in_blocks();
+        let block_proof = self.block_trees[block_index].get_proof(leaf_index).unwrap();
+        let slot_proof = self.tree.get_proof(block_index).unwrap();
+
+        // Combine the paths from the block and slot proofs
+        let mut combined_path = block_proof.path.clone();
+        combined_path.extend(slot_proof.path.clone());
+
+        MerkleProof::<F,D> {
+            index,
+            path: combined_path,
+            nleaves: self.cell_data.len(),
+            zero: block_proof.zero.clone(),
+        }
+    }
+
+    fn get_block_tree(leaves: &Vec<HashOut<F, NUM_HASH_OUT_ELTS>>) -> MerkleTree<F,D> {
+        let zero = HashOut {
+            elements: [F::zero(); 4],
+        };
+        // Build the Merkle tree
+        let block_tree = MerkleTree::<F,D>::new(leaves, zero).unwrap();
+        block_tree
+    }
+}
+
+// -------------- Dataset Tree -------------
+/// Dataset tree containing all slot trees
+#[derive(Clone)]
+pub struct DatasetTree<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+> {
+    pub tree: MerkleTree<F,D>,          // dataset tree
+    pub slot_trees: Vec<SlotTree<F, D>>, // vec of slot trees
+    pub params: InputParams,               // parameters
+}
+
+/// Dataset Merkle proof struct, containing the dataset proof and sampled proofs.
+#[derive(Clone)]
+pub struct DatasetProof<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+> {
+    pub slot_index: F,
+    pub entropy: HashOut<F, NUM_HASH_OUT_ELTS>,
+    pub dataset_proof: MerkleProof<F,D>,    // proof for dataset level tree
+    pub slot_proofs: Vec<MerkleProof<F,D>>, // proofs for sampled slot
+    pub cell_data: Vec<Cell<F,D>>,
+}
+
+impl<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+> DatasetTree<F, D> {
+    /// Dataset tree with fake data, for testing only
+    pub fn new_default(params: &InputParams) -> Self {
+        let mut slot_trees = vec![];
+        let n_slots = 1 << params.dataset_depth_test();
+        for _ in 0..n_slots {
+            slot_trees.push(SlotTree::<F, D>::new_default(params));
+        }
+        Self::new(slot_trees, params.clone())
+    }
+
+    /// Create data for only the specified slot index in params
+    pub fn new_for_testing(params: &InputParams) -> Self {
+        let mut slot_trees = vec![];
+        // let n_slots = 1 << params.dataset_depth();
+        let n_slots = params.n_slots;
+        // zero hash
+        let zero = HashOut {
+            elements: [F::zero(); 4],
+        };
+        let zero_slot = SlotTree::<F, D> {
+            tree: MerkleTree::<F,D>::new(&[zero.clone()], zero.clone()).unwrap(),
+            block_trees: vec![],
+            cell_data: vec![],
+            params: params.clone(),
+        };
+        for i in 0..n_slots {
+            if i == params.testing_slot_index {
+                slot_trees.push(SlotTree::<F, D>::new_default(params));
+            } else {
+                slot_trees.push(zero_slot.clone());
+            }
+        }
+        // get the roots of slot trees
+        let slot_roots = slot_trees
+            .iter()
+            .map(|t| t.tree.root().unwrap())
+            .collect::<Vec<_>>();
+        let dataset_tree = MerkleTree::<F,D>::new(&slot_roots, zero).unwrap();
+        Self {
+            tree: dataset_tree,
+            slot_trees,
+            params: params.clone(),
+        }
+    }
+
+    /// Same as default but with supplied slot trees
+    pub fn new(slot_trees: Vec<SlotTree<F, D>>, params: InputParams) -> Self {
+        // get the roots of slot trees
+        let slot_roots = slot_trees
+            .iter()
+            .map(|t| t.tree.root().unwrap())
+            .collect::<Vec<_>>();
+        // zero hash
+        let zero = HashOut {
+            elements: [F::zero(); 4],
+        };
+        let dataset_tree = MerkleTree::<F,D>::new(&slot_roots, zero).unwrap();
+        Self {
+            tree: dataset_tree,
+            slot_trees,
+            params,
+        }
+    }
+
+    /// Generates a proof for the given slot index
+    /// Also takes entropy so it can use it to sample the slot
+    /// note: proofs are padded based on the params in self
+    pub fn sample_slot(&self, index: usize, entropy: usize) -> DatasetProof<F,D> {
+        let mut dataset_proof = self.tree.get_proof(index).unwrap();
+        Self::pad_proof(&mut dataset_proof, self.params.dataset_max_depth());
+
+        let slot = &self.slot_trees[index];
+        let slot_root = slot.tree.root().unwrap();
+        let mut slot_proofs = vec![];
+        let mut cell_data = vec![];
+        let entropy_field = F::from_canonical_u64(entropy as u64);
+        let mut entropy_as_digest = HashOut::<F, NUM_HASH_OUT_ELTS>::zero();
+        entropy_as_digest.elements[0] = entropy_field;
+
+        // get the index for cell from H(slot_root|counter|entropy)
+        let mask_bits = usize_to_bits_le(self.params.n_cells-1, self.params.max_depth+1);
+        for i in 0..self.params.n_samples {
+            let cell_index_bits = calculate_cell_index_bits(
+                &entropy_as_digest.elements.to_vec(),
+                slot_root,
+                i + 1,
+                self.params.max_depth,
+                mask_bits.clone()
+            );
+            let cell_index = bits_le_padded_to_usize(&cell_index_bits);
+            let mut s_proof = slot.get_proof(cell_index);
+            Self::pad_proof(&mut s_proof, self.params.max_depth);
+            slot_proofs.push(s_proof);
+            let data_i = slot.cell_data[cell_index].data.clone();
+            let cell_i = Cell::<F,D>{
+                data: data_i
+            };
+            cell_data.push(cell_i);
+        }
+
+        DatasetProof {
+            slot_index: F::from_canonical_u64(index as u64),
+            entropy: entropy_as_digest,
+            dataset_proof,
+            slot_proofs,
+            cell_data,
+        }
+    }
+    /// pad the proof with 0s until max_depth
+    pub fn pad_proof(merkle_proof: &mut MerkleProof<F,D>, max_depth: usize){
+        for i in merkle_proof.path.len()..max_depth{
+            merkle_proof.path.push(HashOut::<F, NUM_HASH_OUT_ELTS>::zero());
+        }
+    }
+}
+
+// ------------ helper functions -------------
+
+/// Create a new cell with random data, using the parameters from `Params`
+pub fn new_random_cell<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(params: &InputParams) -> Cell<F,D> {
+    let data = (0..params.n_field_elems_per_cell())
+        .map(|_| F::rand())
+        .collect::<Vec<_>>();
+    Cell::<F,D> {
+        data,
+    }
+}

goldibear_experiments/proof-input/src/gen_input.rs

@@ -0,0 +1,234 @@
+use plonky2::hash::hash_types::RichField;
+use plonky2::plonk::config::{GenericConfig, Hasher};
+use crate::params::{Params,InputParams};
+use crate::utils::{bits_le_padded_to_usize, calculate_cell_index_bits, ceiling_log2, usize_to_bits_le};
+use crate::merkle_tree::merkle_safe::MerkleProof;
+use codex_plonky2_circuits::circuits::sample_cells::{MerklePath, SampleCircuit, SampleCircuitInput, SampleTargets};
+use plonky2::iop::witness::PartialWitness;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData};
+use plonky2::plonk::proof::ProofWithPublicInputs;
+use codex_plonky2_circuits::circuits::params::CircuitParams;
+use plonky2_field::types::HasExtension;
+use crate::data_structs::DatasetTree;
+use crate::sponge::hash_bytes_no_padding;
+use crate::params::{C, D, F, HF};
+use codex_plonky2_circuits::circuits::params::NUM_HASH_OUT_ELTS;
+
+/// generates circuit input (SampleCircuitInput) from fake data for testing
+/// which can be later stored into json see json.rs
+pub fn gen_testing_circuit_input<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(params: &InputParams) -> SampleCircuitInput<F,D>{
+    let dataset_t = DatasetTree::<F, D>::new_for_testing(&params);
+
+    let slot_index = params.testing_slot_index; // samples the specified slot
+    let entropy = params.entropy; // Use the entropy from Params
+
+    let proof = dataset_t.sample_slot(slot_index, entropy);
+    let slot_root = dataset_t.slot_trees[slot_index].tree.root().unwrap();
+
+    let mut slot_paths = vec![];
+    for i in 0..params.n_samples {
+        let path = proof.slot_proofs[i].path.clone();
+        let mp = MerklePath::<F,D>{
+            path,
+        };
+        slot_paths.push(mp);
+    }
+
+    SampleCircuitInput::<F, D> {
+        entropy: proof.entropy,
+        dataset_root: dataset_t.tree.root().unwrap(),
+        slot_index: proof.slot_index.clone(),
+        slot_root,
+        n_cells_per_slot: F::from_canonical_usize(params.n_cells),
+        n_slots_per_dataset: F::from_canonical_usize(params.n_slots),
+        slot_proof: proof.dataset_proof.path.clone(),
+        cell_data: proof.cell_data.clone(),
+        merkle_paths: slot_paths,
+    }
+}
+
+/// verifies the given circuit input.
+/// this is non circuit version for sanity check
+pub fn verify_circuit_input<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(circ_input: SampleCircuitInput<F,D>, params: &InputParams) -> bool{
+    let slot_index = circ_input.slot_index.as_canonical_u64();
+    let slot_root = circ_input.slot_root.clone();
+    // check dataset level proof
+    let slot_proof = circ_input.slot_proof.clone();
+    let dataset_path_bits = usize_to_bits_le(slot_index as usize, params.dataset_max_depth());
+    let (dataset_last_bits, dataset_mask_bits) = ceiling_log2(params.n_slots, params.dataset_max_depth());
+    let reconstructed_slot_root = MerkleProof::<F,D>::reconstruct_root2(
+        slot_root,
+        dataset_path_bits,
+        dataset_last_bits,
+        slot_proof,
+        dataset_mask_bits,
+        params.max_slots.trailing_zeros() as usize,
+    ).unwrap();
+    // assert reconstructed equals dataset root
+    assert_eq!(reconstructed_slot_root, circ_input.dataset_root.clone());
+
+    // check each sampled cell
+    // get the index for cell from H(slot_root|counter|entropy)
+    let mask_bits = usize_to_bits_le(params.n_cells -1, params.max_depth);
+    for i in 0..params.n_samples {
+        let cell_index_bits = calculate_cell_index_bits(
+            &circ_input.entropy.elements.to_vec(),
+            slot_root,
+            i + 1,
+            params.max_depth,
+            mask_bits.clone(),
+        );
+
+        let cell_index = bits_le_padded_to_usize(&cell_index_bits);
+
+        let s_res = verify_cell_proof(&circ_input, &params, cell_index, i);
+        if s_res.unwrap() == false {
+            println!("call {} is false", i);
+            return false;
+        }
+    }
+    true
+}
+
+/// Verify the given proof for slot tree, checks equality with the given root
+pub fn verify_cell_proof<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(circ_input: &SampleCircuitInput<F,D>, params: &InputParams, cell_index: usize, ctr: usize) -> anyhow::Result<bool> {
+    let mut block_path_bits = usize_to_bits_le(cell_index, params.max_depth);
+    let last_index = params.n_cells - 1;
+    let mut block_last_bits = usize_to_bits_le(last_index, params.max_depth);
+
+    let split_point = params.bot_depth();
+
+    let slot_last_bits = block_last_bits.split_off(split_point);
+    let slot_path_bits = block_path_bits.split_off(split_point);
+
+    // pub type HP = <PoseidonHash as Hasher<F>>::Permutation;
+    let leaf_hash = hash_bytes_no_padding::<F,D,HF>(&circ_input.cell_data[ctr].data);
+
+    let mut block_path = circ_input.merkle_paths[ctr].path.clone();
+    let slot_path = block_path.split_off(split_point);
+
+    let mut block_mask_bits = usize_to_bits_le(last_index, params.max_depth+1);
+    let mut slot_mask_bits = block_mask_bits.split_off(split_point);
+
+    block_mask_bits.push(false);
+    slot_mask_bits.push(false);
+
+    let block_res = MerkleProof::<F,D>::reconstruct_root2(
+        leaf_hash,
+        block_path_bits.clone(),
+        block_last_bits.clone(),
+        block_path,
+        block_mask_bits,
+        params.bot_depth(),
+    );
+    let reconstructed_root = MerkleProof::<F,D>::reconstruct_root2(
+        block_res.unwrap(),
+        slot_path_bits,
+        slot_last_bits,
+        slot_path,
+        slot_mask_bits,
+        params.max_depth - params.bot_depth(),
+    );
+
+    Ok(reconstructed_root.unwrap() == circ_input.slot_root)
+}
+
+
+/// returns exactly M default circuit input
+pub fn get_m_default_circ_input<const M: usize>() -> [SampleCircuitInput<F,D>; M]{
+    let params = Params::default().input_params;
+    // let one_circ_input = gen_testing_circuit_input::<F,D>(&params);
+    // let circ_input: [SampleCircuitInput<F,D>; M] = (0..M)
+    //     .map(|_| one_circ_input.clone())
+    //     .collect::<Vec<_>>()
+    //     .try_into().unwrap();
+    // circ_input
+    get_m_circ_input::<M>(params)
+}
+
+/// returns exactly M default circuit input
+pub fn get_m_circ_input<const M: usize>(params: InputParams) -> [SampleCircuitInput<F,D>; M]{
+    // let params = Params::default().input_params;
+    let one_circ_input = gen_testing_circuit_input::<F,D>(&params);
+    let circ_input: [SampleCircuitInput<F,D>; M] = (0..M)
+        .map(|_| one_circ_input.clone())
+        .collect::<Vec<_>>()
+        .try_into().unwrap();
+    circ_input
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::Instant;
+    use super::*;
+    use plonky2::plonk::circuit_data::CircuitConfig;
+    use plonky2::iop::witness::PartialWitness;
+    use plonky2::plonk::circuit_builder::CircuitBuilder;
+    use codex_plonky2_circuits::circuits::params::CircuitParams;
+    use codex_plonky2_circuits::circuits::sample_cells::SampleCircuit;
+    // use crate::params::{C, D, F};
+
+    // Test sample cells (non-circuit)
+    #[test]
+    fn test_gen_verify_proof(){
+        let params = Params::default().input_params;
+        let w = gen_testing_circuit_input::<F,D>(&params);
+        assert!(verify_circuit_input::<F,D>(w, &params));
+    }
+
+    // Test sample cells in-circuit for a selected slot
+    #[test]
+    fn test_proof_in_circuit() -> anyhow::Result<()> {
+        // get input
+        let mut params = Params::default();
+        let mut input_params = params.input_params;
+        input_params.n_samples = 10;
+        let circ_input = gen_testing_circuit_input::<F,D>(&input_params);
+
+        // Create the circuit
+        let config = CircuitConfig::standard_recursion_config_gl();
+        let mut builder = CircuitBuilder::<F, D, NUM_HASH_OUT_ELTS>::new(config);
+
+        let mut circuit_params = params.circuit_params;
+        circuit_params.n_samples = 10;
+
+        // build the circuit
+        let circ = SampleCircuit::<F,D,HF>::new(circuit_params.clone());
+        let mut targets = circ.sample_slot_circuit_with_public_input(&mut builder)?;
+
+        // Create a PartialWitness and assign
+        let mut pw = PartialWitness::new();
+
+        // assign a witness
+        circ.sample_slot_assign_witness(&mut pw, &targets, &circ_input)?;
+
+        // 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(())
+    }
+
+}

goldibear_experiments/proof-input/src/lib.rs

@@ -0,0 +1,7 @@
+
+pub mod gen_input;
+pub mod params;
+pub mod utils;
+pub mod sponge;
+pub mod merkle_tree;
+pub mod data_structs;

goldibear_experiments/proof-input/src/merkle_tree/key_compress.rs

@@ -0,0 +1,27 @@
+use plonky2::hash::hash_types::{HashOut, RichField};
+use plonky2::hash::hashing::PlonkyPermutation;
+use plonky2::plonk::config::Hasher;
+use codex_plonky2_circuits::circuits::params::NUM_HASH_OUT_ELTS;
+use plonky2_field::types::HasExtension;
+
+/// Compression function which takes two 256 bit inputs (HashOut) and u64 key (which is converted to field element in the function)
+/// and returns a 256 bit output (HashOut /  4 Goldilocks field elems).
+pub fn key_compress<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H:Hasher<F>
+>(x: HashOut<F, NUM_HASH_OUT_ELTS>, y: HashOut<F, NUM_HASH_OUT_ELTS>, key: u64) -> HashOut<F, NUM_HASH_OUT_ELTS> {
+
+    let key_field = F::from_canonical_u64(key);
+
+    let mut perm = H::Permutation::new(core::iter::repeat(F::zero()));
+    perm.set_from_slice(&x.elements, 0);
+    perm.set_from_slice(&y.elements, NUM_HASH_OUT_ELTS);
+    perm.set_elt(key_field,NUM_HASH_OUT_ELTS*2);
+
+    perm.permute();
+
+    HashOut {
+        elements: perm.squeeze()[..NUM_HASH_OUT_ELTS].try_into().unwrap(),
+    }
+}

goldibear_experiments/proof-input/src/merkle_tree/merkle_safe.rs

@@ -0,0 +1,247 @@
+// Implementation of "safe" merkle tree
+// consistent with the one in codex:
+// https://github.com/codex-storage/nim-codex/blob/master/codex/merkletree/merkletree.nim
+
+use anyhow::{ensure, Result};
+// use plonky2::field::goldilocks_field::GoldilocksField;
+use plonky2::hash::hash_types::{HashOut, RichField};
+// use plonky2::hash::poseidon::PoseidonHash;
+use plonky2::plonk::config::Hasher;
+use std::ops::Shr;
+use plonky2_field::types::{ HasExtension};
+use crate::merkle_tree::key_compress::key_compress;
+use crate::params::HF;
+use codex_plonky2_circuits::circuits::params::NUM_HASH_OUT_ELTS;
+
+// Constants for the keys used in compression
+pub const KEY_NONE: u64 = 0x0;
+pub const KEY_BOTTOM_LAYER: u64 = 0x1;
+pub const KEY_ODD: u64 = 0x2;
+pub const KEY_ODD_AND_BOTTOM_LAYER: u64 = 0x3;
+
+/// Merkle tree struct, containing the layers, compression function, and zero hash.
+#[derive(Clone)]
+pub struct MerkleTree<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+> {
+    pub layers: Vec<Vec<HashOut<F, NUM_HASH_OUT_ELTS>>>,
+    pub zero: HashOut<F, NUM_HASH_OUT_ELTS>,
+}
+
+impl<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+> MerkleTree<F, D> {
+    /// Constructs a new Merkle tree from the given leaves.
+    pub fn new(
+        leaves: &[HashOut<F, NUM_HASH_OUT_ELTS>],
+        zero: HashOut<F, NUM_HASH_OUT_ELTS>,
+    ) -> Result<Self> {
+        let layers = merkle_tree_worker::<F, D>(leaves, zero, true)?;
+        Ok(Self {
+            layers,
+            zero,
+        })
+    }
+
+    /// Returns the depth of the Merkle tree.
+    pub fn depth(&self) -> usize {
+        self.layers.len() - 1
+    }
+
+    /// Returns the number of leaves in the Merkle tree.
+    pub fn leaves_count(&self) -> usize {
+        self.layers[0].len()
+    }
+
+    /// Returns the root hash of the Merkle tree.
+    pub fn root(&self) -> Result<HashOut<F, NUM_HASH_OUT_ELTS>> {
+        let last_layer = self.layers.last().ok_or_else(|| anyhow::anyhow!("Empty tree"))?;
+        ensure!(last_layer.len() == 1, "Invalid Merkle tree");
+        Ok(last_layer[0])
+    }
+
+    /// Generates a Merkle proof for a given leaf index.
+    pub fn get_proof(&self, index: usize) -> Result<MerkleProof<F, D>> {
+        let depth = self.depth();
+        let nleaves = self.leaves_count();
+
+        ensure!(index < nleaves, "Index out of bounds");
+
+        let mut path = Vec::with_capacity(depth);
+        let mut k = index;
+        let mut m = nleaves;
+
+        for i in 0..depth {
+            let j = k ^ 1;
+            let sibling = if j < m {
+                self.layers[i][j]
+            } else {
+                self.zero
+            };
+            path.push(sibling);
+            k = k >> 1;
+            m = (m + 1) >> 1;
+        }
+
+        Ok(MerkleProof {
+            index,
+            path,
+            nleaves,
+            zero: self.zero,
+        })
+    }
+}
+
+/// Build the Merkle tree layers.
+fn merkle_tree_worker<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+>(
+    xs: &[HashOut<F, NUM_HASH_OUT_ELTS>],
+    zero: HashOut<F, NUM_HASH_OUT_ELTS>,
+    is_bottom_layer: bool,
+) -> Result<Vec<Vec<HashOut<F, NUM_HASH_OUT_ELTS>>>> {
+    let m = xs.len();
+    if !is_bottom_layer && m == 1 {
+        return Ok(vec![xs.to_vec()]);
+    }
+
+    let halfn = m / 2;
+    let n = 2 * halfn;
+    let is_odd = n != m;
+
+    let mut ys = Vec::with_capacity(halfn + if is_odd { 1 } else { 0 });
+
+    for i in 0..halfn {
+        let key = if is_bottom_layer { KEY_BOTTOM_LAYER } else { KEY_NONE };
+        let h = key_compress::<F, D, HF>(xs[2 * i], xs[2 * i + 1], key);
+        ys.push(h);
+    }
+
+    if is_odd {
+        let key = if is_bottom_layer {
+            KEY_ODD_AND_BOTTOM_LAYER
+        } else {
+            KEY_ODD
+        };
+        let h = key_compress::<F, D, HF>(xs[n], zero, key);
+        ys.push(h);
+    }
+
+    let mut layers = vec![xs.to_vec()];
+    let mut upper_layers = merkle_tree_worker::<F, D>(&ys, zero, false)?;
+    layers.append(&mut upper_layers);
+
+    Ok(layers)
+}
+
+/// Merkle proof struct, containing the index, path, and other necessary data.
+#[derive(Clone)]
+pub struct MerkleProof<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+> {
+    pub index: usize,       // Index of the leaf
+    pub path: Vec<HashOut<F, NUM_HASH_OUT_ELTS>>, // Sibling hashes from the leaf to the root
+    pub nleaves: usize,     // Total number of leaves
+    pub zero: HashOut<F, NUM_HASH_OUT_ELTS>,
+}
+
+impl<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+> MerkleProof<F, D> {
+    /// Reconstructs the root hash from the proof and the given leaf.
+    pub fn reconstruct_root(&self, leaf: HashOut<F, NUM_HASH_OUT_ELTS>) -> Result<HashOut<F, NUM_HASH_OUT_ELTS>> {
+        let mut m = self.nleaves;
+        let mut j = self.index;
+        let mut h = leaf;
+        let mut bottom_flag = KEY_BOTTOM_LAYER;
+
+        for p in &self.path {
+            let odd_index = (j & 1) != 0;
+            if odd_index {
+                // The index of the child is odd
+                h = key_compress::<F, D, HF>(*p, h, bottom_flag);
+            } else {
+                if j == m - 1 {
+                    // Single child -> so odd node
+                    h = key_compress::<F, D, HF>(h, *p, bottom_flag + 2);
+                } else {
+                    // Even node
+                    h = key_compress::<F, D, HF>(h, *p, bottom_flag);
+                }
+            }
+            bottom_flag = KEY_NONE;
+            j = j.shr(1);
+            m = (m + 1).shr(1);
+        }
+
+        Ok(h)
+    }
+
+    /// reconstruct the root using path_bits and last_bits in similar way as the circuit
+    /// this is used for testing - sanity check
+    pub fn reconstruct_root2(leaf: HashOut<F, NUM_HASH_OUT_ELTS>, path_bits: Vec<bool>, last_bits:Vec<bool>, path: Vec<HashOut<F, NUM_HASH_OUT_ELTS>>, mask_bits:Vec<bool>, depth: usize) -> Result<HashOut<F, NUM_HASH_OUT_ELTS>> {
+        let is_last = compute_is_last(path_bits.clone(),last_bits);
+
+        let mut h = vec![];
+        h.push(leaf);
+        let mut i = 0;
+
+        for p in &path {
+            let bottom = if(i==0){
+                KEY_BOTTOM_LAYER
+            }else{
+                KEY_NONE
+            };
+
+            let odd = (is_last[i] as usize) * (1-(path_bits[i] as usize));
+
+            let key = bottom + (2 * (odd as u64));
+            let odd_index = path_bits[i];
+            if odd_index {
+                h.push(key_compress::<F, D, HF>(*p, h[i], key));
+            } else {
+                h.push(key_compress::<F,D,HF>(h[i], *p, key));
+            }
+            i += 1;
+        }
+
+        let mut reconstructed_root = HashOut::<F, NUM_HASH_OUT_ELTS>::zero();
+        for k in 0..depth{
+            let diff = (mask_bits[k] as u64) - (mask_bits[k+1] as u64);
+            let mul_res: Vec<F> = h[k+1].elements.iter().map(|e| e.mul(F::from_canonical_u64(diff))).collect();
+            reconstructed_root = HashOut::<F, NUM_HASH_OUT_ELTS>::from_vec(
+                mul_res.iter().zip(reconstructed_root.elements).map(|(e1,e2)| e1.add(e2)).collect()
+            );
+        }
+
+        Ok(reconstructed_root)
+    }
+
+    /// Verifies the proof against a given root and leaf.
+    pub fn verify(&self, leaf: HashOut<F, NUM_HASH_OUT_ELTS>, root: HashOut<F, NUM_HASH_OUT_ELTS>) -> Result<bool> {
+        let reconstructed_root = self.reconstruct_root(leaf)?;
+        Ok(reconstructed_root == root)
+    }
+}
+
+///helper function to compute is_last
+fn compute_is_last(path_bits: Vec<bool>, last_bits: Vec<bool>) -> Vec<bool> {
+    let max_depth = path_bits.len();
+
+    // Initialize isLast vector
+    let mut is_last = vec![false; max_depth + 1];
+    is_last[max_depth] = true; // Set isLast[max_depth] to 1 (true)
+
+    // Iterate over eq and isLast in reverse order
+    for i in (0..max_depth).rev() {
+        let eq_out = path_bits[i] == last_bits[i]; // eq[i].out
+        is_last[i] = is_last[i + 1] && eq_out; // isLast[i] = isLast[i+1] * eq[i].out
+    }
+
+    is_last
+}

goldibear_experiments/proof-input/src/merkle_tree/mod.rs

@@ -0,0 +1,2 @@
+pub mod merkle_safe;
+pub mod key_compress;

goldibear_experiments/proof-input/src/params.rs

@@ -0,0 +1,224 @@
+// params for generating input for proof circuit
+
+use plonky2::hash::poseidon_goldilocks::Poseidon64Hash;
+use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+use std::env;
+use anyhow::{Result, Context};
+use codex_plonky2_circuits::circuits::params::{CircuitParams, NUM_HASH_OUT_ELTS};
+
+// test types
+pub const D: usize = 2;
+pub type C = PoseidonGoldilocksConfig;
+pub type F = <C as GenericConfig<D, NUM_HASH_OUT_ELTS>>::F; // this is the goldilocks field
+pub type HF = Poseidon64Hash;
+
+// hardcoded default params for generating proof input
+const DEFAULT_MAX_DEPTH: usize = 32; // depth of big tree (slot tree depth, includes block tree depth)
+const DEFAULT_MAX_SLOTS: usize = 256; // maximum number of slots
+const DEFAULT_CELL_SIZE: usize = 2048; // cell size in bytes
+const DEFAULT_BLOCK_SIZE: usize = 65536; // block size in bytes
+const DEFAULT_N_SAMPLES: usize = 5; // number of samples to prove
+
+const DEFAULT_ENTROPY: usize = 1234567; // external randomness
+const DEFAULT_SEED: usize = 12345; // seed for creating fake data TODO: not used now
+
+const DEFAULT_N_SLOTS: usize = 11; // number of slots in the dataset
+const DEFAULT_SLOT_INDEX: usize = 3; // the index of the slot to be sampled
+const DEFAULT_N_CELLS: usize = 512; // number of cells in each slot
+
+/// Params struct
+#[derive(Clone)]
+pub struct Params {
+    pub circuit_params: CircuitParams,
+    pub input_params: InputParams,
+}
+
+/// test params
+#[derive(Clone)]
+pub struct InputParams{
+    pub max_depth: usize,
+    pub max_slots: usize,
+    pub cell_size: usize,
+    pub block_size: usize,
+    pub n_samples: usize,
+    pub entropy: usize,
+    pub seed: usize,
+    pub n_slots: usize,
+    pub testing_slot_index: usize,
+    pub n_cells: usize,
+}
+
+/// Implement the Default trait for Params using the hardcoded constants
+impl Default for Params {
+    fn default() -> Self {
+        let input_params = InputParams {
+            max_depth: DEFAULT_MAX_DEPTH,
+            max_slots: DEFAULT_MAX_SLOTS,
+            cell_size: DEFAULT_CELL_SIZE,
+            block_size: DEFAULT_BLOCK_SIZE,
+            n_samples: DEFAULT_N_SAMPLES,
+            entropy: DEFAULT_ENTROPY,
+            seed: DEFAULT_SEED,
+            n_slots: DEFAULT_N_SLOTS,
+            testing_slot_index: DEFAULT_SLOT_INDEX,
+            n_cells: DEFAULT_N_CELLS,
+        };
+        let circuit_params = input_params.get_circuit_params();
+
+        Params{
+            circuit_params,
+            input_params,
+        }
+    }
+}
+
+/// Implement a new function to create Params with custom values
+impl InputParams {
+    // GOLDILOCKS_F_SIZE
+    pub fn goldilocks_f_size(&self) -> usize {
+        64
+    }
+
+    // N_FIELD_ELEMS_PER_CELL
+    pub fn n_field_elems_per_cell(&self) -> usize {
+        (self.cell_size + 62) / 62 * 8
+    }
+
+    // BOT_DEPTH
+    pub fn bot_depth(&self) -> usize {
+        log2(self.block_size / self.cell_size)
+    }
+
+    // N_CELLS_IN_BLOCKS
+    pub fn n_cells_in_blocks(&self) -> usize {
+        1 << self.bot_depth()
+    }
+
+    // N_BLOCKS
+    pub fn n_blocks(&self) -> usize {
+        1 << (self.max_depth - self.bot_depth())
+    }
+
+    // Depth of test input
+    pub fn depth_test(&self) -> usize {
+        self.n_cells.trailing_zeros() as usize
+    }
+
+    // N_BLOCKS for the test input
+    pub fn n_blocks_test(&self) -> usize {
+        1 << (self.depth_test() - self.bot_depth())
+    }
+
+    // DATASET_DEPTH
+    pub fn dataset_max_depth(&self) -> usize {
+        ceiling_log2(self.max_slots)
+    }
+
+    // DATASET_DEPTH for test
+    pub fn dataset_depth_test(&self) -> usize {
+        ceiling_log2(self.n_slots)
+    }
+
+    pub fn get_circuit_params(&self) -> CircuitParams{
+        CircuitParams{
+            max_depth: self.max_depth,
+            max_log2_n_slots: self.dataset_max_depth(),
+            block_tree_depth: self.bot_depth(),
+            n_field_elems_per_cell: self.n_field_elems_per_cell(),
+            n_samples:self.n_samples,
+        }
+    }
+}
+
+pub fn log2(x: usize) -> usize {
+    assert!(x.is_power_of_two(), "Input must be a power of 2.");
+    x.trailing_zeros() as usize
+}
+
+pub fn ceiling_log2(x: usize) -> usize {
+    if x <= 1 {
+        return 0;
+    }
+    usize::BITS as usize - x.saturating_sub(1).leading_zeros() as usize
+}
+
+/// load test params from env
+impl InputParams {
+    pub fn from_env() -> Result<Self> {
+        let max_depth = env::var("MAXDEPTH")
+            .context("MAXDEPTH not set")?
+            .parse::<usize>()
+            .context("Invalid MAXDEPTH")?;
+
+        let max_slots = env::var("MAXSLOTS")
+            .context("MAXSLOTS not set")?
+            .parse::<usize>()
+            .context("Invalid MAXSLOTS")?;
+
+        let cell_size = env::var("CELLSIZE")
+            .context("CELLSIZE not set")?
+            .parse::<usize>()
+            .context("Invalid CELLSIZE")?;
+
+        let block_size = env::var("BLOCKSIZE")
+            .context("BLOCKSIZE not set")?
+            .parse::<usize>()
+            .context("Invalid BLOCKSIZE")?;
+
+        let n_samples = env::var("NSAMPLES")
+            .context("NSAMPLES not set")?
+            .parse::<usize>()
+            .context("Invalid NSAMPLES")?;
+
+        let entropy = env::var("ENTROPY")
+            .context("ENTROPY not set")?
+            .parse::<usize>()
+            .context("Invalid ENTROPY")?;
+
+        let seed = env::var("SEED")
+            .context("SEED not set")?
+            .parse::<usize>()
+            .context("Invalid SEED")?;
+
+        let n_slots = env::var("NSLOTS")
+            .context("NSLOTS not set")?
+            .parse::<usize>()
+            .context("Invalid NSLOTS")?;
+
+        let testing_slot_index = env::var("SLOTINDEX")
+            .context("SLOTINDEX not set")?
+            .parse::<usize>()
+            .context("Invalid SLOTINDEX")?;
+
+        let n_cells = env::var("NCELLS")
+            .context("NCELLS not set")?
+            .parse::<usize>()
+            .context("Invalid NCELLS")?;
+
+        Ok(InputParams {
+            max_depth,
+            max_slots,
+            cell_size,
+            block_size,
+            n_samples,
+            entropy,
+            seed,
+            n_slots,
+            testing_slot_index,
+            n_cells,
+        })
+    }
+}
+
+/// load params from env
+impl Params {
+    pub fn from_env() -> Result<Self> {
+        let input_params = InputParams::from_env()?;
+        let circuit_params = input_params.get_circuit_params();
+
+        Ok(Params{
+            circuit_params,
+            input_params,
+        })
+    }
+}

goldibear_experiments/proof-input/src/sponge.rs

@@ -0,0 +1,292 @@
+use plonky2::hash::hash_types::{HashOut, RichField};
+use plonky2::plonk::config::Hasher;
+use plonky2::hash::hashing::PlonkyPermutation;
+use plonky2_field::types::{HasExtension};
+use codex_plonky2_circuits::circuits::params::NUM_HASH_OUT_ELTS;
+
+/// sponge function similar to the in-circuit one
+/// used here for testing / sanity check
+pub fn hash_n_with_padding<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: Hasher<F>
+>(
+    inputs: &[F],
+) -> HashOut<F, NUM_HASH_OUT_ELTS>{
+    HashOut::<F, NUM_HASH_OUT_ELTS>::from_vec(hash_n_to_m_with_padding::<F,D,H::Permutation>(inputs, NUM_HASH_OUT_ELTS))
+}
+
+pub fn hash_n_to_m_with_padding<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    P: PlonkyPermutation<F>
+>(
+    inputs: &[F],
+    num_outputs: usize,
+) -> Vec<F> {
+    let rate = P::RATE;
+    let width = P::WIDTH; // rate + capacity
+    let zero = F::zero();
+    let one = F::one();
+    let mut perm = P::new(core::iter::repeat(zero).take(width));
+
+    // Set the domain separator at index 8
+    let domsep_value = F::from_canonical_u64(rate as u64 + 256 * 12 + 65536 * 63);
+    perm.set_elt(domsep_value, 8);
+
+    let N = inputs.len();
+    let num_chunks = (N + rate) / rate; // Calculate number of chunks with 10* padding
+    let mut input_iter = inputs.iter();
+
+    // Process all chunks except the last one
+    for _ in 0..(num_chunks - 1) {
+        let mut chunk = Vec::with_capacity(rate);
+        for _ in 0..rate {
+            if let Some(&input) = input_iter.next() {
+                chunk.push(input);
+            } else {
+                // should not happen here
+                panic!("Insufficient input elements for chunk; expected more elements.");
+            }
+        }
+        // Add the chunk to the state
+        for j in 0..rate {
+            perm.set_elt(perm.as_ref()[j] + chunk[j],j);
+        }
+        // Apply permutation
+        perm.permute();
+    }
+
+    // Process the last chunk with 10* padding
+    let rem = num_chunks * rate - N; // Number of padding elements (0 < rem <= rate)
+    let ofs = rate - rem;            // Offset where padding starts
+
+    let mut last_chunk = Vec::with_capacity(rate);
+    // Absorb remaining inputs
+    for _ in 0..ofs {
+        if let Some(&input) = input_iter.next() {
+            last_chunk.push(input);
+        } else {
+            last_chunk.push(zero);
+        }
+    }
+    // Add the '1' padding bit
+    last_chunk.push(one);
+    // Pad with zeros to reach the full rate
+    while last_chunk.len() < rate {
+        last_chunk.push(zero);
+    }
+
+    // Add the last chunk to the state
+    for j in 0..rate {
+        perm.set_elt(perm.as_ref()[j] + last_chunk[j],j);
+    }
+    // Apply permutation
+    perm.permute();
+
+    // Squeeze outputs until we have the desired number
+    let mut outputs = Vec::with_capacity(num_outputs);
+    loop {
+        for &item in perm.squeeze() {
+            outputs.push(item);
+            if outputs.len() == num_outputs {
+                return outputs;
+            }
+        }
+        perm.permute();
+    }
+}
+
+/// sponge function for bytes with no padding
+/// expects the input to be divisible by rate
+/// note: rate is fixed at 8 for now
+/// used here for testing / sanity check
+pub fn hash_bytes_no_padding<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    H: Hasher<F>
+>(
+    inputs: &[F],
+) -> HashOut<F, NUM_HASH_OUT_ELTS>{
+    HashOut::<F, NUM_HASH_OUT_ELTS>::from_vec(hash_bytes_to_m_no_padding::<F, D, H::Permutation>(inputs, NUM_HASH_OUT_ELTS))
+}
+
+pub fn hash_bytes_to_m_no_padding<
+    F: RichField + HasExtension<D>,
+    const D: usize,
+    P: PlonkyPermutation<F>
+>(
+    inputs: &[F],
+    num_outputs: usize,
+) -> Vec<F> {
+    let rate = P::RATE;
+    let width = P::WIDTH; // rate + capacity
+    let zero = F::zero();
+    let one = F::one();
+    let mut perm = P::new(core::iter::repeat(zero).take(width));
+
+    // Set the domain separator at index 8
+    let domsep_value = F::from_canonical_u64(rate as u64 + 256 * 12 + 65536 * 8);
+    perm.set_elt(domsep_value, 8);
+
+    let n = inputs.len();
+    assert_eq!(n % rate, 0, "Input length ({}) must be divisible by rate ({})", n, rate);
+    let num_chunks = n / rate; // Calculate number of chunks
+    let mut input_iter = inputs.iter();
+
+    // Process all chunks
+    for _ in 0..num_chunks {
+        let mut chunk = Vec::with_capacity(rate);
+        for _ in 0..rate {
+            if let Some(&input) = input_iter.next() {
+                chunk.push(input);
+            } else {
+                // should not happen here
+                panic!("Insufficient input elements for chunk; expected more elements.");
+            }
+        }
+        // Add the chunk to the state
+        for j in 0..rate {
+            perm.set_elt(perm.as_ref()[j] + chunk[j],j);
+        }
+        // Apply permutation
+        perm.permute();
+    }
+
+    // Squeeze outputs until we have the desired number
+    let mut outputs = Vec::with_capacity(num_outputs);
+    loop {
+        for &item in perm.squeeze() {
+            outputs.push(item);
+            if outputs.len() == num_outputs {
+                return outputs;
+            }
+        }
+        perm.permute();
+    }
+}
+
+// #[cfg(test)]
+// mod tests {
+//     use crate::sponge::hash_n_with_padding;
+//     use crate::params::{D, F, HF};
+//
+//
+//     #[test]
+//     fn test_sponge_hash_rate_8() {
+//
+//         struct TestCase {
+//             n: usize,
+//             digest: [u64; 4],
+//         }
+//
+//         let test_cases: Vec<TestCase> = vec![
+//             TestCase { n: 0, digest: [0x509f3a747e4a6fca, 0xd6f21d91afb92eb3, 0xf65ef4075dcfb169, 0xbceaf22e0cd21b3d] },
+//             TestCase { n: 1, digest: [0xfa286adad207c7ea, 0x97d864ff2e89415e, 0xcf002b28585bd945, 0x95ec163fbdd0792e] },
+//             TestCase { n: 2, digest: [0xe4b779622cbb574f, 0x1fe4b1bc9a0c9fc7, 0x40051ada5252de9b, 0xb351345b1894a59f] },
+//             TestCase { n: 3, digest: [0x133a5a2fd0cae006, 0x072a7769ca9a550d, 0x92134dad95d394c6, 0x22234de7d7270aab] },
+//             TestCase { n: 4, digest: [0x78269e830f2a824a, 0x76f8b00469a8fa81, 0x6793369b1d75ebf5, 0xfba1a89dc21d9b30] },
+//             TestCase { n: 5, digest: [0x263994efd2cd5c57, 0x7c37a93fd48fc98b, 0xa081b26a68767d13, 0x16af92d6e1e4d7f8] },
+//             TestCase { n: 6, digest: [0x0b0b0f1d64f8d58c, 0x2946089b2eb949fc, 0xf68bcf08b69a95e7, 0x814d6eb4b2df848c] },
+//             TestCase { n: 7, digest: [0xae0c900a194ee051, 0x4555257fba7a500b, 0x1713fd448cc82c3a, 0xaf8f2e895e2136f3] },
+//             TestCase { n: 8, digest: [0x100351f04fc470b7, 0x79d3c3c416087158, 0x113bb1c70a6e84ee, 0x3eab2507cdc254d3] },
+//             TestCase { n: 9, digest: [0xbab284d7f11855d6, 0xe1b53d108f308a1c, 0x971fea7184337830, 0x6d674ae321cfb9ba] },
+//             TestCase { n: 10, digest: [0x68c00dbe0ed03a8f, 0xab5ba3617eb6f76b, 0x5d735bb89418cc0b, 0xff4101076f3f3c70] },
+//             TestCase { n: 11, digest: [0xaecce2fa7de4f97d, 0x07cee3dc720812e0, 0x4155bf667391a9e8, 0xbf8a49a12f40e746] },
+//             TestCase { n: 12, digest: [0xd3f43f06fc7affd2, 0xee9a8ac5ef44071a, 0xe00ec9e7f468d0e2, 0x944e34913a974233] },
+//             TestCase { n: 13, digest: [0xcd50fe6ab5e3de54, 0x9b2093adaeac949c, 0xa176a2a9e2c82787, 0xd35f0635a1ec333f] },
+//             TestCase { n: 14, digest: [0x8f5188d26ca0368c, 0x0116bf587e5cc970, 0x30654ee52a3c66d8, 0xe8ded60382c44b04] },
+//             TestCase { n: 15, digest: [0xc7f020f910327951, 0x13a468945463870d, 0xbcf8ca584edb30f3, 0x7e7234f0b8954e7e] },
+//             TestCase { n: 16, digest: [0xf8a9aef7392048e7, 0x6124715a2c5343eb, 0x1b7f17ebec4a5b13, 0xdf61d868051dad75] },
+//             TestCase { n: 17, digest: [0x44d1fb6822c7f3fa, 0x2623cc2240022e42, 0xc90ce9259c9e1160, 0x7a42bc611acacc12] },
+//             TestCase { n: 18, digest: [0x85dab5b06ef2d176, 0x24a587b13a4e3b30, 0xf547a00373299873, 0xb298a6ef846d64a1] },
+//             TestCase { n: 19, digest: [0x7cc060a3f2a74260, 0xa07dc76e73335eb0, 0xf8ed9acbcf8a242e, 0xd32eaf3150005e49] },
+//             TestCase { n: 20, digest: [0x3e961c84e53106f9, 0x63d9a807f9cfd88c, 0x7031e8834a17821a, 0xf2e1c79698798fa9] },
+//             TestCase { n: 21, digest: [0x8a0ab00081c9828f, 0xa5f7aadaf3af046e, 0xada8b4c6220b3420, 0x80ebc8c91a65518c] },
+//             TestCase { n: 22, digest: [0x39505fc00f052122, 0xb13edc24a35665c7, 0xa7b164fffe37ec64, 0x8f7eeb42c068e19f] },
+//             TestCase { n: 23, digest: [0x1f49d6f25f39522b, 0x879377d8df727784, 0x00f1461600d09cdd, 0xd2c7946a44e1aa66] },
+//             TestCase { n: 24, digest: [0x1c6f7a68537f7dc7, 0x64e6e09714dc0854, 0x9abfed111e51bd96, 0x65061b2bc484ed8b] },
+//             TestCase { n: 25, digest: [0x95fd5cc6bc02ab29, 0xe2e3c96d9b1b8b5d, 0xadcf491caa16549e, 0x97d91e370da3c0b4] },
+//             TestCase { n: 26, digest: [0x7599c5052ba67767, 0x3fe4a05f44e96ed6, 0xbbfe6874aa53808c, 0xd6771e162cc9f0ff] },
+//             TestCase { n: 27, digest: [0xdff28121d822093c, 0x7313ea03b57bb436, 0x10ed29b28a77d8c3, 0x6ee304be541fe36f] },
+//             TestCase { n: 28, digest: [0xce2b7f232b504b48, 0x02c638c398c12cb0, 0x4f1d416215377a86, 0x2d43ff6c5dd88f8c] },
+//             TestCase { n: 29, digest: [0xa60cb008de647e9a, 0x502e2e740f68e2d1, 0xe983eb54e4052013, 0xe76e59c5e5dbcca2] },
+//             TestCase { n: 30, digest: [0x7735e3ac5e08fa00, 0x211a86449207c30d, 0x9d80ddd40e7760b2, 0xe60f32f28597a188] },
+//             TestCase { n: 31, digest: [0x6fab3f12496f0691, 0x5116ad81bedd7d84, 0xaa8a7713a80b323b, 0xce6d94533fc40b88] },
+//             TestCase { n: 32, digest: [0xce51cdbd641d57c0, 0xf638202a88ee7f9c, 0x26c291ecc5162b45, 0x04a0a62b949c236f] },
+//             TestCase { n: 33, digest: [0x923391e4a4cde9e2, 0xdcb3acccba80597d, 0x247bb4b67044a0e1, 0x65bbac92e096d1ec] },
+//             TestCase { n: 34, digest: [0x1550d0234ae35f05, 0x16f4d1708923d4f1, 0x232319cb4090ea4e, 0x8354e1aed093070c] },
+//             TestCase { n: 35, digest: [0xc7dd24e6db4ea70f, 0x80bc3d2aac952cb1, 0xabbd1a878bc50565, 0xf1ebc3b8d513c591] },
+//             TestCase { n: 36, digest: [0xba9c4b1ce906efb1, 0xa332d0daccc62979, 0xfb658fcad0b5fbbd, 0x62d21407f34a35ee] },
+//             TestCase { n: 37, digest: [0xcb2973d44f2b589d, 0x01708b32c4556317, 0x3ad51597c12b8564, 0x28d3a5d7de72cfd5] },
+//             TestCase { n: 38, digest: [0x1dcf1f4ab7338296, 0xb88c661141b5aabb, 0x7e546b6e9b31bc90, 0xf26f7e6ffabb4e69] },
+//             TestCase { n: 39, digest: [0x2e139ff910c0f410, 0xba3d2c0a92ec3845, 0x2860e475933a7108, 0x8f2a6c6d13bedc7a] },
+//             TestCase { n: 40, digest: [0xc18a53c17c360ef4, 0x5e56ea9228988c68, 0xee0bd138436e996d, 0x06afd46a753f8257] },
+//             TestCase { n: 41, digest: [0x2c992403c5277dc5, 0xba8770bc3a54b043, 0x51b882882a7b7864, 0xf75e179a53e7948e] },
+//             TestCase { n: 42, digest: [0xde855183965741c3, 0x93520eac77a8f98d, 0x6412ae8cf0522d78, 0x9db49c6b455a83b4] },
+//             TestCase { n: 43, digest: [0x552e357ddb7e1ef6, 0x5fa779e9c7373b56, 0x18f7c445e27e5dcf, 0x2664ecee5e7bc6c2] },
+//             TestCase { n: 44, digest: [0x37b8a716c87e5489, 0x1201fcd31e407152, 0x0979d7887c42e1ca, 0x902e8b2bf748b356] },
+//             TestCase { n: 45, digest: [0xa48bdd1d464960ed, 0x8e92c1af0cf258bc, 0x7c5b447524b92ba9, 0xac63902e613e4ef0] },
+//             TestCase { n: 46, digest: [0x542e62f9317fe11d, 0xc23ba113a3f3c810, 0x2bda30c42a89cc7e, 0x35616e9f1a00aa8f] },
+//             TestCase { n: 47, digest: [0x1c9194a0acfa97d7, 0x60d536ac106dd774, 0x8855b4a40e110080, 0xc2c408114e8c20d6] },
+//             TestCase { n: 48, digest: [0x0e90b1cc3ac49e0c, 0x1b73aa8e0decbf88, 0x0ca9ef7070e0513f, 0x25cfb975571b6139] },
+//             TestCase { n: 49, digest: [0xba6d6f7aa664f2e7, 0x4b9af896093937b9, 0x115b9aeb6c5f563e, 0x41cb5f42c6d3b115] },
+//             TestCase { n: 50, digest: [0xdc3bdc491154caf6, 0xb95159bae61b2035, 0x98bd384fb3d0100b, 0xd70226f2b71ea465] },
+//             TestCase { n: 51, digest: [0x57f31da51bcd2eab, 0x4a3b3945a8662b5c, 0x44dffaa325530b19, 0x47f4e41c2c1474cf] },
+//             TestCase { n: 52, digest: [0xc3f518f6cf3b43bf, 0x1214790ff48554e4, 0x99c1eabc61b218fd, 0xf90b03954d7937f8] },
+//             TestCase { n: 53, digest: [0x6357b3cdcbc1283a, 0x6acc0c2d5aac9261, 0xdf11e7ad14d432d1, 0x2242b26bdcc8a380] },
+//             TestCase { n: 54, digest: [0x1946dc4471f8c502, 0x6be7d72499e0b4a5, 0x6e11de349239ff90, 0xfca78044256b8b54] },
+//             TestCase { n: 55, digest: [0x302b38fb3df623dd, 0x69b362f7932fd7af, 0x2b47156f9135508b, 0xfe6c574f0a102e92] },
+//             TestCase { n: 56, digest: [0xfdc9bd08a0416122, 0x063ebf4767fc7914, 0x330f36279d94050e, 0x79c61f80746893ec] },
+//             TestCase { n: 57, digest: [0x7b5d8384b67af5c0, 0xa705e0163fa4d839, 0x1e203432e872104e, 0xe0e7699f20a291f4] },
+//             TestCase { n: 58, digest: [0xb0aa74a52fe04366, 0x194b0d4afcdc03d9, 0x5134dc604b5d9f2a, 0x53c6bf9d5a1d502b] },
+//             TestCase { n: 59, digest: [0xd5c8258f6fc80e2b, 0x82bac373eb051b48, 0x5ef620241420462d, 0x58635db0134fb97a] },
+//             TestCase { n: 60, digest: [0x42ebb974ac5dd0ef, 0x676d0c6b3dde78c3, 0x14ed5eda2c9cb9de, 0x0f78a26badaa447c] },
+//             TestCase { n: 61, digest: [0x2b3ca7711db999d5, 0xb74bd29abcb6179a, 0x8ba196525e6adb25, 0x86cb9464ae269a43] },
+//             TestCase { n: 62, digest: [0x3d0e61a2ca7a65a2, 0x31f77852d41a6c8d, 0x2e4ceaa39763a53d, 0x5232ff5a3d78755e] },
+//             TestCase { n: 63, digest: [0xb2ed789e88c1f525, 0x1592c1a1eafd2a9b, 0x98700c512f8c9a5d, 0xf96837b5d99a4eb4] },
+//             TestCase { n: 64, digest: [0xe4b7d14e11de2fa9, 0xe81afff2cee68e14, 0xc58abb080bf37dd3, 0x36ae8b2196b5ae88] },
+//             TestCase { n: 65, digest: [0xa1df9ff199c41d63, 0xd02c067d3d12edc1, 0xc9b598130fa60794, 0x5afe82d34c3fc8fa] },
+//             TestCase { n: 66, digest: [0x0bc0094a1f07256d, 0x33c5b4c2a171d5bd, 0x1f38f1b1dc92aa54, 0x4610d21f276faa11] },
+//             TestCase { n: 67, digest: [0x8072f00df8f7e44f, 0x42f0c2b8fe81d8a0, 0x2b5caf9e7c0ff611, 0x92b0b3a4a4bebe1a] },
+//             TestCase { n: 68, digest: [0x6539f06fab064b57, 0xdb298b91f6c4f44f, 0x5d8f8eec6b7e8c86, 0x848a447123f39006] },
+//             TestCase { n: 69, digest: [0x87f32efc9eaa65f6, 0xc5699d4ab6443852, 0x61008286bc651f4a, 0xcbcf714354843da3] },
+//             TestCase { n: 70, digest: [0xffb8ad2258107315, 0xf7d6a58eb54f2745, 0xaecf888211821114, 0x7e0ea33b4d56976e] },
+//             TestCase { n: 71, digest: [0xa9e5b6d70f67db2b, 0x072fd05840040322, 0x40ffcc86e3909dec, 0x3d80f61616a9e6d7] },
+//             TestCase { n: 72, digest: [0xa77dd95d9ff4d7b8, 0x3a0e0502f74c091a, 0x1fa83de1e7dc716d, 0xe01ae447cc3a0e40] },
+//             TestCase { n: 73, digest: [0xc4a29dc875a308eb, 0xd2ed0da7aab24b0c, 0x4c2aaaed0bc4f059, 0xaea772c635ea901a] },
+//             TestCase { n: 74, digest: [0xaad59bf06c151ecf, 0x5e0f45e55df36692, 0x4798afb8b944a01e, 0xd7152cd819bbd7f8] },
+//             TestCase { n: 75, digest: [0x89ae5b2b35ba07c7, 0x129f4ff59afaa1a3, 0x4275f3f797112650, 0xea3b4baaf7190a19] },
+//             TestCase { n: 76, digest: [0xab068e43be297604, 0x17bd1c3cf4afec96, 0xaa84a8098dba4516, 0xa6e487ceafb02c49] },
+//             TestCase { n: 77, digest: [0x2c85080ef895bb4a, 0xbd280690a789c124, 0xca4f8423b50de8a5, 0xec809bb8c30de95b] },
+//             TestCase { n: 78, digest: [0x51c3d13543e4922b, 0xff9c11d5b93268db, 0xd9cf911cc5326948, 0x4b7bb11eafe7fd44] },
+//             TestCase { n: 79, digest: [0xb435274d75678586, 0x8600e7f2db687493, 0x282873a3600a38da, 0x727791507d1b600e] },
+//             TestCase { n: 80, digest: [0x23ae45602324f628, 0x0dc16b33f43209c5, 0x2455376f83b1aeff, 0xd5470f22ec2113bc] },
+//         ];
+//
+//         for test_case in test_cases {
+//             let n = test_case.n;
+//             let expected_digest = test_case.digest;
+//
+//             // Generate inputs
+//             let inputs: Vec<F> = (0..n)
+//                 .map(|i| F::from_canonical_u64(i as u64 + 1))
+//                 .collect();
+//
+//             // Call the sponge function
+//             let output = hash_n_with_padding::<F,D,HF>(&inputs);
+//
+//             // Compare the outputs
+//             for (i, &out_elem) in output.elements.iter().enumerate() {
+//                 let expected_elem = F::from_canonical_u64(expected_digest[i]);
+//                 assert_eq!(
+//                     out_elem,
+//                     expected_elem,
+//                     "Mismatch at test case n={}, output element {}",
+//                     n,
+//                     i
+//                 );
+//             }
+//         }
+//     }
+// }

goldibear_experiments/proof-input/src/utils.rs

@@ -0,0 +1,103 @@
+use plonky2::hash::hash_types::{HashOut, RichField};
+use crate::params::HF;
+use plonky2::hash::hashing::PlonkyPermutation;
+use plonky2_field::types::HasExtension;
+use crate::sponge::hash_n_with_padding;
+use codex_plonky2_circuits::circuits::params::NUM_HASH_OUT_ELTS;
+
+// --------- helper functions ---------
+
+/// Converts an index to a vector of bits (LSB first) no padding.
+pub fn usize_to_bits_le(index: usize, bit_length: usize) -> Vec<bool> {
+    // Assert that the index can fit within the given bit length.
+    assert!(
+        index < (1 << bit_length),
+        "Index ({}) does not fit in {} bits",
+        index,
+        bit_length
+    );
+
+    let mut bits = Vec::with_capacity(bit_length);
+    for i in 0..bit_length {
+        bits.push(((index >> i) & 1) == 1);
+    }
+
+    // No padding
+    bits
+}
+
+/// returns the first bit_length bits of index
+pub fn low_bits(index: usize, bit_length: usize) -> Vec<bool> {
+
+    let mut bits = Vec::with_capacity(bit_length);
+
+    for i in 0..bit_length {
+        // get the i-th bit and push its bool value
+        bits.push(((index >> i) & 1) == 1);
+    }
+
+    bits
+}
+
+/// calculate the sampled cell index from entropy, slot root, and counter
+/// this is the non-circuit version for testing
+pub fn calculate_cell_index_bits<
+    F: RichField + HasExtension<D>,
+    const D: usize
+>(entropy: &Vec<F>, slot_root: HashOut<F, NUM_HASH_OUT_ELTS>, ctr: usize, depth: usize, mask_bits: Vec<bool>) -> Vec<bool> {
+    let ctr_field = F::from_canonical_u64(ctr as u64);
+    let mut ctr_as_digest = HashOut::<F, NUM_HASH_OUT_ELTS>::zero();
+    ctr_as_digest.elements[0] = ctr_field;
+    let mut hash_inputs = Vec::new();
+    hash_inputs.extend_from_slice(&entropy);
+    hash_inputs.extend_from_slice(&slot_root.elements);
+    hash_inputs.extend_from_slice(&ctr_as_digest.elements);
+    let hash_output = hash_n_with_padding::<F,D,HF>(&hash_inputs);
+    let cell_index_bytes = hash_output.elements[0].as_canonical_u64();
+
+    let cell_index_bits = low_bits(cell_index_bytes as usize, depth);
+
+    let mut masked_cell_index_bits = vec![];
+
+    for i in 0..depth{
+        masked_cell_index_bits.push(cell_index_bits[i] && mask_bits[i]);
+    }
+
+    masked_cell_index_bits
+}
+
+/// Converts a vector of bits (LSB first) into an index (usize).
+pub fn bits_le_padded_to_usize(bits: &[bool]) -> usize {
+    bits.iter().enumerate().fold(0usize, |acc, (i, &bit)| {
+        if bit {
+            acc | (1 << i)
+        } else {
+            acc
+        }
+    })
+}
+
+/// computes the `last_index` (the binary decomposition of `inp-1`) and the `mask_bits`
+pub fn ceiling_log2(
+    inp: usize,
+    n: usize,
+) -> (Vec<bool>, Vec<bool>) {
+    // Handle the case when inp is 0
+    let last_index = if inp == 0 { panic!("input to ceiling_log2 is 0") } else { inp - 1 };
+    let last_bits = usize_to_bits_le(last_index, n);
+
+    // Initialize aux, all false
+    let mut aux = vec![false; n+1];
+    aux[n] = true; // aux[n] = 1
+
+    // Initialize mask vector
+    let mut mask = vec![false; n+1];
+
+    // Compute aux and mask bits
+    for i in (0..n).rev() {
+        aux[i] = aux[i + 1] && !last_bits[i];
+        mask[i] = !aux[i];
+    }
+
+    (last_bits, mask)
+}