feat: add container update ops (#390)

* Add container update ops

* Update src/middleware/operation.rs

Co-authored-by: Eduard S. <eduardsanou@posteo.net>

* Update src/backends/plonky2/mainpod/mod.rs

Co-authored-by: Eduard S. <eduardsanou@posteo.net>

* Code review

---------

Co-authored-by: Eduard S. <eduardsanou@posteo.net>

src/backends/plonky2/circuits/common.rs

@@ -27,7 +27,7 @@ use crate::{
         circuits::mainpod::CustomPredicateVerification,
         error::Result,
         mainpod::{Operation, OperationArg, OperationAux, Statement},
-        primitives::merkletree::MerkleClaimAndProofTarget,
+        primitives::merkletree::{MerkleClaimAndProofTarget, MerkleTreeStateTransitionProofTarget},
     },
     middleware::{
         CustomPredicate, CustomPredicateBatch, CustomPredicateRef, NativeOperation,
@@ -743,6 +743,32 @@ impl From<MerkleClaimAndProofTarget> for MerkleClaimTarget {
     }
 }
 
+/// For the purpose of op verification, we need only look up the
+/// Merkle state transition claim rather than the Merkle state
+/// transition proof since it is verified elsewhere.
+#[derive(Copy, Clone)]
+pub struct MerkleTreeStateTransitionClaimTarget {
+    pub(crate) enabled: BoolTarget,
+    pub(crate) op: Target,
+    pub(crate) old_root: HashOutTarget,
+    pub(crate) new_root: HashOutTarget,
+    pub(crate) op_key: ValueTarget,
+    pub(crate) op_value: ValueTarget,
+}
+
+impl From<MerkleTreeStateTransitionProofTarget> for MerkleTreeStateTransitionClaimTarget {
+    fn from(pf: MerkleTreeStateTransitionProofTarget) -> Self {
+        Self {
+            enabled: pf.enabled,
+            op: pf.op,
+            old_root: pf.old_root,
+            new_root: pf.new_root,
+            op_key: pf.op_key,
+            op_value: pf.op_value,
+        }
+    }
+}
+
 impl Flattenable for HashOutTarget {
     fn flatten(&self) -> Vec<Target> {
         self.elements.to_vec()
@@ -803,6 +829,42 @@ impl Flattenable for MerkleClaimTarget {
     }
 }
 
+impl Flattenable for MerkleTreeStateTransitionClaimTarget {
+    fn flatten(&self) -> Vec<Target> {
+        [
+            vec![self.enabled.target, self.op],
+            self.old_root.elements.to_vec(),
+            self.new_root.elements.to_vec(),
+            self.op_key.elements.to_vec(),
+            self.op_value.elements.to_vec(),
+        ]
+        .concat()
+    }
+
+    fn from_flattened(params: &Params, vs: &[Target]) -> Self {
+        assert_eq!(vs.len(), Self::size(params));
+        Self {
+            enabled: BoolTarget::new_unsafe(vs[0]),
+            op: vs[1],
+            old_root: HashOutTarget::from_vec(vs[2..2 + NUM_HASH_OUT_ELTS].to_vec()),
+            new_root: HashOutTarget::from_vec(
+                vs[2 + NUM_HASH_OUT_ELTS..2 * (1 + NUM_HASH_OUT_ELTS)].to_vec(),
+            ),
+            op_key: ValueTarget::from_slice(
+                &vs[2 * (1 + NUM_HASH_OUT_ELTS)..2 * (1 + NUM_HASH_OUT_ELTS) + VALUE_SIZE],
+            ),
+            op_value: ValueTarget::from_slice(
+                &vs[2 * (1 + NUM_HASH_OUT_ELTS) + VALUE_SIZE
+                    ..2 * (1 + NUM_HASH_OUT_ELTS) + 2 * VALUE_SIZE],
+            ),
+        }
+    }
+
+    fn size(params: &Params) -> usize {
+        2 * (1 + HashOutTarget::size(params)) + 2 * ValueTarget::size(params)
+    }
+}
+
 impl Flattenable for PredicateTarget {
     fn flatten(&self) -> Vec<Target> {
         self.elements.to_vec()

src/backends/plonky2/circuits/mainpod.rs

@@ -23,9 +23,10 @@ use crate::{
             common::{
                 CircuitBuilderPod, CustomPredicateBatchTarget, CustomPredicateEntryTarget,
                 CustomPredicateTarget, CustomPredicateVerifyEntryTarget,
-                CustomPredicateVerifyQueryTarget, Flattenable, MerkleClaimTarget, OperationTarget,
+                CustomPredicateVerifyQueryTarget, Flattenable, MerkleClaimTarget,
-                OperationTypeTarget, PredicateTarget, StatementArgTarget, StatementTarget,
+                MerkleTreeStateTransitionClaimTarget, OperationTarget, OperationTypeTarget,
-                StatementTmplArgTarget, StatementTmplTarget, ValueTarget,
+                PredicateTarget, StatementArgTarget, StatementTarget, StatementTmplArgTarget,
+                StatementTmplTarget, ValueTarget,
             },
             hash::{hash_from_state_circuit, precompute_hash_state},
             mux_table::{MuxTableTarget, TableEntryTarget},
@@ -41,7 +42,9 @@ use crate::{
                 schnorr::SecretKey,
             },
             merkletree::{
-                verify_merkle_proof_circuit, MerkleClaimAndProof, MerkleClaimAndProofTarget,
+                verify_merkle_proof_circuit, verify_merkle_state_transition_circuit,
+                MerkleClaimAndProof, MerkleClaimAndProofTarget, MerkleTreeOp,
+                MerkleTreeStateTransitionProof, MerkleTreeStateTransitionProofTarget,
             },
         },
         recursion::{InnerCircuit, VerifiedProofTarget},
@@ -65,7 +68,7 @@ pub const PI_OFFSET_VDSROOT: usize = 4;
 
 pub const NUM_PUBLIC_INPUTS: usize = 8;
 
-const MAX_VALUE_ARGS: usize = 3;
+const MAX_VALUE_ARGS: usize = 4;
 
 struct StatementArgCache {
     rhs: ValueTarget,
@@ -99,8 +102,8 @@ impl StatementCache {
                 .map(|i| builder.vec_ref(params, prev_statements, i))
                 .collect::<Vec<_>>()
         };
-        assert!(params.max_operation_args >= 3);
+        assert!(params.max_operation_args >= MAX_VALUE_ARGS);
-        assert!(params.max_statement_args >= 3);
+        assert!(params.max_statement_args >= MAX_VALUE_ARGS);
         let equations = array::from_fn(|i| {
             let pred_is_none = op_args[i].has_native_type(builder, params, NativePredicate::None);
             let arg_is_value = builder.statement_arg_is_value(&st.args[i]);
@@ -189,13 +192,16 @@ enum OperationAuxTableTag {
     None = 0,
     MerkleProof = 1,
     PublicKeyOf = 2,
-    CustomPredVerify = 3,
+    MerkleTreeStateTransitionProof = 3,
+    CustomPredVerify = 4,
 }
 
 fn max_operation_aux_entry_len(params: &Params) -> usize {
     [
         (params.max_merkle_proofs_containers > 0).then(|| MerkleClaimTarget::size(params)),
         (params.max_public_key_of > 0).then(|| KeyPairTarget::size(params)),
+        (params.max_merkle_tree_state_transition_proofs_containers > 0)
+            .then(|| MerkleTreeStateTransitionClaimTarget::size(params)),
         (params.max_custom_predicate_verifications > 0)
             .then(|| CustomPredicateVerifyQueryTarget::size(params)),
     ]
@@ -236,6 +242,7 @@ fn build_operation_aux_table_circuit(
     builder: &mut CircuitBuilder,
     merkle_proofs: &[MerkleClaimAndProofTarget],
     public_key_of_sks: &[BigUInt320Target],
+    merkle_tree_state_transition_proofs: &[MerkleTreeStateTransitionProofTarget],
     custom_predicate_verifications: &[CustomPredicateVerifyEntryTarget],
     custom_predicate_table: &[HashOutTarget],
 ) -> Result<MuxTableTarget> {
@@ -285,6 +292,19 @@ fn build_operation_aux_table_circuit(
         measure_gates_end!(builder, measure);
     }
 
+    // Merkle state transition proofs: verify op proof (insert/update/delete)
+    for merkle_tree_state_transition_proof in merkle_tree_state_transition_proofs {
+        verify_merkle_state_transition_circuit(builder, merkle_tree_state_transition_proof);
+        let entry =
+            MerkleTreeStateTransitionClaimTarget::from(merkle_tree_state_transition_proof.clone());
+
+        table.push(
+            builder,
+            OperationAuxTableTag::MerkleTreeStateTransitionProof as u32,
+            &entry,
+        );
+    }
+
     // CustomPredVerify: verify custom predicate statements verification against operations
     for entry in custom_predicate_verifications {
         let measure = measure_gates_begin!(builder, "CustomPredVerify");
@@ -413,6 +433,34 @@ fn verify_operation_circuit(
                 &cache,
             ));
         }
+        if params.max_merkle_tree_state_transition_proofs_containers > 0 {
+            op_checks.extend_from_slice(&[
+                verify_merkle_insert_circuit(
+                    params,
+                    builder,
+                    st,
+                    &op.op_type,
+                    &resolved_aux,
+                    &cache,
+                ),
+                verify_merkle_update_circuit(
+                    params,
+                    builder,
+                    st,
+                    &op.op_type,
+                    &resolved_aux,
+                    &cache,
+                ),
+                verify_merkle_delete_circuit(
+                    params,
+                    builder,
+                    st,
+                    &op.op_type,
+                    &resolved_aux,
+                    &cache,
+                ),
+            ]);
+        }
         if params.max_custom_predicate_verifications > 0 {
             op_checks.push(verify_custom_circuit(
                 builder,
@@ -532,6 +580,224 @@ fn verify_not_contains_from_entries_circuit(
     ok
 }
 
+fn verify_merkle_insert_circuit(
+    params: &Params,
+    builder: &mut CircuitBuilder,
+    st: &StatementTarget,
+    op_type: &OperationTypeTarget,
+    aux: &TableEntryTarget,
+    cache: &StatementCache,
+) -> BoolTarget {
+    let measure = measure_gates_begin!(builder, "MerkleInsertOp");
+    let (aux_tag_ok, resolved_merkle_tree_state_transition_claim) =
+        aux.as_type::<MerkleTreeStateTransitionClaimTarget>(
+            builder,
+            OperationAuxTableTag::MerkleTreeStateTransitionProof as u32,
+        );
+    let op_code_ok = op_type.has_native(builder, NativeOperation::ContainerInsertFromEntries);
+
+    let (arg_types_ok, [new_root_value, old_root_value, op_key_value, op_value_value]) =
+        cache.first_n_args_as_values();
+
+    let expected_merkle_op = builder.constant(F::from_canonical_u8(MerkleTreeOp::Insert as u8));
+
+    // Check Merkle proof (verified elsewhere) against op args.
+    let merkle_proof_checks = [
+        /* The supplied Merkle transition proof must be enabled. */
+        resolved_merkle_tree_state_transition_claim.enabled,
+        /* ...and it must be an insertion proof. */
+        builder.is_equal(
+            resolved_merkle_tree_state_transition_claim.op,
+            expected_merkle_op,
+        ),
+        /* ...for the root-key-value combination in the resolved op args. */
+        builder.is_equal_slice(
+            &old_root_value.elements,
+            &resolved_merkle_tree_state_transition_claim
+                .old_root
+                .elements,
+        ),
+        builder.is_equal_slice(
+            &new_root_value.elements,
+            &resolved_merkle_tree_state_transition_claim
+                .new_root
+                .elements,
+        ),
+        builder.is_equal_slice(
+            &op_key_value.elements,
+            &resolved_merkle_tree_state_transition_claim.op_key.elements,
+        ),
+        builder.is_equal_slice(
+            &op_value_value.elements,
+            &resolved_merkle_tree_state_transition_claim
+                .op_value
+                .elements,
+        ),
+    ];
+
+    let merkle_proof_ok = builder.all(merkle_proof_checks);
+
+    // Check output statement
+    let arg1_expected = cache.equations[0].lhs.clone();
+    let arg2_expected = cache.equations[1].lhs.clone();
+    let arg3_expected = cache.equations[2].lhs.clone();
+    let arg4_expected = cache.equations[3].lhs.clone();
+    let expected_statement = StatementTarget::new_native(
+        builder,
+        params,
+        NativePredicate::ContainerInsert,
+        &[arg1_expected, arg2_expected, arg3_expected, arg4_expected],
+    );
+    let st_ok = builder.is_equal_flattenable(st, &expected_statement);
+
+    let ok = builder.all([op_code_ok, aux_tag_ok, arg_types_ok, merkle_proof_ok, st_ok]);
+    measure_gates_end!(builder, measure);
+    ok
+}
+
+fn verify_merkle_update_circuit(
+    params: &Params,
+    builder: &mut CircuitBuilder,
+    st: &StatementTarget,
+    op_type: &OperationTypeTarget,
+    aux: &TableEntryTarget,
+    cache: &StatementCache,
+) -> BoolTarget {
+    let measure = measure_gates_begin!(builder, "MerkleUpdateOp");
+    let (aux_tag_ok, resolved_merkle_tree_state_transition_claim) =
+        aux.as_type::<MerkleTreeStateTransitionClaimTarget>(
+            builder,
+            OperationAuxTableTag::MerkleTreeStateTransitionProof as u32,
+        );
+    let op_code_ok = op_type.has_native(builder, NativeOperation::ContainerUpdateFromEntries);
+
+    let (arg_types_ok, [new_root_value, old_root_value, op_key_value, op_value_value]) =
+        cache.first_n_args_as_values();
+
+    let expected_merkle_op = builder.constant(F::from_canonical_u8(MerkleTreeOp::Update as u8));
+
+    // Check Merkle proof (verified elsewhere) against op args.
+    let merkle_proof_checks = [
+        /* The supplied Merkle transition proof must be enabled. */
+        resolved_merkle_tree_state_transition_claim.enabled,
+        /* ...and it must be an update proof. */
+        builder.is_equal(
+            resolved_merkle_tree_state_transition_claim.op,
+            expected_merkle_op,
+        ),
+        /* ...for the root-key-value combination in the resolved op args. */
+        builder.is_equal_slice(
+            &old_root_value.elements,
+            &resolved_merkle_tree_state_transition_claim
+                .old_root
+                .elements,
+        ),
+        builder.is_equal_slice(
+            &new_root_value.elements,
+            &resolved_merkle_tree_state_transition_claim
+                .new_root
+                .elements,
+        ),
+        builder.is_equal_slice(
+            &op_key_value.elements,
+            &resolved_merkle_tree_state_transition_claim.op_key.elements,
+        ),
+        builder.is_equal_slice(
+            &op_value_value.elements,
+            &resolved_merkle_tree_state_transition_claim
+                .op_value
+                .elements,
+        ),
+    ];
+
+    let merkle_proof_ok = builder.all(merkle_proof_checks);
+
+    // Check output statement
+    let arg1_expected = cache.equations[0].lhs.clone();
+    let arg2_expected = cache.equations[1].lhs.clone();
+    let arg3_expected = cache.equations[2].lhs.clone();
+    let arg4_expected = cache.equations[3].lhs.clone();
+    let expected_statement = StatementTarget::new_native(
+        builder,
+        params,
+        NativePredicate::ContainerUpdate,
+        &[arg1_expected, arg2_expected, arg3_expected, arg4_expected],
+    );
+    let st_ok = builder.is_equal_flattenable(st, &expected_statement);
+
+    let ok = builder.all([op_code_ok, aux_tag_ok, arg_types_ok, merkle_proof_ok, st_ok]);
+    measure_gates_end!(builder, measure);
+    ok
+}
+
+fn verify_merkle_delete_circuit(
+    params: &Params,
+    builder: &mut CircuitBuilder,
+    st: &StatementTarget,
+    op_type: &OperationTypeTarget,
+    aux: &TableEntryTarget,
+    cache: &StatementCache,
+) -> BoolTarget {
+    let measure = measure_gates_begin!(builder, "MerkleDeleteOp");
+    let (aux_tag_ok, resolved_merkle_tree_state_transition_claim) =
+        aux.as_type::<MerkleTreeStateTransitionClaimTarget>(
+            builder,
+            OperationAuxTableTag::MerkleTreeStateTransitionProof as u32,
+        );
+    let op_code_ok = op_type.has_native(builder, NativeOperation::ContainerDeleteFromEntries);
+
+    let (arg_types_ok, [new_root_value, old_root_value, op_key_value]) =
+        cache.first_n_args_as_values();
+
+    let expected_merkle_op = builder.constant(F::from_canonical_u8(MerkleTreeOp::Delete as u8));
+
+    // Check Merkle proof (verified elsewhere) against op args.
+    let merkle_proof_checks = [
+        /* The supplied Merkle transition proof must be enabled. */
+        resolved_merkle_tree_state_transition_claim.enabled,
+        /* ...and it must be a deletion proof. */
+        builder.is_equal(
+            resolved_merkle_tree_state_transition_claim.op,
+            expected_merkle_op,
+        ),
+        /* ...for the root-key combination in the resolved op args. */
+        builder.is_equal_slice(
+            &old_root_value.elements,
+            &resolved_merkle_tree_state_transition_claim
+                .old_root
+                .elements,
+        ),
+        builder.is_equal_slice(
+            &new_root_value.elements,
+            &resolved_merkle_tree_state_transition_claim
+                .new_root
+                .elements,
+        ),
+        builder.is_equal_slice(
+            &op_key_value.elements,
+            &resolved_merkle_tree_state_transition_claim.op_key.elements,
+        ),
+    ];
+
+    let merkle_proof_ok = builder.all(merkle_proof_checks);
+
+    // Check output statement
+    let arg1_expected = cache.equations[0].lhs.clone();
+    let arg2_expected = cache.equations[1].lhs.clone();
+    let arg3_expected = cache.equations[2].lhs.clone();
+    let expected_statement = StatementTarget::new_native(
+        builder,
+        params,
+        NativePredicate::ContainerDelete,
+        &[arg1_expected, arg2_expected, arg3_expected],
+    );
+    let st_ok = builder.is_equal_flattenable(st, &expected_statement);
+
+    let ok = builder.all([op_code_ok, aux_tag_ok, arg_types_ok, merkle_proof_ok, st_ok]);
+    measure_gates_end!(builder, measure);
+    ok
+}
+
 fn verify_custom_circuit(
     builder: &mut CircuitBuilder,
     st: &StatementTarget,
@@ -1117,9 +1383,9 @@ fn make_custom_statement_circuit(
         params.max_statement_args,
         custom_predicate.predicate.args_len,
     );
-    let st_args = (0..params.max_statement_args)
+    let st_args = std::iter::zip(lt_mask, args)
-        .map(|i| {
+        .map(|(mask, arg)| {
-            let v = builder.select_flattenable(params, lt_mask[i], &args[i], &arg_none);
+            let v = builder.select_flattenable(params, mask, arg, &arg_none);
             StatementArgTarget::wildcard_literal(builder, &v)
         })
         .collect();
@@ -1380,6 +1646,7 @@ fn verify_main_pod_circuit(
         builder,
         &main_pod.merkle_proofs,
         &main_pod.public_key_of_sks,
+        &main_pod.merkle_tree_state_transition_proofs,
         &main_pod.custom_predicate_verifications,
         &custom_predicate_table,
     )?;
@@ -1446,6 +1713,7 @@ pub struct MainPodVerifyTarget {
     operations: Vec<OperationTarget>,
     merkle_proofs: Vec<MerkleClaimAndProofTarget>,
     public_key_of_sks: Vec<BigUInt320Target>,
+    merkle_tree_state_transition_proofs: Vec<MerkleTreeStateTransitionProofTarget>,
     custom_predicate_batches: Vec<CustomPredicateBatchTarget>,
     custom_predicate_verifications: Vec<CustomPredicateVerifyEntryTarget>,
 }
@@ -1482,6 +1750,15 @@ impl MainPodVerifyTarget {
             public_key_of_sks: (0..params.max_public_key_of)
                 .map(|_| builder.add_virtual_biguint320_target())
                 .collect(),
+            merkle_tree_state_transition_proofs: (0..params
+                .max_merkle_tree_state_transition_proofs_containers)
+                .map(|_| {
+                    MerkleTreeStateTransitionProofTarget::new_virtual(
+                        params.max_depth_mt_containers,
+                        builder,
+                    )
+                })
+                .collect(),
             custom_predicate_batches: (0..params.max_custom_predicate_batches)
                 .map(|_| builder.add_virtual_custom_predicate_batch(params))
                 .collect(),
@@ -1511,6 +1788,7 @@ pub struct MainPodVerifyInput {
     pub operations: Vec<mainpod::Operation>,
     pub merkle_proofs: Vec<MerkleClaimAndProof>,
     pub public_key_of_sks: Vec<SecretKey>,
+    pub merkle_tree_state_transition_proofs: Vec<MerkleTreeStateTransitionProof>,
     pub custom_predicate_batches: Vec<Arc<CustomPredicateBatch>>,
     pub custom_predicate_verifications: Vec<CustomPredicateVerification>,
 }
@@ -1641,6 +1919,25 @@ impl InnerCircuit for MainPodVerifyTarget {
             pw.set_biguint320_target(&self.public_key_of_sks[i], &pad_sk)?;
         }
 
+        assert!(
+            input.merkle_tree_state_transition_proofs.len()
+                <= self
+                    .params
+                    .max_merkle_tree_state_transition_proofs_containers
+        );
+        for (i, mtp) in input.merkle_tree_state_transition_proofs.iter().enumerate() {
+            self.merkle_tree_state_transition_proofs[i].set_targets(pw, true, mtp)?;
+        }
+        // Padding
+        let pad_mtp = MerkleTreeStateTransitionProof::empty();
+        for i in input.merkle_tree_state_transition_proofs.len()
+            ..self
+                .params
+                .max_merkle_tree_state_transition_proofs_containers
+        {
+            self.merkle_tree_state_transition_proofs[i].set_targets(pw, false, &pad_mtp)?;
+        }
+
         assert!(input.custom_predicate_batches.len() <= self.params.max_custom_predicate_batches);
         for (i, cpb) in input.custom_predicate_batches.iter().enumerate() {
             self.custom_predicate_batches[i].set_targets(pw, &self.params, cpb)?;
@@ -1703,7 +2000,7 @@ mod tests {
             mainpod::{calculate_id, OperationArg, OperationAux},
             primitives::{
                 ec::schnorr::SecretKey,
-                merkletree::{MerkleClaimAndProof, MerkleTree},
+                merkletree::{MerkleClaimAndProof, MerkleTree, MerkleTreeStateTransitionProof},
             },
         },
         frontend::{self, literal, CustomPredicateBatchBuilder, StatementTmplBuilder},
@@ -1719,6 +2016,7 @@ mod tests {
         prev_statements: Vec<mainpod::Statement>,
         merkle_proofs: Vec<MerkleClaimAndProof>,
         secret_keys: Vec<SecretKey>,
+        merkle_tree_state_transition_proofs: Vec<MerkleTreeStateTransitionProof>,
     ) -> Result<()> {
         let params = Params {
             max_custom_predicate_batches: 0,
@@ -1749,11 +2047,23 @@ mod tests {
             .map(|sk| builder.constant_biguint320(&sk.0))
             .collect();
 
+        let merkle_tree_state_transition_proofs_target: Vec<_> =
+            merkle_tree_state_transition_proofs
+                .iter()
+                .map(|_| {
+                    MerkleTreeStateTransitionProofTarget::new_virtual(
+                        params.max_depth_mt_containers,
+                        &mut builder,
+                    )
+                })
+                .collect();
+
         let aux_table = build_operation_aux_table_circuit(
             &params,
             &mut builder,
             &merkle_proofs_target,
             &secret_keys_target,
+            &merkle_tree_state_transition_proofs_target,
             &[],
             &[],
         )?;
@@ -1781,6 +2091,17 @@ mod tests {
         {
             merkle_proof_target.set_targets(&mut pw, true, merkle_proof)?
         }
+        for (merkle_tree_state_transition_proof_target, merkle_tree_state_transition_proof) in
+            merkle_tree_state_transition_proofs_target
+                .iter()
+                .zip(merkle_tree_state_transition_proofs.iter())
+        {
+            merkle_tree_state_transition_proof_target.set_targets(
+                &mut pw,
+                true,
+                merkle_tree_state_transition_proof,
+            )?
+        }
 
         // generate & verify proof
         let data = builder.build::<C>();
@@ -1927,7 +2248,7 @@ mod tests {
         .into_iter()
         .for_each(|(op, st)| {
             let check = std::panic::catch_unwind(|| {
-                operation_verify(st, op, prev_statements.to_vec(), vec![], vec![])
+                operation_verify(st, op, prev_statements.to_vec(), vec![], vec![], vec![])
             });
             match check {
                 Err(e) => {
@@ -1992,7 +2313,9 @@ mod tests {
         ]
         .into_iter()
         .for_each(|(op, st)| {
-            assert!(operation_verify(st, op, prev_statements.to_vec(), vec![], vec![]).is_err())
+            assert!(
+                operation_verify(st, op, prev_statements.to_vec(), vec![], vec![], vec![]).is_err()
+            )
         });
     }
 
@@ -2005,7 +2328,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![Statement::None.into()];
-        operation_verify(st, op, prev_statements, vec![], vec![])
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])
     }
 
     #[test]
@@ -2023,7 +2346,7 @@ mod tests {
             vec![],
             OperationAux::None,
         );
-        operation_verify(st1, op, prev_statements, vec![], vec![])
+        operation_verify(st1, op, prev_statements, vec![], vec![], vec![])
     }
 
     #[test]
@@ -2035,7 +2358,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![Statement::None.into()];
-        operation_verify(st, op, prev_statements, vec![], vec![])
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])
     }
 
     #[test]
@@ -2058,7 +2381,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st1, st2];
-        operation_verify(st, op, prev_statements, vec![], vec![])
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])
     }
 
     #[test]
@@ -2081,7 +2404,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st1, st2];
-        operation_verify(st, op, prev_statements, vec![], vec![])
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])
     }
 
     #[test]
@@ -2104,7 +2427,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st1, st2.clone()];
-        operation_verify(st, op, prev_statements, vec![], vec![])?;
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])?;
 
         // Also check negative < negative
         let st3: mainpod::Statement = Statement::equal(
@@ -2128,7 +2451,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st3.clone(), st4];
-        operation_verify(st, op, prev_statements, vec![], vec![])?;
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])?;
 
         // Also check negative < positive
         let st: mainpod::Statement = Statement::lt(
@@ -2142,7 +2465,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st3, st2];
-        operation_verify(st, op, prev_statements, vec![], vec![])
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])
     }
 
     #[test]
@@ -2165,7 +2488,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st1, st2.clone()];
-        operation_verify(st, op, prev_statements, vec![], vec![])?;
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])?;
 
         // Also check negative <= negative
         let st3: mainpod::Statement = Statement::equal(
@@ -2189,7 +2512,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st3.clone(), st4];
-        operation_verify(st, op, prev_statements, vec![], vec![])?;
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])?;
 
         // Also check negative <= positive
         let st: mainpod::Statement = Statement::lt_eq(
@@ -2203,7 +2526,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st3, st2];
-        operation_verify(st, op, prev_statements.clone(), vec![], vec![])?;
+        operation_verify(st, op, prev_statements.clone(), vec![], vec![], vec![])?;
 
         // Also check equality, both positive and negative.
         let st: mainpod::Statement = Statement::lt_eq(
@@ -2216,7 +2539,7 @@ mod tests {
             vec![OperationArg::Index(0), OperationArg::Index(0)],
             OperationAux::None,
         );
-        operation_verify(st, op, prev_statements.clone(), vec![], vec![])?;
+        operation_verify(st, op, prev_statements.clone(), vec![], vec![], vec![])?;
         let st: mainpod::Statement = Statement::lt_eq(
             AnchoredKey::from((PodId(RawValue::from(88).into()), "hello")),
             AnchoredKey::from((PodId(RawValue::from(88).into()), "hello")),
@@ -2227,7 +2550,7 @@ mod tests {
             vec![OperationArg::Index(1), OperationArg::Index(1)],
             OperationAux::None,
         );
-        operation_verify(st, op, prev_statements, vec![], vec![])
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])
     }
 
     #[test]
