add proving cells circuit and clean up

codex-plonky2-circuits/Cargo.toml

@@ -21,3 +21,11 @@ rand = "0.8.5"
 [dev-dependencies]
 criterion = { version = "0.5.1", default-features = false }
 tynm = { version = "0.1.6", default-features = false }
+
+[[bench]]
+name = "safe_circuit"
+harness = false
+
+[[bench]]
+name = "prove_cells"
+harness = false

codex-plonky2-circuits/benches/prove_cells.rs

@@ -0,0 +1,156 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use anyhow::Result;
+use std::time::{Duration, Instant};
+
+use codex_plonky2_circuits::{
+    merkle_tree::merkle_safe::MerkleProof,
+    circuits::safe_tree_circuit::MerkleTreeCircuit,
+};
+use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData};
+use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, Hasher, PoseidonGoldilocksConfig};
+use plonky2::iop::witness::PartialWitness;
+use plonky2::hash::poseidon::PoseidonHash;
+use plonky2::field::extension::Extendable;
+use plonky2::hash::hash_types::RichField;
+use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
+use std::marker::PhantomData;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use codex_plonky2_circuits::circuits::prove_single_cell::SlotTree;
+
+macro_rules! pretty_print {
+    ($($arg:tt)*) => {
+        print!("\x1b[0;36mINFO ===========>\x1b[0m ");
+        println!($($arg)*);
+    }
+}
+
+// Hash function used
+type HF = PoseidonHash;
+
+fn prepare_data<F, H>(N: usize) -> Result<(
+    SlotTree<F, H>,
+    Vec<usize>,
+    Vec<MerkleProof<F, H>>,
+)>
+where
+    F: RichField + Extendable<2> + Poseidon2,
+    H: Hasher<F> + AlgebraicHasher<F> + Hasher<F>,
+{
+    // Initialize the slot tree with default data
+    let slot_tree = SlotTree::<F, H>::default();
+
+    // Select N leaf indices to prove
+    let leaf_indices: Vec<usize> = (0..N).collect();
+
+    // Get the Merkle proofs for the selected leaves
+    let proofs: Vec<_> = leaf_indices
+        .iter()
+        .map(|&leaf_index| slot_tree.get_proof(leaf_index))
+        .collect();
+
+    Ok((slot_tree, leaf_indices, proofs))
+}
+
+fn build_circuit<F, C, const D: usize, H>(
+    slot_tree: &SlotTree<F, H>,
+    leaf_indices: &[usize],
+    proofs: &[MerkleProof<F, H>],
+) -> Result<(CircuitData<F, C, D>, PartialWitness<F>)>
+where
+    F: RichField + Extendable<D> + Poseidon2,
+    C: GenericConfig<D, F = F>,
+    H: Hasher<F> + AlgebraicHasher<F> + Hasher<F>,
+{
+    // Create the circuit
+    let config = CircuitConfig::standard_recursion_config();
+    let mut builder = CircuitBuilder::<F, D>::new(config);
+
+    // Create a PartialWitness
+    let mut pw = PartialWitness::new();
+
+    // Initialize the circuit instance
+    let mut circuit_instance = MerkleTreeCircuit::<F, C, D, H> {
+        tree: slot_tree.tree.clone(),
+        _phantom: PhantomData,
+    };
+
+    // For each proof, create targets, add constraints, and assign witnesses
+    for (i, &leaf_index) in leaf_indices.iter().enumerate() {
+        // Build the circuit for each proof
+        let mut targets = circuit_instance.prove_single_cell2(&mut builder);
+
+        // Assign witnesses for each proof
+        circuit_instance.single_cell_assign_witness(
+            &mut pw,
+            &mut targets,
+            leaf_index,
+            &slot_tree.cell_data[leaf_index],
+            proofs[i].clone(),
+        )?;
+    }
+
+    // Build the circuit
+    let data = builder.build::<C>();
+
+    Ok((data, pw))
+}
+
+fn single_cell_proof_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Single Cell Proof Benchmark");
+
+    // Circuit parameters
+    const D: usize = 2;
+    type C = PoseidonGoldilocksConfig;
+    type F = <C as GenericConfig<D>>::F;
+    type H = PoseidonHash;
+
+    // Prepare the data that will be used in all steps
+    let N = 5; // Number of leaves to prove
+    let (slot_tree, leaf_indices, proofs) = prepare_data::<F, H>(N).unwrap();
+
+    // Benchmark the circuit building
+    group.bench_function("Single Cell Proof Build", |b| {
+        b.iter(|| {
+            build_circuit::<F, C, D, H>(&slot_tree, &leaf_indices, &proofs).unwrap();
+        })
+    });
+
+    // Build the circuit
+    let (data, pw) = build_circuit::<F, C, D, H>(&slot_tree, &leaf_indices, &proofs).unwrap();
+
+    pretty_print!(
+        "Circuit size: 2^{} gates",
+        data.common.degree_bits()
+    );
+
+    let start_time = Instant::now();
+    let proof_with_pis = data.prove(pw.clone()).unwrap();
+    println!("prove_time = {:?}", start_time.elapsed());
+
+    // Benchmark the proving time
+    group.bench_function("Single Cell Proof Prove", |b| {
+        b.iter(|| {
+            let _proof_with_pis = data.prove(pw.clone()).unwrap();
+        })
+    });
+
+    // Generate the proof
+    let proof_with_pis = data.prove(pw.clone()).unwrap();
+    let verifier_data = data.verifier_data();
+
+    pretty_print!("Proof size: {} bytes", proof_with_pis.to_bytes().len());
+
+    // Benchmark the verification time
+    group.bench_function("Single Cell Proof Verify", |b| {
+        b.iter(|| {
+            verifier_data.verify(proof_with_pis.clone()).unwrap();
+        })
+    });
+
+    group.finish();
+}
+
+criterion_group!(name = benches;
+    config = Criterion::default().sample_size(10);
+    targets = single_cell_proof_benchmark);
+criterion_main!(benches);

codex-plonky2-circuits/benches/safe_circuit.rs

