Merkle tree for custom predicate batches (#471)

Resolve https://github.com/0xPARC/pod2/issues/466

Now batches are identified by the root of a merkle tree that contains all the predicates (using sequential indices as keys).  This means that the format to identify a custom predicate reference is still a hash + index, but the calculation of the hash is different.
The MainPod circuit now isn't limited by number of batches but instead number of custom predicates; and for each one we verify a merkle proof to verify the batch id.

I've removed a bunch of tests from lang that were testing splitting into multiple batches because there's no longer any need for that.  In a future PR we'll remove the code that handles batch splitting.

Each custom predicate needs 148.2 gates (which is very close to my estimate of 142.7 in https://github.com/0xPARC/pod2/issues/466#issuecomment-3823531286 where I actually made a mistake and considered 5 predicates per batch instead of 4 in the previous Params).

src/backends/plonky2/circuits/common.rs

@@ -28,14 +28,17 @@ use crate::{
         circuits::mainpod::CustomPredicateVerification,
         error::Result,
         mainpod::{Operation, OperationArg, OperationAux, Statement},
-        primitives::merkletree::{MerkleClaimAndProofTarget, MerkleTreeStateTransitionProofTarget},
+        primitives::merkletree::{
+            verify_merkle_proof_circuit, MerkleClaimAndProof, MerkleClaimAndProofTarget,
+            MerkleProof, MerkleTreeStateTransitionProofTarget,
+        },
     },
     middleware::{
-        CustomPredicate, CustomPredicateBatch, CustomPredicateRef, NativeOperation,
-        NativePredicate, OperationType, Params, Predicate, PredicateOrWildcard,
-        PredicateOrWildcardPrefix, PredicatePrefix, RawValue, StatementArg, StatementTmpl,
-        StatementTmplArg, StatementTmplArgPrefix, ToFields, Value, EMPTY_VALUE, F, HASH_SIZE,
-        STATEMENT_ARG_F_LEN, VALUE_SIZE,
+        hash_fields, CustomPredicate, CustomPredicateRef, NativeOperation, NativePredicate,
+        OperationType, Params, Predicate, PredicateOrWildcard, PredicateOrWildcardPrefix,
+        PredicatePrefix, RawValue, StatementArg, StatementTmpl, StatementTmplArg,
+        StatementTmplArgPrefix, ToFields, Value, EMPTY_VALUE, F, HASH_SIZE, STATEMENT_ARG_F_LEN,
+        VALUE_SIZE,
     },
 };
 
@@ -688,34 +691,65 @@ impl CustomPredicateTarget {
     }
 }
 
-/// This type is used to build the custom predicate table, which exposes the custom predicates with
-/// normalized statement templates indexed by batch_id and custom_predicate_index.
+/// Custom predicate structure that can be verified to belong to a batch id at a particular index
 #[derive(Clone, Serialize, Deserialize)]