@@ -2276,7 +2599,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st1, st2, st3];
-        operation_verify(st, op, prev_statements, vec![], vec![])
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])
     }
 
     #[test]
@@ -2322,7 +2645,7 @@ mod tests {
                     OperationAux::None,
                 );
                 let prev_statements = vec![st1, st2, st3];
-                operation_verify(st, op, prev_statements, vec![], vec![])
+                operation_verify(st, op, prev_statements, vec![], vec![], vec![])
             })
     }
 
@@ -2369,7 +2692,7 @@ mod tests {
                     OperationAux::None,
                 );
                 let prev_statements = vec![st1, st2, st3];
-                operation_verify(st, op, prev_statements, vec![], vec![])
+                operation_verify(st, op, prev_statements, vec![], vec![], vec![])
             })
     }
 
@@ -2411,7 +2734,7 @@ mod tests {
                 OperationAux::None,
             );
             let prev_statements = vec![st1, st2, st3];
-            operation_verify(st, op, prev_statements, vec![], vec![])
+            operation_verify(st, op, prev_statements, vec![], vec![], vec![])
         })
     }
 
@@ -2456,7 +2779,7 @@ mod tests {
                 let prev_statements = [st1, st2, st3];
 
                 let check = std::panic::catch_unwind(|| {
-                    operation_verify(st, op, prev_statements.to_vec(), vec![], vec![])
+                    operation_verify(st, op, prev_statements.to_vec(), vec![], vec![], vec![])
                 });
                 match check {
                     Err(e) => {
@@ -2489,7 +2812,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st1];
-        operation_verify(st, op, prev_statements, vec![], vec![])
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])
     }
 
     #[test]