@@ -0,0 +1,164 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use anyhow::Result;
+
+use codex_plonky2_circuits::{merkle_tree::merkle_safe::MerkleTree, circuits::safe_tree_circuit::MerkleTreeCircuit};
+use plonky2::field::types::Field;
+use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData};
+use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, Hasher, PoseidonGoldilocksConfig};
+use plonky2::iop::witness::PartialWitness;
+use plonky2::hash::hash_types::HashOut;
+use plonky2::hash::poseidon::PoseidonHash;
+use plonky2::field::extension::Extendable;
+use plonky2::hash::hash_types::RichField;
+use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
+use std::marker::PhantomData;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use codex_plonky2_circuits::merkle_tree::merkle_safe::MerkleProof;
+
+macro_rules! pretty_print {
+    ($($arg:tt)*) => {
+        print!("\x1b[0;36mINFO ===========>\x1b[0m ");
+        println!($($arg)*);
+    }
+}
+
+fn prepare_data<F, H>(N: usize) -> Result<(
+    MerkleTree<F, H>,
+    Vec<HashOut<F>>,
+    Vec<usize>,
+    Vec<MerkleProof<F, H>>,
+    HashOut<F>,
+)>
+    where
+        F: RichField + Extendable<2> + Poseidon2,
+        H: Hasher<F> + AlgebraicHasher<F> + Hasher<F>,
+{
+    // Total number of leaves in the Merkle tree
+    let nleaves = 1u64 << 16;
+
+    // Generate leaf data
+    let data = (0..nleaves)
+        .map(|i| F::from_canonical_u64(i as u64))
+        .collect::<Vec<_>>();
+
+    // Hash the data to obtain leaf hashes
+    let leaves: Vec<HashOut<F>> = data
+        .iter()
+        .map(|&element| {
+            PoseidonHash::hash_no_pad(&[element])
+        })
+        .collect();
+
+    let zero_hash = HashOut {
+        elements: [F::ZERO; 4],
+    };
+    let tree = MerkleTree::<F, H>::new(&leaves, zero_hash)?;
+
+    // Select N leaf indices to prove
+    let leaf_indices: Vec<usize> = (0..N).collect();
+
+    // Get the Merkle proofs for the selected leaves
+    let proofs: Vec<_> = leaf_indices
+        .iter()
+        .map(|&leaf_index| tree.get_proof(leaf_index))
+        .collect::<Result<Vec<_>, _>>()?;
+
+    // Expected Merkle root
+    let expected_root = tree.root()?;
+
+    Ok((tree, leaves, leaf_indices, proofs, expected_root))
+}
+
+fn build_circuit<F, C, const D: usize, H>(
+    tree: &MerkleTree<F, H>,
+    leaf_indices: &[usize],
+) -> Result<(CircuitData<F, C, D>, PartialWitness<F>)>
+    where
+        F: RichField + Extendable<D> + Poseidon2,
+        C: GenericConfig<D, F = F>,
+        H: Hasher<F> + AlgebraicHasher<F> + Hasher<F>,
+{
+    // Create the circuit
+    let config = CircuitConfig::standard_recursion_config();
+    let mut builder = CircuitBuilder::<F, D>::new(config);
+
+    // Create a PartialWitness
+    let mut pw = PartialWitness::new();
+
+    // Initialize the circuit instance
+    let mut circuit_instance = MerkleTreeCircuit::<F, C, D, H> {
+        tree: tree.clone(),
+        _phantom: PhantomData,
+    };
+
+    // For each proof, create targets, add constraints, and assign witnesses
+    for &leaf_index in leaf_indices.iter() {
+        // Build the circuit for each proof
+        let mut targets = circuit_instance.build_circuit(&mut builder);
+
+        // Assign witnesses for each proof
+        circuit_instance.assign_witness(&mut pw, &mut targets, leaf_index)?;
+    }
+
+    // Build the circuit
+    let data = builder.build::<C>();
+
+    Ok((data, pw))
+}
+
+fn merkle_proof_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Merkle Proof Benchmark");
+
+    // Circuit parameters
+    const D: usize = 2;
+    type C = PoseidonGoldilocksConfig;
+    type F = <C as GenericConfig<D>>::F;
+    type H = PoseidonHash;
+
+    // Prepare the data that will be used in all steps
+    let N = 5; // Number of leaves to prove
+    let (tree, _leaves, leaf_indices, _proofs, _expected_root) = prepare_data::<F, H>(N).unwrap();
+
+    // Benchmark the circuit building
+    group.bench_function("Merkle Proof Build", |b| {
+        b.iter(|| {
+            build_circuit::<F, C, D, H>(&tree, &leaf_indices).unwrap();
+        })
+    });
+
+    // Build the circuit once to get the data for the proving and verifying steps
+    let (data, pw) = build_circuit::<F, C, D, H>(&tree, &leaf_indices).unwrap();
+
+    pretty_print!(
+        "circuit size: 2^{} gates",
+        data.common.degree_bits()
+    );
+
+    // Benchmark the proving time
+    group.bench_function("Merkle Proof Prove", |b| {
+        b.iter(|| {
+            let _proof_with_pis = data.prove(pw.clone()).unwrap();
+        })
+    });
+
+    // Generate the proof once for verification
+    let proof_with_pis = data.prove(pw.clone()).unwrap();
+    let verifier_data = data.verifier_data();
+
+    pretty_print!("proof size: {}", proof_with_pis.to_bytes().len());
+
+    // Benchmark the verification time
+    group.bench_function("Merkle Proof Verify", |b| {
+        b.iter(|| {
+            verifier_data.verify(proof_with_pis.clone()).unwrap();
+        })
+    });
+
+    group.finish();
+}
+
+// criterion_group!(benches, merkle_proof_benchmark);
+criterion_group!(name = benches;
+    config = Criterion::default().sample_size(10);
+    targets = merkle_proof_benchmark);
+criterion_main!(benches);

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

@@ -1,7 +1,10 @@
+// circuit for regular merkle tree implementation (non-safe version)
+// the circuit uses caps in similar way as in Plonky2 Merkle tree implementation
+// NOTE: this might be deleted at later time, since we don't use it for codex
+
 use anyhow::Result;
 use plonky2::field::extension::Extendable;
 use plonky2::hash::hash_types::RichField;
-use plonky2::hash::hashing::hash_n_to_m_no_pad;
 use plonky2::iop::target::{BoolTarget, Target};
 use plonky2::iop::witness::{PartialWitness, WitnessWrite, Witness};
 use plonky2::plonk::circuit_builder::CircuitBuilder;
@@ -15,16 +18,14 @@ use crate::merkle_tree::capped_tree::MerkleTree;
 use plonky2::hash::poseidon::PoseidonHash;
 
 use plonky2::hash::hash_types::{HashOutTarget, MerkleCapTarget, NUM_HASH_OUT_ELTS};
-use crate::merkle_tree::capped_tree::{MerkleProof, MerkleProofTarget};
+use crate::merkle_tree::capped_tree::MerkleProofTarget;
-use plonky2_poseidon2::poseidon2_hash::poseidon2::{Poseidon2, Poseidon2Hash};
+use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
 
-use plonky2::field::goldilocks_field::GoldilocksField;
 use plonky2::plonk::config::PoseidonGoldilocksConfig;
 use plonky2::plonk::proof::Proof;
 
 use plonky2::hash::hashing::PlonkyPermutation;
 use plonky2::plonk::circuit_data::VerifierCircuitTarget;
-use crate::merkle_tree::capped_tree::MerkleCap;
 
 // size of leaf data (in number of field elements)
 pub const LEAF_LEN: usize = 4;
@@ -66,7 +67,6 @@ impl<
     }
 
     // build the circuit and returns the circuit data
