move MainPod layouting to mainpod (#196)

* move MainPod layouting to mainpod

* wip

* use MerkleClaimAndProof in merkle circuit set_targets

* fix empty mt proof

src/backends/plonky2/circuits/common.rs

@@ -19,7 +19,7 @@ use plonky2::{
 use crate::{
     backends::plonky2::{
         basetypes::D,
-        mock::mainpod::{Operation, OperationArg, Statement},
+        mainpod::{Operation, OperationArg, Statement},
         primitives::merkletree::MerkleClaimAndProofTarget,
     },
     middleware::{

src/backends/plonky2/circuits/mainpod.rs

@@ -16,10 +16,9 @@ use crate::{
             },
             signedpod::{SignedPodVerifyGadget, SignedPodVerifyTarget},
         },
-        mock::{mainpod, mainpod::MerkleClaimAndProof},
+        mainpod,
-        primitives::{
+        primitives::merkletree::{
-            merkletree,
+            MerkleClaimAndProof, MerkleClaimAndProofTarget, MerkleProofGadget,
-            merkletree::{MerkleClaimAndProofTarget, MerkleProofGadget},
         },
         signedpod::SignedPod,
     },
@@ -495,20 +494,8 @@ impl MainPodVerifyTarget {
         }
         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);
+            assert_eq!(mp.proof.siblings.len(), self.params.max_depth_mt_gadget);
-            self.merkle_proofs[i].set_targets(
+            self.merkle_proofs[i].set_targets(pw, mp)?;
-                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(())
     }
@@ -531,26 +518,23 @@ impl MainPodVerifyCircuit {
 
 #[cfg(test)]
 mod tests {
-    use merkletree::MerkleTree;
     use plonky2::plonk::{circuit_builder::CircuitBuilder, circuit_data::CircuitConfig};
 
     use super::*;
     use crate::{
         backends::plonky2::{
             basetypes::C,
-            mock::{
-                mainpod,
             mainpod::{OperationArg, OperationAux},
+            primitives::merkletree::{MerkleClaimAndProof, MerkleTree},
         },
-        },
+        middleware::{Hash, OperationType, PodId, RawValue},
-        middleware::{OperationType, PodId, RawValue},
     };
 
     fn operation_verify(
         st: mainpod::Statement,
         op: mainpod::Operation,
         prev_statements: Vec<mainpod::Statement>,
-        merkle_proofs: Vec<mainpod::MerkleClaimAndProof>,
+        merkle_proofs: Vec<MerkleClaimAndProof>,
     ) -> Result<()> {
         let params = Params::default();
         let mp_gadget = MerkleProofGadget {
@@ -595,15 +579,7 @@ mod tests {
         for (merkle_proof_target, merkle_proof) in
             merkle_proofs_target.iter().zip(merkle_proofs.iter())
         {
-            merkle_proof_target.set_targets(
+            merkle_proof_target.set_targets(&mut pw, &merkle_proof)?
-                &mut pw,
-                merkle_proof.enabled,
-                merkle_proof.existence,
-                merkle_proof.root,
-                merkle_proof.clone().try_into()?,
-                merkle_proof.key,
-                merkle_proof.value,
-            )?
         }
 
         // generate & verify proof
@@ -730,10 +706,10 @@ mod tests {
             OperationAux::MerkleProofIndex(0),
         );
 
-        let merkle_proofs = vec![mainpod::MerkleClaimAndProof::try_from_middleware(
+        let merkle_proofs = vec![MerkleClaimAndProof::new(
-            &params,
+            params.max_depth_mt_gadget,
-            &root.raw(),
+            Hash::from(root.raw()),
-            &key,
+            key,
             None,
             &no_key_pf,
         )?];

src/backends/plonky2/circuits/signedpod.rs

@@ -16,14 +16,16 @@ use crate::{
         basetypes::D,
         circuits::common::{CircuitBuilderPod, StatementArgTarget, StatementTarget, ValueTarget},
         primitives::{
-            merkletree::{MerkleProof, MerkleProofExistenceGadget, MerkleProofExistenceTarget},
+            merkletree::{
+                MerkleClaimAndProof, MerkleProofExistenceGadget, MerkleProofExistenceTarget,
+            },
             signature::{PublicKey, SignatureVerifyGadget, SignatureVerifyTarget},
         },
         signedpod::SignedPod,
     },
     middleware::{
-        hash_str, Key, NativePredicate, Params, PodType, Predicate, RawValue, ToFields, Value,
+        hash_str, Key, NativePredicate, Params, PodType, Predicate, RawValue, ToFields, Value, F,
-        EMPTY_VALUE, F, KEY_SIGNER, KEY_TYPE, SELF,
+        KEY_SIGNER, KEY_TYPE, SELF,
     },
 };
 
@@ -132,11 +134,13 @@ impl SignedPodVerifyTarget {
                 let (v, proof) = pod.dict.prove(k)?;
                 self.mt_proofs[i].set_targets(
                     pw,
-                    true,
+                    &MerkleClaimAndProof::new(
+                        self.params.max_depth_mt_gadget,
                         pod.dict.commitment(),
-                    proof,
                         k.raw(),
-                    v.raw(),
+                        Some(v.raw()),
+                        &proof,
+                    )?,
                 )?;
                 Ok(v)
             })
@@ -156,11 +160,13 @@ impl SignedPodVerifyTarget {
 
             self.mt_proofs[curr].set_targets(
                 pw,
-                true,
+                &MerkleClaimAndProof::new(
+                    self.params.max_depth_mt_gadget,
                     pod.dict.commitment(),
-                proof,
                     k.raw(),
-                v.raw(),
+                    Some(v.raw()),
+                    &proof,
+                )?,
             )?;
             curr += 1;
         }
@@ -168,20 +174,10 @@ impl SignedPodVerifyTarget {
         assert!(curr <= self.params.max_signed_pod_values);
 
         // add the proofs of empty leaves (if needed), till the max_signed_pod_values
+        let mut mp = MerkleClaimAndProof::empty(self.params.max_depth_mt_gadget);
+        mp.root = pod.dict.commitment();
         for i in curr..self.params.max_signed_pod_values {
-            self.mt_proofs[i].set_targets(
+            self.mt_proofs[i].set_targets(pw, &mp)?;
-                pw,
-                false, // disable verification
-                pod.dict.commitment(),
-                // use an empty proof:
-                MerkleProof {
-                    existence: true,
-                    siblings: vec![],
-                    other_leaf: None,
-                },
-                EMPTY_VALUE,
-                EMPTY_VALUE,
-            )?;
         }
 
         // get the signer pk

src/backends/plonky2/mainpod.rs

@@ -1,244 +0,0 @@
-use std::any::Any;
-
-use anyhow::{anyhow, Result};
-use itertools::Itertools;
-use plonky2::{
-    iop::witness::PartialWitness,
-    plonk::{
-        circuit_builder::CircuitBuilder, circuit_data::CircuitConfig, proof::ProofWithPublicInputs,
-    },
-};
-
-use crate::{
-    backends::plonky2::{
-        basetypes::{C, D},
-        circuits::mainpod::{MainPodVerifyCircuit, MainPodVerifyInput},
-        mock::mainpod::{hash_statements, MockMainPod, Statement},
-        signedpod::SignedPod,
-    },
-    middleware::{
-        self, AnchoredKey, MainPodInputs, Params, Pod, PodId, PodProver, StatementArg, F, SELF,
-    },
-};
-// TODO: Move the shared components between MockMainPod and MainPod to a common place.
-
-pub struct Prover {}
-
-impl PodProver for Prover {
-    // TODO: Be consistent on where we apply the padding, here, or in the set_targets?
-    fn prove(&mut self, params: &Params, inputs: MainPodInputs) -> Result<Box<dyn Pod>> {
-        let config = CircuitConfig::standard_recursion_config();
-        let mut builder = CircuitBuilder::<F, D>::new(config);
-        let main_pod = MainPodVerifyCircuit {
-            params: params.clone(),
-        }
-        .eval(&mut builder)?;
-
-        let mut pw = PartialWitness::<F>::new();
-        let signed_pods_input: Vec<SignedPod> = inputs
-            .signed_pods
-            .iter()
-            .map(|p| {
-                let p = p
-                    .as_any()
-                    .downcast_ref::<SignedPod>()
-                    .expect("type SignedPod");
-                p.clone()
-            })
-            .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,
-            &merkle_proofs,
-            inputs.operations,
-        )?;
-        let operations =
-            MockMainPod::process_public_statements_operations(params, &statements, operations)?;
-
-        let public_statements =
-            statements[statements.len() - params.max_public_statements..].to_vec();
-        // get the id out of the public statements
-        let id: PodId = PodId(hash_statements(&public_statements, params));
-
-        let input = MainPodVerifyInput {
-            signed_pods: signed_pods_input,
-            statements: statements[statements.len() - params.max_statements..].to_vec(),
-            operations,
-            merkle_proofs,
-        };
-        main_pod.set_targets(&mut pw, &input)?;
-
-        // generate & verify proof
-        let data = builder.build::<C>();
-        let proof = data.prove(pw)?;
-
-        Ok(Box::new(MainPod {
-            params: params.clone(),
-            id,
-            public_statements,
-            proof,
-        }))
-    }
-}
-
-#[derive(Clone, Debug)]
-pub struct MainPod {
-    params: Params,
-    id: PodId,
-    public_statements: Vec<Statement>,
-    proof: ProofWithPublicInputs<F, C, 2>,
-}
-
-impl Pod for MainPod {
-    fn verify(&self) -> Result<()> {
-        // 2. get the id out of the public statements
-        let id: PodId = PodId(hash_statements(&self.public_statements, &self.params));
-        if id != self.id {
-            return Err(anyhow!(
-                "id does not match, expected {}, computed {}",
-                self.id,
-                id
-            ));
-        }
-
-        // 1, 3, 4, 5 verification via the zkSNARK proof
-        // TODO: cache these artefacts
-        let config = CircuitConfig::standard_recursion_config();
-        let mut builder = CircuitBuilder::<F, D>::new(config);
-        let _main_pod = MainPodVerifyCircuit {
-            params: self.params.clone(),
-        }
-        .eval(&mut builder)?;
-
-        let data = builder.build::<C>();
-        data.verify(self.proof.clone())
-            .map_err(|e| anyhow!("MainPod proof verification failure: {:?}", e))
-    }
-
-    fn id(&self) -> PodId {
-        self.id
-    }
-
-    fn pub_statements(&self) -> Vec<middleware::Statement> {
-        // return the public statements, where when origin=SELF is replaced by origin=self.id()
-        self.public_statements
-            .iter()
-            .cloned()
-            .map(|statement| {
-                Statement(
-                    statement.0.clone(),
-                    statement
-                        .1
-                        .iter()
-                        .map(|sa| match &sa {
-                            StatementArg::Key(AnchoredKey { pod_id, key }) if *pod_id == SELF => {
-                                StatementArg::Key(AnchoredKey::new(self.id(), key.clone()))
-                            }
-                            _ => sa.clone(),
-                        })
-                        .collect(),
-                )
-                .try_into()
-                .unwrap()
-            })
-            .collect()
-    }
-
-    fn into_any(self: Box<Self>) -> Box<dyn Any> {
-        self
-    }
-    fn as_any(&self) -> &dyn Any {
-        self
-    }
-
-    fn serialized_proof(&self) -> String {
-        todo!()
-    }
-}
-
-#[cfg(test)]
-pub mod tests {
-    use super::*;
-    use crate::{
-        backends::plonky2::{
-            mock::mainpod::MockProver, primitives::signature::SecretKey, signedpod::Signer,
-        },
-        examples::{zu_kyc_pod_builder, zu_kyc_sign_pod_builders},
-        frontend, middleware,
-        middleware::RawValue,
-        op,
-    };
-
-    #[test]
-    fn test_main_zu_kyc() -> Result<()> {
-        let params = middleware::Params {
-            // Currently the circuit uses random access that only supports vectors of length 64.
-            // With max_input_main_pods=3 we need random access to a vector of length 73.
-            max_input_main_pods: 1,
-            ..Default::default()
-        };
-
-        let (gov_id_builder, pay_stub_builder, sanction_list_builder) =
-            zu_kyc_sign_pod_builders(&params);
-        let mut signer = Signer(SecretKey(RawValue::from(1)));
-        let gov_id_pod = gov_id_builder.sign(&mut signer)?;
-        let mut signer = Signer(SecretKey(RawValue::from(2)));
-        let pay_stub_pod = pay_stub_builder.sign(&mut signer)?;
-        let mut signer = Signer(SecretKey(RawValue::from(3)));
-        let sanction_list_pod = sanction_list_builder.sign(&mut signer)?;
-        let kyc_builder =
-            zu_kyc_pod_builder(&params, &gov_id_pod, &pay_stub_pod, &sanction_list_pod)?;
-
-        let mut prover = Prover {};
-        let kyc_pod = kyc_builder.prove(&mut prover, &params)?;
-        let pod = kyc_pod.pod.into_any().downcast::<MainPod>().unwrap();
-
-        pod.verify()
-    }
-
-    #[test]
-    fn test_mini_0() {
-        let params = middleware::Params {
-            max_input_signed_pods: 1,
-            max_input_main_pods: 1,
-            max_signed_pod_values: 6,
-            max_statements: 8,
-            max_public_statements: 4,
-            ..Default::default()
-        };
-
-        let mut gov_id_builder = frontend::SignedPodBuilder::new(&params);
-        gov_id_builder.insert("idNumber", "4242424242");
-        gov_id_builder.insert("dateOfBirth", 1169909384);
-        gov_id_builder.insert("socialSecurityNumber", "G2121210");
-        let mut signer = Signer(SecretKey(RawValue::from(42)));
-        let gov_id = gov_id_builder.sign(&mut signer).unwrap();
-        let now_minus_18y: i64 = 1169909388;
-        let mut kyc_builder = frontend::MainPodBuilder::new(&params);
-        kyc_builder.add_signed_pod(&gov_id);
-        kyc_builder
-            .pub_op(op!(lt, (&gov_id, "dateOfBirth"), now_minus_18y))
-            .unwrap();
-
-        println!("{}", kyc_builder);
-        println!();
-
-        // Mock
-        let mut prover = MockProver {};
-        let kyc_pod = kyc_builder.prove(&mut prover, &params).unwrap();
-        let pod = kyc_pod.pod.into_any().downcast::<MockMainPod>().unwrap();
-        pod.verify().unwrap();
-        println!("{:#}", pod);
-
-        // Real
-        let mut prover = Prover {};
-        let kyc_pod = kyc_builder.prove(&mut prover, &params).unwrap();
-        let pod = kyc_pod.pod.into_any().downcast::<MainPod>().unwrap();
-        pod.verify().unwrap()
-    }
-}

src/backends/plonky2/mainpod/mod.rs

@@ -0,0 +1,515 @@
+pub mod operation;
+pub mod statement;
+use std::any::Any;
+
+use anyhow::{anyhow, Result};
+use itertools::Itertools;
+pub use operation::*;
+use plonky2::{
+    hash::poseidon::PoseidonHash,
+    iop::witness::PartialWitness,
+    plonk::{
+        circuit_builder::CircuitBuilder, circuit_data::CircuitConfig, config::Hasher,
+        proof::ProofWithPublicInputs,
+    },
+};
+pub use statement::*;
+
+use crate::{
+    backends::plonky2::{
+        basetypes::{C, D},
+        circuits::mainpod::{MainPodVerifyCircuit, MainPodVerifyInput},
+        primitives::{merkletree, merkletree::MerkleClaimAndProof},
+        signedpod::SignedPod,
+    },
+    middleware::{
+        self, AnchoredKey, Hash, MainPodInputs, NativeOperation, NonePod, OperationType, Params,
+        Pod, PodId, PodProver, PodType, StatementArg, ToFields, F, KEY_TYPE, SELF,
+    },
+};
+
+/// Hash a list of public statements to derive the PodId
+pub(crate) fn hash_statements(statements: &[Statement], _params: &Params) -> middleware::Hash {
+    let field_elems = statements
+        .iter()
+        .flat_map(|statement| statement.clone().to_fields(_params))
+        .collect::<Vec<_>>();
+    Hash(PoseidonHash::hash_no_pad(&field_elems).elements)
+}
+
+/// 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::new(
+                params.max_depth_mt_gadget,
+                Hash::from(root.raw()),
+                key.raw(),
+                Some(value.raw()),
+                pf,
+            )),
+            middleware::Operation::NotContainsFromEntries(
+                middleware::Statement::ValueOf(_, root),
+                middleware::Statement::ValueOf(_, key),
+                pf,
+            ) => Some(MerkleClaimAndProof::new(
+                params.max_depth_mt_gadget,
+                Hash::from(root.raw()),
+                key.raw(),
+                None,
+                pf,
+            )),
+            _ => None,
+        })
+        .collect::<Result<Vec<_>>>()?;
+    if merkle_proofs.len() > params.max_merkle_proofs {
+        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)
+    }
+}
+
+/// Find the operation argument statement in the list of previous statements and return the index.
+fn find_op_arg(statements: &[Statement], op_arg: &middleware::Statement) -> Result<OperationArg> {
+    match op_arg {
+        middleware::Statement::None => Ok(OperationArg::None),
+        _ => statements
+            .iter()
+            .enumerate()
+            .find_map(|(i, s)| {
+                (&middleware::Statement::try_from(s.clone()).ok()? == op_arg).then_some(i)
+            })
+            .map(OperationArg::Index)
+            .ok_or(anyhow!(
+                "Statement corresponding to op arg {} not found",
+                op_arg
+            )),
+    }
+}
+
+/// Find the operation auxiliary data in the list of auxiliary data and return the index.
+fn find_op_aux(
+    merkle_proofs: &[MerkleClaimAndProof],
+    op_aux: &middleware::OperationAux,
+) -> Result<OperationAux> {
+    match op_aux {
+        middleware::OperationAux::None => Ok(OperationAux::None),
+        middleware::OperationAux::MerkleProof(pf_arg) => merkle_proofs
+            .iter()
+            .enumerate()
+            .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",
+                op_aux
+            )),
+    }
+}
+
+fn fill_pad<T: Clone>(v: &mut Vec<T>, pad_value: T, len: usize) {
+    if v.len() > len {
+        panic!("length exceeded");
+    }
+    while v.len() < len {
+        v.push(pad_value.clone());
+    }
+}
+
+fn pad_statement(params: &Params, s: &mut Statement) {
+    fill_pad(&mut s.1, StatementArg::None, params.max_statement_args)
+}
+
+fn pad_operation_args(params: &Params, args: &mut Vec<OperationArg>) {
+    fill_pad(args, OperationArg::None, params.max_operation_args)
+}
+
+/// Returns the statements from the given MainPodInputs, padding to the
+/// respective max lengths defined at the given Params.
+pub(crate) fn layout_statements(params: &Params, inputs: &MainPodInputs) -> Vec<Statement> {
+    let mut statements = Vec::new();
+
+    // Input signed pods region
+    let none_sig_pod_box: Box<dyn Pod> = Box::new(NonePod {});
+    let none_sig_pod = none_sig_pod_box.as_ref();
+    assert!(inputs.signed_pods.len() <= params.max_input_signed_pods);
+    for i in 0..params.max_input_signed_pods {
+        let pod = inputs.signed_pods.get(i).unwrap_or(&none_sig_pod);
+        let sts = pod.pub_statements();
+        assert!(sts.len() <= params.max_signed_pod_values);
+        for j in 0..params.max_signed_pod_values {
+            let mut st = sts
+                .get(j)
+                .unwrap_or(&middleware::Statement::None)
+                .clone()
+                .into();
+            pad_statement(params, &mut st);
+            statements.push(st);
+        }
+    }
+
+    // Input main pods region
+    let none_main_pod_box: Box<dyn Pod> = Box::new(NonePod {});
+    let none_main_pod = none_main_pod_box.as_ref();
+    assert!(inputs.main_pods.len() <= params.max_input_main_pods);
+    for i in 0..params.max_input_main_pods {
+        let pod = inputs.main_pods.get(i).copied().unwrap_or(none_main_pod);
+        let sts = pod.pub_statements();
+        assert!(sts.len() <= params.max_public_statements);
+        for j in 0..params.max_public_statements {
+            let mut st = sts
+                .get(j)
+                .unwrap_or(&middleware::Statement::None)
+                .clone()
+                .into();
+            pad_statement(params, &mut st);
+            statements.push(st);
+        }
+    }
+
+    // Input statements
+    assert!(inputs.statements.len() <= params.max_priv_statements());
+    for i in 0..params.max_priv_statements() {
+        let mut st = inputs
+            .statements
+            .get(i)
+            .unwrap_or(&middleware::Statement::None)
+            .clone()
+            .into();
+        pad_statement(params, &mut st);
+        statements.push(st);
+    }
+
+    // Public statements
+    assert!(inputs.public_statements.len() < params.max_public_statements);
+    let mut type_st = middleware::Statement::ValueOf(
+        AnchoredKey::from((SELF, KEY_TYPE)),
+        middleware::Value::from(PodType::MockMain),
+    )
+    .into();
+    pad_statement(params, &mut type_st);
+    statements.push(type_st);
+
+    for i in 0..(params.max_public_statements - 1) {
+        let mut st = inputs
+            .public_statements
+            .get(i)
+            .unwrap_or(&middleware::Statement::None)
+            .clone()
+            .into();
+        pad_statement(params, &mut st);
+        statements.push(st);
+    }
+
+    statements
+}
+
+pub(crate) fn process_private_statements_operations(
+    params: &Params,
+    statements: &[Statement],
+    merkle_proofs: &[MerkleClaimAndProof],
+    input_operations: &[middleware::Operation],
+) -> Result<Vec<Operation>> {
+    let mut operations = Vec::new();
+    for i in 0..params.max_priv_statements() {
+        let op = input_operations
+            .get(i)
+            .unwrap_or(&middleware::Operation::None)
+            .clone();
+        let mid_args = op.args();
+        let mut args = mid_args
+            .iter()
+            .map(|mid_arg| find_op_arg(statements, mid_arg))
+            .collect::<Result<Vec<_>>>()?;
+
+        let mid_aux = op.aux();
+        let aux = find_op_aux(merkle_proofs, &mid_aux)?;
+
+        pad_operation_args(params, &mut args);
+        operations.push(Operation(op.op_type(), args, aux));
+    }
+    Ok(operations)
+}
+
+// NOTE: In this implementation public statements are always copies from
+// previous statements, so we fill in the operations accordingly.
+/// This method assumes that the given `statements` array has been padded to
+/// `params.max_statements`.
+pub(crate) fn process_public_statements_operations(
+    params: &Params,
+    statements: &[Statement],
+    mut operations: Vec<Operation>,
+) -> Result<Vec<Operation>> {
+    let offset_public_statements = statements.len() - params.max_public_statements;
+    operations.push(Operation(
+        OperationType::Native(NativeOperation::NewEntry),
+        vec![],
+        OperationAux::None,
+    ));
+    for i in 0..(params.max_public_statements - 1) {
+        let st = &statements[offset_public_statements + i + 1];
+        let mut op = if st.is_none() {
+            Operation(
+                OperationType::Native(NativeOperation::None),
+                vec![],
+                OperationAux::None,
+            )
+        } else {
+            let mid_arg = st.clone();
+            Operation(
+                OperationType::Native(NativeOperation::CopyStatement),
+                vec![find_op_arg(statements, &mid_arg.try_into()?)?],
+                OperationAux::None,
+            )
+        };
+        fill_pad(&mut op.1, OperationArg::None, params.max_operation_args);
+        operations.push(op);
+    }
+    Ok(operations)
+}
+
+pub struct Prover {}
+
+impl PodProver for Prover {
+    // TODO: Be consistent on where we apply the padding, here, or in the set_targets?
+    fn prove(&mut self, params: &Params, inputs: MainPodInputs) -> Result<Box<dyn Pod>> {
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+        let main_pod = MainPodVerifyCircuit {
+            params: params.clone(),
+        }
+        .eval(&mut builder)?;
+
+        let mut pw = PartialWitness::<F>::new();
+        let signed_pods_input: Vec<SignedPod> = inputs
+            .signed_pods
+            .iter()
+            .map(|p| {
+                let p = p
+                    .as_any()
+                    .downcast_ref::<SignedPod>()
+                    .expect("type SignedPod");
+                p.clone()
+            })
+            .collect_vec();
+
+        let merkle_proofs = extract_merkle_proofs(params, inputs.operations)?;
+
+        let statements = layout_statements(params, &inputs);
+        let operations = process_private_statements_operations(
+            params,
+            &statements,
+            &merkle_proofs,
+            inputs.operations,
+        )?;
+        let operations = process_public_statements_operations(params, &statements, operations)?;
+
+        let public_statements =
+            statements[statements.len() - params.max_public_statements..].to_vec();
+        // get the id out of the public statements
+        let id: PodId = PodId(hash_statements(&public_statements, params));
+
+        let input = MainPodVerifyInput {
+            signed_pods: signed_pods_input,
+            statements: statements[statements.len() - params.max_statements..].to_vec(),
+            operations,
+            merkle_proofs,
+        };
+        main_pod.set_targets(&mut pw, &input)?;
+
+        // generate & verify proof
+        let data = builder.build::<C>();
+        let proof = data.prove(pw)?;
+
+        Ok(Box::new(MainPod {
+            params: params.clone(),
+            id,
+            public_statements,
+            proof,
+        }))
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct MainPod {
+    params: Params,
+    id: PodId,
+    public_statements: Vec<Statement>,
+    proof: ProofWithPublicInputs<F, C, 2>,
+}
+
+/// Convert a Statement into middleware::Statement and replace references to SELF by `self_id`.
+pub(crate) fn normalize_statement(statement: &Statement, self_id: PodId) -> middleware::Statement {
+    Statement(
+        statement.0.clone(),
+        statement
+            .1
+            .iter()
+            .map(|sa| match &sa {
+                StatementArg::Key(AnchoredKey { pod_id, key }) if *pod_id == SELF => {
+                    StatementArg::Key(AnchoredKey::new(self_id, key.clone()))
+                }
+                _ => sa.clone(),
+            })
+            .collect(),
+    )
+    .try_into()
+    .unwrap()
+}
+
+impl Pod for MainPod {
+    fn verify(&self) -> Result<()> {
+        // 2. get the id out of the public statements
+        let id: PodId = PodId(hash_statements(&self.public_statements, &self.params));
+        if id != self.id {
+            return Err(anyhow!(
+                "id does not match, expected {}, computed {}",
+                self.id,
+                id
+            ));
+        }
+
+        // 1, 3, 4, 5 verification via the zkSNARK proof
+        // TODO: cache these artefacts
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+        let _main_pod = MainPodVerifyCircuit {
+            params: self.params.clone(),
+        }
+        .eval(&mut builder)?;
+
+        let data = builder.build::<C>();
+        data.verify(self.proof.clone())
+            .map_err(|e| anyhow!("MainPod proof verification failure: {:?}", e))
+    }
+
+    fn id(&self) -> PodId {
+        self.id
+    }
+
+    fn pub_statements(&self) -> Vec<middleware::Statement> {
+        // return the public statements, where when origin=SELF is replaced by origin=self.id()
+        self.public_statements
+            .iter()
+            .cloned()
+            .map(|statement| normalize_statement(&statement, self.id()))
+            .collect()
+    }
+
+    fn into_any(self: Box<Self>) -> Box<dyn Any> {
+        self
+    }
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn serialized_proof(&self) -> String {
+        todo!()
+    }
+}
+
+#[cfg(test)]
+pub mod tests {
+    use super::*;
+    use crate::{
+        backends::plonky2::{
+            mock::mainpod::{MockMainPod, MockProver},
+            primitives::signature::SecretKey,
+            signedpod::Signer,
+        },
+        examples::{zu_kyc_pod_builder, zu_kyc_sign_pod_builders},
+        frontend, middleware,
+        middleware::RawValue,
+        op,
+    };
+
+    #[test]
+    fn test_main_zu_kyc() -> Result<()> {
+        let params = middleware::Params {
+            // Currently the circuit uses random access that only supports vectors of length 64.
+            // With max_input_main_pods=3 we need random access to a vector of length 73.
+            max_input_main_pods: 1,
+            ..Default::default()
+        };
+
+        let (gov_id_builder, pay_stub_builder, sanction_list_builder) =
+            zu_kyc_sign_pod_builders(&params);
+        let mut signer = Signer(SecretKey(RawValue::from(1)));
+        let gov_id_pod = gov_id_builder.sign(&mut signer)?;
+        let mut signer = Signer(SecretKey(RawValue::from(2)));
+        let pay_stub_pod = pay_stub_builder.sign(&mut signer)?;
+        let mut signer = Signer(SecretKey(RawValue::from(3)));
+        let sanction_list_pod = sanction_list_builder.sign(&mut signer)?;
+        let kyc_builder =
+            zu_kyc_pod_builder(&params, &gov_id_pod, &pay_stub_pod, &sanction_list_pod)?;
+
+        let mut prover = Prover {};
+        let kyc_pod = kyc_builder.prove(&mut prover, &params)?;
+        let pod = kyc_pod.pod.into_any().downcast::<MainPod>().unwrap();
+
+        pod.verify()
+    }
+
+    #[test]
+    fn test_mini_0() {
+        let params = middleware::Params {
+            max_input_signed_pods: 1,
+            max_input_main_pods: 1,
+            max_signed_pod_values: 6,
+            max_statements: 8,
+            max_public_statements: 4,
+            ..Default::default()
+        };
+
+        let mut gov_id_builder = frontend::SignedPodBuilder::new(&params);
+        gov_id_builder.insert("idNumber", "4242424242");
+        gov_id_builder.insert("dateOfBirth", 1169909384);
+        gov_id_builder.insert("socialSecurityNumber", "G2121210");
+        let mut signer = Signer(SecretKey(RawValue::from(42)));
+        let gov_id = gov_id_builder.sign(&mut signer).unwrap();
+        let now_minus_18y: i64 = 1169909388;
+        let mut kyc_builder = frontend::MainPodBuilder::new(&params);
+        kyc_builder.add_signed_pod(&gov_id);
+        kyc_builder
+            .pub_op(op!(lt, (&gov_id, "dateOfBirth"), now_minus_18y))
+            .unwrap();
+
+        println!("{}", kyc_builder);
+        println!();
+
+        // Mock
+        let mut prover = MockProver {};
+        let kyc_pod = kyc_builder.prove(&mut prover, &params).unwrap();
+        let pod = kyc_pod.pod.into_any().downcast::<MockMainPod>().unwrap();
+        pod.verify().unwrap();
+        println!("{:#}", pod);
+
+        // Real
+        let mut prover = Prover {};
+        let kyc_pod = kyc_builder.prove(&mut prover, &params).unwrap();
+        let pod = kyc_pod.pod.into_any().downcast::<MainPod>().unwrap();
+        pod.verify().unwrap()
+    }
+}

src/backends/plonky2/mainpod/operation.rs

@@ -0,0 +1,111 @@
+use std::fmt;
+
+use anyhow::{anyhow, Result};
+use plonky2::field::types::Field;
+
+// use serde::{Deserialize, Serialize};
+use crate::{
+    backends::plonky2::{mainpod::Statement, primitives::merkletree::MerkleClaimAndProof},
+    middleware::{self, OperationType, Params, ToFields, F},
+};
+
+#[derive(Clone, Debug, PartialEq)]
+pub enum OperationArg {
+    None,
+    Index(usize),
+}
+
+impl ToFields for OperationArg {
+    fn to_fields(&self, _params: &Params) -> Vec<F> {
+        let f = match self {
+            Self::None => F::ZERO,
+            Self::Index(i) => F::from_canonical_usize(*i),
+        };
+        vec![f]
+    }
+}
+
+impl OperationArg {
+    pub fn is_none(&self) -> bool {
+        matches!(self, OperationArg::None)
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub enum OperationAux {
+    None,
+    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)]
+pub struct Operation(pub OperationType, pub Vec<OperationArg>, pub OperationAux);
+
+impl Operation {
+    pub fn op_type(&self) -> OperationType {
+        self.0.clone()
+    }
+    pub fn args(&self) -> &[OperationArg] {
+        &self.1
+    }
+    pub fn aux(&self) -> &OperationAux {
+        &self.2
+    }
+    pub fn deref(
+        &self,
+        statements: &[Statement],
+        merkle_proofs: &[MerkleClaimAndProof],
+    ) -> Result<crate::middleware::Operation> {
+        let deref_args = self
+            .1
+            .iter()
+            .flat_map(|arg| match arg {
+                OperationArg::None => None,
+                OperationArg::Index(i) => Some(statements[*i].clone().try_into()),
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let deref_aux = match self.2 {
+            OperationAux::None => Ok(crate::middleware::OperationAux::None),
+            OperationAux::MerkleProofIndex(i) => merkle_proofs
+                .get(i)
+                .cloned()
+                .ok_or(anyhow!("Missing Merkle proof index {}", i))
+                .and_then(|mp| {
+                    mp.try_into()
+                        .map(crate::middleware::OperationAux::MerkleProof)
+                }),
+        }?;
+        middleware::Operation::op(self.0.clone(), &deref_args, &deref_aux)
+    }
+}
+
+impl fmt::Display for Operation {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{:?} ", self.0)?;
+        for (i, arg) in self.1.iter().enumerate() {
+            if f.alternate() || !arg.is_none() {
+                if i != 0 {
+                    write!(f, " ")?;
+                }
+                match arg {
+                    OperationArg::None => write!(f, "none")?,
+                    OperationArg::Index(i) => write!(f, "{:02}", i)?,
+                }
+            }
+        }
+        match self.2 {
+            OperationAux::None => (),
+            OperationAux::MerkleProofIndex(i) => write!(f, " merkle_proof_{:02}", i)?,
+        }
+        Ok(())
+    }
+}

src/backends/plonky2/mock/mainpod.rs

@@ -1,24 +1,28 @@
+//
+// MainPod
+//
+
 use std::{any::Any, fmt};
 
 use anyhow::{anyhow, Result};
-// use base64::prelude::*;
-use plonky2::{hash::poseidon::PoseidonHash, plonk::config::Hasher};
 
+// use base64::prelude::*;
 // use serde::{Deserialize, Serialize};
+use crate::backends::plonky2::mainpod::process_private_statements_operations;
 use crate::{
-    backends::plonky2::primitives::merkletree,
+    backends::plonky2::{
+        mainpod::{
+            extract_merkle_proofs, hash_statements, layout_statements, normalize_statement,
+            process_public_statements_operations, Operation, Statement,
+        },
+        primitives::merkletree::MerkleClaimAndProof,
+    },
     middleware::{
-        self, hash_str, AnchoredKey, Hash, MainPodInputs, NativeOperation, NativePredicate,
+        self, hash_str, AnchoredKey, MainPodInputs, NativePredicate, Params, Pod, PodId, PodProver,
-        NonePod, OperationType, Params, Pod, PodId, PodProver, PodType, Predicate, StatementArg,
+        Predicate, StatementArg, KEY_TYPE, SELF,
-        ToFields, KEY_TYPE, SELF,
     },
 };
 
-mod operation;
-mod statement;
-pub use operation::*;
-pub use statement::*;
-
 pub struct MockProver {}
 
 impl PodProver for MockProver {
@@ -111,15 +115,6 @@ fn fmt_statement_index(
     Ok(())
 }
 
-pub fn fill_pad<T: Clone>(v: &mut Vec<T>, pad_value: T, len: usize) {
-    if v.len() > len {
-        panic!("length exceeded");
-    }
-    while v.len() < len {
-        v.push(pad_value.clone());
-    }
-}
-
 /// Inputs are sorted as:
 /// - SignedPods
 /// - MainPods
@@ -136,267 +131,21 @@ impl MockMainPod {
     fn offset_public_statements(&self) -> usize {
         self.offset_input_statements() + self.params.max_priv_statements()
     }
-    fn pad_statement(params: &Params, s: &mut Statement) {
-        fill_pad(&mut s.1, StatementArg::None, params.max_statement_args)
-    }
-
-    /// Returns the statements from the given MainPodInputs, padding to the
-    /// respective max lengths defined at the given Params.
-    pub(crate) fn layout_statements(params: &Params, inputs: &MainPodInputs) -> Vec<Statement> {
-        let mut statements = Vec::new();
-
-        // Input signed pods region
-        let none_sig_pod_box: Box<dyn Pod> = Box::new(NonePod {});
-        let none_sig_pod = none_sig_pod_box.as_ref();
-        assert!(inputs.signed_pods.len() <= params.max_input_signed_pods);
-        for i in 0..params.max_input_signed_pods {
-            let pod = inputs.signed_pods.get(i).unwrap_or(&none_sig_pod);
-            let sts = pod.pub_statements();
-            assert!(sts.len() <= params.max_signed_pod_values);
-            for j in 0..params.max_signed_pod_values {
-                let mut st = sts
-                    .get(j)
-                    .unwrap_or(&middleware::Statement::None)
-                    .clone()
-                    .into();
-                Self::pad_statement(params, &mut st);
-                statements.push(st);
-            }
-        }
-
-        // Input main pods region
-        let none_main_pod_box: Box<dyn Pod> = Box::new(NonePod {});
-        let none_main_pod = none_main_pod_box.as_ref();
-        assert!(inputs.main_pods.len() <= params.max_input_main_pods);
-        for i in 0..params.max_input_main_pods {
-            let pod = inputs.main_pods.get(i).copied().unwrap_or(none_main_pod);
-            let sts = pod.pub_statements();
-            assert!(sts.len() <= params.max_public_statements);
-            for j in 0..params.max_public_statements {
-                let mut st = sts
-                    .get(j)
-                    .unwrap_or(&middleware::Statement::None)
-                    .clone()
-                    .into();
-                Self::pad_statement(params, &mut st);
-                statements.push(st);
-            }
-        }
-
-        // Input statements
-        assert!(inputs.statements.len() <= params.max_priv_statements());
-        for i in 0..params.max_priv_statements() {
-            let mut st = inputs
-                .statements
-                .get(i)
-                .unwrap_or(&middleware::Statement::None)
-                .clone()
-                .into();
-            Self::pad_statement(params, &mut st);
-            statements.push(st);
-        }
-
-        // Public statements
-        assert!(inputs.public_statements.len() < params.max_public_statements);
-        let mut type_st = middleware::Statement::ValueOf(
-            AnchoredKey::from((SELF, KEY_TYPE)),
-            middleware::Value::from(PodType::MockMain),
-        )
-        .into();
-        Self::pad_statement(params, &mut type_st);
-        statements.push(type_st);
-
-        for i in 0..(params.max_public_statements - 1) {
-            let mut st = inputs
-                .public_statements
-                .get(i)
-                .unwrap_or(&middleware::Statement::None)
-                .clone()
-                .into();
-            Self::pad_statement(params, &mut st);
-            statements.push(st);
-        }
-
-        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.raw(),
-                    &key.raw(),
-                    Some(&value.raw()),
-                    pf,
-                )),
-                middleware::Operation::NotContainsFromEntries(
-                    middleware::Statement::ValueOf(_, root),
-                    middleware::Statement::ValueOf(_, key),
-                    pf,
-                ) => Some(MerkleClaimAndProof::try_from_middleware(
-                    params,
-                    &root.raw(),
-                    &key.raw(),
-                    None,
-                    pf,
-                )),
-                _ => None,
-            })
-            .collect::<Result<Vec<_>>>()?;
-        if merkle_proofs.len() > params.max_merkle_proofs {
-            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,
-    ) -> Result<OperationArg> {
-        match op_arg {
-            middleware::Statement::None => Ok(OperationArg::None),
-            _ => statements
-                .iter()
-                .enumerate()
-                .find_map(|(i, s)| {
-                    (&middleware::Statement::try_from(s.clone()).ok()? == op_arg).then_some(i)
-                })
-                .map(OperationArg::Index)
-                .ok_or(anyhow!(
-                    "Statement corresponding to op arg {} not found",
-                    op_arg
-                )),
-        }
-    }
-
-    fn find_op_aux(
-        merkle_proofs: &[MerkleClaimAndProof],
-        op_aux: &middleware::OperationAux,
-    ) -> Result<OperationAux> {
-        match op_aux {
-            middleware::OperationAux::None => Ok(OperationAux::None),
-            middleware::OperationAux::MerkleProof(pf_arg) => merkle_proofs
-                .iter()
-                .enumerate()
-                .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",
-                    op_aux
-                )),
-        }
-    }
-
-    pub(crate) fn process_private_statements_operations(
-        params: &Params,
-        statements: &[Statement],
-        merkle_proofs: &[MerkleClaimAndProof],
-        input_operations: &[middleware::Operation],
-    ) -> Result<Vec<Operation>> {
-        let mut operations = Vec::new();
-        for i in 0..params.max_priv_statements() {
-            let op = input_operations
-                .get(i)
-                .unwrap_or(&middleware::Operation::None)
-                .clone();
-            let mid_args = op.args();
-            let mut args = mid_args
-                .iter()
-                .map(|mid_arg| Self::find_op_arg(statements, mid_arg))
-                .collect::<Result<Vec<_>>>()?;
-
-            let mid_aux = op.aux();
-            let aux = Self::find_op_aux(merkle_proofs, &mid_aux)?;
-
-            Self::pad_operation_args(params, &mut args);
-            operations.push(Operation(op.op_type(), args, aux));
-        }
-        Ok(operations)
-    }
-
-    // NOTE: In this implementation public statements are always copies from
-    // previous statements, so we fill in the operations accordingly.
-    /// This method assumes that the given `statements` array has been padded to
-    /// `params.max_statements`.
-    pub(crate) fn process_public_statements_operations(
-        params: &Params,
-        statements: &[Statement],
-        mut operations: Vec<Operation>,
-    ) -> Result<Vec<Operation>> {
-        let offset_public_statements = statements.len() - params.max_public_statements;
-        operations.push(Operation(
-            OperationType::Native(NativeOperation::NewEntry),
-            vec![],
-            OperationAux::None,
-        ));
-        for i in 0..(params.max_public_statements - 1) {
-            let st = &statements[offset_public_statements + i + 1];
-            let mut op = if st.is_none() {
-                Operation(
-                    OperationType::Native(NativeOperation::None),
-                    vec![],
-                    OperationAux::None,
-                )
-            } else {
-                let mid_arg = st.clone();
-                Operation(
-                    OperationType::Native(NativeOperation::CopyStatement),
-                    vec![Self::find_op_arg(statements, &mid_arg.try_into()?)?],
-                    OperationAux::None,
-                )
-            };
-            fill_pad(&mut op.1, OperationArg::None, params.max_operation_args);
-            operations.push(op);
-        }
-        Ok(operations)
-    }
 
     pub fn new(params: &Params, inputs: MainPodInputs) -> Result<Self> {
-        // TODO: Figure out a way to handle public statements.  For example, in the public slots
-        // use copy operations taking the private statements that need to be public.  We may change
-        // the MainPodInputs type to accommodate for that.
         // TODO: Insert a new public statement of ValueOf with `key=KEY_TYPE,
         // value=PodType::MockMainPod`
-        let statements = Self::layout_statements(params, &inputs);
+        let statements = layout_statements(params, &inputs);
         // Extract Merkle proofs and pad.
-        let merkle_proofs = Self::extract_merkle_proofs(params, inputs.operations)?;
+        let merkle_proofs = extract_merkle_proofs(params, inputs.operations)?;
 
-        let operations = Self::process_private_statements_operations(
+        let operations = process_private_statements_operations(
             params,
             &statements,
             &merkle_proofs,
             inputs.operations,
         )?;
-        let operations =
+        let operations = process_public_statements_operations(params, &statements, operations)?;
-            Self::process_public_statements_operations(params, &statements, operations)?;
 
         let public_statements =
             statements[statements.len() - params.max_public_statements..].to_vec();
@@ -417,10 +166,6 @@ impl MockMainPod {
         })
     }
 
-    fn pad_operation_args(params: &Params, args: &mut Vec<OperationArg>) {
-        fill_pad(args, OperationArg::None, params.max_operation_args)
-    }
-
     // pub fn deserialize(serialized: String) -> Result<Self> {
     //     let proof = String::from_utf8(BASE64_STANDARD.decode(&serialized)?)
     //         .map_err(|e| anyhow::anyhow!("Invalid base64 encoding: {}", e))?;
@@ -431,14 +176,6 @@ impl MockMainPod {
     // }
 }
 
-pub fn hash_statements(statements: &[Statement], _params: &Params) -> middleware::Hash {
-    let field_elems = statements
-        .iter()
-        .flat_map(|statement| statement.clone().to_fields(_params))
-        .collect::<Vec<_>>();
-    Hash(PoseidonHash::hash_no_pad(&field_elems).elements)
-}
-
 impl Pod for MockMainPod {
     fn verify(&self) -> Result<()> {
         // 1. TODO: Verify input pods
@@ -533,23 +270,7 @@ impl Pod for MockMainPod {
             .iter()
             .skip(self.offset_public_statements())
             .cloned()
-            .map(|statement| {
+            .map(|statement| normalize_statement(&statement, self.id()))
-                Statement(
-                    statement.0.clone(),
-                    statement
-                        .1
-                        .iter()
-                        .map(|sa| match &sa {
-                            StatementArg::Key(AnchoredKey { pod_id, key }) if *pod_id == SELF => {
-                                StatementArg::Key(AnchoredKey::new(self.id(), key.clone()))
-                            }
-                            _ => sa.clone(),
-                        })
-                        .collect(),
-                )
-                .try_into()
-                .unwrap()
-            })
             .collect()
     }
 

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

@@ -1,218 +0,0 @@
-use std::{fmt, iter};
-
-use anyhow::{anyhow, Result};
-use plonky2::field::types::Field;
-
-// use serde::{Deserialize, Serialize};
-use crate::{
-    backends::plonky2::{
-        mock::mainpod::Statement,
-        primitives::merkletree::{self},
-    },
-    middleware::{
-        self, Hash, OperationType, Params, RawValue, ToFields, EMPTY_HASH, EMPTY_VALUE, F,
-    },
-};
-
-#[derive(Clone, Debug, PartialEq)]
-pub enum OperationArg {
-    None,
-    Index(usize),
-}
-
-impl ToFields for OperationArg {
-    fn to_fields(&self, _params: &Params) -> Vec<F> {
-        let f = match self {
-            Self::None => F::ZERO,
-            Self::Index(i) => F::from_canonical_usize(*i),
-        };
-        vec![f]
-    }
-}
-
-impl OperationArg {
-    pub fn is_none(&self) -> bool {
-        matches!(self, OperationArg::None)
-    }
-}
-
-#[derive(Clone, Debug, PartialEq)]
-pub enum OperationAux {
-    None,
-    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)]
-pub struct MerkleClaimAndProof {
-    pub enabled: bool,
-    pub root: Hash,
-    pub key: RawValue,
-    pub value: RawValue,
-    pub existence: bool,
-    pub siblings: Vec<Hash>,
-    pub case_ii_selector: bool,
-    pub other_key: RawValue,
-    pub other_value: RawValue,
-}
-
-impl MerkleClaimAndProof {
-    pub fn empty(max_depth: usize) -> Self {
-        Self {
-            enabled: false,
-            root: EMPTY_HASH,
-            key: RawValue::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: &RawValue,
-        key: &RawValue,
-        value: Option<&RawValue>,
-        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).into(),
-                key: *key,
-                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)]
-pub struct Operation(pub OperationType, pub Vec<OperationArg>, pub OperationAux);
-
-impl Operation {
-    pub fn op_type(&self) -> OperationType {
-        self.0.clone()
-    }
-    pub fn args(&self) -> &[OperationArg] {
-        &self.1
-    }
-    pub fn aux(&self) -> &OperationAux {
-        &self.2
-    }
-    pub fn deref(
-        &self,
-        statements: &[Statement],
-        merkle_proofs: &[MerkleClaimAndProof],
-    ) -> Result<crate::middleware::Operation> {
-        let deref_args = self
-            .1
-            .iter()
-            .flat_map(|arg| match arg {
-                OperationArg::None => None,
-                OperationArg::Index(i) => Some(statements[*i].clone().try_into()),
-            })
-            .collect::<Result<Vec<_>>>()?;
-        let deref_aux = match self.2 {
-            OperationAux::None => Ok(crate::middleware::OperationAux::None),
-            OperationAux::MerkleProofIndex(i) => merkle_proofs
-                .get(i)
-                .cloned()
-                .ok_or(anyhow!("Missing Merkle proof index {}", i))
-                .and_then(|mp| {
-                    mp.try_into()
-                        .map(crate::middleware::OperationAux::MerkleProof)
-                }),
-        }?;
-        middleware::Operation::op(self.0.clone(), &deref_args, &deref_aux)
-    }
-}
-
-impl fmt::Display for Operation {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{:?} ", self.0)?;
-        for (i, arg) in self.1.iter().enumerate() {
-            if f.alternate() || !arg.is_none() {
-                if i != 0 {
-                    write!(f, " ")?;
-                }
-                match arg {
-                    OperationArg::None => write!(f, "none")?,
-                    OperationArg::Index(i) => write!(f, "{:02}", i)?,
-                }
-            }
-        }
-        match self.2 {
-            OperationAux::None => (),
-            OperationAux::MerkleProofIndex(i) => write!(f, " merkle_proof_{:02}", i)?,
-        }
-        Ok(())
-    }
-}

src/backends/plonky2/primitives/merkletree.rs

@@ -1,13 +1,17 @@
 //! Module that implements the MerkleTree specified at
 //! https://0xparc.github.io/pod2/merkletree.html .
-use std::{collections::HashMap, fmt, iter::IntoIterator};
+use std::{
+    collections::HashMap,
+    fmt,
+    iter::{self, IntoIterator},
+};
 
 use anyhow::{anyhow, Result};
 use plonky2::field::types::Field;
 
 // use serde::{Deserialize, Serialize};
 pub use super::merkletree_circuit::*;
-use crate::middleware::{hash_fields, Hash, RawValue, EMPTY_HASH, F};
+use crate::middleware::{hash_fields, Hash, RawValue, EMPTY_HASH, EMPTY_VALUE, F};
 
 /// Implements the MerkleTree specified at
 /// https://0xparc.github.io/pod2/merkletree.html
@@ -260,6 +264,99 @@ impl MerkleProof {
     }
 }
 