@@ -2515,7 +2838,7 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![st1, st2];
-        operation_verify(st, op, prev_statements, vec![], vec![])
+        operation_verify(st, op, prev_statements, vec![], vec![], vec![])
     }
 
     #[test]
@@ -2555,7 +2878,7 @@ mod tests {
             no_key_pf,
         )];
         let prev_statements = vec![root_st, key_st];
-        operation_verify(st, op, prev_statements, merkle_proofs, vec![])
+        operation_verify(st, op, prev_statements, merkle_proofs, vec![], vec![])
     }
 
     #[test]
@@ -2602,7 +2925,163 @@ mod tests {
             key_pf,
         )];
         let prev_statements = vec![root_st, key_st, value_st];
-        operation_verify(st, op, prev_statements, merkle_proofs, vec![])
+        operation_verify(st, op, prev_statements, merkle_proofs, vec![], vec![])
+    }
+
+    #[test]
+    fn test_operation_verify_merkle_insert() -> Result<()> {
+        let params = Params::default();
+
+        let mut tree = MerkleTree::new(params.max_depth_mt_containers, &[].into())?;
+
+        let key = 175.into();
+        let key_ak = AnchoredKey::from((PodId(RawValue::from(70).into()), "key"));
+
+        let value = 0.into();
+        let value_ak = AnchoredKey::from((PodId(RawValue::from(72).into()), "value"));
+
+        let state_transition_proof = tree.insert(&key, &value)?;
+
+        let old_root = Value::from(state_transition_proof.old_root);
+        let old_root_ak = AnchoredKey::from((PodId(RawValue::from(73).into()), "old_root"));
+
+        let new_root = Value::from(state_transition_proof.new_root);
+        let new_root_ak = AnchoredKey::from((PodId(RawValue::from(74).into()), "new_root"));
+
+        let new_root_st: mainpod::Statement =
+            Statement::equal(new_root_ak.clone(), new_root.clone()).into();
+        let old_root_st: mainpod::Statement =
+            Statement::equal(old_root_ak.clone(), old_root.clone()).into();
+        let key_st: mainpod::Statement = Statement::equal(key_ak.clone(), key).into();
+        let value_st: mainpod::Statement = Statement::equal(value_ak.clone(), value).into();
+
+        let st: mainpod::Statement =
+            Statement::insert(new_root_ak, old_root_ak, key_ak, value_ak).into();
+        let op = mainpod::Operation(
+            OperationType::Native(NativeOperation::ContainerInsertFromEntries),
+            vec![
+                OperationArg::Index(0),
+                OperationArg::Index(1),
+                OperationArg::Index(2),
+                OperationArg::Index(3),
+            ],
+            OperationAux::MerkleTreeStateTransitionProofIndex(0),
+        );
+
+        let merkle_tree_state_transition_proofs = vec![state_transition_proof];
+        let prev_statements = vec![new_root_st, old_root_st, key_st, value_st];
+        operation_verify(
+            st,
+            op,
+            prev_statements,
+            vec![],
+            vec![],
+            merkle_tree_state_transition_proofs,
+        )
+    }
+
+    #[test]
+    fn test_operation_verify_merkle_update() -> Result<()> {
+        let params = Params::default();
+
+        let mut tree = MerkleTree::new(
+            params.max_depth_mt_containers,
+            &[(175.into(), 55.into())].into(),
+        )?;
+
+        let key = 175.into();
+        let key_ak = AnchoredKey::from((PodId(RawValue::from(70).into()), "key"));
+
+        let value = 0.into();
+        let value_ak = AnchoredKey::from((PodId(RawValue::from(72).into()), "value"));
+
+        let state_transition_proof = tree.update(&key, &value)?;
+
+        let old_root = Value::from(state_transition_proof.old_root);
+        let old_root_ak = AnchoredKey::from((PodId(RawValue::from(73).into()), "old_root"));
+
+        let new_root = Value::from(state_transition_proof.new_root);
+        let new_root_ak = AnchoredKey::from((PodId(RawValue::from(74).into()), "new_root"));
+
+        let new_root_st: mainpod::Statement =
+            Statement::equal(new_root_ak.clone(), new_root.clone()).into();
+        let old_root_st: mainpod::Statement =
+            Statement::equal(old_root_ak.clone(), old_root.clone()).into();
+        let key_st: mainpod::Statement = Statement::equal(key_ak.clone(), key).into();
+        let value_st: mainpod::Statement = Statement::equal(value_ak.clone(), value).into();
+
+        let st: mainpod::Statement =
+            Statement::update(new_root_ak, old_root_ak, key_ak, value_ak).into();
+        let op = mainpod::Operation(
+            OperationType::Native(NativeOperation::ContainerUpdateFromEntries),
+            vec![
+                OperationArg::Index(0),
+                OperationArg::Index(1),
+                OperationArg::Index(2),
+                OperationArg::Index(3),
+            ],
+            OperationAux::MerkleTreeStateTransitionProofIndex(0),
+        );
+
+        let merkle_tree_state_transition_proofs = vec![state_transition_proof];
+        let prev_statements = vec![new_root_st, old_root_st, key_st, value_st];
+        operation_verify(
+            st,
+            op,
+            prev_statements,
+            vec![],
+            vec![],
+            merkle_tree_state_transition_proofs,
+        )
+    }
+
+    #[test]
+    fn test_operation_verify_merkle_delete() -> Result<()> {
+        let params = Params::default();
+
+        let mut tree = MerkleTree::new(
+            params.max_depth_mt_containers,
+            &[(175.into(), 55.into())].into(),
+        )?;
+
+        let key = 175.into();
+        let key_ak = AnchoredKey::from((PodId(RawValue::from(70).into()), "key"));
+
+        let state_transition_proof = tree.delete(&key)?;
+
+        let old_root = Value::from(state_transition_proof.old_root);
+        let old_root_ak = AnchoredKey::from((PodId(RawValue::from(73).into()), "old_root"));
+
+        let new_root = Value::from(state_transition_proof.new_root);
+        let new_root_ak = AnchoredKey::from((PodId(RawValue::from(74).into()), "new_root"));
+
+        let new_root_st: mainpod::Statement =
+            Statement::equal(new_root_ak.clone(), new_root.clone()).into();
+        let old_root_st: mainpod::Statement =
+            Statement::equal(old_root_ak.clone(), old_root.clone()).into();
+        let key_st: mainpod::Statement = Statement::equal(key_ak.clone(), key).into();
+
+        let st: mainpod::Statement = Statement::delete(new_root_ak, old_root_ak, key_ak).into();
+        let op = mainpod::Operation(
+            OperationType::Native(NativeOperation::ContainerDeleteFromEntries),
+            vec![
+                OperationArg::Index(0),
+                OperationArg::Index(1),
+                OperationArg::Index(2),
+            ],
+            OperationAux::MerkleTreeStateTransitionProofIndex(0),
+        );
+
+        let merkle_tree_state_transition_proofs = vec![state_transition_proof];
+        let prev_statements = vec![new_root_st, old_root_st, key_st];
+        operation_verify(
+            st,
+            op,
+            prev_statements,
+            vec![],
+            vec![],
+            merkle_tree_state_transition_proofs,
+        )
     }
 
     #[test]