-    // note, this fn generate circuit data with
     pub fn build_circuit(&mut self, builder: &mut CircuitBuilder::<F, D>) -> MerkleTreeTargets<F, C, D, H>{
 
         let proof_t = MerkleProofTarget {
@@ -89,8 +89,6 @@ impl<
         );
 
         MerkleTreeTargets{
-            // depth: 0,
-            // cap_height: 0,
             proof_target: proof_t,
             cap_target: cap_t,
             leaf: leaf_t.to_vec(),
@@ -360,7 +358,7 @@ pub mod tests {
     use super::*;
     use plonky2::field::types::Field;
     use crate::merkle_tree::capped_tree::MerkleTree;
-    use plonky2::iop::witness::{PartialWitness, WitnessWrite};
+    use plonky2::iop::witness::PartialWitness;
     use plonky2::plonk::circuit_data::CircuitConfig;
     use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
 

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

@@ -1,2 +1,3 @@
 pub mod capped_tree_circuit;
-pub mod safe_tree_circuit;
+pub mod safe_tree_circuit;
+pub mod prove_single_cell;

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

@@ -0,0 +1,443 @@
+// prove single cell
+// consistent with:
+// https://github.com/codex-storage/codex-storage-proofs-circuits/blob/master/circuit/codex/single_cell.circom
+// circuit consists of:
+// - reconstruct the block merkle root
+// - use merkle root as leaf and reconstruct slot root
+// - check equality with given slot root
+
+use anyhow::Result;
+use plonky2::field::extension::Extendable;
+use plonky2::hash::hash_types::{HashOut, RichField};
+use plonky2::iop::target::{BoolTarget, Target};
+use plonky2::iop::witness::{PartialWitness, WitnessWrite, Witness};
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::CircuitConfig;
+use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, Hasher, GenericHashOut};
+use std::marker::PhantomData;
+use itertools::Itertools;
+
+use crate::merkle_tree::merkle_safe::MerkleTree;
+use plonky2::hash::poseidon::PoseidonHash;
+
+use plonky2::hash::hash_types::{HashOutTarget, NUM_HASH_OUT_ELTS};
+use crate::merkle_tree::merkle_safe::{MerkleProof, MerkleProofTarget};
+use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
+
+use plonky2::field::goldilocks_field::GoldilocksField;
+use plonky2::plonk::config::PoseidonGoldilocksConfig;
+
+use plonky2::hash::hashing::PlonkyPermutation;
+use crate::circuits::safe_tree_circuit::{MerkleTreeCircuit, MerkleTreeTargets};
+
+// constants and types used throughout the circuit
+pub const N_FIELD_ELEMS_PER_CELL: usize = 4;
+pub const BOT_DEPTH: usize = 5; // block depth - depth of the block merkle tree
+pub const MAX_DEPTH: usize = 16; // depth of big tree (slot tree depth + block tree depth)
+const N_CELLS_IN_BLOCKS: usize = 1<<BOT_DEPTH; //2^BOT_DEPTH
+const N_BLOCKS: usize = 1<<(MAX_DEPTH - BOT_DEPTH); // 2^(MAX_DEPTH - BOT_DEPTH)
+const N_CELLS: usize = N_CELLS_IN_BLOCKS * N_BLOCKS;
+// hash function used. this is hackish way of doing it because
+// H::Hash is not consistent with HashOut<F> and causing a lot of headache
+// will look into this later.
+type HF = PoseidonHash;
+
+// ------ Slot Tree --------
+
+#[derive(Clone)]
+pub struct SlotTree<F: RichField, H: Hasher<F>> {
+    pub tree: MerkleTree<F,H>, // slot tree
+    pub block_trees: Vec<MerkleTree<F,H>>, // vec of block trees
+    pub cell_data: Vec<Vec<F>>, // cell data as field elements
+    pub cell_hash: Vec<HashOut<F>>, // hash of above
+}
+
+impl<F: RichField, H: Hasher<F>> Default for SlotTree<F,H>{
+    /// slot tree with fake data, for testing only
+    fn default() -> Self {
+        // generate fake cell data
+        let mut cell_data = (0..N_CELLS)
+            .map(|i|{
+                (0..N_FIELD_ELEMS_PER_CELL)
+                    .map(|j| F::from_canonical_u64((j+i) as u64))
+                    .collect::<Vec<_>>()
+            })
+            .collect::<Vec<_>>();
+        // hash it
+        let leaves: Vec<HashOut<F>> = cell_data
+            .iter()
+            .map(|element| {
+                HF::hash_no_pad(&element)
+            })
+            .collect();
+        // zero hash
+        let zero = HashOut {
+            elements: [F::ZERO; 4],
+        };
+        // create block tree
+        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()) // use helper function
+            })
+            .collect::<Vec<_>>();
+        // get the roots or block trees
+        let block_roots = block_trees.iter()
+            .map(|t| {
+                t.root().unwrap()
+            })
+            .collect::<Vec<_>>();
+        // create slot tree
+        let slot_tree = MerkleTree::<F, H>::new(&block_roots, zero).unwrap();
+        Self{
+            tree: slot_tree,
+            block_trees,
+            cell_data,
+            cell_hash: leaves,
+        }
+    }
+}
+
+impl<F: RichField, H: Hasher<F>> SlotTree<F, H> {
+
+    /// same as default but with supplied cell data
+    pub fn new(cell_data: Vec<Vec<F>>) -> Self{
+        let leaves: Vec<HashOut<F>> = cell_data
+            .iter()
+            .map(|element| {
+                HF::hash_no_pad(element)
+            })
+            .collect();
+        let zero = HashOut {
+            elements: [F::ZERO; 4],
+        };
+        let block_trees = (0..N_BLOCKS as usize)
+            .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, H>::new(&block_roots, zero).unwrap();
+        Self{
+            tree: slot_tree,
+            block_trees,
+            cell_data,
+            cell_hash: leaves,
+        }
+    }
+
+    /// generates a proof for 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, H> {
+        let block_index = index/ N_CELLS_IN_BLOCKS;
+        let leaf_index = index % 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, H> {
+            index: index,
+            path: combined_path,
+            nleaves: self.cell_hash.len(),
+            zero: block_proof.zero.clone(),
+            phantom_data: Default::default(),
+        }
+
+    }
+
+    /// verify the given proof for slot tree, checks equality with given root
+    pub fn verify_cell_proof(&self, proof: MerkleProof<F, H>, root: HashOut<F>) -> Result<bool>{
+        let mut block_path_bits = self.usize_to_bits_le_padded(proof.index, MAX_DEPTH);
+        let last_index = N_CELLS - 1;
+        let mut block_last_bits = self.usize_to_bits_le_padded(last_index, MAX_DEPTH);
+
+        let split_point = BOT_DEPTH;
+
+        let slot_last_bits = block_last_bits.split_off(split_point);
+        let slot_path_bits = block_path_bits.split_off(split_point);
+
+        let leaf_hash = self.cell_hash[proof.index];
+
+        let mut block_path = proof.path;
+        let slot_path = block_path.split_off(split_point);
+
+        let block_res = MerkleProof::<F,H>::reconstruct_root2(leaf_hash,block_path_bits.clone(),block_last_bits.clone(),block_path);
+        let reconstructed_root = MerkleProof::<F,H>::reconstruct_root2(block_res.unwrap(),slot_path_bits,slot_last_bits,slot_path);
+
+        Ok(reconstructed_root.unwrap() == root)
+    }
+
+    fn get_block_tree(leaves: &Vec<HashOut<F>>) -> MerkleTree<F, H> {
+        let zero = HashOut {
+            elements: [F::ZERO; 4],
+        };
+        // Build the Merkle tree
+        let block_tree = MerkleTree::<F, H>::new(leaves, zero).unwrap();
+        return block_tree;
+    }
+
+    /// Converts an index to a vector of bits (LSB first) with padding.
+    pub(crate) fn usize_to_bits_le_padded(&self, index: usize, bit_length: usize) -> Vec<bool> {
+        let mut bits = Vec::with_capacity(bit_length);
+        for i in 0..bit_length {
+            bits.push(((index >> i) & 1) == 1);
+        }
+        // If index requires fewer bits, pad with `false`
+        while bits.len() < bit_length {
+            bits.push(false);
+        }
+        bits
+    }
+}
+
+//------- single cell struct ------
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct SingleCellTargets<
+    F: RichField + Extendable<D> + Poseidon2,
+    C: GenericConfig<D, F = F>,
+    const D: usize,
+    H: Hasher<F> + AlgebraicHasher<F>,
+> {
+    pub expected_slot_root_target: HashOutTarget,
+    pub proof_target: MerkleProofTarget,
+    pub leaf_target: Vec<Target>,
+    pub path_bits: Vec<BoolTarget>,
+    pub last_bits: Vec<BoolTarget>,
+    _phantom: PhantomData<(C,H)>,
+}
+
+//------- circuit impl --------
+
+impl<
+    F: RichField + Extendable<D> + Poseidon2,
+    C: GenericConfig<D, F=F>,
+    const D: usize,
+    H: Hasher<F> + AlgebraicHasher<F> + Hasher<F>,
+> MerkleTreeCircuit<F, C, D, H> {
+
+    pub fn prove_single_cell2(
+        &mut self,
+        builder: &mut CircuitBuilder::<F, D>
+    ) -> SingleCellTargets<F, C, D, H> {
+
+        // Retrieve tree depth
+        let depth = MAX_DEPTH;
+
+        // Create virtual targets
+        let mut leaf = (0..N_FIELD_ELEMS_PER_CELL).map(|_| builder.add_virtual_target()).collect::<Vec<_>>();
+
+        let mut perm_inputs:Vec<Target>= Vec::new();
+        perm_inputs.extend_from_slice(&leaf);
+        let leaf_hash = builder.hash_n_to_hash_no_pad::<H>(perm_inputs);
+
+        // path bits (binary decomposition of leaf_index)
+        let mut block_path_bits = (0..BOT_DEPTH).map(|_| builder.add_virtual_bool_target_safe()).collect::<Vec<_>>();
+        let mut slot_path_bits = (0..(depth - BOT_DEPTH)).map(|_| builder.add_virtual_bool_target_safe()).collect::<Vec<_>>();
+
+        // last bits (binary decomposition of last_index = nleaves - 1)
+        let block_last_bits = (0..BOT_DEPTH).map(|_| builder.add_virtual_bool_target_safe()).collect::<Vec<_>>();
+        let slot_last_bits = (0..(depth-BOT_DEPTH)).map(|_| builder.add_virtual_bool_target_safe()).collect::<Vec<_>>();
+
+        // Merkle path (sibling hashes from leaf to root)
+        let mut block_merkle_path = MerkleProofTarget {
+            path: (0..BOT_DEPTH).map(|_| builder.add_virtual_hash()).collect(),
+        };
+        let mut slot_merkle_path = MerkleProofTarget {
+            path: (0..(depth - BOT_DEPTH)).map(|_| builder.add_virtual_hash()).collect(),
+        };
+
+        // expected Merkle root
+        let slot_expected_root = builder.add_virtual_hash();
+
+        let mut block_targets = MerkleTreeTargets {
+            leaf: leaf_hash,
+            path_bits:block_path_bits,
+            last_bits: block_last_bits,
+            merkle_path: block_merkle_path,
+            _phantom: PhantomData,
+        };
+
+        // reconstruct block root
+        let block_root = self.reconstruct_merkle_root_circuit(builder, &mut block_targets);
+
+        // create MerkleTreeTargets struct
+        let mut slot_targets = MerkleTreeTargets {
+            leaf: block_root,
+            path_bits:slot_path_bits,
+            last_bits:slot_last_bits,
+            merkle_path:slot_merkle_path,
+            _phantom: PhantomData,
+        };
+
+        // reconstruct slot root with block root as leaf
+        let slot_root = self.reconstruct_merkle_root_circuit(builder, &mut slot_targets);
+
+        // check equality with expected root
+        for i in 0..NUM_HASH_OUT_ELTS {
+            builder.connect(slot_expected_root.elements[i], slot_root.elements[i]);
+        }
+
+        let mut proof_target = MerkleProofTarget{
+            path: block_targets.merkle_path.path,
+        };
+        proof_target.path.extend_from_slice(&slot_targets.merkle_path.path);
+
+        let mut path_bits = block_targets.path_bits;
+        path_bits.extend_from_slice(&slot_targets.path_bits);
+
+        let mut last_bits = block_targets.last_bits;
+        last_bits.extend_from_slice(&slot_targets.last_bits);
+
+        let mut cell_targets = SingleCellTargets {
+            expected_slot_root_target: slot_expected_root,
+            proof_target,
+            leaf_target: leaf,
+            path_bits,
+            last_bits,
+            _phantom: Default::default(),
+        };
+
+        // Return MerkleTreeTargets
+        cell_targets
+    }
+
+    /// assign the witness values in the circuit targets
+    /// this takes leaf_index, leaf, and proof (generated from slot_tree)
+    /// and fills all required circuit targets(circuit inputs)
+    pub fn single_cell_assign_witness(
+        &mut self,
+        pw: &mut PartialWitness<F>,
+        targets: &mut SingleCellTargets<F, C, D, H>,
+        leaf_index: usize,
+        leaf: &Vec<F>,
+        proof: MerkleProof<F,H>,
+    )-> Result<()> {
+
+        // Assign the leaf to the leaf target
+        for i in 0..targets.leaf_target.len(){
+            pw.set_target(targets.leaf_target[i], leaf[i]);
+        }
+
+        // Convert `leaf_index` to binary bits and assign as path_bits
+        let path_bits = self.usize_to_bits_le_padded(leaf_index, MAX_DEPTH);
+        for (i, bit) in path_bits.iter().enumerate() {
+            pw.set_bool_target(targets.path_bits[i], *bit);
+        }
+
+        // get `last_index` (nleaves - 1) in binary bits and assign
+        let last_index = N_CELLS - 1;
+        let last_bits = self.usize_to_bits_le_padded(last_index, MAX_DEPTH);
+        for (i, bit) in last_bits.iter().enumerate() {
+            pw.set_bool_target(targets.last_bits[i], *bit);
+        }
+
+        // assign the Merkle path (sibling hashes) to the targets
+        for (i, sibling_hash) in proof.path.iter().enumerate() {
+            // This is a bit hacky because it should be HashOutTarget, but it is H:Hash
+            // pw.set_hash_target(targets.merkle_path.path[i],sibling_hash);
+            // TODO: fix this HashOutTarget later
+            let sibling_hash_out = sibling_hash.to_vec();
+            for j in 0..sibling_hash_out.len() {
+                pw.set_target(targets.proof_target.path[i].elements[j], sibling_hash_out[j]);
+            }
+        }
+
+        // assign the expected Merkle root to the target
+        let expected_root = self.tree.root()?;
+        // TODO: fix this HashOutTarget later same issue as above
+        let expected_root_hash_out = expected_root.to_vec();
+        for j in 0..expected_root_hash_out.len() {
+            pw.set_target(targets.expected_slot_root_target.elements[j], expected_root_hash_out[j]);
+        }
+
+        Ok(())
+    }
+
+    fn hash_leaf(builder: &mut CircuitBuilder<F, D >, leaf: &mut Vec<Target>){
+        builder.hash_n_to_hash_no_pad::<H>(leaf.to_owned());
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::Instant;
+    use super::*;
+    use plonky2::plonk::circuit_data::CircuitConfig;
+    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+    use plonky2::iop::witness::PartialWitness;
+
+    //types for tests
+    type F = GoldilocksField;
+    type H = PoseidonHash;
+
+    #[test]
+    fn test_prove_single_cell(){
+        let slot_t = SlotTree::<F,H>::default();
+        let index = 8;
+        let proof = slot_t.get_proof(index);
+        let res = slot_t.verify_cell_proof(proof,slot_t.tree.root().unwrap()).unwrap();
+        assert_eq!(res, true);
+    }
+
+    #[test]
+    fn test_cell_build_circuit() -> Result<()> {
+        // circuit params
+        const D: usize = 2;
+        type C = PoseidonGoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+        type H = PoseidonHash;
+
+        let slot_t = SlotTree::<F,H>::default();
+
+        // select leaf index to prove
+        let leaf_index: usize = 8;
+
+        let proof = slot_t.get_proof(leaf_index);
+        // get the expected Merkle root
+        let expected_root = slot_t.tree.root().unwrap();
+        let res = slot_t.verify_cell_proof(proof.clone(),expected_root).unwrap();
+        assert_eq!(res, true);
+
+        // create the circuit
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+        let mut circuit_instance = MerkleTreeCircuit::<F, C, D, H> {
+            tree: slot_t.tree.clone(),
+            _phantom: PhantomData,
+        };
+        let mut targets = circuit_instance.prove_single_cell2(&mut builder);
+
+        // create a PartialWitness and assign
+        let mut pw = PartialWitness::new();
+        circuit_instance.single_cell_assign_witness(&mut pw, &mut targets, leaf_index, &slot_t.cell_data[leaf_index], proof)?;
+
+        // build the circuit
+        let data = builder.build::<C>();
+        println!("circuit size = {:?}", data.common.degree_bits());
+
+        // Prove the circuit with the assigned witness
+        let start_time = Instant::now();
+        let proof_with_pis = data.prove(pw)?;
+        println!("prove_time = {:?}", start_time.elapsed());
+
+        // verify the proof
+        let verifier_data = data.verifier_data();
+        assert!(
+            verifier_data.verify(proof_with_pis).is_ok(),
+            "Merkle proof verification failed"
+        );
+
+        Ok(())
+    }
+}
+

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

@@ -6,21 +6,21 @@ use anyhow::Result;
 use plonky2::field::extension::Extendable;
 use plonky2::field::goldilocks_field::GoldilocksField;
 use plonky2::field::types::Field;
-use plonky2::hash::hash_types::{HashOut, HashOutTarget, MerkleCapTarget, RichField, NUM_HASH_OUT_ELTS};
+use plonky2::hash::hash_types::{HashOut, HashOutTarget, RichField, NUM_HASH_OUT_ELTS};
-use plonky2::hash::hashing::{hash_n_to_m_no_pad, PlonkyPermutation};
+use plonky2::hash::hashing::PlonkyPermutation;
 use plonky2::hash::poseidon::PoseidonHash;
 use plonky2::iop::target::{BoolTarget, Target};
 use plonky2::iop::witness::{PartialWitness, Witness, WitnessWrite};
 use plonky2::plonk::circuit_builder::CircuitBuilder;
-use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, VerifierCircuitData, VerifierCircuitTarget};
+use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, VerifierCircuitData};
 use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, GenericHashOut, Hasher, PoseidonGoldilocksConfig};
 use plonky2::plonk::proof::{Proof, ProofWithPublicInputs};
 use std::marker::PhantomData;
 use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
 use serde::Serialize;
 
