chore(backend): implement more circuit op logic (#173)

* Add backend MerkleProof type

* Add eval_not_contains

* Remove print statement

* Handle some edge cases

* Add test

* Add missing ?

* Optimisation and stylistic changes

* Code review

src/backends/plonky2/circuits/common.rs

@@ -3,14 +3,16 @@
 use crate::backends::plonky2::basetypes::D;
 use crate::backends::plonky2::mock::mainpod::Statement;
 use crate::backends::plonky2::mock::mainpod::{Operation, OperationArg};
+use crate::backends::plonky2::primitives::merkletree::MerkleClaimAndProofTarget;
 use crate::middleware::{
     NativeOperation, NativePredicate, Params, Predicate, StatementArg, ToFields, Value,
-    EMPTY_VALUE, F, HASH_SIZE, OPERATION_ARG_F_LEN, STATEMENT_ARG_F_LEN, VALUE_SIZE,
+    EMPTY_VALUE, F, HASH_SIZE, OPERATION_ARG_F_LEN, OPERATION_AUX_F_LEN, STATEMENT_ARG_F_LEN,
+    VALUE_SIZE,
 };
 use anyhow::Result;
 use plonky2::field::extension::Extendable;
 use plonky2::field::types::{Field, PrimeField64};
-use plonky2::hash::hash_types::RichField;
+use plonky2::hash::hash_types::{HashOutTarget, RichField, NUM_HASH_OUT_ELTS};
 use plonky2::iop::target::{BoolTarget, Target};
 use plonky2::iop::witness::{PartialWitness, WitnessWrite};
 use plonky2::plonk::circuit_builder::CircuitBuilder;
@@ -155,6 +157,7 @@ impl StatementTarget {
 pub struct OperationTarget {
     pub op_type: [Target; Params::operation_type_size()],
     pub args: Vec<[Target; OPERATION_ARG_F_LEN]>,
+    pub aux: [Target; OPERATION_AUX_F_LEN],
 }
 
 impl OperationTarget {
@@ -174,6 +177,7 @@ impl OperationTarget {
         {
             pw.set_target_arr(&self.args[i], &arg.to_fields(params))?;
         }
+        pw.set_target_arr(&self.aux, &op.aux().to_fields(params))?;
         Ok(())
     }
 
@@ -197,6 +201,56 @@ pub trait Flattenable {
     fn from_flattened(vs: &[Target]) -> Self;
 }
 
+/// For the purpose of op verification, we need only look up the
+/// Merkle claim rather than the Merkle proof since it is verified
+/// elsewhere.
+pub struct MerkleClaimTarget {
+    pub(crate) enabled: BoolTarget,
+    pub(crate) root: HashOutTarget,
+    pub(crate) key: ValueTarget,
+    pub(crate) value: ValueTarget,
+    pub(crate) existence: BoolTarget,
+}
+
+impl From<MerkleClaimAndProofTarget> for MerkleClaimTarget {
+    fn from(pf: MerkleClaimAndProofTarget) -> Self {
+        Self {
+            enabled: pf.enabled,
+            root: pf.root,
+            key: pf.key,
+            value: pf.value,
+            existence: pf.existence,
+        }
+    }
+}
+
+impl Flattenable for MerkleClaimTarget {
+    fn flatten(&self) -> Vec<Target> {
+        [
+            vec![self.enabled.target],
+            self.root.elements.to_vec(),
+            self.key.elements.to_vec(),
+            self.value.elements.to_vec(),
+            vec![self.existence.target],
+        ]
+        .concat()
+    }
+
+    fn from_flattened(vs: &[Target]) -> Self {
+        Self {
+            enabled: BoolTarget::new_unsafe(vs[0]),
+            root: HashOutTarget::from_vec((&vs[1..1 + NUM_HASH_OUT_ELTS]).to_vec()),
+            key: ValueTarget::from_slice(
+                &vs[1 + NUM_HASH_OUT_ELTS..1 + NUM_HASH_OUT_ELTS + VALUE_SIZE],
+            ),
+            value: ValueTarget::from_slice(
+                &vs[1 + NUM_HASH_OUT_ELTS + VALUE_SIZE..1 + NUM_HASH_OUT_ELTS + 2 * VALUE_SIZE],
+            ),
+            existence: BoolTarget::new_unsafe(vs[1 + NUM_HASH_OUT_ELTS + 2 * VALUE_SIZE]),
+        }
+    }
+}
+
 impl Flattenable for StatementTarget {
     fn flatten(&self) -> Vec<Target> {
         self.predicate
@@ -290,6 +344,7 @@ impl CircuitBuilderPod<F, D> for CircuitBuilder<F, D> {
             args: (0..params.max_operation_args)
                 .map(|_| self.add_virtual_target_arr())
                 .collect(),
+            aux: self.add_virtual_target_arr(),
         }
     }
 

src/backends/plonky2/circuits/mainpod.rs

@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use itertools::zip_eq;
 use plonky2::{
     hash::{hash_types::HashOutTarget, poseidon::PoseidonHash},
@@ -6,19 +6,27 @@ use plonky2::{
     plonk::circuit_builder::CircuitBuilder,
 };
 
-use crate::backends::plonky2::basetypes::{Value, D, EMPTY_HASH, F, VALUE_SIZE};
-use crate::backends::plonky2::circuits::common::{
-    CircuitBuilderPod, OperationTarget, StatementTarget, ValueTarget,
-};
 use crate::backends::plonky2::mock::mainpod;
 use crate::backends::plonky2::signedpod::SignedPod;
+use crate::backends::plonky2::{
+    basetypes::{Value, D, EMPTY_HASH, F, VALUE_SIZE},
+    mock::mainpod::MerkleClaimAndProof,
+    primitives::merkletree::{MerkleClaimAndProofTarget, MerkleProofGadget},
+};
 use crate::middleware::{
     hash_str, AnchoredKey, NativeOperation, NativePredicate, Params, PodType, Statement,
     StatementArg, ToFields, KEY_TYPE, SELF,
 };
+use crate::{
+    backends::plonky2::{
+        circuits::common::{CircuitBuilderPod, OperationTarget, StatementTarget, ValueTarget},
+        primitives::merkletree,
+    },
+    middleware,
+};
 
 use super::{
-    common::Flattenable,
+    common::{Flattenable, MerkleClaimTarget},
     signedpod::{SignedPodVerifyGadget, SignedPodVerifyTarget},
 };
 
@@ -37,6 +45,7 @@ impl OperationVerifyGadget {
         st: &StatementTarget,
         op: &OperationTarget,
         prev_statements: &[StatementTarget],
+        merkle_claims: &[MerkleClaimTarget],
     ) -> Result<OperationVerifyTarget> {
         let _true = builder._true();
         let _false = builder._false();
@@ -54,6 +63,12 @@ impl OperationVerifyGadget {
                 .collect::<Vec<_>>()
         };
 
+        // Certain operations (Contains/NotContains) will refer to one
+        // of the provided Merkle proofs (if any). These proofs have already
+        // been verified, so we need only look up the claim.
+        let resolved_merkle_claim =
+            (merkle_claims.len() > 0).then(|| builder.vec_ref(merkle_claims, op.aux[0]));
+
         // The verification may require aux data which needs to be stored in the
         // `OperationVerifyTarget` so that we can set during witness generation.
 
@@ -77,6 +92,18 @@ impl OperationVerifyGadget {
                     self.eval_lt_from_entries(builder, st, op, &resolved_op_args),
                 ]
             },
+            // Skip these if there are no resolved Merkle claims
+            if let Some(resolved_merkle_claim) = resolved_merkle_claim {
+                vec![self.eval_not_contains_from_entries(
+                    builder,
+                    st,
+                    op,
+                    resolved_merkle_claim,
+                    &resolved_op_args,
+                )]
+            } else {
+                vec![]
+            },
         ]
         .concat();
 
@@ -87,6 +114,74 @@ impl OperationVerifyGadget {
         Ok(OperationVerifyTarget {})
     }
 
+    fn eval_not_contains_from_entries(
+        &self,
+        builder: &mut CircuitBuilder<F, D>,
+        st: &StatementTarget,
+        op: &OperationTarget,
+        resolved_merkle_claim: MerkleClaimTarget,
+        resolved_op_args: &[StatementTarget],
+    ) -> BoolTarget {
+        let op_code_ok = op.has_native_type(builder, NativeOperation::NotContainsFromEntries);
+
+        // Expect 2 op args of type `ValueOf`.
+        let op_arg_type_checks = resolved_op_args
+            .iter()
+            .take(2)
+            .map(|op_arg| op_arg.has_native_type(builder, &self.params, NativePredicate::ValueOf))
+            .collect::<Vec<_>>();
+        let op_arg_types_ok = builder.all(op_arg_type_checks);
+
+        // The values embedded in the op args must be values, i.e. the
+        // last `STATEMENT_ARG_F_LEN - VALUE_SIZE` slots of each being
+        // 0.
+        let merkle_root_arg = &resolved_op_args[0].args[1];
+        let key_arg = &resolved_op_args[1].args[1];
+        let op_arg_range_checks = [
+            builder.statement_arg_is_value(merkle_root_arg),
+            builder.statement_arg_is_value(key_arg),
+        ];
+        let op_arg_range_ok = builder.all(op_arg_range_checks);
+
+        // Check Merkle proof (verified elsewhere) against op args.
+        let merkle_proof_checks = [
+            /* The supplied Merkle proof must be enabled. */
+            resolved_merkle_claim.enabled,
+            /* ...and it must be a nonexistence proof. */
+            builder.not(resolved_merkle_claim.existence),
+            /* ...for the root-key pair in the resolved op args. */
+            builder.is_equal_slice(
+                &merkle_root_arg.elements[..VALUE_SIZE],
+                &resolved_merkle_claim.root.elements,
+            ),
+            builder.is_equal_slice(
+                &key_arg.elements[..VALUE_SIZE],
+                &resolved_merkle_claim.key.elements,
+            ),
+        ];
+
+        let merkle_proof_ok = builder.all(merkle_proof_checks);
+
+        // Check output statement
+        let arg1_key = resolved_op_args[0].args[0].clone();
+        let arg2_key = resolved_op_args[1].args[0].clone();
+        let expected_statement = StatementTarget::new_native(
+            builder,
+            &self.params,
+            NativePredicate::NotContains,
+            &[arg1_key, arg2_key],
+        );
+        let st_ok = builder.is_equal_flattenable(st, &expected_statement);
+
+        builder.all([
+            op_code_ok,
+            op_arg_types_ok,
+            op_arg_range_ok,
+            merkle_proof_ok,
+            st_ok,
+        ])
+    }
+
     fn eval_eq_from_entries(
         &self,
         builder: &mut CircuitBuilder<F, D>,
@@ -315,6 +410,19 @@ impl MainPodVerifyGadget {
         let pub_statements =
             &input_statements[input_statements.len() - params.max_public_statements..];
 
+        // Add Merkle claim/proof targets
+        let mp_gadget = MerkleProofGadget {
+            max_depth: params.max_depth_mt_gadget,
+        };
+        let merkle_proofs: Vec<_> = (0..params.max_merkle_proofs)
+            .map(|_| mp_gadget.eval(builder))
+            .collect::<Result<_>>()?;
+        let merkle_claims: Vec<_> = merkle_proofs
+            .clone()
+            .into_iter()
+            .map(|pf| pf.into())
+            .collect();
+
         // 2. Calculate the Pod Id from the public statements
         let pub_statements_flattened = pub_statements
             .iter()
@@ -350,7 +458,7 @@ impl MainPodVerifyGadget {
             let op_verification = OperationVerifyGadget {
                 params: params.clone(),
             }
-            .eval(builder, st, op, prev_statements)?;
+            .eval(builder, st, op, prev_statements, &merkle_claims)?;
             op_verifications.push(op_verification);
         }
 
@@ -360,6 +468,7 @@ impl MainPodVerifyGadget {
             signed_pods,
             statements: input_statements.to_vec(),
             operations,
+            merkle_proofs,
             op_verifications,
         })
     }
@@ -372,6 +481,7 @@ pub struct MainPodVerifyTarget {
     // The KEY_TYPE statement must be the first public one
     statements: Vec<StatementTarget>,
     operations: Vec<OperationTarget>,
+    merkle_proofs: Vec<MerkleClaimAndProofTarget>,
     op_verifications: Vec<OperationVerifyTarget>,
 }
 
@@ -379,6 +489,7 @@ pub struct MainPodVerifyInput {
     pub signed_pods: Vec<SignedPod>,
     pub statements: Vec<mainpod::Statement>,
     pub operations: Vec<mainpod::Operation>,
+    pub merkle_proofs: Vec<MerkleClaimAndProof>,
 }
 
 impl MainPodVerifyTarget {
@@ -402,6 +513,23 @@ impl MainPodVerifyTarget {
             self.statements[i].set_targets(pw, &self.params, st)?;
             self.operations[i].set_targets(pw, &self.params, op)?;
         }
+        assert_eq!(input.merkle_proofs.len(), self.params.max_merkle_proofs);
+        for (i, mp) in input.merkle_proofs.iter().enumerate() {
+            assert_eq!(mp.siblings.len(), self.params.max_depth_mt_gadget);
+            self.merkle_proofs[i].set_targets(
+                pw,
+                mp.enabled,
+                mp.existence,
+                mp.root,
+                mp.clone().try_into().unwrap_or(merkletree::MerkleProof {
+                    existence: mp.existence,
+                    siblings: mp.siblings.clone(),
+                    other_leaf: None,
+                }),
+                mp.key,
+                mp.value,
+            )?;
+        }
         Ok(())
     }
 }
@@ -423,6 +551,7 @@ impl MainPodVerifyCircuit {
 
 #[cfg(test)]
 mod tests {
+    use merkletree::MerkleTree;
     use plonky2::plonk::{circuit_builder::CircuitBuilder, circuit_data::CircuitConfig};
 
     use super::*;
@@ -437,8 +566,12 @@ mod tests {
         st: mainpod::Statement,
         op: mainpod::Operation,
         prev_statements: Vec<mainpod::Statement>,
+        merkle_proofs: Vec<mainpod::MerkleClaimAndProof>,
     ) -> Result<()> {
         let params = Params::default();
+        let mp_gadget = MerkleProofGadget {
+            max_depth: params.max_depth_mt_gadget,
+        };
 
         let config = CircuitConfig::standard_recursion_config();
         let mut builder = CircuitBuilder::<F, D>::new(config);
@@ -448,6 +581,15 @@ mod tests {
         let prev_statements_target: Vec<_> = (0..prev_statements.len())
             .map(|_| builder.add_virtual_statement(&params))
             .collect();
+        let merkle_proofs_target: Vec<_> = merkle_proofs
+            .iter()
+            .map(|_| mp_gadget.eval(&mut builder))
+            .collect::<Result<_>>()?;
+        let merkle_claims_target: Vec<_> = merkle_proofs_target
+            .clone()
+            .into_iter()
+            .map(|pf| pf.into())
+            .collect();
 
         let operation_verify = OperationVerifyGadget {
             params: params.clone(),
@@ -457,6 +599,7 @@ mod tests {
             &st_target,
             &op_target,
             &prev_statements_target,
+            &merkle_claims_target,
         )?;
 
         let mut pw = PartialWitness::<F>::new();
@@ -465,6 +608,19 @@ mod tests {
         for (prev_st_target, prev_st) in prev_statements_target.iter().zip(prev_statements.iter()) {
             prev_st_target.set_targets(&mut pw, &params, prev_st)?;
         }
+        for (merkle_proof_target, merkle_proof) in
+            merkle_proofs_target.iter().zip(merkle_proofs.iter())
+        {
+            merkle_proof_target.set_targets(
+                &mut pw,
+                merkle_proof.enabled,
+                merkle_proof.existence,
+                merkle_proof.root,
+                merkle_proof.clone().try_into()?,
+                merkle_proof.key,
+                merkle_proof.value,
+            )?
+        }
         let input = OperationVerifyInput {};
         operation_verify.set_targets(&mut pw, &input)?;
 
@@ -478,6 +634,8 @@ mod tests {
 
     #[test]
     fn test_operation_verify() -> Result<()> {
+        let params = Params::default();
+
         // None
         let st: mainpod::Statement = Statement::None.into();
         let op = mainpod::Operation(
@@ -486,7 +644,13 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![Statement::None.into()];
-        operation_verify(st.clone(), op, prev_statements.clone())?;
+        let merkle_proofs = vec![];
+        operation_verify(
+            st.clone(),
+            op,
+            prev_statements.clone(),
+            merkle_proofs.clone(),
+        )?;
 
         // NewEntry
         let st1: mainpod::Statement =
@@ -502,7 +666,12 @@ mod tests {
             vec![],
             OperationAux::None,
         );
-        operation_verify(st1.clone(), op, prev_statements.clone())?;
+        operation_verify(
+            st1.clone(),
+            op,
+            prev_statements.clone(),
+            merkle_proofs.clone(),
+        )?;
 
         // Copy
         let st: mainpod::Statement = Statement::None.into();
@@ -512,7 +681,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![Statement::None.into()];
-        operation_verify(st, op, prev_statements)?;
+        operation_verify(st, op, prev_statements, merkle_proofs.clone())?;
 
         // Eq
         let st2: mainpod::Statement = Statement::ValueOf(
@@ -531,7 +700,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st1.clone(), st2];
-        operation_verify(st, op, prev_statements)?;
+        operation_verify(st, op, prev_statements, merkle_proofs.clone())?;
 
         // Lt
         let st2: mainpod::Statement = Statement::ValueOf(
@@ -550,7 +719,40 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st1.clone(), st2];
-        operation_verify(st, op, prev_statements)?;
+        operation_verify(st, op, prev_statements, merkle_proofs.clone())?;
+
+        // NotContainsFromEntries
+        let kvs = [
+            (1.into(), 55.into()),
+            (2.into(), 88.into()),
+            (175.into(), 0.into()),
+        ]
+        .into_iter()
+        .collect();
+        let mt = MerkleTree::new(params.max_depth_mt_gadget, &kvs)?;
+
+        let root = mt.root().into();
+        let root_ak = AnchoredKey(PodId(Value::from(88).into()), "merkle root".into());
+
+        let key = 5.into();
+        let key_ak = AnchoredKey(PodId(Value::from(88).into()), "key".into());
+
+        let no_key_pf = mt.prove_nonexistence(&key)?;
+
+        let root_st: mainpod::Statement = Statement::ValueOf(root_ak, root).into();
+        let key_st: mainpod::Statement = Statement::ValueOf(key_ak, key).into();
+        let st: mainpod::Statement = Statement::NotContains(root_ak, key_ak).into();
+        let op = mainpod::Operation(
+            OperationType::Native(NativeOperation::NotContainsFromEntries),
+            vec![OperationArg::Index(0), OperationArg::Index(1)],
+            OperationAux::MerkleProofIndex(0),
+        );
+
+        let merkle_proofs = vec![mainpod::MerkleClaimAndProof::try_from_middleware(
+            &params, &root, &key, None, &no_key_pf,
+        )?];
+        let prev_statements = vec![root_st, key_st];
+        operation_verify(st, op, prev_statements, merkle_proofs.clone())?;
 
         Ok(())
     }

src/backends/plonky2/mainpod.rs

@@ -23,6 +23,8 @@ use crate::middleware::{
     SELF,
 };
 
+use super::mock::mainpod::MerkleClaimAndProof;
+
 pub struct Prover {}
 
 impl PodProver for Prover {
@@ -47,12 +49,14 @@ impl PodProver for Prover {
             })
             .collect_vec();
 
+        let merkle_proofs = MockMainPod::extract_merkle_proofs(params, &inputs.operations)?;
+
         // TODO: Move these methods from the mock main pod to a common place
         let statements = MockMainPod::layout_statements(params, &inputs);
         let operations = MockMainPod::process_private_statements_operations(
             params,
             &statements,
-            &[], // TODO: fill in the merkle proofs for Contains/NotContains ops
+            &merkle_proofs,
             inputs.operations,
         )?;
         let operations =
@@ -67,6 +71,7 @@ impl PodProver for Prover {
             signed_pods: signed_pods_input,
             statements: statements[statements.len() - params.max_statements..].to_vec(),
             operations,
+            merkle_proofs,
         };
         main_pod.set_targets(&mut pw, &input)?;
 
@@ -173,19 +178,28 @@ pub mod tests {
         pay_stub: &frontend::SignedPod,
         sanction_list: &frontend::SignedPod,
     ) -> Result<frontend::MainPodBuilder> {
+        let sanction_set = match sanction_list.kvs.get("sanctionList") {
+            Some(frontend::Value::Set(s)) => Ok(s),
+            _ => Err(anyhow!("Missing sanction list!")),
+        }?;
         let now_minus_18y: i64 = 1169909388;
         let now_minus_1y: i64 = 1706367566;
 
+        let gov_id_kvs = gov_id.kvs();
+        let id_number_value = gov_id_kvs.get(&"idNumber".into()).unwrap();
+
         let mut kyc = frontend::MainPodBuilder::new(params);
         kyc.add_signed_pod(gov_id);
         kyc.add_signed_pod(pay_stub);
         kyc.add_signed_pod(sanction_list);
-        // NOTE: Unimplemented in the circuit
-        // kyc.pub_op(op!(
-        //     not_contains,
-        //     (sanction_list, "sanctionList"),
-        //     (gov_id, "idNumber")
-        // ))?;
+        kyc.pub_op(op!(
+            set_not_contains,
+            (sanction_list, "sanctionList"),
+            (gov_id, "idNumber"),
+            sanction_set
+                .middleware_set()
+                .prove_nonexistence(id_number_value)?
+        ))?;
         kyc.pub_op(op!(lt, (gov_id, "dateOfBirth"), now_minus_18y))?;
         kyc.pub_op(op!(
             eq,

src/backends/plonky2/mock/mainpod/mod.rs

@@ -9,7 +9,7 @@ use std::any::Any;
 use std::fmt;
 
 use crate::{
-    backends::plonky2::primitives::merkletree::MerkleProof,
+    backends::plonky2::primitives::merkletree,
     middleware::{
         self, hash_str, AnchoredKey, Hash, MainPodInputs, NativeOperation, NativePredicate,
         NonePod, OperationType, Params, Pod, PodId, PodProver, PodType, Predicate, StatementArg,
@@ -44,7 +44,7 @@ pub struct MockMainPod {
     statements: Vec<Statement>,
     // All Merkle proofs
     // TODO: Use a backend-specific representation
-    merkle_proofs: Vec<MerkleProof>,
+    merkle_proofs: Vec<MerkleClaimAndProof>,
 }
 
 impl fmt::Display for MockMainPod {
@@ -227,6 +227,52 @@ impl MockMainPod {
         statements
     }
 
+    /// Extracts and pads Merkle proofs from Contains/NotContains ops.
+    pub(crate) fn extract_merkle_proofs(
+        params: &Params,
+        operations: &[middleware::Operation],
+    ) -> Result<Vec<MerkleClaimAndProof>> {
+        let mut merkle_proofs = operations
+            .iter()
+            .flat_map(|op| match op {
+                middleware::Operation::ContainsFromEntries(
+                    middleware::Statement::ValueOf(_, root),
+                    middleware::Statement::ValueOf(_, key),
+                    middleware::Statement::ValueOf(_, value),
+                    pf,
+                ) => Some(MerkleClaimAndProof::try_from_middleware(
+                    params,
+                    root,
+                    key,
+                    Some(value),
+                    pf,
+                )),
+                middleware::Operation::NotContainsFromEntries(
+                    middleware::Statement::ValueOf(_, root),
+                    middleware::Statement::ValueOf(_, key),
+                    pf,
+                ) => Some(MerkleClaimAndProof::try_from_middleware(
+                    params, root, key, None, pf,
+                )),
+                _ => None,
+            })
+            .collect::<Result<Vec<_>>>()?;
+        if merkle_proofs.len() > params.max_merkle_proofs {
+            return Err(anyhow!(
+                "The number of required Merkle proofs ({}) exceeds the maximum number ({}).",
+                merkle_proofs.len(),
+                params.max_merkle_proofs
+            ));
+        } else {
+            fill_pad(
+                &mut merkle_proofs,
+                MerkleClaimAndProof::empty(params.max_depth_mt_gadget),
+                params.max_merkle_proofs,
+            );
+            Ok(merkle_proofs)
+        }
+    }
+
     fn find_op_arg(
         statements: &[Statement],
         op_arg: &middleware::Statement,
@@ -248,7 +294,7 @@ impl MockMainPod {
     }
 
     fn find_op_aux(
-        merkle_proofs: &[MerkleProof],
+        merkle_proofs: &[MerkleClaimAndProof],
         op_aux: &middleware::OperationAux,
     ) -> Result<OperationAux> {
         match op_aux {
@@ -256,7 +302,14 @@ impl MockMainPod {
             middleware::OperationAux::MerkleProof(pf_arg) => merkle_proofs
                 .iter()
                 .enumerate()
-                .find_map(|(i, pf)| (pf == pf_arg).then_some(i))
+                .find_map(|(i, pf)| {
+                    pf.clone()
+                        .try_into()
+                        .ok()
+                        .and_then(|mid_pf: merkletree::MerkleProof| {
+                            (&mid_pf == pf_arg).then_some(i)
+                        })
+                })
                 .map(OperationAux::MerkleProofIndex)
                 .ok_or(anyhow!(
                     "Merkle proof corresponding to op arg {} not found",
@@ -268,7 +321,7 @@ impl MockMainPod {
     pub(crate) fn process_private_statements_operations(
         params: &Params,
         statements: &[Statement],
-        merkle_proofs: &[MerkleProof],
+        merkle_proofs: &[MerkleClaimAndProof],
         input_operations: &[middleware::Operation],
     ) -> Result<Vec<Operation>> {
         let mut operations = Vec::new();
@@ -336,15 +389,9 @@ impl MockMainPod {
         // TODO: Insert a new public statement of ValueOf with `key=KEY_TYPE,
         // value=PodType::MockMainPod`
         let statements = Self::layout_statements(params, &inputs);
-        let merkle_proofs = inputs
-            .operations
-            .iter()
-            .flat_map(|op| match op {
-                middleware::Operation::ContainsFromEntries(_, _, _, pf) => Some(pf.clone()),
-                middleware::Operation::NotContainsFromEntries(_, _, pf) => Some(pf.clone()),
-                _ => None,
-            })
-            .collect::<Vec<_>>();
+        // Extract Merkle proofs and pad.
+        let merkle_proofs = Self::extract_merkle_proofs(params, &inputs.operations)?;
+
         let operations = Self::process_private_statements_operations(
             params,
             &statements,

src/backends/plonky2/mock/mainpod/operation.rs

@@ -1,12 +1,12 @@
 use super::Statement;
 use crate::{
-    backends::plonky2::primitives::merkletree::MerkleProof,
-    middleware::{self, OperationType, Params, ToFields, F},
+    backends::plonky2::primitives::merkletree::{self, kv_hash},
+    middleware::{self, Hash, OperationType, Params, ToFields, Value, EMPTY_HASH, EMPTY_VALUE, F},
 };
 use anyhow::{anyhow, Result};
 use plonky2::field::types::{Field, PrimeField64};
 use serde::{Deserialize, Serialize};
-use std::fmt;
+use std::{fmt, iter};
 
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub enum OperationArg {
@@ -36,6 +36,118 @@ pub enum OperationAux {
     MerkleProofIndex(usize),
 }
 
+impl ToFields for OperationAux {
+    fn to_fields(&self, _params: &Params) -> Vec<F> {
+        let f = match self {
+            Self::None => F::ZERO,
+            Self::MerkleProofIndex(i) => F::from_canonical_usize(*i),
+        };
+        vec![f]
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct MerkleClaimAndProof {
+    pub enabled: bool,
+    pub root: Hash,
+    pub key: Value,
+    pub value: Value,
+    pub existence: bool,
+    pub siblings: Vec<Hash>,
+    pub case_ii_selector: bool,
+    pub other_key: Value,
+    pub other_value: Value,
+}
+
+impl MerkleClaimAndProof {
+    pub fn empty(max_depth: usize) -> Self {
+        Self {
+            enabled: false,
+            root: EMPTY_HASH,
+            key: Value::from(1),
+            value: EMPTY_VALUE,
+            existence: false,
+            siblings: iter::repeat(EMPTY_HASH).take(max_depth).collect(),
+            case_ii_selector: false,
+            other_key: EMPTY_VALUE,
+            other_value: EMPTY_VALUE,
+        }
+    }
+    pub fn try_from_middleware(
+        params: &Params,
+        root: &Value,
+        key: &Value,
+        value: Option<&Value>,
+        mid_mp: &merkletree::MerkleProof,
+    ) -> Result<Self> {
+        if mid_mp.siblings.len() > params.max_depth_mt_gadget {
+            Err(anyhow!(
+                "Number of siblings ({}) exceeds maximum depth ({})",
+                mid_mp.siblings.len(),
+                params.max_depth_mt_gadget
+            ))
+        } else {
+            let (other_key, other_value) = mid_mp.other_leaf.unwrap_or((EMPTY_VALUE, EMPTY_VALUE));
+            Ok(Self {
+                enabled: true,
+                root: root.clone().into(),
+                key: key.clone(),
+                value: value.cloned().unwrap_or(EMPTY_VALUE),
+                existence: mid_mp.existence,
+                siblings: mid_mp
+                    .siblings
+                    .iter()
+                    .cloned()
+                    .chain(iter::repeat(EMPTY_HASH))
+                    .take(params.max_depth_mt_gadget)
+                    .collect(),
+                case_ii_selector: mid_mp.other_leaf.is_some(),
+                other_key,
+                other_value,
+            })
+        }
+    }
+}
+
+impl TryFrom<MerkleClaimAndProof> for merkletree::MerkleProof {
+    type Error = anyhow::Error;
+    fn try_from(mp: MerkleClaimAndProof) -> Result<Self> {
+        if !mp.enabled {
+            return Err(anyhow!("Not a valid Merkle proof."));
+        }
+        let existence = mp.existence;
+        let other_leaf = if mp.case_ii_selector {
+            Some((mp.other_key, mp.other_value))
+        } else {
+            None
+        };
+        // Trim padding (if any).
+        let siblings = mp
+            .siblings
+            .into_iter()
+            .rev()
+            .skip_while(|s| s == &EMPTY_HASH)
+            .collect::<Vec<_>>()
+            .into_iter()
+            .rev()
+            .collect();
+        Ok(merkletree::MerkleProof {
+            existence,
+            siblings,
+            other_leaf,
+        })
+    }
+}
+
+impl fmt::Display for MerkleClaimAndProof {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match merkletree::MerkleProof::try_from(self.clone()) {
+            Err(_) => write!(f, "∅"),
+            Ok(mp) => mp.fmt(f),
+        }
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Operation(pub OperationType, pub Vec<OperationArg>, pub OperationAux);
 
@@ -46,10 +158,13 @@ impl Operation {
     pub fn args(&self) -> &[OperationArg] {
         &self.1
     }
+    pub fn aux(&self) -> &OperationAux {
+        &self.2
+    }
     pub fn deref(
         &self,
         statements: &[Statement],
-        merkle_proofs: &[MerkleProof],
+        merkle_proofs: &[MerkleClaimAndProof],
     ) -> Result<crate::middleware::Operation> {
         let deref_args = self
             .1
@@ -65,7 +180,10 @@ impl Operation {
                 .get(i)
                 .cloned()
                 .ok_or(anyhow!("Missing Merkle proof index {}", i))
-                .map(crate::middleware::OperationAux::MerkleProof),
+                .and_then(|mp| {
+                    mp.try_into()
+                        .map(crate::middleware::OperationAux::MerkleProof)
+                }),
         }?;
         middleware::Operation::op(self.0.clone(), &deref_args, &deref_aux)
     }

src/backends/plonky2/primitives/merkletree.rs

@@ -23,13 +23,21 @@ pub struct MerkleTree {
 impl MerkleTree {
     /// builds a new `MerkleTree` where the leaves contain the given key-values
     pub fn new(max_depth: usize, kvs: &HashMap<Value, Value>) -> Result<Self> {
-        let mut root = Node::Intermediate(Intermediate::empty());
+        // Construct leaves.
+        let mut leaves: Vec<_> = kvs
+            .iter()
+            .map(|(k, v)| Leaf::new(max_depth, *k, *v))
+            .collect::<Result<_>>()?;
 
-        for (k, v) in kvs.iter() {
-            let leaf = Leaf::new(max_depth, *k, *v)?;
+        // Start with a leaf or conclude with an empty node as root.
+        let mut root = leaves.pop().map(|l| Node::Leaf(l)).unwrap_or(Node::None);
+
+        // Iterate over remaining leaves (if any) and add them.
+        for leaf in leaves.into_iter() {
             root.add_leaf(0, max_depth, leaf)?;
         }
 
+        // Fill in hashes.
         let _ = root.compute_hash();
         Ok(Self { max_depth, root })
     }

src/backends/plonky2/primitives/merkletree_circuit.rs

@@ -35,8 +35,9 @@ pub struct MerkleProofGadget {
     pub max_depth: usize,
 }
 
-pub struct MerkleProofTarget {
-    max_depth: usize,
+#[derive(Clone)]
+pub struct MerkleClaimAndProofTarget {
+    pub(crate) max_depth: usize,
     // `enabled` determines if the merkleproof verification is enabled
     pub(crate) enabled: BoolTarget,
     pub(crate) root: HashOutTarget,
@@ -51,7 +52,7 @@ pub struct MerkleProofTarget {
 
 impl MerkleProofGadget {
     /// creates the targets and defines the logic of the circuit
-    pub fn eval(&self, builder: &mut CircuitBuilder<F, D>) -> Result<MerkleProofTarget> {
+    pub fn eval(&self, builder: &mut CircuitBuilder<F, D>) -> Result<MerkleClaimAndProofTarget> {
         let enabled = builder.add_virtual_bool_target_safe();
         let root = builder.add_virtual_hash();
         let key = builder.add_virtual_value();
@@ -133,7 +134,7 @@ impl MerkleProofGadget {
             builder.connect(computed_root[j], expected_root[j]);
         }
 
-        Ok(MerkleProofTarget {
+        Ok(MerkleClaimAndProofTarget {
             max_depth: self.max_depth,
             enabled,
             existence,
@@ -148,7 +149,7 @@ impl MerkleProofGadget {
     }
 }
 
-impl MerkleProofTarget {
+impl MerkleClaimAndProofTarget {
     /// assigns the given values to the targets
     pub fn set_targets(
         &self,