+#[derive(Clone, Debug, PartialEq)]
+pub struct MerkleClaimAndProof {
+    /// `enabled` determines if the merkleproof verification is enabled
+    pub enabled: bool,
+    pub root: Hash,
+    pub key: RawValue,
+    pub value: RawValue,
+    /// The siblings in this proof are padded to max_depth
+    pub proof: MerkleProof,
+}
+
+impl MerkleClaimAndProof {
+    pub fn empty(max_depth: usize) -> Self {
+        Self {
+            enabled: false,
+            root: EMPTY_HASH,
+            key: EMPTY_VALUE,
+            value: EMPTY_VALUE,
+            proof: MerkleProof {
+                existence: true,
+                siblings: iter::repeat(EMPTY_HASH).take(max_depth).collect(),
+                other_leaf: None,
+            },
+        }
+    }
+    pub fn new(
+        max_depth: usize,
+        root: Hash,
+        key: RawValue,
+        value: Option<RawValue>,
+        proof: &MerkleProof,
+    ) -> Result<Self> {
+        if proof.siblings.len() > max_depth {
+            Err(anyhow!(
+                "Number of siblings ({}) exceeds maximum depth ({})",
+                proof.siblings.len(),
+                max_depth
+            ))
+        } else {
+            Ok(Self {
+                enabled: true,
+                root,
+                key,
+                value: value.unwrap_or(EMPTY_VALUE),
+                proof: MerkleProof {
+                    existence: proof.existence,
+                    siblings: proof
+                        .siblings
+                        .iter()
+                        .cloned()
+                        .chain(iter::repeat(EMPTY_HASH))
+                        .take(max_depth)
+                        .collect(),
+                    other_leaf: proof.other_leaf,
+                },
+            })
+        }
+    }
+}
+
+impl TryFrom<MerkleClaimAndProof> for MerkleProof {
+    type Error = anyhow::Error;
+    fn try_from(mp: MerkleClaimAndProof) -> Result<Self> {
+        if !mp.enabled {
+            return Err(anyhow!("Not a valid Merkle proof."));
+        }
+        Ok(MerkleProof {
+            existence: mp.proof.existence,
+            // Trim padding (if any).
+            siblings: mp
+                .proof
+                .siblings
+                .into_iter()
+                .rev()
+                .skip_while(|s| s == &EMPTY_HASH)
+                .collect::<Vec<_>>()
+                .into_iter()
+                .rev()
+                .collect(),
+            other_leaf: mp.proof.other_leaf,
+        })
+    }
+}
+
+impl fmt::Display for MerkleClaimAndProof {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match MerkleProof::try_from(self.clone()) {
+            Err(_) => write!(f, "∅"),
+            Ok(mp) => mp.fmt(f),
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 enum Node {
     None,

src/backends/plonky2/primitives/merkletree_circuit.rs

@@ -28,9 +28,9 @@ use crate::{
     backends::plonky2::{
         basetypes::D,
         circuits::common::{CircuitBuilderPod, ValueTarget},
-        primitives::merkletree::MerkleProof,
+        primitives::merkletree::MerkleClaimAndProof,
     },
-    middleware::{Hash, RawValue, EMPTY_HASH, EMPTY_VALUE, F, HASH_SIZE},
+    middleware::{EMPTY_HASH, EMPTY_VALUE, F, HASH_SIZE},
 };
 
 /// `MerkleProofGadget` allows to verify both proofs of existence and proofs
@@ -158,34 +158,20 @@ impl MerkleProofGadget {
 impl MerkleClaimAndProofTarget {
     /// assigns the given values to the targets
     #[allow(clippy::too_many_arguments)]
-    pub fn set_targets(
+    pub fn set_targets(&self, pw: &mut PartialWitness<F>, mp: &MerkleClaimAndProof) -> Result<()> {
-        &self,
+        pw.set_bool_target(self.enabled, mp.enabled)?;
-        pw: &mut PartialWitness<F>,
+        pw.set_hash_target(self.root, HashOut::from_vec(mp.root.0.to_vec()))?;
-        // `enabled` determines if the merkleproof verification is enabled
+        pw.set_target_arr(&self.key.elements, &mp.key.0)?;
-        enabled: bool,
+        pw.set_target_arr(&self.value.elements, &mp.value.0)?;
-        existence: bool,
+        pw.set_bool_target(self.existence, mp.proof.existence)?;
-        root: Hash,
-        proof: MerkleProof,
-        key: RawValue,
-        value: RawValue,
-    ) -> Result<()> {
-        pw.set_bool_target(self.enabled, enabled)?;
-        pw.set_hash_target(self.root, HashOut::from_vec(root.0.to_vec()))?;
-        pw.set_target_arr(&self.key.elements, &key.0)?;
-        pw.set_target_arr(&self.value.elements, &value.0)?;
-        pw.set_bool_target(self.existence, existence)?;
 
-        // pad siblings with zeros to length max_depth
+        assert_eq!(mp.proof.siblings.len(), self.max_depth);
-        let mut siblings = proof.siblings.clone();
+        for (i, sibling) in mp.proof.siblings.iter().enumerate() {
-        siblings.resize(self.max_depth, EMPTY_HASH);
-        assert_eq!(self.siblings.len(), siblings.len());
-
-        for (i, sibling) in siblings.iter().enumerate() {
             pw.set_hash_target(self.siblings[i], HashOut::from_vec(sibling.0.to_vec()))?;
         }
 
-        match proof.other_leaf {
+        match mp.proof.other_leaf {
-            Some((k, v)) if !existence => {
+            Some((k, v)) if !mp.proof.existence => {
                 // non-existence case ii) expected leaf does exist but it has a different key
                 pw.set_bool_target(self.case_ii_selector, true)?;
                 pw.set_target_arr(&self.other_key.elements, &k.0)?;
@@ -264,29 +250,18 @@ impl MerkleProofExistenceGadget {
 
 impl MerkleProofExistenceTarget {
     /// assigns the given values to the targets
-    pub fn set_targets(
+    pub fn set_targets(&self, pw: &mut PartialWitness<F>, mp: &MerkleClaimAndProof) -> Result<()> {
-        &self,
+        assert!(mp.proof.existence); // sanity check
-        pw: &mut PartialWitness<F>,
-        // `enabled` determines if the merkleproof verification is enabled
-        enabled: bool,
-        root: Hash,
-        proof: MerkleProof,
-        key: RawValue,
-        value: RawValue,
-    ) -> Result<()> {
-        assert!(proof.existence); // sanity check
 
-        pw.set_bool_target(self.enabled, enabled)?;
+        pw.set_bool_target(self.enabled, mp.enabled)?;
-        pw.set_hash_target(self.root, HashOut::from_vec(root.0.to_vec()))?;
+        pw.set_hash_target(self.root, HashOut::from_vec(mp.root.0.to_vec()))?;
-        pw.set_target_arr(&self.key.elements, &key.0)?;
+        pw.set_target_arr(&self.key.elements, &mp.key.0)?;
-        pw.set_target_arr(&self.value.elements, &value.0)?;
+        pw.set_target_arr(&self.value.elements, &mp.value.0)?;
 
         // pad siblings with zeros to length max_depth
-        let mut siblings = proof.siblings.clone();
+        assert_eq!(mp.proof.siblings.len(), self.max_depth);
-        siblings.resize(self.max_depth, EMPTY_HASH);
-        assert_eq!(self.siblings.len(), siblings.len());
 
-        for (i, sibling) in siblings.iter().enumerate() {
+        for (i, sibling) in mp.proof.siblings.iter().enumerate() {
             pw.set_hash_target(self.siblings[i], HashOut::from_vec(sibling.0.to_vec()))?;
         }
 
@@ -540,12 +515,7 @@ pub mod tests {
         let targets = MerkleProofGadget { max_depth }.eval(&mut builder)?;
         targets.set_targets(
             &mut pw,
-            true, // verification enabled
+            &MerkleClaimAndProof::new(max_depth, tree.root(), key, Some(value), &proof)?,
-            existence,
-            tree.root(),
-            proof,
-            key,
-            value,
         )?;
 
         // generate & verify proof
@@ -588,7 +558,10 @@ pub mod tests {
         let mut pw = PartialWitness::<F>::new();
 
         let targets = MerkleProofExistenceGadget { max_depth }.eval(&mut builder)?;
-        targets.set_targets(&mut pw, true, tree.root(), proof, key, value)?;
+        targets.set_targets(
+            &mut pw,
+            &MerkleClaimAndProof::new(max_depth, tree.root(), key, Some(value), &proof)?,
+        )?;
 
         // generate & verify proof
         let data = builder.build::<C>();
@@ -661,12 +634,7 @@ pub mod tests {
         let targets = MerkleProofGadget { max_depth }.eval(&mut builder)?;
         targets.set_targets(
             &mut pw,
-            true, // verification enabled
+            &MerkleClaimAndProof::new(max_depth, tree.root(), key, Some(value), &proof)?,
-            proof.existence,
-            tree.root(),
-            proof,
-            key,
-            value,
         )?;
 
         // generate & verify proof
@@ -708,15 +676,9 @@ pub mod tests {
         let mut pw = PartialWitness::<F>::new();
 
         let targets = MerkleProofGadget { max_depth }.eval(&mut builder)?;
-        targets.set_targets(
+        // verification enabled & proof of existence
-            &mut pw,
+        let mut mp = MerkleClaimAndProof::new(max_depth, tree2.root(), key, Some(value), &proof)?;
-            true, // verification enabled
+        targets.set_targets(&mut pw, &mp)?;
-            true, // proof of existence
-            tree2.root(),
-            proof.clone(),
-            key,
-            value,
-        )?;
 
         // generate proof, expecting it to fail (since we're using the wrong
         // root)
@@ -730,15 +692,9 @@ pub mod tests {
         let mut pw = PartialWitness::<F>::new();
 
         let targets = MerkleProofGadget { max_depth }.eval(&mut builder)?;
-        targets.set_targets(
+        // verification disabled & proof of existence
-            &mut pw,
+        mp.enabled = false;
-            false, // verification disabled
+        targets.set_targets(&mut pw, &mp)?;
-            true,  // proof of existence
-            tree2.root(),
-            proof,
-            key,
-            value,
-        )?;
 
         // generate proof, should pass despite using wrong witness, since the
         // `enabled=false`