-use crate::merkle_tree::merkle_safe::{MerkleTree, MerkleProof, MerkleProofTarget, KeyedHasher};
+use crate::merkle_tree::merkle_safe::{MerkleTree, MerkleProofTarget};
-use crate::merkle_tree::merkle_safe::{KEY_NONE,KEY_BOTTOM_LAYER,KEY_ODD,KEY_ODD_AND_BOTTOM_LAYER};
+use crate::merkle_tree::merkle_safe::{KEY_NONE,KEY_BOTTOM_LAYER};
 
 
 /// Merkle tree targets representing the input to the circuit
@@ -32,14 +32,13 @@ pub struct MerkleTreeTargets<
     F: RichField + Extendable<D> + Poseidon2,
     C: GenericConfig<D, F = F>,
     const D: usize,
-    H: Hasher<F> + AlgebraicHasher<F> + KeyedHasher<F>,
+    H: Hasher<F> + AlgebraicHasher<F>,
 > {
     pub leaf: HashOutTarget,
     pub path_bits: Vec<BoolTarget>,
     pub last_bits: Vec<BoolTarget>,
     pub merkle_path: MerkleProofTarget,
-    pub expected_root: HashOutTarget,
+    pub _phantom: PhantomData<(C, H)>,
-    _phantom: PhantomData<(C, H)>,
 }
 
 /// Merkle tree circuit contains the tree and functions for