-pub struct CustomPredicateBatchTarget {
-    pub predicates: Vec<CustomPredicateTarget>,
+pub struct CustomPredicateInBatchTarget {
+    pub id: HashOutTarget,
+    pub index: Target,
+    /// Predicate that may use references to another predicate of the batch with BatchSelf
+    pub self_predicate: CustomPredicateTarget,
+    pub mtp: MerkleClaimAndProofTarget,
 }
 
-impl CustomPredicateBatchTarget {
-    pub fn id(&self, builder: &mut CircuitBuilder) -> HashOutTarget {
-        let flattened: Vec<_> = self.predicates.iter().flat_map(|cp| cp.flatten()).collect();
-        builder.hash_n_to_hash_no_pad::<PoseidonHash>(flattened)
-    }
+impl CustomPredicateInBatchTarget {
+    /// This constructor connects the merkle proof and claim targets with with the (index,
+    /// self_predicate) and id.
+    pub fn new_virtual(builder: &mut CircuitBuilder) -> CustomPredicateInBatchTarget {
+        let index = builder.add_virtual_target();
+        let self_predicate = builder.add_virtual_custom_predicate(true);
+        // Existence Merkle Tree proof of (index, hash(self_predicate)) -> id
+        let mtp =
+            MerkleClaimAndProofTarget::new_virtual(Params::max_depth_custom_batch_mt(), builder);
+        let _true = builder._true();
+        builder.connect(_true.target, mtp.enabled.target);
+        builder.connect(_true.target, mtp.existence.target);
+        let zero = builder.constant(F(0));
+        let key = ValueTarget {
+            elements: [index, zero, zero, zero],
+        };
+        builder.connect_values(key, mtp.key);
+        let id = mtp.root;
 
+        Self {
+            id,
+            index,
+            mtp,
+            self_predicate,
+        }
+    }
+    /// Hash the predicate, connect it to the merkle proof claim value and verify the merkle proof.
+    pub fn verify_circuit(&self, builder: &mut CircuitBuilder) {
+        let value = builder.hash_n_to_hash_no_pad::<PoseidonHash>(self.self_predicate.flatten());
+        builder.connect_array(value.elements, self.mtp.value.elements);
+        verify_merkle_proof_circuit(builder, &self.mtp);
+    }
     pub fn set_targets(
         &self,
         pw: &mut PartialWitness<F>,
-        custom_predicate_batch: &CustomPredicateBatch,
+        predicate_ref: &CustomPredicateRef,
+        mtp: &MerkleProof,
     ) -> Result<()> {
-        let pad_predicate = CustomPredicate::empty();
-        for (i, predicate) in custom_predicate_batch
-            .predicates()
-            .iter()
-            .chain(iter::repeat(&pad_predicate))
-            .take(Params::max_custom_batch_size())
-            .enumerate()
-        {
-            self.predicates[i].set_targets(pw, predicate)?;
-        }
+        pw.set_target_arr(&self.id.elements, &predicate_ref.batch.id().0)?;
+        pw.set_target(self.index, F::from_canonical_usize(predicate_ref.index))?;
+        let predicate = predicate_ref.predicate();
+        self.self_predicate.set_targets(pw, predicate)?;
+        let mtp_claim = MerkleClaimAndProof {
+            root: predicate_ref.batch.id(),
+            key: Value::from(predicate_ref.index as i64).raw(),
+            value: RawValue::from(hash_fields(&predicate.to_fields())),
+            proof: mtp.clone(),
+        };
+        self.mtp.set_targets(pw, true, &mtp_claim)?;
         Ok(())
     }
 }
@@ -812,11 +846,9 @@ pub struct CustomPredicateVerifyEntryTarget {
 
 impl CustomPredicateVerifyEntryTarget {
     pub fn new_virtual(params: &Params, builder: &mut CircuitBuilder) -> Self {
-        let custom_predicate_table_len =
-            params.max_custom_predicate_batches * Params::max_custom_batch_size();
         CustomPredicateVerifyEntryTarget {
             custom_predicate_table_index: IndexTarget::new_virtual(
-                custom_predicate_table_len,
+                params.max_custom_predicates,
                 builder,
             ),
             custom_predicate: builder.add_virtual_custom_predicate_entry(),
@@ -1245,8 +1277,6 @@ pub trait CircuitBuilderPod<F: RichField + Extendable<D>, const D: usize> {
     fn add_virtual_statement_tmpl_arg(&mut self) -> StatementTmplArgTarget;
     fn add_virtual_statement_tmpl(&mut self, with_pred: bool) -> StatementTmplTarget;
     fn add_virtual_custom_predicate(&mut self, with_pred: bool) -> CustomPredicateTarget;
-    fn add_virtual_custom_predicate_batch(&mut self, with_pred: bool)
-        -> CustomPredicateBatchTarget;
     fn add_virtual_custom_predicate_entry(&mut self) -> CustomPredicateEntryTarget;
     fn select_value(&mut self, b: BoolTarget, x: ValueTarget, y: ValueTarget) -> ValueTarget;
     fn select_statement_arg(
@@ -1435,18 +1465,6 @@ impl CircuitBuilderPod<F, D> for CircuitBuilder {
         }
     }
 
-    /// See `add_virtual_statement_tmpl` for the meaning of `with_pred`.
-    fn add_virtual_custom_predicate_batch(
-        &mut self,
-        with_pred: bool,
-    ) -> CustomPredicateBatchTarget {
-        CustomPredicateBatchTarget {
-            predicates: (0..Params::max_custom_batch_size())
-                .map(|_| self.add_virtual_custom_predicate(with_pred))
-                .collect(),
-        }
-    }
-
     /// See `add_virtual_statement_tmpl` for the meaning of `with_pred`.
     fn add_virtual_custom_predicate_entry(&mut self) -> CustomPredicateEntryTarget {
         CustomPredicateEntryTarget {
@@ -1869,6 +1887,8 @@ impl SimpleGenerator<F, D> for LtMaskGenerator {
 
 #[cfg(test)]
 pub(crate) mod tests {
+    use std::sync::Arc;
+
     use anyhow::anyhow;
     use itertools::Itertools;
     use plonky2::plonk::{
@@ -1878,8 +1898,10 @@ pub(crate) mod tests {
 
     use super::*;
     use crate::{
-        backends::plonky2::basetypes::C, examples::custom::eth_dos_batch, frontend,
-        frontend::CustomPredicateBatchBuilder, middleware::CustomPredicateBatch,
+        backends::plonky2::basetypes::C,
+        examples::custom::eth_dos_batch,
+        frontend::{self, CustomPredicateBatchBuilder},
+        middleware::CustomPredicateBatch,
     };
 
     pub(crate) const I64_TEST_PAIRS: [(i64, i64); 36] = [
@@ -1952,50 +1974,54 @@ pub(crate) mod tests {
         Ok(())
     }
 
-    fn helper_custom_predicate_batch_target_id(
-        custom_predicate_batch: &CustomPredicateBatch,
+    fn helper_custom_predicate_in_batch_target(
+        custom_predicate_batch: &Arc<CustomPredicateBatch>,
     ) -> Result<()> {
-        let config = CircuitConfig::standard_recursion_config();
-        let mut builder = CircuitBuilder::<F, D>::new(config);
+        for index in 0..custom_predicate_batch.predicates().len() {
+            let cpr = custom_predicate_batch
+                .predicate_ref_by_index(index)
+                .unwrap();
 
-        let custom_predicate_batch_target = builder.add_virtual_custom_predicate_batch(false);
+            let config = CircuitConfig::standard_recursion_config();
+            let mut builder = CircuitBuilder::<F, D>::new(config);
 
-        // Calculate the id in constraints and compare it against the id calculated natively
-        let id_target = custom_predicate_batch_target.id(&mut builder);
+            let custom_pred_in_batch_target =
+                CustomPredicateInBatchTarget::new_virtual(&mut builder);
+            custom_pred_in_batch_target.verify_circuit(&mut builder);
 
-        let mut pw = PartialWitness::<F>::new();
-        custom_predicate_batch_target.set_targets(&mut pw, custom_predicate_batch)?;
-        let id = custom_predicate_batch.id();
-        pw.set_target_arr(&id_target.elements, &id.0)?;
+            let mut pw = PartialWitness::<F>::new();
+            let (_, mtp) = custom_predicate_batch
+                .mt()
+                .prove(&Value::from(index as i64).raw())
+                .unwrap();
+            custom_pred_in_batch_target.set_targets(&mut pw, &cpr, &mtp)?;
 
-        // generate & verify proof
-        let data = builder.build::<C>();
-        let proof = data.prove(pw).unwrap();
-        data.verify(proof.clone()).unwrap();
+            // generate & verify proof
+            let data = builder.build::<C>();
+            let proof = data.prove(pw).unwrap();
+            data.verify(proof.clone()).unwrap();
+        }
 
         Ok(())
     }
 
     #[test]
-    fn test_custom_predicate_batch_target_id() -> frontend::Result<()> {
-        let params = Params {
-            max_custom_predicate_wildcards: 12,
-            ..Default::default()
-        };
+    fn test_custom_predicate_in_batch_target() -> frontend::Result<()> {
+        let params = Params::default();
 
         // Empty case
         let mut cpb_builder = CustomPredicateBatchBuilder::new(params.clone(), "empty".into());
         _ = cpb_builder.predicate_and("empty", &[], &[], &[])?;
         let custom_predicate_batch = cpb_builder.finish();
-        helper_custom_predicate_batch_target_id(&custom_predicate_batch).unwrap();
+        helper_custom_predicate_in_batch_target(&custom_predicate_batch).unwrap();
 
         // Some cases from the examples
         let custom_predicate_batch = eth_dos_batch(&params)?;
-        helper_custom_predicate_batch_target_id(&custom_predicate_batch).unwrap();
+        helper_custom_predicate_in_batch_target(&custom_predicate_batch).unwrap();
 
         let custom_predicate_batch =
-            CustomPredicateBatch::new(&params, "empty".to_string(), vec![CustomPredicate::empty()]);
-        helper_custom_predicate_batch_target_id(&custom_predicate_batch).unwrap();
+            CustomPredicateBatch::new("empty".to_string(), vec![CustomPredicate::empty()]);
+        helper_custom_predicate_in_batch_target(&custom_predicate_batch).unwrap();
 
         Ok(())
     }

src/backends/plonky2/circuits/mainpod.rs

@@ -1,4 +1,4 @@
-use std::{array, iter, sync::Arc};
+use std::{array, iter};
 
 use itertools::{izip, zip_eq, Itertools};
 use num::{BigUint, One};
@@ -21,7 +21,7 @@ use crate::{
         basetypes::{CircuitBuilder, VDSet},
         circuits::{
             common::{
-                CircuitBuilderPod, CustomPredicateBatchTarget, CustomPredicateEntryTarget,
+                CircuitBuilderPod, CustomPredicateEntryTarget, CustomPredicateInBatchTarget,
                 CustomPredicateTarget, CustomPredicateVerifyEntryTarget,
                 CustomPredicateVerifyQueryTarget, Flattenable, MerkleClaimTarget,
                 MerkleTreeStateTransitionClaimTarget, OperationTarget, OperationTypeTarget,
@@ -44,7 +44,7 @@ use crate::{
             },
             merkletree::{
                 verify_merkle_proof_circuit, verify_merkle_state_transition_circuit,
-                MerkleClaimAndProof, MerkleClaimAndProofTarget, MerkleTreeOp,
+                MerkleClaimAndProof, MerkleClaimAndProofTarget, MerkleProof, MerkleTreeOp,
                 MerkleTreeStateTransitionProof, MerkleTreeStateTransitionProofTarget,
             },
             signature::{verify_signature_circuit, SignatureVerifyTarget},
@@ -1573,37 +1573,34 @@ fn normalize_st_tmpl_circuit(
 fn build_custom_predicate_table_circuit(
     params: &Params,
     builder: &mut CircuitBuilder,
-    custom_predicate_batches: &[CustomPredicateBatchTarget],
+    custom_predicates: &[CustomPredicateInBatchTarget],
 ) -> Result<Vec<HashOutTarget>> {
     let measure = measure_gates_begin!(builder, "BuildCustomPredTbl");
-    let mut custom_predicate_table =
-        Vec::with_capacity(params.max_custom_predicate_batches * Params::max_custom_batch_size());
-    for cpb in custom_predicate_batches {
-        let measure_cpb = measure_gates_begin!(builder, "CustomPredBatch");
-        let id = cpb.id(builder); // constrain the id
-        for (index, cp) in cpb.predicates.iter().enumerate() {
-            let statements = cp
-                .statements
-                .iter()
-                .map(|st_with_pred_tmpl| {
-                    normalize_st_tmpl_circuit(params, builder, st_with_pred_tmpl, id)
-                })
-                .collect_vec();
-            let cp = CustomPredicateTarget {
-                conjunction: cp.conjunction,
+    let mut custom_predicate_table = Vec::with_capacity(params.max_custom_predicates);
+    for cp in custom_predicates {
+        let measure_cp = measure_gates_begin!(builder, "CustomPred");
+        cp.verify_circuit(builder);
+        let statements = cp
+            .self_predicate
+            .statements
+            .iter()
+            .map(|st_with_pred_tmpl| {
+                normalize_st_tmpl_circuit(params, builder, st_with_pred_tmpl, cp.id)
+            })
+            .collect_vec();
+        let entry = CustomPredicateEntryTarget {
+            id: cp.id,       // output
+            index: cp.index, // input
+            predicate: CustomPredicateTarget {
+                conjunction: cp.self_predicate.conjunction,
                 statements,
-                args_len: cp.args_len,
-            };
-            let entry = CustomPredicateEntryTarget {
-                id,                                                      // output
-                index: builder.constant(F::from_canonical_usize(index)), // constant
-                predicate: cp.clone(),                                   // input
-            };
+                args_len: cp.self_predicate.args_len,
+            }, // input
+        };
 
-            let in_query_hash = entry.hash(builder);
-            custom_predicate_table.push(in_query_hash);
-        }
-        measure_gates_end!(builder, measure_cpb);
+        let in_query_hash = entry.hash(builder);
+        custom_predicate_table.push(in_query_hash);
+        measure_gates_end!(builder, measure_cp);
     }
     measure_gates_end!(builder, measure);
     Ok(custom_predicate_table)
@@ -1711,7 +1708,7 @@ fn verify_main_pod_circuit(
 
     // Table of custom predicate batches with batch_id calculation
     let custom_predicate_table =
-        build_custom_predicate_table_circuit(params, builder, &main_pod.custom_predicate_batches)?;
+        build_custom_predicate_table_circuit(params, builder, &main_pod.custom_predicates)?;
 
     let aux_table = build_operation_aux_table_circuit(
         params,
@@ -1754,7 +1751,7 @@ pub struct MainPodVerifyTarget {
     public_key_of_sks: Vec<BigUInt320Target>,
     signed_bys: Vec<SignedByTarget>,
     merkle_tree_state_transition_proofs: Vec<MerkleTreeStateTransitionProofTarget>,
-    custom_predicate_batches: Vec<CustomPredicateBatchTarget>,
+    custom_predicates: Vec<CustomPredicateInBatchTarget>,
     custom_predicate_verifications: Vec<CustomPredicateVerifyEntryTarget>,
 }
 
@@ -1799,8 +1796,8 @@ impl MainPodVerifyTarget {
                     )
                 })
                 .collect(),
-            custom_predicate_batches: (0..params.max_custom_predicate_batches)
-                .map(|_| builder.add_virtual_custom_predicate_batch(true))
+            custom_predicates: (0..params.max_custom_predicates)
+                .map(|_| CustomPredicateInBatchTarget::new_virtual(builder))
                 .collect(),
             custom_predicate_verifications: (0..params.max_custom_predicate_verifications)
                 .map(|_| CustomPredicateVerifyEntryTarget::new_virtual(params, builder))
@@ -1830,7 +1827,7 @@ pub struct MainPodVerifyInput {
     pub public_key_of_sks: Vec<SecretKey>,
     pub signed_bys: Vec<SignedBy>,
     pub merkle_tree_state_transition_proofs: Vec<MerkleTreeStateTransitionProof>,
-    pub custom_predicate_batches: Vec<Arc<CustomPredicateBatch>>,
+    pub custom_predicates_with_mpt_proofs: Vec<(CustomPredicateRef, MerkleProof)>,
     pub custom_predicate_verifications: Vec<CustomPredicateVerification>,
 }
 
@@ -1972,18 +1969,20 @@ impl InnerCircuit for MainPodVerifyTarget {
             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, cpb)?;
+        assert!(input.custom_predicates_with_mpt_proofs.len() <= self.params.max_custom_predicates);
+        for (i, (cp, mtp)) in input.custom_predicates_with_mpt_proofs.iter().enumerate() {
+            self.custom_predicates[i].set_targets(pw, cp, mtp)?;
         }
         // Padding
-        let pad_cpb = CustomPredicateBatch::new(
-            &self.params,
-            "empty".to_string(),
-            vec![CustomPredicate::empty()],
-        );
-        for i in input.custom_predicate_batches.len()..self.params.max_custom_predicate_batches {
-            self.custom_predicate_batches[i].set_targets(pw, &pad_cpb)?;
+        let pad_cpb =
+            CustomPredicateBatch::new("empty".to_string(), vec![CustomPredicate::empty()]);
+        let pad_cp = pad_cpb.predicate_ref_by_index(0).expect("index 0 exists");
+        let (_, pad_mtp) = pad_cpb
+            .mt()
+            .prove(&Value::from(0i64).raw())
+            .expect("exists");
+        for i in input.custom_predicates_with_mpt_proofs.len()..self.params.max_custom_predicates {
+            self.custom_predicates[i].set_targets(pw, &pad_cp, &pad_mtp)?;
         }
 
         assert!(
@@ -2096,7 +2095,7 @@ mod tests {
                 .merkle_tree_state_transition_proofs
                 .len(),
             max_custom_predicate_verifications: 0,
-            max_custom_predicate_batches: 0,
+            max_custom_predicates: 0,
             ..Default::default()
         };
 

src/backends/plonky2/mainpod/mod.rs

@@ -1,7 +1,7 @@
 pub mod operation;
 use crate::middleware::{wildcard_values_from_op_st, PodType};
 pub mod statement;
-use std::{iter, sync::Arc};
+use std::iter;
 
 use itertools::{zip_eq, Itertools};
 use num_bigint::BigUint;
@@ -37,9 +37,9 @@ use crate::{
         serialize_proof, serialize_verifier_only,
     },
     middleware::{
-        self, value_from_op, CustomPredicateBatch, Error as MiddlewareError, Hash, MainPodInputs,
+        self, value_from_op, CustomPredicateRef, Error as MiddlewareError, Hash, MainPodInputs,
         MainPodProver, NativeOperation, OperationType, Params, Pod, RawValue, StatementArg,
-        ToFields, VDSet,
+        ToFields, VDSet, Value,
     },
     timed,
 };
@@ -68,27 +68,27 @@ pub fn calculate_statements_hash(statements: &[Statement]) -> middleware::Hash {
     Hash(PoseidonHash::hash_no_pad(&field_elems).elements)
 }
 
-/// Extracts unique `CustomPredicateBatch`es from Custom ops.
-pub(crate) fn extract_custom_predicate_batches(
+/// Extracts unique `CustomPredicate`s from Custom ops.
+pub(crate) fn extract_custom_predicates(
     params: &Params,
     operations: &[middleware::Operation],
-) -> Result<Vec<Arc<CustomPredicateBatch>>> {
-    let custom_predicate_batches: Vec<_> = operations
+) -> Result<Vec<CustomPredicateRef>> {
+    let custom_predicates: Vec<_> = operations
         .iter()
         .flat_map(|op| match op {
-            middleware::Operation::Custom(cpr, _) => Some(cpr.batch.clone()),
+            middleware::Operation::Custom(cpr, _) => Some(cpr.clone()),
             _ => None,
         })
-        .unique_by(|cpr| cpr.id())
+        .unique()
         .collect();
-    if custom_predicate_batches.len() > params.max_custom_predicate_batches {
+    if custom_predicates.len() > params.max_custom_predicates {
         return Err(Error::custom(format!(
-            "The number of required `CustomPredicateBatch`es ({}) exceeds the maximum number ({}).",
-            custom_predicate_batches.len(),
-            params.max_custom_predicate_batches
+            "The number of required `CustomPredicate`s ({}) exceeds the maximum number ({}).",
+            custom_predicates.len(),
+            params.max_custom_predicates
         )));
     }
-    Ok(custom_predicate_batches)
+    Ok(custom_predicates)
 }
 
 /// Extracts all custom predicate operations with all the data required to verify them.
@@ -97,7 +97,7 @@ pub(crate) fn extract_custom_predicate_verifications(
     aux_list: &mut [OperationAux],
     operations: &[middleware::Operation],
     statements: &[middleware::Statement],
-    custom_predicate_batches: &[Arc<CustomPredicateBatch>],
+    custom_predicates: &[CustomPredicateRef],
 ) -> Result<Vec<CustomPredicateVerification>> {
     let mut table = Vec::new();
     for (i, (op, st)) in zip_eq(operations.iter(), statements.iter()).enumerate() {
@@ -108,13 +108,11 @@ pub(crate) fn extract_custom_predicate_verifications(
                     wildcard_values_from_op_st(params, cpr.predicate(), sts, st_args)
                         .expect("resolved wildcards");
                 let sts = sts.iter().map(|s| Statement::from(s.clone())).collect();
-                let batch_index = custom_predicate_batches
+                let custom_predicate_table_index = custom_predicates
                     .iter()
                     .enumerate()
-                    .find_map(|(i, cpb)| (cpb.id() == cpr.batch.id()).then_some(i))
+                    .find_map(|(i, table_cpr)| (table_cpr == cpr).then_some(i))
                     .expect("find the custom predicate from the extracted unique list");
-                let custom_predicate_table_index =
-                    batch_index * Params::max_custom_batch_size() + cpr.index;
                 aux_list[i] = OperationAux::CustomPredVerifyIndex(table.len());
                 table.push(CustomPredicateVerification {
                     custom_predicate_table_index,
@@ -497,14 +495,25 @@ impl MainPodProver for Prover {
         let mut aux_list = vec![OperationAux::None; params.max_priv_statements()];
         let merkle_proofs =
             extract_merkle_proofs(params, &mut aux_list, inputs.operations, inputs.statements)?;
-        let custom_predicate_batches = extract_custom_predicate_batches(params, inputs.operations)?;
+        let custom_predicates = extract_custom_predicates(params, inputs.operations)?;
         let custom_predicate_verifications = extract_custom_predicate_verifications(
             params,
             &mut aux_list,
             inputs.operations,
             inputs.statements,
-            &custom_predicate_batches,
+            &custom_predicates,
         )?;
+        let custom_predicates_with_mpt_proofs = custom_predicates
+            .into_iter()
+            .map(|cpr| {
+                let (_, mtp) = cpr
+                    .batch
+                    .mt()
+                    .prove(&Value::from(cpr.index as i64).raw())
+                    .expect("index by construction exists");
+                (cpr, mtp)
+            })
+            .collect_vec();
         let public_key_of_sks =
             extract_public_key_of(params, &mut aux_list, inputs.operations, inputs.statements)?;
         let signed_bys =
@@ -572,7 +581,7 @@ impl MainPodProver for Prover {
             public_key_of_sks,
             signed_bys,
             merkle_tree_state_transition_proofs,
-            custom_predicate_batches,
+            custom_predicates_with_mpt_proofs,
             custom_predicate_verifications,
         };
 
@@ -840,7 +849,7 @@ pub mod tests {
             // Currently the circuit uses random access that only supports vectors of length 64.
             // With max_input_main_pods=3 we need random access to a vector of length 73.
             max_input_pods: 0,
-            max_custom_predicate_batches: 0,
+            max_custom_predicates: 0,
             max_custom_predicate_verifications: 0,
             ..Default::default()
         };
@@ -961,7 +970,7 @@ pub mod tests {
             max_merkle_proofs_containers: 0,
             max_public_key_of: 0,
             max_custom_predicate_verifications: 0,
-            max_custom_predicate_batches: 0,
+            max_custom_predicates: 0,
             ..Default::default()
         };
         let mut vds = DEFAULT_VD_LIST.clone();
@@ -995,7 +1004,7 @@ pub mod tests {
             max_statements: 5,
             max_public_statements: 2,
             max_operation_args: 5,
-            max_custom_predicate_batches: 2,
+            max_custom_predicates: 2,
             max_custom_predicate_verifications: 2,
             max_custom_predicate_wildcards: 3,
             max_merkle_proofs_containers: 2,

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

@@ -23,6 +23,13 @@ pub struct MerkleTree {
     root: Node,
 }
 
+impl PartialEq for MerkleTree {
+    fn eq(&self, other: &Self) -> bool {
+        self.root() == other.root()
+    }
+}
+impl Eq for MerkleTree {}
+
 impl MerkleTree {
     /// builds a new `MerkleTree` where the leaves contain the given key-values
     pub fn new(kvs: &HashMap<RawValue, RawValue>) -> Self {