@@ -2631,7 +3110,7 @@ mod tests {
                 OperationAux::PublicKeyOfIndex(0),
             );
             let prev_statements = vec![public_key_st, secret_key_st];
-            operation_verify(st, op, prev_statements, vec![], vec![secret_key])
+            operation_verify(st, op, prev_statements, vec![], vec![secret_key], vec![])
         })
     }
 
@@ -2654,7 +3133,9 @@ mod tests {
             OperationAux::PublicKeyOfIndex(0),
         );
         let prev_statements = vec![public_key_st, secret_key_st];
-        assert!(operation_verify(st, op, prev_statements, vec![], vec![secret_key]).is_err())
+        assert!(
+            operation_verify(st, op, prev_statements, vec![], vec![secret_key], vec![]).is_err()
+        )
     }
 
     #[test]
@@ -2676,7 +3157,9 @@ mod tests {
             OperationAux::None,
         );
         let prev_statements = vec![public_key_st, secret_key_st];
-        assert!(operation_verify(st, op, prev_statements, vec![], vec![secret_key]).is_err())
+        assert!(
+            operation_verify(st, op, prev_statements, vec![], vec![secret_key], vec![]).is_err()
+        )
     }
 
     #[test]
@@ -2703,7 +3186,8 @@ mod tests {
             op,
             prev_statements,
             vec![],
-            vec![SecretKey(BigUint::from(123u32))]
+            vec![SecretKey(BigUint::from(123u32))],
+            vec![]
         )
         .is_err())
     }
@@ -2727,7 +3211,9 @@ mod tests {
             OperationAux::PublicKeyOfIndex(0),
         );
         let prev_statements = vec![public_key_st, secret_key_st];