@@ -49,7 +48,7 @@ pub struct MerkleTreeCircuit<
     F: RichField + Extendable<D> + Poseidon2,
     C: GenericConfig<D, F = F>,
     const D: usize,
-    H: Hasher<F> + AlgebraicHasher<F> + KeyedHasher<F>,
+    H: Hasher<F> + AlgebraicHasher<F>,
 > {
     pub tree: MerkleTree<F, H>,
     pub _phantom: PhantomData<C>,
@@ -59,9 +58,10 @@ impl<
     F: RichField + Extendable<D> + Poseidon2,
     C: GenericConfig<D, F=F>,
     const D: usize,
-    H: Hasher<F> + AlgebraicHasher<F> + KeyedHasher<F>,
+    H: Hasher<F> + AlgebraicHasher<F>,
 > MerkleTreeCircuit<F, C, D, H> {
 
+    /// defines the computations inside the circuit and returns the targets used
     pub fn build_circuit(
         &mut self,
         builder: &mut CircuitBuilder::<F, D>
@@ -83,21 +83,17 @@ impl<
             path: (0..depth).map(|_| builder.add_virtual_hash()).collect(),
         };
 
-        // expected Merkle root
-        let expected_root = builder.add_virtual_hash();
-
         // create MerkleTreeTargets struct
         let mut targets = MerkleTreeTargets {
             leaf,
             path_bits,
             last_bits,
             merkle_path,
-            expected_root,
             _phantom: PhantomData,
         };
 
         // Add Merkle proof verification constraints to the circuit
-        self.verify_merkle_proof_circuit2(builder, &mut targets);
+        self.reconstruct_merkle_root_circuit(builder, &mut targets);
 
         // Return MerkleTreeTargets
         targets
@@ -171,24 +167,16 @@ impl<
             }
         }
 
-        // assign the expected Merkle root to the target
-        let expected_root = self.tree.root()?;
-        // TODO: fix this HashOutTarget later same issue as above
-        let expected_root_hash_out = expected_root.to_vec();
-        for j in 0..expected_root_hash_out.len() {
-            pw.set_target(targets.expected_root.elements[j], expected_root_hash_out[j]);
-        }
-
         Ok(())
     }
 
-    /// Verifies a Merkle proof within the circuit.
     /// takes the params from the targets struct
-    pub fn verify_merkle_proof_circuit2(
+    /// outputs the reconstructed merkle root
+    pub fn reconstruct_merkle_root_circuit(
         &self,
         builder: &mut CircuitBuilder<F, D>,
         targets: &mut MerkleTreeTargets<F, C, D, H>,
-    ) {
+    ) -> HashOutTarget {
         let max_depth = targets.path_bits.len();
         let mut state: HashOutTarget = targets.leaf;
         let zero = builder.zero();
@@ -239,11 +227,7 @@ impl<
             i += 1;
         }
 
-        // check equality with expected root
+        return state;
-        for i in 0..NUM_HASH_OUT_ELTS {
-            builder.connect(targets.expected_root.elements[i], state.elements[i]);
-        }
-
     }
 
 
@@ -254,10 +238,10 @@ impl<
     F: RichField + Extendable<D> + Poseidon2,
     C: GenericConfig<D, F = F>,
     const D: usize,