-        assert!(operation_verify(st, op, prev_statements, vec![], vec![secret_key]).is_err())
+        assert!(
+            operation_verify(st, op, prev_statements, vec![], vec![secret_key], vec![]).is_err()
+        )
     }
 
     fn helper_statement_arg_from_template(

src/backends/plonky2/mainpod/mod.rs

@@ -19,7 +19,10 @@ use crate::{
         error::{Error, Result},
         hash_common_data,
         mock::emptypod::MockEmptyPod,
-        primitives::{ec::schnorr::SecretKey, merkletree::MerkleClaimAndProof},
+        primitives::{
+            ec::schnorr::SecretKey,
+            merkletree::{MerkleClaimAndProof, MerkleTreeStateTransitionProof},
+        },
         recursion::{
             hash_verifier_data, prove_rec_circuit, RecursiveCircuit, RecursiveCircuitTarget,
         },
@@ -173,6 +176,33 @@ pub(crate) fn extract_merkle_proofs(
     Ok(table)
 }
 
+/// Extracts Merkle state transition proofs from container update ops.
+pub(crate) fn extract_merkle_tree_state_transition_proofs(
+    params: &Params,
+    aux_list: &mut [OperationAux],
+    operations: &[middleware::Operation],
+) -> Result<Vec<MerkleTreeStateTransitionProof>> {
+    let mut table = Vec::new();
+    for (i, op) in operations.iter().enumerate() {
+        let pf = match op {
+            middleware::Operation::ContainerInsertFromEntries(_, _, _, _, pf)
+            | middleware::Operation::ContainerUpdateFromEntries(_, _, _, _, pf)
+            | middleware::Operation::ContainerDeleteFromEntries(_, _, _, pf) => pf.clone(),
+            _ => continue,
+        };
+        aux_list[i] = OperationAux::MerkleTreeStateTransitionProofIndex(table.len());
+        table.push(pf);
+    }
+    if table.len() > params.max_merkle_tree_state_transition_proofs_containers {
+        return Err(Error::custom(format!(
+            "The number of required Merkle proofs ({}) exceeds the maximum number ({}).",
+            table.len(),
+            params.max_merkle_tree_state_transition_proofs_containers
+        )));
+    }
+    Ok(table)
+}
+
 pub(crate) fn extract_public_key_of(
     params: &Params,
     aux_list: &mut [OperationAux],
@@ -471,6 +501,9 @@ impl PodProver for Prover {
         let public_key_of_sks =
             extract_public_key_of(params, &mut aux_list, inputs.operations, inputs.statements)?;
 
+        let merkle_tree_state_transition_proofs =
+            extract_merkle_tree_state_transition_proofs(params, &mut aux_list, inputs.operations)?;
+
         let (statements, public_statements) = layout_statements(params, false, &inputs)?;
         let operations = process_private_statements_operations(
             params,
@@ -513,6 +546,7 @@ impl PodProver for Prover {
             operations,
             merkle_proofs,
             public_key_of_sks,
+            merkle_tree_state_transition_proofs,
             custom_predicate_batches,
             custom_predicate_verifications,
         };
@@ -912,14 +946,15 @@ pub mod tests {
             max_signed_pod_values: 2,
             max_public_statements: 2,
             num_public_statements_id: 4,
-            max_statement_args: 3,
+            max_statement_args: 4,
-            max_operation_args: 3,
+            max_operation_args: 4,
             max_custom_predicate_batches: 2,
             max_custom_predicate_verifications: 2,
             max_custom_predicate_arity: 2,
             max_custom_predicate_wildcards: 3,
             max_custom_batch_size: 2,
             max_merkle_proofs_containers: 2,
+            max_merkle_tree_state_transition_proofs_containers: 2,
             max_public_key_of: 2,
             max_depth_mt_containers: 4,
             max_depth_mt_vds: 6,
@@ -979,13 +1014,14 @@ pub mod tests {
             max_input_recursive_pods: 0,
             max_statements: 9,
             max_public_statements: 4,
-            max_statement_args: 3,
+            max_statement_args: 4,
-            max_operation_args: 3,
+            max_operation_args: 4,
             max_custom_predicate_arity: 3,
             max_custom_batch_size: 3,
             max_custom_predicate_wildcards: 4,
             max_custom_predicate_verifications: 2,
             max_merkle_proofs_containers: 0,
+            max_merkle_tree_state_transition_proofs_containers: 0,
             ..Default::default()
         };
         println!("{:#?}", params);

src/backends/plonky2/mainpod/operation.rs

@@ -6,7 +6,7 @@ use crate::{
     backends::plonky2::{
         error::{Error, Result},
         mainpod::Statement,
-        primitives::merkletree::MerkleClaimAndProof,
+        primitives::merkletree::{MerkleClaimAndProof, MerkleTreeStateTransitionProof},
     },
     middleware::{self, OperationType, Params},
 };
@@ -35,6 +35,7 @@ pub enum OperationAux {
     None,
     MerkleProofIndex(usize),
     PublicKeyOfIndex(usize),
+    MerkleTreeStateTransitionProofIndex(usize),
     CustomPredVerifyIndex(usize),
 }
 
@@ -46,12 +47,17 @@ impl OperationAux {
     fn table_offset_public_key_of(params: &Params) -> usize {
         Self::table_offset_merkle_proof(params) + params.max_merkle_proofs_containers
     }
-    fn table_offset_custom_pred_verify(params: &Params) -> usize {
+    fn table_offset_merkle_tree_state_transition_proof(params: &Params) -> usize {
         Self::table_offset_public_key_of(params) + params.max_public_key_of
     }
+    fn table_offset_custom_pred_verify(params: &Params) -> usize {
+        Self::table_offset_merkle_tree_state_transition_proof(params)
+            + params.max_merkle_tree_state_transition_proofs_containers
+    }
     pub(crate) fn table_size(params: &Params) -> usize {
         1 + params.max_merkle_proofs_containers
             + params.max_public_key_of
+            + params.max_merkle_tree_state_transition_proofs_containers
             + params.max_custom_predicate_verifications
     }
     pub fn table_index(&self, params: &Params) -> usize {
@@ -59,6 +65,9 @@ impl OperationAux {
             Self::None => 0,
             Self::MerkleProofIndex(i) => Self::table_offset_merkle_proof(params) + *i,
             Self::PublicKeyOfIndex(i) => Self::table_offset_public_key_of(params) + *i,
+            Self::MerkleTreeStateTransitionProofIndex(i) => {
+                Self::table_offset_merkle_tree_state_transition_proof(params) + *i
+            }
             Self::CustomPredVerifyIndex(i) => Self::table_offset_custom_pred_verify(params) + *i,
         }
     }
@@ -81,6 +90,7 @@ impl Operation {
         &self,
         statements: &[Statement],
         merkle_proofs: &[MerkleClaimAndProof],
+        merkle_tree_state_transition_proofs: &[MerkleTreeStateTransitionProof],
     ) -> Result<crate::middleware::Operation> {
         let deref_args = self
             .1
@@ -104,6 +114,17 @@ impl Operation {
                     .proof
                     .clone(),
             ),
+            OperationAux::MerkleTreeStateTransitionProofIndex(i) => {
+                crate::middleware::OperationAux::MerkleTreeStateTransitionProof(
+                    merkle_tree_state_transition_proofs
+                        .get(i)
+                        .ok_or(Error::custom(format!(
+                            "Missing Merkle state transition proof index {}",
+                            i
+                        )))?
+                        .clone(),
+                )
+            }
             OperationAux::PublicKeyOfIndex(_) => crate::middleware::OperationAux::None,
         };
         Ok(middleware::Operation::op(
@@ -133,6 +154,9 @@ impl fmt::Display for Operation {
             OperationAux::MerkleProofIndex(i) => write!(f, " merkle_proof_{:02}", i)?,
             OperationAux::CustomPredVerifyIndex(i) => write!(f, " custom_pred_verify_{:02}", i)?,
             OperationAux::PublicKeyOfIndex(i) => write!(f, " public_key_of_{:02}", i)?,
+            OperationAux::MerkleTreeStateTransitionProofIndex(i) => {
+                write!(f, " merkle_tree_state_transition_proof_{:02}", i)?
+            }
         }
         Ok(())
     }

src/backends/plonky2/mainpod/statement.rs

@@ -74,6 +74,21 @@ impl TryFrom<Statement> for middleware::Statement {
                     S::HashOf(a1.try_into()?, a2.try_into()?, a3.try_into()?)
                 }
                 (NP::PublicKeyOf, &[a1, a2]) => S::PublicKeyOf(a1.try_into()?, a2.try_into()?),
+                (NP::ContainerInsert, &[a1, a2, a3, a4]) => S::ContainerInsert(
+                    a1.try_into()?,
+                    a2.try_into()?,
+                    a3.try_into()?,
+                    a4.try_into()?,
+                ),
+                (NP::ContainerUpdate, &[a1, a2, a3, a4]) => S::ContainerUpdate(
+                    a1.try_into()?,
+                    a2.try_into()?,
+                    a3.try_into()?,
+                    a4.try_into()?,
+                ),
+                (NP::ContainerDelete, &[a1, a2, a3]) => {
+                    S::ContainerDelete(a1.try_into()?, a2.try_into()?, a3.try_into()?)
+                }
                 _ => Err(Error::custom(format!(
                     "Ill-formed statement expression {:?}",
                     s

src/backends/plonky2/mock/mainpod.rs

@@ -12,12 +12,12 @@ use crate::{
         basetypes::{Proof, VerifierOnlyCircuitData},
         error::{Error, Result},
         mainpod::{
-            calculate_id, extract_merkle_proofs, layout_statements,
+            calculate_id, extract_merkle_proofs, extract_merkle_tree_state_transition_proofs,
-            process_private_statements_operations, process_public_statements_operations, Operation,
+            layout_statements, process_private_statements_operations,
-            OperationAux, Statement,
+            process_public_statements_operations, Operation, OperationAux, Statement,
         },
         mock::emptypod::MockEmptyPod,
-        primitives::merkletree::MerkleClaimAndProof,
+        primitives::merkletree::{MerkleClaimAndProof, MerkleTreeStateTransitionProof},
         recursion::hash_verifier_data,
         signedpod::SignedPod,
     },
@@ -55,6 +55,8 @@ pub struct MockMainPod {
     public_statements: Vec<Statement>,
     // All Merkle proofs
     merkle_proofs_containers: Vec<MerkleClaimAndProof>,
+    // All Merkle tree state transition proofs
+    merkle_tree_state_transition_proofs_containers: Vec<MerkleTreeStateTransitionProof>,
 }
 
 impl Eq for MockMainPod {}
@@ -145,6 +147,7 @@ struct Data {
     operations: Vec<Operation>,
     statements: Vec<Statement>,
     merkle_proofs: Vec<MerkleClaimAndProof>,
+    merkle_tree_state_transition_proofs: Vec<MerkleTreeStateTransitionProof>,
     input_signed_pods: Vec<(usize, PodId, serde_json::Value)>,
     input_recursive_pods: Vec<(usize, Params, PodId, VDSet, serde_json::Value)>,
 }
@@ -176,6 +179,9 @@ impl MockMainPod {
         // Extract Merkle proofs and pad.
         let merkle_proofs =
             extract_merkle_proofs(params, &mut aux_list, inputs.operations, inputs.statements)?;
+        // Similarly for Merkle state transition proofs.
+        let merkle_tree_state_transition_proofs =
+            extract_merkle_tree_state_transition_proofs(params, &mut aux_list, inputs.operations)?;
 
         let operations = process_private_statements_operations(
             params,
@@ -214,6 +220,7 @@ impl MockMainPod {
             statements,
             operations,
             merkle_proofs_containers: merkle_proofs,
+            merkle_tree_state_transition_proofs_containers: merkle_tree_state_transition_proofs,
         })
     }
 
@@ -313,6 +320,7 @@ impl Pod for MockMainPod {
                     .deref(
                         &self.statements[..input_statement_offset + i],
                         &self.merkle_proofs_containers,
+                        &self.merkle_tree_state_transition_proofs_containers,
                     )?
                     .check_and_log(&self.params, &s.clone().try_into()?)
                     .map_err(|e| e.into())
@@ -362,6 +370,9 @@ impl Pod for MockMainPod {
             operations: self.operations.clone(),
             statements: self.statements.clone(),
             merkle_proofs: self.merkle_proofs_containers.clone(),
+            merkle_tree_state_transition_proofs: self
+                .merkle_tree_state_transition_proofs_containers
+                .clone(),
             input_signed_pods,
             input_recursive_pods,
         })
@@ -396,6 +407,7 @@ impl RecursivePod for MockMainPod {
             operations,
             statements,
             merkle_proofs,
+            merkle_tree_state_transition_proofs,
             input_signed_pods,
             input_recursive_pods,
         } = serde_json::from_value(data)?;
@@ -419,6 +431,7 @@ impl RecursivePod for MockMainPod {
             operations,
             statements,
             merkle_proofs_containers: merkle_proofs,
+            merkle_tree_state_transition_proofs_containers: merkle_tree_state_transition_proofs,
         }))
     }
 }

src/backends/plonky2/primitives/merkletree/circuit.rs

@@ -399,6 +399,7 @@ fn kv_hash_target(
 /// has been done correctly for the given new_key. This will allow verifying
 /// correct new leaf insertion, and leaf edition&deletion (if needed).
 /// See `MerkleTreeStateTransitionProof` struct for an explanation of the fields.
+#[derive(Clone, Serialize, Deserialize)]
 pub struct MerkleTreeStateTransitionProofTarget {
     pub(crate) max_depth: usize,
     // `enabled` determines if the merkleproof state transition verification is enabled

src/backends/plonky2/primitives/merkletree/mod.rs

@@ -344,6 +344,17 @@ impl MerkleTree {
                 // 2) at i==d, if old_siblings[i] != new_siblings[i]:
                 //     old_siblings[i] == EMPTY_HASH
                 //     new_siblings[i] == old_leaf_hash
+
+                // First rule out the case of insertion into empty tree.
+                if new_siblings.is_empty() {
+                    return (old_siblings.is_empty() && proof.old_root == EMPTY_HASH)
+                        .then_some(())
+                        .ok_or(TreeError::state_transition_fail(
+                            "new tree has no siblings yet old tree is not the empty tree"
+                                .to_string(),
+                        ));
+                }
+
                 let d = new_siblings.len() - 1;
                 old_siblings.resize(d + 1, EMPTY_HASH);
                 for i in 0..d {
@@ -550,6 +561,23 @@ pub struct MerkleTreeStateTransitionProof {
     pub(crate) siblings: Vec<Hash>,
 }
 
+impl MerkleTreeStateTransitionProof {
+    /// Value used for padding.
+    pub fn empty() -> Self {
+        let empty_proof_and_claim = MerkleClaimAndProof::empty();
+        Self {
+            op: MerkleTreeOp::Insert,
+            old_root: empty_proof_and_claim.root,
+            op_proof: empty_proof_and_claim.proof,
+            new_root: empty_proof_and_claim.root,
+            op_key: empty_proof_and_claim.key,
+            op_value: empty_proof_and_claim.value,
+            value: None,
+            siblings: vec![],
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 enum Node {
     None,

src/frontend/mod.rs

@@ -240,6 +240,60 @@ impl MainPodBuilder {
                 Operation(Native(LtFromEntries), vec![entry2, entry1], op.2)
             }),
             Native(GtToNotEqual) => Ok(Operation(Native(LtToNotEqual), op.1, op.2)),
+            Native(DictInsertFromEntries) => {
+                <[_; 4]>::try_from(op.1).map(|[new_dict, old_dict, key, value]| {
+                    Operation(
+                        Native(ContainerInsertFromEntries),
+                        vec![new_dict, old_dict, key, value],
+                        op.2,
+                    )
+                })
+            }
+            Native(DictUpdateFromEntries) => {
+                <[_; 4]>::try_from(op.1).map(|[new_dict, old_dict, key, value]| {
+                    Operation(
+                        Native(ContainerUpdateFromEntries),
+                        vec![new_dict, old_dict, key, value],
+                        op.2,
+                    )
+                })
+            }
+            Native(DictDeleteFromEntries) => {
+                <[_; 3]>::try_from(op.1).map(|[new_dict, old_dict, key]| {
+                    Operation(
+                        Native(ContainerDeleteFromEntries),
+                        vec![new_dict, old_dict, key],
+                        op.2,
+                    )
+                })
+            }
+            Native(SetInsertFromEntries) => {
+                <[_; 3]>::try_from(op.1).map(|[new_set, old_set, value]| {
+                    Operation(
+                        Native(ContainerInsertFromEntries),
+                        vec![new_set, old_set, value.clone(), value],
+                        op.2,
+                    )
+                })
+            }
+            Native(SetDeleteFromEntries) => {
+                <[_; 3]>::try_from(op.1).map(|[new_set, old_set, value]| {
+                    Operation(
+                        Native(ContainerDeleteFromEntries),
+                        vec![new_set, old_set, value],
+                        op.2,
+                    )
+                })
+            }
+            Native(ArrayUpdateFromEntries) => {
+                <[_; 4]>::try_from(op.1).map(|[new_arr, old_arr, i, value]| {
+                    Operation(
+                        Native(ContainerUpdateFromEntries),
+                        vec![new_arr, old_arr, i, value],
+                        op.2,
+                    )
+                })
+            }
             _ => Ok(op),
         }
         .map_err(|_| {
@@ -249,7 +303,10 @@ impl MainPodBuilder {
 
     /// Fills in auxiliary data if necessary/possible.
     fn fill_in_aux(op: Operation) -> Result<Operation> {
-        use NativeOperation::{ContainsFromEntries, NotContainsFromEntries};
+        use NativeOperation::{
+            ContainerDeleteFromEntries, ContainerInsertFromEntries, ContainerUpdateFromEntries,
+            ContainsFromEntries, NotContainsFromEntries,
+        };
         use OperationAux as OpAux;
         use OperationType::Native;
 
@@ -279,6 +336,45 @@ impl MainPodBuilder {
                 };
                 Ok(Operation(op_type.clone(), op.1, OpAux::MerkleProof(proof)))
             }
+            (Native(ContainerInsertFromEntries), OpAux::None)
+            | (Native(ContainerUpdateFromEntries), OpAux::None)
+            | (Native(ContainerDeleteFromEntries), OpAux::None) => {
+                let old_container =
+                    op.1.get(1)
+                        .and_then(|arg| arg.value())
+                        .ok_or(Error::custom(format!(
+                            "Invalid container argument for op {}.",
+                            op
+                        )))?;
+                let key =
+                    op.1.get(2)
+                        .and_then(|arg| arg.value())
+                        .ok_or(Error::custom(format!(
+                            "Invalid key argument for op {}.",
+                            op
+                        )))?;
+                let value =
+                    op.1.get(3)
+                        .and_then(|arg| arg.value())
+                        .ok_or(Error::custom(format!(
+                            "Invalid key argument for op {}.",
+                            op
+                        )));
+                let proof = match op_type {
+                    Native(ContainerInsertFromEntries) => {
+                        old_container.prove_insertion(key, value?)?
+                    }
+                    Native(ContainerUpdateFromEntries) => {
+                        old_container.prove_update(key, value?)?
+                    }
+                    _ => old_container.prove_deletion(key)?,
+                };
+                Ok(Operation(
+                    op_type.clone(),
+                    op.1,
+                    OpAux::MerkleTreeStateTransitionProof(proof),
+                ))
+            }
             _ => Ok(op),
         }
     }
@@ -405,6 +501,29 @@ impl MainPodBuilder {
                             return Err(native_arg_error());
                         }
                     }
+                    (ContainerInsertFromEntries, &[a1, a2, a3, a4]) => {
+                        let (r1, _v1) = a1.value_and_ref().ok_or_else(native_arg_error)?;
+                        let (r2, _v2) = a2.value_and_ref().ok_or_else(native_arg_error)?;
+                        let (r3, _v3) = a3.value_and_ref().ok_or_else(native_arg_error)?;
+                        let (r4, _v4) = a4.value_and_ref().ok_or_else(native_arg_error)?;
+                        // TODO: validate proof
+                        Statement::ContainerInsert(r1, r2, r3, r4)
+                    }
+                    (ContainerUpdateFromEntries, &[a1, a2, a3, a4]) => {
+                        let (r1, _v1) = a1.value_and_ref().ok_or_else(native_arg_error)?;
+                        let (r2, _v2) = a2.value_and_ref().ok_or_else(native_arg_error)?;
+                        let (r3, _v3) = a3.value_and_ref().ok_or_else(native_arg_error)?;
+                        let (r4, _v4) = a4.value_and_ref().ok_or_else(native_arg_error)?;
+                        // TODO: validate proof
+                        Statement::ContainerUpdate(r1, r2, r3, r4)
+                    }
+                    (ContainerDeleteFromEntries, &[a1, a2, a3]) => {
+                        let (r1, _v1) = a1.value_and_ref().ok_or_else(native_arg_error)?;
+                        let (r2, _v2) = a2.value_and_ref().ok_or_else(native_arg_error)?;
+                        let (r3, _v3) = a3.value_and_ref().ok_or_else(native_arg_error)?;
+                        // TODO: validate proof
+                        Statement::ContainerDelete(r1, r2, r3)
+                    }
                     (t, _) => {
                         if t.is_syntactic_sugar() {
                             return Err(Error::custom(format!(
@@ -741,7 +860,10 @@ pub mod tests {
             tickets_pod_full_flow, zu_kyc_pod_builder, zu_kyc_pod_request,
             zu_kyc_sign_pod_builders, EthDosHelper, MOCK_VD_SET,
         },
-        middleware::{containers::Dictionary, Value},
+        middleware::{
+            containers::{Array, Dictionary, Set},
+            Value,
+        },
     };
 
     // Check that frontend public statements agree with those
@@ -987,19 +1109,140 @@ pub mod tests {
         let st1 = builder.op(true, Operation::new_entry("key", "a")).unwrap();
         let st2 = builder.literal(false, Value::from(1)).unwrap();
 
-        builder
+        builder.pub_op(Operation(
-            .pub_op(Operation(
             // OperationType
             OperationType::Native(NativeOperation::DictContainsFromEntries),
             // Vec<OperationArg>
             vec![
-                    OperationArg::Statement(st0),
+                OperationArg::Statement(st0.clone()),
                 OperationArg::Statement(st1),
                 OperationArg::Statement(st2),
             ],
             OperationAux::MerkleProof(dict.prove(&Key::from("a")).unwrap().1),
-            ))
+        ))?;
-            .unwrap();
+
+        let mut new_dict = dict.clone();
+        new_dict.insert(&Key::from("d"), &Value::from(4))?;
+
+        builder.pub_op(Operation(
+            OperationType::Native(NativeOperation::DictInsertFromEntries),
+            vec![
+                Value::from(new_dict.clone()).into(),
+                OperationArg::Statement(st0.clone()),
+                "d".into(),
+                4.into(),
+            ],
+            OperationAux::None,
+        ))?;
+
+        let mut new_old_dict = new_dict.clone();
+        new_old_dict.delete(&Key::from("d"))?;
+
+        assert_eq!(new_old_dict, dict);
+
+        builder.pub_op(Operation(
+            OperationType::Native(NativeOperation::DictDeleteFromEntries),
+            vec![
+                OperationArg::Statement(st0.clone()),
+                Value::from(new_dict).into(),
+                "d".into(),
+            ],
+            OperationAux::None,
+        ))?;
+
+        new_old_dict.update(&Key::from("c"), &55.into())?;
+
+        builder.pub_op(Operation(
+            OperationType::Native(NativeOperation::DictUpdateFromEntries),
+            vec![
+                Value::from(new_old_dict).into(),
+                OperationArg::Statement(st0.clone()),
+                "c".into(),
+                55.into(),
+            ],
+            OperationAux::None,
+        ))?;
+
+        let main_prover = MockProver {};
+        let main_pod = builder.prove(&main_prover).unwrap();
+
+        println!("{}", main_pod);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_sets() -> Result<()> {
+        let params = Params::default();
+        let vd_set = &*MOCK_VD_SET;
+        let mut builder = MainPodBuilder::new(&params, vd_set);
+
+        let empty_set = Set::new(params.max_depth_mt_containers, [].into())?;
+
+        let mut set1 = empty_set.clone();
+        set1.insert(&1.into())?;
+
+        let mut set2 = set1.clone();
+        set2.delete(&1.into())?;
+
+        assert_eq!(set2, empty_set);
+
+        builder.pub_op(Operation(
+            // OperationType
+            OperationType::Native(NativeOperation::SetInsertFromEntries),
+            // Vec<OperationArg>
+            vec![
+                Value::from(set1.clone()).into(),
+                Value::from(empty_set.clone()).into(),
+                1.into(),
+            ],
+            OperationAux::None,
+        ))?;
+
+        builder.pub_op(Operation(
+            // OperationType
+            OperationType::Native(NativeOperation::SetDeleteFromEntries),
+            // Vec<OperationArg>
+            vec![
+                Value::from(empty_set.clone()).into(),
+                Value::from(set1.clone()).into(),
+                1.into(),
+            ],
+            OperationAux::None,
+        ))?;
+
+        let main_prover = MockProver {};
+        let main_pod = builder.prove(&main_prover).unwrap();
+
+        println!("{}", main_pod);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_arrays() -> Result<()> {
+        let params = Params::default();
+        let vd_set = &*MOCK_VD_SET;
+        let mut builder = MainPodBuilder::new(&params, vd_set);
+
+        let array1 = Array::new(params.max_depth_mt_containers, [1.into()].into())?;
+
+        let mut array2 = array1.clone();
+        array2.update(0, &5.into())?;
+
+        builder.pub_op(Operation(
+            // OperationType
+            OperationType::Native(NativeOperation::ArrayUpdateFromEntries),
+            // Vec<OperationArg>
+            vec![
+                Value::from(array2.clone()).into(),
+                Value::from(array1.clone()).into(),
+                0.into(),
+                5.into(),
+            ],
+            OperationAux::None,
+        ))?;
+
         let main_prover = MockProver {};
         let main_pod = builder.prove(&main_prover).unwrap();
 

src/frontend/operation.rs

@@ -139,6 +139,21 @@ macro_rules! op_impl_oa {
             )
         }
     };
+
+    ($fn_name: ident, $op_name: ident, 4) => {
+        pub fn $fn_name(
+            a1: impl Into<OperationArg>,
+            a2: impl Into<OperationArg>,
+            a3: impl Into<OperationArg>,
+            a4: impl Into<OperationArg>,
+        ) -> Self {
+            Self(
+                OperationType::Native(NativeOperation::$op_name),
+                vec![a1.into(), a2.into(), a3.into(), a4.into()],
+                OperationAux::None,
+            )
+        }
+    };
 }
 
 macro_rules! op_impl_st {
@@ -201,4 +216,10 @@ impl Operation {
     op_impl_oa!(set_not_contains, SetNotContainsFromEntries, 2);
     op_impl_oa!(array_contains, ArrayContainsFromEntries, 3);
     op_impl_oa!(public_key_of, PublicKeyOf, 2);
+    op_impl_oa!(dict_insert, DictInsertFromEntries, 4);
+    op_impl_oa!(dict_update, DictUpdateFromEntries, 4);
+    op_impl_oa!(dict_delete, DictDeleteFromEntries, 3);
+    op_impl_oa!(set_insert, SetInsertFromEntries, 3);
+    op_impl_oa!(set_delete, SetDeleteFromEntries, 3);
+    op_impl_oa!(array_update, ArrayUpdateFromEntries, 4);
 }

src/lang/processor.rs

@@ -381,7 +381,16 @@ fn validate_and_build_statement_template(
                 | NativePredicate::SumOf
                 | NativePredicate::ProductOf
                 | NativePredicate::MaxOf
-                | NativePredicate::HashOf => 3,
+                | NativePredicate::HashOf
+                | NativePredicate::ContainerDelete
+                | NativePredicate::DictDelete
+                | NativePredicate::SetInsert
+                | NativePredicate::SetDelete => 3,
+                NativePredicate::ContainerInsert
+                | NativePredicate::ContainerUpdate
+                | NativePredicate::DictInsert
+                | NativePredicate::DictUpdate
+                | NativePredicate::ArrayUpdate => 4,
                 NativePredicate::None | NativePredicate::False => 0,
             };
 

src/middleware/containers.rs

@@ -9,7 +9,10 @@ use serde::{Deserialize, Deserializer, Serialize};
 use super::serialization::{ordered_map, ordered_set};
 #[cfg(feature = "backend_plonky2")]
 use crate::backends::plonky2::primitives::merkletree::{MerkleProof, MerkleTree};
-use crate::middleware::{Error, Hash, Key, RawValue, Result, Value};
+use crate::{
+    backends::plonky2::primitives::merkletree::MerkleTreeStateTransitionProof,
+    middleware::{Error, Hash, Key, RawValue, Result, Value},
+};
 
 /// Dictionary: the user original keys and values are hashed to be used in the leaf.
 ///    leaf.key=hash(original_key)
@@ -52,6 +55,21 @@ impl Dictionary {
     pub fn prove_nonexistence(&self, key: &Key) -> Result<MerkleProof> {
         Ok(self.mt.prove_nonexistence(&key.raw())?)
     }
+    pub fn insert(&mut self, key: &Key, value: &Value) -> Result<MerkleTreeStateTransitionProof> {
+        let mtp = self.mt.insert(&key.raw(), &value.raw())?;
+        self.kvs.insert(key.clone(), value.clone());
+        Ok(mtp)
+    }
+    pub fn update(&mut self, key: &Key, value: &Value) -> Result<MerkleTreeStateTransitionProof> {
+        let mtp = self.mt.update(&key.raw(), &value.raw())?;
+        self.kvs.insert(key.clone(), value.clone());
+        Ok(mtp)
+    }
+    pub fn delete(&mut self, key: &Key) -> Result<MerkleTreeStateTransitionProof> {
+        let mtp = self.mt.delete(&key.raw())?;
+        self.kvs.remove(key);
+        Ok(mtp)
+    }
     pub fn verify(
         max_depth: usize,
         root: Hash,
@@ -79,6 +97,12 @@ impl Dictionary {
             max_depth, root, proof, &key,
         )?)
     }
+    pub fn verify_state_transition(
+        max_depth: usize,
+        proof: &MerkleTreeStateTransitionProof,
+    ) -> Result<()> {
+        MerkleTree::verify_state_transition(max_depth, proof).map_err(|e| e.into())
+    }
     // TODO: Rename to dict to be consistent maybe?
     pub fn kvs(&self) -> &HashMap<Key, Value> {
         &self.kvs
@@ -156,6 +180,17 @@ impl Set {
         let rv = value.raw();
         Ok(self.mt.prove_nonexistence(&rv)?)
     }
+    pub fn insert(&mut self, value: &Value) -> Result<MerkleTreeStateTransitionProof> {
+        let raw_value = value.raw();
+        let mtp = self.mt.insert(&raw_value, &raw_value)?;
+        self.set.insert(value.clone());
+        Ok(mtp)
+    }
+    pub fn delete(&mut self, value: &Value) -> Result<MerkleTreeStateTransitionProof> {
+        let mtp = self.mt.delete(&value.raw())?;
+        self.set.remove(value);
+        Ok(mtp)
+    }
     pub fn verify(max_depth: usize, root: Hash, proof: &MerkleProof, value: &Value) -> Result<()> {
         let rv = value.raw();
         Ok(MerkleTree::verify(max_depth, root, proof, &rv, &rv)?)
@@ -171,6 +206,12 @@ impl Set {
             max_depth, root, proof, &rv,
         )?)
     }
+    pub fn verify_state_transition(
+        max_depth: usize,
+        proof: &MerkleTreeStateTransitionProof,
+    ) -> Result<()> {
+        MerkleTree::verify_state_transition(max_depth, proof).map_err(|e| e.into())
+    }
     pub fn set(&self) -> &HashSet<Value> {
         &self.set
     }
@@ -244,6 +285,11 @@ impl Array {
         let value = self.array.get(i).expect("valid index");
         Ok((value, mtp))
     }
+    pub fn update(&mut self, i: usize, value: &Value) -> Result<MerkleTreeStateTransitionProof> {
+        let mtp = self.mt.update(&(i as i64).into(), &value.raw())?;
+        self.array[i] = value.clone();
+        Ok(mtp)
+    }
     pub fn verify(
         max_depth: usize,
         root: Hash,
@@ -259,6 +305,12 @@ impl Array {
             &value.raw(),
         )?)
     }
+    pub fn verify_state_transition(
+        max_depth: usize,
+        proof: &MerkleTreeStateTransitionProof,
+    ) -> Result<()> {
+        MerkleTree::verify_state_transition(max_depth, proof).map_err(|e| e.into())
+    }
     pub fn array(&self) -> &[Value] {
         &self.array
     }

src/middleware/mod.rs

@@ -35,7 +35,7 @@ pub use statement::*;
 
 use crate::backends::plonky2::primitives::{
     ec::{curve::Point as PublicKey, schnorr::SecretKey},
-    merkletree::MerkleProof,
+    merkletree::{MerkleProof, MerkleTreeStateTransitionProof},
 };
 
 pub const SELF: PodId = PodId(SELF_ID_HASH);
@@ -535,6 +535,59 @@ impl Value {
             ))),
         }
     }
+    /// Returns a Merkle state transition proof for inserting a
+    /// key-value pair (if applicable).
+    pub(crate) fn prove_insertion(
+        &self,
+        key: &Value,
+        value: &Value,
+    ) -> Result<MerkleTreeStateTransitionProof> {
+        let container = self.typed().clone();
+        match container {
+            TypedValue::Dictionary(mut d) => d.insert(&key.typed().clone().try_into()?, value),
+            TypedValue::Set(mut s) => s.insert(value),
+            _ => Err(Error::custom(format!(
+                "Invalid container value {}",
+                self.typed()
+            ))),
+        }
+    }
+    /// Returns a Merkle state transition proof for updating a
+    /// key-value pair (if applicable).
+    pub(crate) fn prove_update(
+        &self,
+        key: &Value,
+        value: &Value,
+    ) -> Result<MerkleTreeStateTransitionProof> {
+        let container = self.typed().clone();
+        match container {
+            TypedValue::Array(mut a) => match key.typed() {
+                TypedValue::Int(i) if i >= &0 => a.update(*i as usize, value),
+                _ => Err(Error::custom(format!(
+                    "Invalid key {} for container {}.",
+                    key, self
+                )))?,
+            },
+            TypedValue::Dictionary(mut d) => d.update(&key.typed().clone().try_into()?, value),
+            _ => Err(Error::custom(format!(
+                "Invalid container value {} for update op",
+                self.typed()
+            ))),
+        }
+    }
+    /// Returns a Merkle state transition proof for deleting a
+    /// key (if applicable).
+    pub(crate) fn prove_deletion(&self, key: &Value) -> Result<MerkleTreeStateTransitionProof> {
+        let container = self.typed().clone();
+        match container {
+            TypedValue::Dictionary(mut d) => d.delete(&key.typed().clone().try_into()?),
+            TypedValue::Set(mut s) => s.delete(key),
+            _ => Err(Error::custom(format!(
+                "Invalid container value {}",
+                self.typed()
+            ))),
+        }
+    }
 }
 
 // A Value can be created from any type Into<TypedValue> type: bool, string-like, i64, ...
@@ -760,6 +813,8 @@ pub struct Params {
     pub max_custom_predicate_wildcards: usize,
     // maximum number of merkle proofs used for container operations
     pub max_merkle_proofs_containers: usize,
+    // maximum number of merkle tree state transition proofs used for container update operations
+    pub max_merkle_tree_state_transition_proofs_containers: usize,
     // maximum depth for merkle tree gadget used for container operations
     pub max_depth_mt_containers: usize,
     // maximum depth of the merkle tree gadget used for verifier_data membership
@@ -804,6 +859,7 @@ impl Default for Params {
             max_custom_predicate_wildcards: 10,
             max_custom_batch_size: 5, // TODO: Move down to 4?
             max_merkle_proofs_containers: 5,
+            max_merkle_tree_state_transition_proofs_containers: 5,
             max_depth_mt_containers: 32,
             max_depth_mt_vds: 6, // up to 64 (2^6) different pod circuits
             max_public_key_of: 2,

src/middleware/operation.rs

@@ -10,7 +10,7 @@ use crate::{
             curve::{Point as PublicKey, GROUP_ORDER},
             schnorr::SecretKey,
         },
-        merkletree::{MerkleProof, MerkleTree},
+        merkletree::{MerkleProof, MerkleTree, MerkleTreeOp, MerkleTreeStateTransitionProof},
     },
     middleware::{
         hash_values, AnchoredKey, CustomPredicate, CustomPredicateRef, Error, NativePredicate,
@@ -29,6 +29,7 @@ pub enum OperationType {
 pub enum OperationAux {
     None,
     MerkleProof(MerkleProof),
+    MerkleTreeStateTransitionProof(MerkleTreeStateTransitionProof),
 }
 
 impl fmt::Display for OperationAux {
@@ -36,6 +37,10 @@ impl fmt::Display for OperationAux {
         match self {
             Self::None => write!(f, "<no aux>")?,
             Self::MerkleProof(pf) => write!(f, "merkle_proof({})", pf)?,
+            // TODO: Make this look nicer.
+            Self::MerkleTreeStateTransitionProof(pf) => {
+                write!(f, "merkle_tree_state_transition_proof({:?})", pf)?
+            }
         }
         Ok(())
     }
@@ -80,6 +85,9 @@ pub enum NativeOperation {
     MaxOf = 13,
     HashOf = 14,
     PublicKeyOf = 15,
+    ContainerInsertFromEntries = 16,
+    ContainerUpdateFromEntries = 17,
+    ContainerDeleteFromEntries = 18,
 
     // Syntactic sugar operations.  These operations are not supported by the backend.  The
     // frontend compiler is responsible of translating these operations into the operations above.
@@ -91,6 +99,12 @@ pub enum NativeOperation {
     GtEqFromEntries = 1006,
     GtFromEntries = 1007,
     GtToNotEqual = 1008,
+    DictInsertFromEntries = 1009,
+    DictUpdateFromEntries = 1010,
+    DictDeleteFromEntries = 1011,
+    SetInsertFromEntries = 1012,
+    SetDeleteFromEntries = 1013,
+    ArrayUpdateFromEntries = 1014,
 }
 
 impl NativeOperation {
@@ -140,6 +154,15 @@ impl OperationType {
                 NativeOperation::PublicKeyOf => {
                     Some(Predicate::Native(NativePredicate::PublicKeyOf))
                 }
+                NativeOperation::ContainerInsertFromEntries => {
+                    Some(Predicate::Native(NativePredicate::ContainerInsert))
+                }
+                NativeOperation::ContainerUpdateFromEntries => {
+                    Some(Predicate::Native(NativePredicate::ContainerUpdate))
+                }
+                NativeOperation::ContainerDeleteFromEntries => {
+                    Some(Predicate::Native(NativePredicate::ContainerDelete))
+                }
                 no => unreachable!("Unexpected syntactic sugar op {:?}", no),
             },
             OperationType::Custom(cpr) => Some(Predicate::Custom(cpr.clone())),
@@ -175,6 +198,26 @@ pub enum Operation {
     MaxOf(Statement, Statement, Statement),
     HashOf(Statement, Statement, Statement),
     PublicKeyOf(Statement, Statement),
+    ContainerInsertFromEntries(
+        /* new_root */ Statement,
+        /* old_root */ Statement,
+        /*  key    */ Statement,
+        /*  value  */ Statement,
+        /*  proof  */ MerkleTreeStateTransitionProof,
+    ),
+    ContainerUpdateFromEntries(
+        /* new_root */ Statement,
+        /* old_root */ Statement,
+        /*  key    */ Statement,
+        /*  value  */ Statement,
+        /*  proof  */ MerkleTreeStateTransitionProof,
+    ),
+    ContainerDeleteFromEntries(
+        /* new_root */ Statement,
+        /* old_root */ Statement,
+        /*  key    */ Statement,
+        /*  proof  */ MerkleTreeStateTransitionProof,
+    ),
     Custom(CustomPredicateRef, Vec<Statement>),
 }
 
@@ -215,6 +258,13 @@ impl Operation {
             Self::MaxOf(_, _, _) => OT::Native(MaxOf),
             Self::HashOf(_, _, _) => OT::Native(HashOf),
             Self::PublicKeyOf(_, _) => OT::Native(PublicKeyOf),
+            Self::ContainerInsertFromEntries(_, _, _, _, _) => {
+                OT::Native(ContainerInsertFromEntries)
+            }
+            Self::ContainerUpdateFromEntries(_, _, _, _, _) => {
+                OT::Native(ContainerUpdateFromEntries)
+            }
+            Self::ContainerDeleteFromEntries(_, _, _, _) => OT::Native(ContainerDeleteFromEntries),
             Self::Custom(cpr, _) => OT::Custom(cpr.clone()),
         }
     }
@@ -237,6 +287,9 @@ impl Operation {
             Self::MaxOf(s1, s2, s3) => vec![s1, s2, s3],
             Self::HashOf(s1, s2, s3) => vec![s1, s2, s3],
             Self::PublicKeyOf(s1, s2) => vec![s1, s2],
+            Self::ContainerInsertFromEntries(s1, s2, s3, s4, _pf) => vec![s1, s2, s3, s4],
+            Self::ContainerUpdateFromEntries(s1, s2, s3, s4, _pf) => vec![s1, s2, s3, s4],
+            Self::ContainerDeleteFromEntries(s1, s2, s3, _pf) => vec![s1, s2, s3],
             Self::Custom(_, args) => args,
         }
     }
@@ -290,6 +343,33 @@ impl Operation {
                     Self::HashOf(s1.clone(), s2.clone(), s3.clone())
                 }
                 (NO::PublicKeyOf, &[s1, s2], OA::None) => Self::PublicKeyOf(s1.clone(), s2.clone()),
+                (
+                    NO::ContainerInsertFromEntries,
+                    &[s1, s2, s3, s4],
+                    OA::MerkleTreeStateTransitionProof(pf),
+                ) => Self::ContainerInsertFromEntries(
+                    s1.clone(),
+                    s2.clone(),
+                    s3.clone(),
+                    s4.clone(),
+                    pf,
+                ),
+                (
+                    NO::ContainerUpdateFromEntries,
+                    &[s1, s2, s3, s4],
+                    OA::MerkleTreeStateTransitionProof(pf),
+                ) => Self::ContainerUpdateFromEntries(
+                    s1.clone(),
+                    s2.clone(),
+                    s3.clone(),
+                    s4.clone(),
+                    pf,
+                ),
+                (
+                    NO::ContainerDeleteFromEntries,
+                    &[s1, s2, s3],
+                    OA::MerkleTreeStateTransitionProof(pf),
+                ) => Self::ContainerDeleteFromEntries(s1.clone(), s2.clone(), s3.clone(), pf),
                 _ => Err(Error::custom(format!(
                     "Ill-formed operation {:?} with {} arguments {:?} and aux {:?}.",
                     op_code,
@@ -392,6 +472,67 @@ impl Operation {
             (Self::PublicKeyOf(s1, s2), PublicKeyOf(v3, v4)) => {
                 Self::check_public_key(&val(v3, s1)?, &val(v4, s2)?)?
             }
+            (
+                Self::ContainerInsertFromEntries(new_root_s, old_root_s, key_s, val_s, pf),
+                ContainerInsert(new_root_v, old_root_v, key_v, val_v),
+            ) => {
+                let old_root = val(old_root_v, old_root_s)?;
+                let new_root = val(new_root_v, new_root_s)?;
+                let key = val(key_v, key_s)?;
+                let value = val(val_v, val_s)?;
+                (pf.op == MerkleTreeOp::Insert
+                    && Value::from(pf.old_root) == old_root
+                    && Value::from(pf.new_root) == new_root
+                    && pf.op_key == key.raw()
+                    && pf.op_value == value.raw())
+                .then_some(())
+                .ok_or(Error::custom(
+                    "The provided Merkle tree state transition proof does not match the claim."
+                        .into(),
+                ))?;
+                MerkleTree::verify_state_transition(params.max_depth_mt_containers, pf)?;
+                true
+            }
+            (
+                Self::ContainerUpdateFromEntries(new_root_s, old_root_s, key_s, val_s, pf),
+                ContainerUpdate(new_root_v, old_root_v, key_v, val_v),
+            ) => {
+                let old_root = val(old_root_v, old_root_s)?;
+                let new_root = val(new_root_v, new_root_s)?;
+                let key = val(key_v, key_s)?;
+                let value = val(val_v, val_s)?;
+                (pf.op == MerkleTreeOp::Update
+                    && Value::from(pf.old_root) == old_root
+                    && Value::from(pf.new_root) == new_root
+                    && pf.op_key == key.raw()
+                    && pf.op_value == value.raw())
+                .then_some(())
+                .ok_or(Error::custom(
+                    "The provided Merkle tree state transition proof does not match the claim."
+                        .into(),
+                ))?;
+                MerkleTree::verify_state_transition(params.max_depth_mt_containers, pf)?;
+                true
+            }
+            (
+                Self::ContainerDeleteFromEntries(new_root_s, old_root_s, key_s, pf),
+                ContainerDelete(new_root_v, old_root_v, key_v),
+            ) => {
+                let old_root = val(old_root_v, old_root_s)?;
+                let new_root = val(new_root_v, new_root_s)?;
+                let key = val(key_v, key_s)?;
+                (pf.op == MerkleTreeOp::Delete
+                    && Value::from(pf.old_root) == old_root
+                    && Value::from(pf.new_root) == new_root
+                    && pf.op_key == key.raw())
+                .then_some(())
+                .ok_or(Error::custom(
+                    "The provided Merkle tree state transition proof does not match the claim."
+                        .into(),
+                ))?;
+                MerkleTree::verify_state_transition(params.max_depth_mt_containers, pf)?;
+                true
+            }
             (Self::Custom(CustomPredicateRef { batch, index }, args), Custom(cpr, s_args))
                 if batch == &cpr.batch && index == &cpr.index =>
             {
@@ -687,6 +828,126 @@ mod tests {
         })
     }
 
+    #[test]
+    fn check_container_update_ops() -> Result<()> {
+        let params = Params::default();
+        let pod_id = PodId::default();
+        let new_root_ak = AnchoredKey::new(pod_id, Key::new("new_root".into()));
+        let old_root_ak = AnchoredKey::new(pod_id, Key::new("new_root".into()));
+        let key_ak = AnchoredKey::new(pod_id, Key::new("key".into()));
+        let val_ak = AnchoredKey::new(pod_id, Key::new("value".into()));
+
+        // Form Merkle tree
+        let kvs = (0..10)
+            .map(|i| (hash_value(&i.into()).into(), i.into()))
+            .collect::<HashMap<_, _>>();
+        let mut mt = MerkleTree::new(params.max_depth_mt_containers, &kvs)?;
+
+        // Check insertion proofs
+        (11..20)
+            .map(|i| (hash_value(&i.into()).into(), i.into()))
+            .try_for_each(|(k, v)| {
+                let old_root_s = Statement::Equal(old_root_ak.clone().into(), mt.root().into());
+                let mtp = mt.insert(&k, &v)?;
+                // Form op args
+                let new_root_s = Statement::Equal(new_root_ak.clone().into(), mt.root().into());
+                let key_s = Statement::Equal(key_ak.clone().into(), k.into());
+                let value_s = Statement::Equal(val_ak.clone().into(), v.into());
+
+                // Form op
+                let op = Operation::ContainerInsertFromEntries(
+                    new_root_s, old_root_s, key_s, value_s, mtp,
+                );
+                // Form output statement
+                let st = Statement::ContainerInsert(
+                    new_root_ak.clone().into(),
+                    old_root_ak.clone().into(),
+                    key_ak.clone().into(),
+                    val_ak.clone().into(),
+                );
+
+                // Check op against output statement
+                op.check(&params, &st).and_then(|ind| {
+                    if ind {
+                        Ok(())
+                    } else {
+                        Err(Error::custom(format!(
+                            "Insertion op check failed for pair ({},{})",
+                            k, v
+                        )))
+                    }
+                })
+            })?;
+
+        // Check update proofs
+        (11..20)
+            .map(|i| (hash_value(&i.into()).into(), (i + 1).into()))
+            .try_for_each(|(k, v)| {
+                let old_root_s = Statement::Equal(old_root_ak.clone().into(), mt.root().into());
+                let mtp = mt.update(&k, &v)?;
+                // Form op args
+                let new_root_s = Statement::Equal(new_root_ak.clone().into(), mt.root().into());
+                let key_s = Statement::Equal(key_ak.clone().into(), k.into());
+                let value_s = Statement::Equal(val_ak.clone().into(), v.into());
+
+                // Form op
+                let op = Operation::ContainerUpdateFromEntries(
+                    new_root_s, old_root_s, key_s, value_s, mtp,
+                );
+                // Form output statement
+                let st = Statement::ContainerUpdate(
+                    new_root_ak.clone().into(),
+                    old_root_ak.clone().into(),
+                    key_ak.clone().into(),
+                    val_ak.clone().into(),
+                );
+
+                // Check op against output statement
+                op.check(&params, &st).and_then(|ind| {
+                    if ind {
+                        Ok(())
+                    } else {
+                        Err(Error::custom(format!(
+                            "Update op check failed for pair ({},{})",
+                            k, v
+                        )))
+                    }
+                })
+            })?;
+
+        // Check deletion proofs
+        (11..20)
+            .map(|i| hash_value(&i.into()).into())
+            .try_for_each(|k| {
+                let old_root_s = Statement::Equal(old_root_ak.clone().into(), mt.root().into());
+                let mtp = mt.delete(&k)?;
+                // Form op args
+                let new_root_s = Statement::Equal(new_root_ak.clone().into(), mt.root().into());
+                let key_s = Statement::Equal(key_ak.clone().into(), k.into());
+
+                // Form op
+                let op = Operation::ContainerDeleteFromEntries(new_root_s, old_root_s, key_s, mtp);
+                // Form output statement
+                let st = Statement::ContainerDelete(
+                    new_root_ak.clone().into(),
+                    old_root_ak.clone().into(),
+                    key_ak.clone().into(),
+                );
+
+                // Check op against output statement
+                op.check(&params, &st).and_then(|ind| {
+                    if ind {
+                        Ok(())
+                    } else {
+                        Err(Error::custom(format!(
+                            "Deletion op check failed for key {}",
+                            k
+                        )))
+                    }
+                })
+            })
+    }
+
     #[test]
     fn check_public_key_of_op() -> Result<()> {
         let fixed_sk = SecretKey(BigUint::from(0x1234567890abcdefu64));

src/middleware/statement.rs

@@ -34,6 +34,9 @@ pub enum NativePredicate {
     MaxOf = 10,
     HashOf = 11,
     PublicKeyOf = 12,
+    ContainerInsert = 13,
+    ContainerUpdate = 14,
+    ContainerDelete = 15,
 
     // Syntactic sugar predicates.  These predicates are not supported by the backend.  The
     // frontend compiler is responsible of translating these predicates into the predicates above.
@@ -44,6 +47,12 @@ pub enum NativePredicate {
     ArrayContains = 1004, // there is no ArrayNotContains
     GtEq = 1005,
     Gt = 1006,
+    DictInsert = 1009,
+    DictUpdate = 1010,
+    DictDelete = 1011,
+    SetInsert = 1012,
+    SetDelete = 1013,
+    ArrayUpdate = 1014,
 }
 
 impl Display for NativePredicate {
@@ -64,11 +73,20 @@ impl Display for NativePredicate {
             NativePredicate::MaxOf => "MaxOf",
             NativePredicate::HashOf => "HashOf",
             NativePredicate::PublicKeyOf => "PublicKeyOf",
+            NativePredicate::ContainerInsert => "ContainerInsert",
+            NativePredicate::ContainerUpdate => "ContainerUpdate",
+            NativePredicate::ContainerDelete => "ContainerDelete",
             NativePredicate::DictContains => "DictContains",
             NativePredicate::DictNotContains => "DictNotContains",
             NativePredicate::ArrayContains => "ArrayContains",
             NativePredicate::SetContains => "SetContains",
             NativePredicate::SetNotContains => "SetNotContains",
+            NativePredicate::DictInsert => "DictInsert",
+            NativePredicate::DictUpdate => "DictUpdate",
+            NativePredicate::DictDelete => "DictDelete",
+            NativePredicate::SetInsert => "SetInsert",
+            NativePredicate::SetDelete => "SetDelete",
+            NativePredicate::ArrayUpdate => "ArrayUpdate",
         };
         write!(f, "{}", s)
     }
@@ -178,6 +196,23 @@ pub enum Statement {
     MaxOf(ValueRef, ValueRef, ValueRef),
     HashOf(ValueRef, ValueRef, ValueRef),
     PublicKeyOf(ValueRef, ValueRef),
+    ContainerInsert(
+        /* new_root */ ValueRef,
+        /* old_root */ ValueRef,
+        /*   key    */ ValueRef,
+        /*  value   */ ValueRef,
+    ),
+    ContainerUpdate(
+        /* new_root */ ValueRef,
+        /* old_root */ ValueRef,
+        /*   key    */ ValueRef,
+        /*  value   */ ValueRef,
+    ),
+    ContainerDelete(
+        /* new_root */ ValueRef,
+        /* old_root */ ValueRef,
+        /*   key    */ ValueRef,
+    ),
     Custom(CustomPredicateRef, Vec<Value>),
 }
 
@@ -196,6 +231,16 @@ macro_rules! statement_constructor {
             Self::$cons_name(v1.into(), v2.into(), v3.into())
         }
     };
+    ($var_name: ident, $cons_name: ident, 4) => {
+        pub fn $var_name(
+            v1: impl Into<ValueRef>,
+            v2: impl Into<ValueRef>,
+            v3: impl Into<ValueRef>,
+            v4: impl Into<ValueRef>,
+        ) -> Self {
+            Self::$cons_name(v1.into(), v2.into(), v3.into(), v4.into())
+        }
+    };
 }
 
 impl Statement {
@@ -213,6 +258,9 @@ impl Statement {
     statement_constructor!(max_of, MaxOf, 3);
     statement_constructor!(hash_of, HashOf, 3);
     statement_constructor!(public_key_of, PublicKeyOf, 2);
+    statement_constructor!(insert, ContainerInsert, 4);
+    statement_constructor!(update, ContainerUpdate, 4);
+    statement_constructor!(delete, ContainerDelete, 3);
     pub fn predicate(&self) -> Predicate {
         use Predicate::*;
         match self {
@@ -228,6 +276,9 @@ impl Statement {
             Self::MaxOf(_, _, _) => Native(NativePredicate::MaxOf),
             Self::HashOf(_, _, _) => Native(NativePredicate::HashOf),
             Self::PublicKeyOf(_, _) => Native(NativePredicate::PublicKeyOf),
+            Self::ContainerInsert(_, _, _, _) => Native(NativePredicate::ContainerInsert),
+            Self::ContainerUpdate(_, _, _, _) => Native(NativePredicate::ContainerUpdate),
+            Self::ContainerDelete(_, _, _) => Native(NativePredicate::ContainerDelete),
             Self::Custom(cpr, _) => Custom(cpr.clone()),
         }
     }
@@ -246,6 +297,13 @@ impl Statement {
             Self::MaxOf(ak1, ak2, ak3) => vec![ak1.into(), ak2.into(), ak3.into()],
             Self::HashOf(ak1, ak2, ak3) => vec![ak1.into(), ak2.into(), ak3.into()],
             Self::PublicKeyOf(ak1, ak2) => vec![ak1.into(), ak2.into()],
+            Self::ContainerInsert(ak1, ak2, ak3, ak4) => {
+                vec![ak1.into(), ak2.into(), ak3.into(), ak4.into()]
+            }
+            Self::ContainerUpdate(ak1, ak2, ak3, ak4) => {
+                vec![ak1.into(), ak2.into(), ak3.into(), ak4.into()]
+            }
+            Self::ContainerDelete(ak1, ak2, ak3) => vec![ak1.into(), ak2.into(), ak3.into()],
             Self::Custom(_, args) => Vec::from_iter(args.into_iter().map(Literal)),
         }
     }
@@ -293,6 +351,21 @@ impl Statement {
             (Native(NativePredicate::PublicKeyOf), &[a1, a2]) => {
                 Self::PublicKeyOf(a1.try_into()?, a2.try_into()?)
             }
+            (Native(NativePredicate::ContainerInsert), &[a1, a2, a3, a4]) => Self::ContainerInsert(
+                a1.try_into()?,
+                a2.try_into()?,
+                a3.try_into()?,
+                a4.try_into()?,
+            ),
+            (Native(NativePredicate::ContainerUpdate), &[a1, a2, a3, a4]) => Self::ContainerUpdate(
+                a1.try_into()?,
+                a2.try_into()?,
+                a3.try_into()?,
+                a4.try_into()?,
+            ),
+            (Native(NativePredicate::ContainerDelete), &[a1, a2, a3]) => {
+                Self::ContainerDelete(a1.try_into()?, a2.try_into()?, a3.try_into()?)
+            }
             (Native(np), _) => {
                 return Err(Error::custom(format!("Predicate {:?} is syntax sugar", np)))
             }