-    H: Hasher<F> + AlgebraicHasher<F> + KeyedHasher<F>,
+    H: Hasher<F> + AlgebraicHasher<F>,
 > MerkleTreeCircuit<F, C, D, H> {
     /// Converts an index to a vector of bits (LSB first) with padding.
-    fn usize_to_bits_le_padded(&self, index: usize, bit_length: usize) -> Vec<bool> {
+    pub(crate) fn usize_to_bits_le_padded(&self, index: usize, bit_length: usize) -> Vec<bool> {
         let mut bits = Vec::with_capacity(bit_length);
         for i in 0..bit_length {
             bits.push(((index >> i) & 1) == 1);
@@ -270,6 +254,9 @@ impl<
     }
 }
 
+
+// NOTE: for now these tests don't check the reconstructed root is equal to expected_root
+// will be fixed later, but for that test check the prove_single_cell tests
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -277,7 +264,6 @@ mod tests {
     use plonky2::plonk::circuit_data::CircuitConfig;
     use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
     use plonky2::iop::witness::PartialWitness;
-    use rand::Rng;
 
     #[test]
     fn test_build_circuit() -> Result<()> {
@@ -305,7 +291,7 @@ mod tests {
         let zero_hash = HashOut {
             elements: [GoldilocksField::ZERO; 4],
         };
-        let tree = MerkleTree::<F, H>::new(&leaves, zero_hash, H::compress)?;
+        let tree = MerkleTree::<F, H>::new(&leaves, zero_hash)?;
 
         // select leaf index to prove
         let leaf_index: usize = 8;
@@ -372,7 +358,7 @@ mod tests {
         let zero_hash = HashOut {
             elements: [GoldilocksField::ZERO; 4],
         };
-        let tree = MerkleTree::<F, H>::new(&leaves, zero_hash, H::compress)?;
+        let tree = MerkleTree::<F, H>::new(&leaves, zero_hash)?;
 
         let expected_root = tree.root()?;
 

codex-plonky2-circuits/src/merkle_tree/merkle_safe.rs

@@ -2,6 +2,7 @@
 // consistent with the one in codex:
 // https://github.com/codex-storage/nim-codex/blob/master/codex/merkletree/merkletree.nim
 
+use std::marker::PhantomData;
 use anyhow::{ensure, Result};
 use plonky2::field::goldilocks_field::GoldilocksField;
 use plonky2::hash::hash_types::{HashOut, HashOutTarget, RichField};
@@ -10,48 +11,38 @@ use plonky2::plonk::config::Hasher;
 use std::ops::Shr;
 use plonky2_field::types::Field;
 
+
 // 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;
 
-/// Trait for a hash function that supports keyed compression.
+// hash function used. this is hackish way of doing it because
-pub trait KeyedHasher<F: RichField>: Hasher<F> {
+// H::Hash is not consistent with HashOut<F> and causing a lot of headache
-    fn compress(x: Self::Hash, y: Self::Hash, key: u64) -> Self::Hash;
+// will look into this later.
-}
+type HF = PoseidonHash;
-
-impl KeyedHasher<GoldilocksField> for PoseidonHash {
-    fn compress(x: Self::Hash, y: Self::Hash, key: u64) -> Self::Hash {
-        let key_field = GoldilocksField::from_canonical_u64(key);
-        let mut inputs = Vec::new();
-        inputs.extend_from_slice(&x.elements);
-        inputs.extend_from_slice(&y.elements);
-        inputs.push(key_field);
-        PoseidonHash::hash_no_pad(&inputs) // TODO: double-check this function
-    }
-}
 
 /// Merkle tree struct, containing the layers, compression function, and zero hash.
 #[derive(Clone)]
-pub struct MerkleTree<F: RichField, H: KeyedHasher<F>> {
+pub struct MerkleTree<F: RichField, H: Hasher<F>> {
-    pub layers: Vec<Vec<H::Hash>>,
+    pub layers: Vec<Vec<HashOut<F>>>,
-    pub compress: fn(H::Hash, H::Hash, u64) -> H::Hash,
+    pub zero: HashOut<F>,
-    pub zero: H::Hash,
+    phantom_data: PhantomData<H>
 }
 
-impl<F: RichField, H: KeyedHasher<F>> MerkleTree<F, H> {
+impl<F: RichField, H: Hasher<F>> MerkleTree<F, H> {
     /// Constructs a new Merkle tree from the given leaves.
     pub fn new(
-        leaves: &[H::Hash],
+        leaves: &[HashOut<F>],
-        zero: H::Hash,
+        zero: HashOut<F>,
-        compress: fn(H::Hash, H::Hash, u64) -> H::Hash,
     ) -> Result<Self> {
-        let layers = merkle_tree_worker::<F,H>(leaves, zero, compress, true)?;
+        let layers = merkle_tree_worker::<F,H>(leaves, zero, true)?;
         Ok(Self {
             layers,
-            compress,
+            // compress,
             zero,
+            phantom_data: Default::default(),
         })
     }
 
@@ -66,7 +57,7 @@ impl<F: RichField, H: KeyedHasher<F>> MerkleTree<F, H> {
     }
 
     /// Returns the root hash of the Merkle tree.
-    pub fn root(&self) -> Result<H::Hash> {
+    pub fn root(&self) -> Result<HashOut<F>> {
         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])
@@ -99,19 +90,28 @@ impl<F: RichField, H: KeyedHasher<F>> MerkleTree<F, H> {
             index,
             path,
             nleaves,
-            compress: self.compress,
             zero: self.zero,
+            phantom_data: Default::default(),
         })
     }
 }
 
+/// compress input (x and y) with key using the define HF hash function
+fn key_compress<F: RichField>(x: HashOut<F>, y: HashOut<F>, key: u64) -> HashOut<F> {
+    let key_field = F::from_canonical_u64(key);
+    let mut inputs = Vec::new();
+    inputs.extend_from_slice(&x.elements);
+    inputs.extend_from_slice(&y.elements);
+    inputs.push(key_field);
+    HF::hash_no_pad(&inputs) // TODO: double-check this function
+}
+
 /// Build the Merkle tree layers.
-fn merkle_tree_worker<F: RichField, H: KeyedHasher<F>>(
+fn merkle_tree_worker<F: RichField, H: Hasher<F>>(
-    xs: &[H::Hash],
+    xs: &[HashOut<F>],
-    zero: H::Hash,
+    zero: HashOut<F>,
-    compress: fn(H::Hash, H::Hash, u64) -> H::Hash,
     is_bottom_layer: bool,
-) -> Result<Vec<Vec<H::Hash>>> {
+) -> Result<Vec<Vec<HashOut<F>>>> {
     let m = xs.len();
     if !is_bottom_layer && m == 1 {
         return Ok(vec![xs.to_vec()]);
@@ -125,7 +125,7 @@ fn merkle_tree_worker<F: RichField, H: KeyedHasher<F>>(
 
     for i in 0..halfn {
         let key = if is_bottom_layer { KEY_BOTTOM_LAYER } else { KEY_NONE };
-        let h = compress(xs[2 * i], xs[2 * i + 1], key);
+        let h = key_compress::<F>(xs[2 * i], xs[2 * i + 1], key);
         ys.push(h);
     }
 
@@ -135,12 +135,12 @@ fn merkle_tree_worker<F: RichField, H: KeyedHasher<F>>(
         } else {
             KEY_ODD
         };
-        let h = compress(xs[n], zero, key);
+        let h = key_compress::<F>(xs[n], zero, key);
         ys.push(h);
     }
 
     let mut layers = vec![xs.to_vec()];
-    let mut upper_layers = merkle_tree_worker::<F,H>(&ys, zero, compress, false)?;
+    let mut upper_layers = merkle_tree_worker::<F,H>(&ys, zero, false)?;
     layers.append(&mut upper_layers);
 
     Ok(layers)
@@ -148,12 +148,12 @@ fn merkle_tree_worker<F: RichField, H: KeyedHasher<F>>(
 
 /// Merkle proof struct, containing the index, path, and other necessary data.
 #[derive(Clone)]
-pub struct MerkleProof<F: RichField, H: KeyedHasher<F>> {
+pub struct MerkleProof<F: RichField, H: Hasher<F>> {
     pub index: usize,       // Index of the leaf
-    pub path: Vec<H::Hash>, // Sibling hashes from the leaf to the root
+    pub path: Vec<HashOut<F>>, // Sibling hashes from the leaf to the root
     pub nleaves: usize,     // Total number of leaves
-    pub compress: fn(H::Hash, H::Hash, u64) -> H::Hash, // compression function - TODO: make it generic instead
+    pub zero: HashOut<F>,
-    pub zero: H::Hash,
+    pub(crate) phantom_data: PhantomData<H>
 }
 
 #[derive(Clone, Debug, Eq, PartialEq)]
@@ -162,9 +162,9 @@ pub struct MerkleProofTarget {
     pub path: Vec<HashOutTarget>,
 }
 
-impl<F: RichField, H: KeyedHasher<F>> MerkleProof<F, H> {
+impl<F: RichField, H: Hasher<F>> MerkleProof<F, H> {
     /// Reconstructs the root hash from the proof and the given leaf.
-    pub fn reconstruct_root(&self, leaf: H::Hash) -> Result<H::Hash> {
+    pub fn reconstruct_root(&self, leaf: HashOut<F>) -> Result<HashOut<F>> {
         let mut m = self.nleaves;
         let mut j = self.index;
         let mut h = leaf;
@@ -174,14 +174,14 @@ impl<F: RichField, H: KeyedHasher<F>> MerkleProof<F, H> {
             let odd_index = (j & 1) != 0;
             if odd_index {
                 // The index of the child is odd
-                h = (self.compress)(*p, h, bottom_flag);
+                h = key_compress::<F>(*p, h, bottom_flag);
             } else {
                 if j == m - 1 {
                     // Single child -> so odd node
-                    h = (self.compress)(h, *p, bottom_flag + 2);
+                    h = key_compress::<F>(h, *p, bottom_flag + 2);
                 } else {
                     // Even node
-                    h = (self.compress)(h, *p, bottom_flag);
+                    h = key_compress::<F>(h, *p, bottom_flag);
                 }
             }
             bottom_flag = KEY_NONE;
@@ -192,30 +192,75 @@ impl<F: RichField, H: KeyedHasher<F>> MerkleProof<F, H> {
         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>, path_bits: Vec<bool>, last_bits:Vec<bool>,  path: Vec<HashOut<F>>) -> Result<HashOut<F>> {
+        let is_last = compute_is_last(path_bits.clone(),last_bits);
+
+        let mut h = 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 = key_compress::<F>(*p, h, key);
+            } else {
+                h = key_compress::<F>(h, *p, key);
+            }
+            i += 1;
+        }
+
+        Ok(h)
+    }
+
     /// Verifies the proof against a given root and leaf.
-    pub fn verify(&self, leaf: H::Hash, root: H::Hash) -> Result<bool> {
+    pub fn verify(&self, leaf: HashOut<F>, root: HashOut<F>) -> 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
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use plonky2::field::types::Field;
 
-    // Constants for the keys used in compression
+    // types used in all tests
-    // const KEY_NONE: u64 = 0x0;
+    type F = GoldilocksField;
-    // const KEY_BOTTOM_LAYER: u64 = 0x1;
+    type H = PoseidonHash;
-    // const KEY_ODD: u64 = 0x2;
-    // const KEY_ODD_AND_BOTTOM_LAYER: u64 = 0x3;
 
     fn compress(
-        x: HashOut<GoldilocksField>,
+        x: HashOut<F>,
-        y: HashOut<GoldilocksField>,
+        y: HashOut<F>,
         key: u64,
-    ) -> HashOut<GoldilocksField> {
+    ) -> HashOut<F> {
-        let key_field = GoldilocksField::from_canonical_u64(key);
+        let key_field = F::from_canonical_u64(key);
         let mut inputs = Vec::new();
         inputs.extend_from_slice(&x.elements);
         inputs.extend_from_slice(&y.elements);
@@ -224,46 +269,42 @@ mod tests {
     }
 
     fn make_tree(
-        data: &[GoldilocksField],
+        data: &[F],
-        zero: HashOut<GoldilocksField>,
+        zero: HashOut<F>,
-    ) -> Result<MerkleTree<GoldilocksField, PoseidonHash>> {
+    ) -> Result<MerkleTree<F, H>> {
-        let compress_fn = PoseidonHash::compress;
-
         // Hash the data to obtain leaf hashes
         let leaves: Vec<HashOut<GoldilocksField>> = data
             .iter()
             .map(|&element| {
                 // Hash each field element to get the leaf hash
-                PoseidonHash::hash_no_pad(&[element])
+                H::hash_no_pad(&[element])
             })
             .collect();
 
-        MerkleTree::<GoldilocksField, PoseidonHash>::new(&leaves, zero, compress_fn)
+        MerkleTree::<F, H>::new(&leaves, zero)
     }
 
     #[test]
     fn single_proof_test() -> Result<()> {
         let data = (1u64..=8)
-            .map(|i| GoldilocksField::from_canonical_u64(i))
+            .map(|i| F::from_canonical_u64(i))
             .collect::<Vec<_>>();
 
         // Hash the data to obtain leaf hashes
-        let leaves: Vec<HashOut<GoldilocksField>> = data
+        let leaves: Vec<HashOut<F>> = data
             .iter()
             .map(|&element| {
                 // Hash each field element to get the leaf hash
-                PoseidonHash::hash_no_pad(&[element])
+                H::hash_no_pad(&[element])
             })
             .collect();
 
         let zero = HashOut {
-            elements: [GoldilocksField::ZERO; 4],
+            elements: [F::ZERO; 4],
         };
 
-        let compress_fn = PoseidonHash::compress;
-
         // Build the Merkle tree
-        let tree = MerkleTree::<GoldilocksField, PoseidonHash>::new(&leaves, zero, compress_fn)?;
+        let tree = MerkleTree::<F, H>::new(&leaves, zero)?;
 
         // Get the root
         let root = tree.root()?;
@@ -282,18 +323,18 @@ mod tests {
     fn test_correctness_even_bottom_layer() -> Result<()> {
         // Data for the test (field elements)
         let data = (1u64..=8)
-            .map(|i| GoldilocksField::from_canonical_u64(i))
+            .map(|i| F::from_canonical_u64(i))
             .collect::<Vec<_>>();
 
         // Hash the data to get leaf hashes
-        let leaf_hashes: Vec<HashOut<GoldilocksField>> = data
+        let leaf_hashes: Vec<HashOut<F>> = data
             .iter()
-            .map(|&element| PoseidonHash::hash_no_pad(&[element]))
+            .map(|&element| H::hash_no_pad(&[element]))
             .collect();
 
         // zero hash
         let zero = HashOut {
-            elements: [GoldilocksField::ZERO; 4],
+            elements: [F::ZERO; 4],
         };
 
         let expected_root =
@@ -343,18 +384,18 @@ mod tests {
     fn test_correctness_odd_bottom_layer() -> Result<()> {
         // Data for the test (field elements)
         let data = (1u64..=7)
-            .map(|i| GoldilocksField::from_canonical_u64(i))
+            .map(|i| F::from_canonical_u64(i))
             .collect::<Vec<_>>();
 
         // Hash the data to get leaf hashes
-        let leaf_hashes: Vec<HashOut<GoldilocksField>> = data
+        let leaf_hashes: Vec<HashOut<F>> = data
             .iter()
-            .map(|&element| PoseidonHash::hash_no_pad(&[element]))
+            .map(|&element| H::hash_no_pad(&[element]))
             .collect();
 
         // zero hash
         let zero = HashOut {
-            elements: [GoldilocksField::ZERO; 4],
+            elements: [F::ZERO; 4],
         };
 
         let expected_root =
@@ -404,18 +445,18 @@ mod tests {
     fn test_correctness_even_bottom_odd_upper_layers() -> Result<()> {
         // Data for the test (field elements)
         let data = (1u64..=10)
-            .map(|i| GoldilocksField::from_canonical_u64(i))
+            .map(|i| F::from_canonical_u64(i))
             .collect::<Vec<_>>();
 
         // Hash the data to get leaf hashes
-        let leaf_hashes: Vec<HashOut<GoldilocksField>> = data
+        let leaf_hashes: Vec<HashOut<F>> = data
             .iter()
-            .map(|&element| PoseidonHash::hash_no_pad(&[element]))
+            .map(|&element| H::hash_no_pad(&[element]))
             .collect();
 
         // zero hash
         let zero = HashOut {
-            elements: [GoldilocksField::ZERO; 4],
+            elements: [F::ZERO; 4],
         };
 
         let expected_root = compress(
@@ -480,24 +521,22 @@ mod tests {
     fn test_proofs() -> Result<()> {
         // Data for the test (field elements)
         let data = (1u64..=10)
-            .map(|i| GoldilocksField::from_canonical_u64(i))
+            .map(|i| F::from_canonical_u64(i))
             .collect::<Vec<_>>();
 
         // Hash the data to get leaf hashes
-        let leaf_hashes: Vec<HashOut<GoldilocksField>> = data
+        let leaf_hashes: Vec<HashOut<F>> = data
             .iter()
-            .map(|&element| PoseidonHash::hash_no_pad(&[element]))
+            .map(|&element| H::hash_no_pad(&[element]))
             .collect();
 
         // zero hash
         let zero = HashOut {
-            elements: [GoldilocksField::ZERO; 4],
+            elements: [F::ZERO; 4],
         };
 
-        let compress_fn = PoseidonHash::compress;
-
         // Build the tree
-        let tree = MerkleTree::<GoldilocksField, PoseidonHash>::new(&leaf_hashes, zero, compress_fn)?;
+        let tree = MerkleTree::<F, H>::new(&leaf_hashes, zero)?;
 
         // Get the root
         let expected_root = tree.root()?;