Add table multiplexer (and use it for container, custom pred & PublicKeyOf ops) (#376)

- Extend the `Flattenable` trait to include a `size` method that returns the number of `Target`s the type requires.  This is used in the table to figure out the max length of an array that must fit all entry types.
- Move the circuit methods to precalculate hash states and do hashes started from a precomputed state to a new module
- Introduce `MuxTableTarget` which allows easy multiplexing of tables where each sub-table may have entries of different lengths.  The table access is done via hashing + unhashing automatically (via use of a generator)
- Use the `MuxTableTarget` to access merkle tree claims and custom predicate verification, which where previously in different tables and accessed with independent random accesses each
- Move the public key derivation for the PublicKeyOf operation check to the same multiplexed table.  Now we can choose how many of those operations a circuit supports.

Resolve https://github.com/0xPARC/pod2/issues/357
Resolve https://github.com/0xPARC/pod2/issues/361

Cargo.toml

@@ -41,7 +41,7 @@ serde_arrays = "0.2.0"
 sha2 = { version = "0.10.9" }
 
 # Uncomment for debugging with https://github.com/ed255/plonky2/ at branch `feat/debug`.  The repo directory needs to be checked out next to the pod2 repo directory.
-# [patch."https://github.com/0xPolygonZero/plonky2"]
+# [patch."https://github.com/0xPARC/plonky2"]
 # plonky2 = { path = "../plonky2/plonky2" }
 
 [dev-dependencies]
@@ -61,3 +61,7 @@ time = []
 examples = []
 disk_cache = ["directories", "minicbor-serde"]
 mem_cache = []
+
+# Uncomment in order to enable debug information in the release builds.  This allows getting panic backtraces with a performance similar to regular release.
+# [profile.release]
+# debug = true

src/backends/plonky2/circuits/common.rs

@@ -26,7 +26,7 @@ use crate::{
         basetypes::{CircuitBuilder, CommonCircuitData, D},
         circuits::mainpod::CustomPredicateVerification,
         error::Result,
-        mainpod::{Operation, OperationArg, Statement},
+        mainpod::{Operation, OperationArg, OperationAux, Statement},
         primitives::merkletree::MerkleClaimAndProofTarget,
     },
     middleware::{
@@ -128,6 +128,10 @@ impl StatementArgTarget {
     pub fn as_value(&self) -> ValueTarget {
         ValueTarget::from_slice(&self.elements[..VALUE_SIZE])
     }
+
+    fn size(_params: &Params) -> usize {
+        STATEMENT_ARG_F_LEN
+    }
 }
 
 #[derive(Clone, Serialize, Deserialize)]
@@ -249,6 +253,10 @@ impl OperationTypeTarget {
     ) -> Result<()> {
         Ok(pw.set_target_arr(&self.elements, &op_type.to_fields(params))?)
     }
+
+    fn size(_params: &Params) -> usize {
+        Params::operation_type_size()
+    }
 }
 
 // TODO: Implement Operation::to_field to determine the size of each element
@@ -256,8 +264,7 @@ impl OperationTypeTarget {
 pub struct OperationTarget {
     pub op_type: OperationTypeTarget,
     pub args: Vec<IndexTarget>,
-    #[serde(with = "serde_arrays")]
+    pub aux_index: IndexTarget,
-    pub aux: [IndexTarget; 2],
 }
 
 impl OperationTarget {
@@ -277,11 +284,13 @@ impl OperationTarget {
         {
             self.args[i].set_targets(pw, arg.as_usize())?;
         }
-        let indexes = op.aux().as_usizes();
+        self.aux_index.set_targets(pw, op.aux().table_index(params))
-        for (index_target, index) in self.aux.iter().zip_eq(indexes.iter()) {
+    }
-            index_target.set_targets(pw, *index)?;
+
-        }
+    fn size(params: &Params) -> usize {
-        Ok(())
+        OperationTypeTarget::size(params)
+            + params.max_operation_args * IndexTarget::size(params)
+            + IndexTarget::size(params)
     }
 }
 
@@ -570,12 +579,16 @@ impl Flattenable for CustomPredicateEntryTarget {
             .collect()
     }
     fn from_flattened(params: &Params, vs: &[Target]) -> Self {
+        assert_eq!(vs.len(), Self::size(params));
         Self {
             id: HashOutTarget::from_flattened(params, &vs[0..4]),
             index: vs[4],
             predicate: CustomPredicateTarget::from_flattened(params, &vs[5..]),
         }
     }
+    fn size(params: &Params) -> usize {
+        HashOutTarget::size(params) + 1 + CustomPredicateTarget::size(params)
+    }
 }
 
 impl CustomPredicateEntryTarget {
@@ -669,15 +682,16 @@ impl Flattenable for CustomPredicateVerifyQueryTarget {
             .collect()
     }
     fn from_flattened(params: &Params, vs: &[Target]) -> Self {
-        let (pos, size) = (0, params.statement_size());
+        assert_eq!(vs.len(), Self::size(params));
+        let (pos, size) = (0, StatementTarget::size(params));
         let statement = StatementTarget::from_flattened(params, &vs[pos..pos + size]);
-        let (pos, size) = (pos + size, params.operation_size(IndexTarget::f_len()));
+        let (pos, size) = (pos + size, OperationTypeTarget::size(params));
         let op_type = OperationTypeTarget {
             elements: vs[pos..pos + size]
                 .try_into()
                 .expect("len = operation_type_size"),
         };
-        let (pos, size) = (pos + size, params.statement_size());
+        let (pos, size) = (pos + size, StatementTarget::size(params));
         let op_args = (0..params.max_operation_args)
             .map(|i| {
                 StatementTarget::from_flattened(params, &vs[pos + i * size..pos + (1 + i) * size])
@@ -689,6 +703,10 @@ impl Flattenable for CustomPredicateVerifyQueryTarget {
             op_args,
         }
     }
+    fn size(params: &Params) -> usize {
+        StatementTarget::size(params) * (1 + params.max_operation_args)
+            + OperationTarget::size(params)
+    }
 }
 
 /// Trait for target structs that may be converted to and from vectors
@@ -696,8 +714,11 @@ impl Flattenable for CustomPredicateVerifyQueryTarget {
 pub trait Flattenable {
     fn flatten(&self) -> Vec<Target>;
     fn from_flattened(params: &Params, vs: &[Target]) -> Self;
+    /// Size in number of `Target`s
+    fn size(params: &Params) -> usize;
 }
 
+// TODO: Figure out why this is defined in common and not in the merkletree directory
 /// For the purpose of op verification, we need only look up the
 /// Merkle claim rather than the Merkle proof since it is verified
 /// elsewhere.
@@ -726,21 +747,28 @@ impl Flattenable for HashOutTarget {
     fn flatten(&self) -> Vec<Target> {
         self.elements.to_vec()
     }
-    fn from_flattened(_params: &Params, vs: &[Target]) -> Self {
+    fn from_flattened(params: &Params, vs: &[Target]) -> Self {
-        assert_eq!(vs.len(), HASH_SIZE);
+        assert_eq!(vs.len(), Self::size(params));
         Self {
             elements: array::from_fn(|i| vs[i]),
         }
     }
+    fn size(_params: &Params) -> usize {
+        4
+    }
 }
 
 impl Flattenable for ValueTarget {
     fn flatten(&self) -> Vec<Target> {
         self.elements.to_vec()
     }
-    fn from_flattened(_params: &Params, vs: &[Target]) -> Self {
+    fn from_flattened(params: &Params, vs: &[Target]) -> Self {
+        assert_eq!(vs.len(), Self::size(params));
         Self::from_slice(vs)
     }
+    fn size(_params: &Params) -> usize {
+        4
+    }
 }
 
 impl Flattenable for MerkleClaimTarget {
@@ -755,7 +783,8 @@ impl Flattenable for MerkleClaimTarget {
         .concat()
     }
 
-    fn from_flattened(_params: &Params, vs: &[Target]) -> Self {
+    fn from_flattened(params: &Params, vs: &[Target]) -> Self {
+        assert_eq!(vs.len(), Self::size(params));
         Self {
             enabled: BoolTarget::new_unsafe(vs[0]),
             root: HashOutTarget::from_vec(vs[1..1 + NUM_HASH_OUT_ELTS].to_vec()),
@@ -768,6 +797,10 @@ impl Flattenable for MerkleClaimTarget {
             existence: BoolTarget::new_unsafe(vs[1 + NUM_HASH_OUT_ELTS + 2 * VALUE_SIZE]),
         }
     }
+
+    fn size(params: &Params) -> usize {
+        2 + HashOutTarget::size(params) + 2 * ValueTarget::size(params)
+    }
 }
 
 impl Flattenable for PredicateTarget {
@@ -775,11 +808,15 @@ impl Flattenable for PredicateTarget {
         self.elements.to_vec()
     }
 
-    fn from_flattened(_params: &Params, v: &[Target]) -> Self {
+    fn from_flattened(params: &Params, v: &[Target]) -> Self {
+        assert_eq!(v.len(), Self::size(params));
         Self {
             elements: v.try_into().expect("len is predicate_size"),
         }
     }
+    fn size(_params: &Params) -> usize {
+        Params::predicate_size()
+    }
 }
 
 impl Flattenable for StatementTarget {
@@ -792,13 +829,9 @@ impl Flattenable for StatementTarget {
     }
 
     fn from_flattened(params: &Params, v: &[Target]) -> Self {
-        let num_args = (v.len() - Params::predicate_size()) / STATEMENT_ARG_F_LEN;
+        assert_eq!(v.len(), Self::size(params));
-        assert_eq!(
-            v.len(),
-            Params::predicate_size() + num_args * STATEMENT_ARG_F_LEN
-        );
         let predicate = PredicateTarget::from_flattened(params, &v[..Params::predicate_size()]);
-        let args = (0..num_args)
+        let args = (0..params.max_statement_args)
             .map(|i| StatementArgTarget {
                 elements: array::from_fn(|j| {
                     v[Params::predicate_size() + i * STATEMENT_ARG_F_LEN + j]
@@ -808,6 +841,10 @@ impl Flattenable for StatementTarget {
 
         Self { predicate, args }
     }
+
+    fn size(params: &Params) -> usize {
+        PredicateTarget::size(params) + params.max_statement_args * StatementArgTarget::size(params)
+    }
 }
 
 impl Flattenable for CustomPredicateTarget {
@@ -819,6 +856,7 @@ impl Flattenable for CustomPredicateTarget {
     }
 
     fn from_flattened(params: &Params, v: &[Target]) -> Self {
+        assert_eq!(v.len(), Self::size(params));
         // We assume that `from_flattened` is always called with the output of `flattened`, so
         // this `BoolTarget` should actually safe.
         let conjunction = BoolTarget::new_unsafe(v[0]);
@@ -836,6 +874,9 @@ impl Flattenable for CustomPredicateTarget {
             args_len,
         }
     }
+    fn size(params: &Params) -> usize {
+        2 + params.max_custom_predicate_arity * StatementTmplTarget::size(params)
+    }
 }
 
 impl Flattenable for StatementTmplTarget {
@@ -848,6 +889,7 @@ impl Flattenable for StatementTmplTarget {
     }
 
     fn from_flattened(params: &Params, v: &[Target]) -> Self {
+        assert_eq!(v.len(), Self::size(params));
         let pred_end = Params::predicate_size();
         let pred = PredicateTarget::from_flattened(params, &v[..pred_end]);
         let sta_size = Params::statement_tmpl_arg_size();
@@ -859,6 +901,11 @@ impl Flattenable for StatementTmplTarget {
             .collect();
         Self { pred, args }
     }
+
+    fn size(params: &Params) -> usize {
+        PredicateTarget::size(params)
+            + params.max_statement_args * StatementTmplArgTarget::size(params)
+    }
 }
 
 impl Flattenable for StatementTmplArgTarget {
@@ -866,24 +913,28 @@ impl Flattenable for StatementTmplArgTarget {
         self.elements.to_vec()
     }
 
-    fn from_flattened(_params: &Params, v: &[Target]) -> Self {
+    fn from_flattened(params: &Params, v: &[Target]) -> Self {
+        assert_eq!(v.len(), Self::size(params));
         Self {
             elements: v.try_into().expect("len is statement_tmpl_arg_size"),
         }
     }
+    fn size(_params: &Params) -> usize {
+        Params::statement_tmpl_arg_size()
+    }
 }
 
 /// Index to an array for random access
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
 pub struct IndexTarget {
-    max_array_len: usize,
+    pub max_array_len: usize,
-    low: Target,
+    pub low: Target,
-    high: Target,
+    pub high: Target,
 }
 
 impl IndexTarget {
     // Length in field elements
-    pub const fn f_len() -> usize {
+    pub fn size(_params: &Params) -> usize {
         2
     }
     pub fn new_virtual(max_array_len: usize, builder: &mut CircuitBuilder) -> Self {
@@ -1051,10 +1102,7 @@ impl CircuitBuilderPod<F, D> for CircuitBuilder {
             args: (0..params.max_operation_args)
                 .map(|_| IndexTarget::new_virtual(params.statement_table_size(), self))
                 .collect(),
-            aux: [
+            aux_index: IndexTarget::new_virtual(OperationAux::table_size(params), self),
-                IndexTarget::new_virtual(params.max_merkle_proofs_containers, self),
-                IndexTarget::new_virtual(params.max_custom_predicate_verifications, self),
-            ],
         }
     }
 

src/backends/plonky2/circuits/hash.rs

@@ -0,0 +1,57 @@
+use plonky2::{
+    hash::{
+        hash_types::{HashOutTarget, RichField, NUM_HASH_OUT_ELTS},
+        hashing::PlonkyPermutation,
+    },
+    iop::target::Target,
+    plonk::config::AlgebraicHasher,
+};
+
+use crate::{backends::plonky2::basetypes::CircuitBuilder, middleware::F};
+
+/// Precompute the hash state by absorbing all full chunks from `inputs` and return the reminder
+/// elements that didn't fit into a chunk.
+pub fn precompute_hash_state<F: RichField, P: PlonkyPermutation<F>>(inputs: &[F]) -> (P, &[F]) {
+    let (inputs, inputs_rem) = inputs.split_at((inputs.len() / P::RATE) * P::RATE);
+    let mut perm = P::new(core::iter::repeat(F::ZERO));
+
+    // Absorb all inputs up to the biggest multiple of RATE.
+    for input_chunk in inputs.chunks(P::RATE) {
+        perm.set_from_slice(input_chunk, 0);
+        perm.permute();
+    }
+
+    (perm, inputs_rem)
+}
+
+/// Hash `inputs` starting from a circuit-constant `perm` state.
+pub fn hash_from_state_circuit<H: AlgebraicHasher<F>, P: PlonkyPermutation<F>>(
+    builder: &mut CircuitBuilder,
+    perm: P,
+    inputs: &[Target],
+) -> HashOutTarget {
+    let mut state =
+        H::AlgebraicPermutation::new(perm.as_ref().iter().map(|v| builder.constant(*v)));
+
+    // Absorb all input chunks.
+    for input_chunk in inputs.chunks(H::AlgebraicPermutation::RATE) {
+        // Overwrite the first r elements with the inputs. This differs from a standard sponge,
+        // where we would xor or add in the inputs. This is a well-known variant, though,
+        // sometimes called "overwrite mode".
+        state.set_from_slice(input_chunk, 0);
+        state = builder.permute::<H>(state);
+    }
+
+    let num_outputs = NUM_HASH_OUT_ELTS;
+    // Squeeze until we have the desired number of outputs.
+    let mut outputs = Vec::with_capacity(num_outputs);
+    loop {
+        for &s in state.squeeze() {
+            outputs.push(s);
+            if outputs.len() == num_outputs {
+                return HashOutTarget::from_vec(outputs);
+            }
+        }
+        state = builder.permute::<H>(state);
+    }
+}

src/backends/plonky2/circuits/mod.rs

@@ -1,5 +1,7 @@
 pub mod common;
+pub mod hash;
 pub mod mainpod;
 pub mod metrics;
+pub mod mux_table;
 pub mod signedpod;
 pub mod utils;

src/backends/plonky2/circuits/mux_table.rs

@@ -0,0 +1,216 @@
+use std::iter;
+
+use itertools::Itertools;
+use plonky2::{
+    field::{extension::Extendable, types::Field},
+    hash::{
+        hash_types::{HashOutTarget, RichField},
+        poseidon::{PoseidonHash, PoseidonPermutation},
+    },
+    iop::{
+        generator::{GeneratedValues, SimpleGenerator},
+        target::{BoolTarget, Target},
+        witness::{PartitionWitness, Witness, WitnessWrite},
+    },
+    plonk::circuit_data::CommonCircuitData,
+    util::serialization::{Buffer, IoResult, Read, Write},
+};
+
+use crate::{
+    backends::plonky2::{
+        basetypes::CircuitBuilder,
+        circuits::{
+            common::{CircuitBuilderPod, Flattenable, IndexTarget},
+            hash::{hash_from_state_circuit, precompute_hash_state},
+        },
+    },
+    measure_gates_begin, measure_gates_end,
+    middleware::{Params, F},
+};
+
+// This structure allows multiplexing multiple tables into one by using tags.  The table entries
+// are computed by hashing the concatenation of the tag with the flattened target, with zero
+// padding to normalize the size of all flattened entries.  We use zero-padding on then reverse the
+// array so that smaller entries can skip the initial hashes by using the precomputed hash state of
+// the prefixed zeroes.
+// The table offers an indexing API that returns a flattened entry that includes the "unhashing",
+// this allows doing a single lookup for different possible tagged entries at the same time.
+pub struct MuxTableTarget {
+    params: Params,
+    max_flattened_entry_len: usize,
+    hashed_tagged_entries: Vec<HashOutTarget>,
+    tagged_entries: Vec<Vec<Target>>,
+}
+
+impl MuxTableTarget {
+    pub fn new(params: &Params, max_flattened_entry_len: usize) -> Self {
+        Self {
+            params: params.clone(),
+            max_flattened_entry_len,
+            hashed_tagged_entries: Vec::new(),
+            tagged_entries: Vec::new(),
+        }
+    }
+
+    #[allow(clippy::len_without_is_empty)]
+    pub fn len(&self) -> usize {
+        self.hashed_tagged_entries.len()
+    }
+
+    pub fn push<T: Flattenable>(&mut self, builder: &mut CircuitBuilder, tag: u32, entry: &T) {
+        let flattened_entry = entry.flatten();
+        self.push_flattened(builder, tag, &flattened_entry);
+    }
+
+    pub fn push_flattened(
+        &mut self,
+        builder: &mut CircuitBuilder,
+        tag: u32,
+        flattened_entry: &[Target],
+    ) {
+        let measure = measure_gates_begin!(builder, "HashTaggedTblEntry");
+        assert!(flattened_entry.len() <= self.max_flattened_entry_len);
+        let flattened = [&[builder.constant(F(tag as u64))], flattened_entry].concat();
+        self.tagged_entries.push(flattened.clone());
+
+        let tagged_entry_max_len = 1 + self.max_flattened_entry_len;
+        let front_pad_elts = iter::repeat(F::ZERO)
+            .take(tagged_entry_max_len - flattened.len())
+            .collect_vec();
+
+        let (perm, front_pad_elts_rem) =
+            precompute_hash_state::<F, PoseidonPermutation<F>>(&front_pad_elts);
+
+        let rev_flattened = flattened.iter().rev().copied();
+        // Precompute the Poseidon state for the initial padding chunks
+        let inputs = front_pad_elts_rem
+            .iter()
+            .map(|v| builder.constant(*v))
+            .chain(rev_flattened)
+            .collect_vec();
+        let hash =
+            hash_from_state_circuit::<PoseidonHash, PoseidonPermutation<F>>(builder, perm, &inputs);
+
+        measure_gates_end!(builder, measure);
+        self.hashed_tagged_entries.push(hash);
+    }
+
+    pub fn get(&self, builder: &mut CircuitBuilder, index: &IndexTarget) -> TableEntryTarget {
+        let measure = measure_gates_begin!(builder, "GetTaggedTblEntry");
+        let entry_hash = builder.vec_ref(&self.params, &self.hashed_tagged_entries, index);
+
+        let mut rev_resolved_tagged_flattened =
+            builder.add_virtual_targets(1 + self.max_flattened_entry_len);
+        let query_hash =
+            builder.hash_n_to_hash_no_pad::<PoseidonHash>(rev_resolved_tagged_flattened.clone());
+        builder.connect_flattenable(&entry_hash, &query_hash);
+        rev_resolved_tagged_flattened.reverse();
+        let resolved_tagged_flattened = rev_resolved_tagged_flattened;
+
+        builder.add_simple_generator(TableGetGenerator {
+            index: index.clone(),
+            tagged_entries: self.tagged_entries.clone(),
+            get_tagged_entry: resolved_tagged_flattened.clone(),
+        });
+        measure_gates_end!(builder, measure);
+        TableEntryTarget {
+            params: self.params.clone(),
+            tagged_flattened_entry: resolved_tagged_flattened,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct TableGetGenerator {
+    index: IndexTarget,
+    tagged_entries: Vec<Vec<Target>>,
+    get_tagged_entry: Vec<Target>,
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F, D> for TableGetGenerator {
+    fn id(&self) -> String {
+        "TableGetGenerator".to_string()
+    }
+
+    fn dependencies(&self) -> Vec<Target> {
+        [self.index.low, self.index.high]
+            .into_iter()
+            .chain(self.tagged_entries.iter().flatten().copied())
+            .collect()
+    }
+
+    fn run_once(
+        &self,
+        witness: &PartitionWitness<F>,
+        out_buffer: &mut GeneratedValues<F>,
+    ) -> anyhow::Result<()> {
+        let index_low = witness.get_target(self.index.low);
+        let index_high = witness.get_target(self.index.high);
+        let index = (index_low + index_high * F::from_canonical_usize(1 << 6)).to_canonical_u64();
+
+        let entry = witness.get_targets(&self.tagged_entries[index as usize]);
+
+        for (target, value) in self.get_tagged_entry.iter().zip(
+            entry
+                .iter()
+                .chain(iter::repeat(&F::ZERO).take(self.get_tagged_entry.len())),
+        ) {
+            out_buffer.set_target(*target, *value)?;
+        }
+
+        Ok(())
+    }
+
+    fn serialize(&self, dst: &mut Vec<u8>, _common_data: &CommonCircuitData<F, D>) -> IoResult<()> {
+        dst.write_usize(self.index.max_array_len)?;
+        dst.write_target(self.index.low)?;
+        dst.write_target(self.index.high)?;
+
+        dst.write_usize(self.tagged_entries.len())?;
+        for tagged_entry in &self.tagged_entries {
+            dst.write_target_vec(tagged_entry)?;
+        }
+
+        dst.write_target_vec(&self.get_tagged_entry)
+    }
+
+    fn deserialize(src: &mut Buffer, _common_data: &CommonCircuitData<F, D>) -> IoResult<Self> {
+        let index = IndexTarget {
+            max_array_len: src.read_usize()?,
+            low: src.read_target()?,
+            high: src.read_target()?,
+        };
+        let len = src.read_usize()?;
+        let mut tagged_entries = Vec::with_capacity(len);
+        for _ in 0..len {
+            tagged_entries.push(src.read_target_vec()?);
+        }
+        let get_tagged_entry = src.read_target_vec()?;
+
+        Ok(Self {
+            index,
+            tagged_entries,
+            get_tagged_entry,
+        })
+    }
+}
+
+pub struct TableEntryTarget {
+    params: Params,
+    tagged_flattened_entry: Vec<Target>,
+}
+
+impl TableEntryTarget {
+    pub fn as_type<T: Flattenable>(
+        &self,
+        builder: &mut CircuitBuilder,
+        tag: u32,
+    ) -> (BoolTarget, T) {
+        let tag_target = self.tagged_flattened_entry[0];
+        let flattened_entry = &self.tagged_flattened_entry[1..];
+        let entry = T::from_flattened(&self.params, &flattened_entry[..T::size(&self.params)]);
+        let tag_expect = builder.constant(F(tag as u64));
+        let tag_ok = builder.is_equal(tag_expect, tag_target);
+        (tag_ok, entry)
+    }
+}

src/backends/plonky2/mainpod/mod.rs

@@ -18,7 +18,7 @@ use crate::{
         emptypod::EmptyPod,
         error::{Error, Result},
         mock::emptypod::MockEmptyPod,
-        primitives::merkletree::MerkleClaimAndProof,
+        primitives::{ec::schnorr::SecretKey, merkletree::MerkleClaimAndProof},
         recursion::{
             hash_verifier_data, prove_rec_circuit, RecursiveCircuit, RecursiveCircuitTarget,
         },
@@ -29,9 +29,9 @@ use crate::{
         signedpod::SignedPod,
     },
     middleware::{
-        self, resolve_wildcard_values, value_from_op, AnchoredKey, CustomPredicateBatch, Hash,
+        self, resolve_wildcard_values, value_from_op, AnchoredKey, CustomPredicateBatch,
-        MainPodInputs, NativeOperation, OperationType, Params, Pod, PodId, PodProver, PodType,
+        Error as MiddlewareError, Hash, MainPodInputs, NativeOperation, OperationType, Params, Pod,
-        RecursivePod, StatementArg, ToFields, VDSet, KEY_TYPE, SELF,
+        PodId, PodProver, PodType, RecursivePod, StatementArg, ToFields, VDSet, KEY_TYPE, SELF,
     },
     timed,
 };
@@ -87,16 +87,13 @@ pub(crate) fn extract_custom_predicate_batches(
 /// Extracts all custom predicate operations with all the data required to verify them.
 pub(crate) fn extract_custom_predicate_verifications(
     params: &Params,
+    aux_list: &mut [OperationAux],
     operations: &[middleware::Operation],
     custom_predicate_batches: &[Arc<CustomPredicateBatch>],
 ) -> Result<Vec<CustomPredicateVerification>> {
-    let custom_predicate_data: Vec<_> = operations
+    let mut table = Vec::new();
-        .iter()
+    for (i, op) in operations.iter().enumerate() {
-        .flat_map(|op| match op {
+        if let middleware::Operation::Custom(cpr, sts) = op {
-            middleware::Operation::Custom(cpr, sts) => Some((cpr, sts)),
-            _ => None,
-        })
-        .map(|(cpr, sts)| {
             let wildcard_values =
                 resolve_wildcard_values(params, cpr.predicate(), sts).expect("resolved wildcards");
             let sts = sts.iter().map(|s| Statement::from(s.clone())).collect();
@@ -107,73 +104,105 @@ pub(crate) fn extract_custom_predicate_verifications(
                 .expect("find the custom predicate from the extracted unique list");
             let custom_predicate_table_index =
                 batch_index * params.max_custom_batch_size + cpr.index;
-            CustomPredicateVerification {
+            aux_list[i] = OperationAux::CustomPredVerifyIndex(table.len());
+            table.push(CustomPredicateVerification {
                 custom_predicate_table_index,
                 custom_predicate: cpr.clone(),
                 args: wildcard_values,
                 op_args: sts,
-            }
+            });
-        })
+        }
-        .collect();
+    }
-    if custom_predicate_data.len() > params.max_custom_predicate_verifications {
+
+    if table.len() > params.max_custom_predicate_verifications {
         return Err(Error::custom(format!(
             "The number of required custom predicate verifications ({}) exceeds the maximum number ({}).",
-            custom_predicate_data.len(),
+            table.len(),
             params.max_custom_predicate_verifications
         )));
     }
-    Ok(custom_predicate_data)
+    Ok(table)
 }
 
 /// Extracts Merkle proofs from Contains/NotContains ops.
 pub(crate) fn extract_merkle_proofs(
     params: &Params,
+    aux_list: &mut [OperationAux],
     operations: &[middleware::Operation],
     statements: &[middleware::Statement],
 ) -> Result<Vec<MerkleClaimAndProof>> {
-    assert_eq!(operations.len(), statements.len());
+    let mut table = Vec::new();
-    let merkle_proofs: Vec<_> = operations
+    for (i, (op, st)) in operations.iter().zip(statements.iter()).enumerate() {
-        .iter()
+        let deduction_err = || MiddlewareError::invalid_deduction(op.clone(), st.clone());
-        .zip(statements.iter())
+        let (root, key, value, pf) = match (op, st) {
-        .flat_map(|(op, st)| match (op, st) {
             (
                 middleware::Operation::ContainsFromEntries(root_s, key_s, value_s, pf),
                 middleware::Statement::Contains(root_ref, key_ref, value_ref),
             ) => {
-                let root = value_from_op(root_s, root_ref)?;
+                let root = value_from_op(root_s, root_ref).ok_or_else(deduction_err)?;
-                let key = value_from_op(key_s, key_ref)?;
+                let key = value_from_op(key_s, key_ref).ok_or_else(deduction_err)?;
-                let value = value_from_op(value_s, value_ref)?;
+                let value = value_from_op(value_s, value_ref).ok_or_else(deduction_err)?;
-                Some(MerkleClaimAndProof::new(
+                (root.raw(), key.raw(), Some(value.raw()), pf)
-                    Hash::from(root.raw()),
-                    key.raw(),
-                    Some(value.raw()),
-                    pf.clone(),
-                ))
             }
             (
                 middleware::Operation::NotContainsFromEntries(root_s, key_s, pf),
                 middleware::Statement::NotContains(root_ref, key_ref),
             ) => {
-                let root = value_from_op(root_s, root_ref)?;
+                let root = value_from_op(root_s, root_ref).ok_or_else(deduction_err)?;
-                let key = value_from_op(key_s, key_ref)?;
+                let key = value_from_op(key_s, key_ref).ok_or_else(deduction_err)?;
-                Some(MerkleClaimAndProof::new(
+                (root.raw(), key.raw(), None, pf)
-                    Hash::from(root.raw()),
-                    key.raw(),
-                    None,
-                    pf.clone(),
-                ))
             }
-            _ => None,
+            _ => continue,
-        })
+        };
-        .collect();
+        aux_list[i] = OperationAux::MerkleProofIndex(table.len());
-    if merkle_proofs.len() > params.max_merkle_proofs_containers {
+        table.push(MerkleClaimAndProof::new(
+            Hash::from(root),
+            key,
+            value,
+            pf.clone(),
+        ));
+    }
+    if table.len() > params.max_merkle_proofs_containers {
         return Err(Error::custom(format!(
             "The number of required Merkle proofs ({}) exceeds the maximum number ({}).",
-            merkle_proofs.len(),
+            table.len(),
             params.max_merkle_proofs_containers
         )));
     }
-    Ok(merkle_proofs)
+    Ok(table)
+}
+
+pub(crate) fn extract_public_key_of(
+    params: &Params,
+    aux_list: &mut [OperationAux],
+    operations: &[middleware::Operation],
+    statements: &[middleware::Statement],
+) -> Result<Vec<SecretKey>> {
+    let mut table = Vec::new();
+    for (i, (op, st)) in operations.iter().zip(statements.iter()).enumerate() {
+        if let (
+            middleware::Operation::PublicKeyOf(_, sk_s),
+            middleware::Statement::PublicKeyOf(_, sk_ref),
+        ) = (op, st)
+        {
+            let deduction_err = || MiddlewareError::invalid_deduction(op.clone(), st.clone());
+            let sk = SecretKey::try_from(
+                value_from_op(sk_s, sk_ref)
+                    .ok_or_else(deduction_err)?
+                    .typed(),
+            )?;
+            aux_list[i] = OperationAux::PublicKeyOfIndex(table.len());
+            table.push(sk);
+        }
+    }
+    if table.len() > params.max_public_key_of {
+        return Err(Error::custom(format!(
+            "The number of required PublicKeyOf verifications ({}) exceeds the maximum number ({}).",
+            table.len(),
+            params.max_public_statements
+        )));
+    }
+    Ok(table)
 }
 
 /// Find the operation argument statement in the list of previous statements and return the index.
@@ -192,52 +221,6 @@ fn find_op_arg(statements: &[Statement], op_arg: &middleware::Statement) -> Resu
         )))
 }
 
-/// Find the operation auxiliary data in the list of auxiliary data and return the index.
-// NOTE: The `custom_predicate_verifications` is optional because in the MainPod we want to store
-// the index of a custom predicate verification in the aux data, but in the MockMainPod we don't
-// need that because we keep a reference to the custom predicate in the operation type, which
-// removes the need for indexing.  We could change the OperationType and Predicate for the backend
-// to not keep a reference to the custom predicate and instead just keep the id and index and then
-// do the same double indexing that the MainPod does to verify custom predicates.
-fn find_op_aux(
-    merkle_proofs: &[MerkleClaimAndProof],
-    custom_predicate_verifications: Option<&[CustomPredicateVerification]>,
-    op: &middleware::Operation,
-) -> Result<OperationAux> {
-    let op_aux = op.aux();
-    if let (middleware::Operation::Custom(cpr, op_args), Some(cpvs)) =
-        (op, custom_predicate_verifications)
-    {
-        return Ok(cpvs
-            .iter()
-            .enumerate()
-            .find_map(|(i, cpv)| {
-                (cpv.custom_predicate.batch.id() == cpr.batch.id()
-                    && cpv.custom_predicate.index == cpr.index
-                    && cpv
-                        .op_args
-                        .iter()
-                        .zip_eq(op_args.iter())
-                        .all(|(a0, a1)| a0.0 == a1.predicate() && a0.1 == a1.args()))
-                .then_some(i)
-            })
-            .map(OperationAux::CustomPredVerifyIndex)
-            .expect("custom predicate verification in the list"));
-    }
-    match &op_aux {
-        middleware::OperationAux::None => Ok(OperationAux::None),
-        middleware::OperationAux::MerkleProof(pf_arg) => merkle_proofs
-            .iter()
-            .enumerate()
-            .find_map(|(i, pf)| (pf.proof == *pf_arg).then_some(i))
-            .map(OperationAux::MerkleProofIndex)
-            .ok_or(Error::custom(format!(
-                "Merkle proof corresponding to op arg {} not found",
-                op_aux
-            ))),
-    }
-}
-
 fn fill_pad<T: Clone>(v: &mut Vec<T>, pad_value: T, len: usize) {
     if v.len() > len {
         panic!("length exceeded");
@@ -367,12 +350,12 @@ pub(crate) fn layout_statements(
 pub(crate) fn process_private_statements_operations(
     params: &Params,
     statements: &[Statement],
-    merkle_proofs: &[MerkleClaimAndProof],
+    aux_list: &[OperationAux],
-    custom_predicate_verifications: Option<&[CustomPredicateVerification]>,
     input_operations: &[middleware::Operation],
 ) -> Result<Vec<Operation>> {
+    assert_eq!(params.max_priv_statements(), aux_list.len());
     let mut operations = Vec::new();
-    for i in 0..params.max_priv_statements() {
+    for (i, aux) in aux_list.iter().enumerate() {
         let op = input_operations
             .get(i)
             .unwrap_or(&middleware::Operation::None)
@@ -383,10 +366,8 @@ pub(crate) fn process_private_statements_operations(
             .map(|mid_arg| find_op_arg(statements, mid_arg))
             .collect::<Result<Vec<_>>>()?;
 
-        let aux = find_op_aux(merkle_proofs, custom_predicate_verifications, &op)?;
-
         pad_operation_args(params, &mut args);
-        operations.push(Operation(op.op_type(), args, aux));
+        operations.push(Operation(op.op_type(), args, *aux));
     }
     Ok(operations)
 }
@@ -475,20 +456,25 @@ impl PodProver for Prover {
             })
             .collect_vec();
 
-        let merkle_proofs = extract_merkle_proofs(params, inputs.operations, inputs.statements)?;
+        // Aux values for backend::Operation
+        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_predicate_verifications = extract_custom_predicate_verifications(
             params,
+            &mut aux_list,
             inputs.operations,
             &custom_predicate_batches,
         )?;
+        let public_key_of_sks =
+            extract_public_key_of(params, &mut aux_list, inputs.operations, inputs.statements)?;
 
         let (statements, public_statements) = layout_statements(params, false, &inputs)?;
         let operations = process_private_statements_operations(
             params,
             &statements,
-            &merkle_proofs,
+            &aux_list,
-            Some(&custom_predicate_verifications),
             inputs.operations,
         )?;
         let operations = process_public_statements_operations(params, &statements, operations)?;
@@ -523,6 +509,7 @@ impl PodProver for Prover {
             statements: statements[statements.len() - params.max_statements..].to_vec(),
             operations,
             merkle_proofs,
+            public_key_of_sks,
             custom_predicate_batches,
             custom_predicate_verifications,
         };
@@ -845,6 +832,45 @@ pub mod tests {
         pod.verify().unwrap()
     }
 
+    // This pod does nothing but it's useful for debugging to keep things small.
+    #[ignore]
+    #[test]
+    fn test_mini_1() {
+        let params = middleware::Params {
+            max_input_signed_pods: 0,
+            max_input_recursive_pods: 0,
+            max_signed_pod_values: 0,
+            max_statements: 2,
+            max_public_statements: 1,
+            max_input_pods_public_statements: 0,
+            max_merkle_proofs_containers: 0,
+            max_public_key_of: 0,
+            max_custom_predicate_verifications: 0,
+            max_custom_predicate_batches: 0,
+            ..Default::default()
+        };
+        let mut vds = DEFAULT_VD_LIST.clone();
+        vds.push(rec_main_pod_circuit_data(&params).1.verifier_only.clone());
+        let vd_set = VDSet::new(params.max_depth_mt_vds, &vds).unwrap();
+
+        let builder = frontend::MainPodBuilder::new(&params, &vd_set);
+        println!("{}", builder);
+        println!();
+
+        // Mock
+        let prover = MockProver {};
+        let pod = builder.prove(&prover).unwrap();
+        let pod = (pod.pod as Box<dyn Any>).downcast::<MockMainPod>().unwrap();
+        pod.verify().unwrap();
+        println!("{:#}", pod);
+
+        // Real
+        let prover = Prover {};
+        let pod = builder.prove(&prover).unwrap();
+        let pod = (pod.pod as Box<dyn Any>).downcast::<MainPod>().unwrap();
+        pod.verify().unwrap()
+    }
+
     #[test]
     fn test_mainpod_small_empty() {
         let params = middleware::Params {
@@ -863,6 +889,7 @@ pub mod tests {
             max_custom_predicate_wildcards: 3,
             max_custom_batch_size: 2,
             max_merkle_proofs_containers: 2,
+            max_public_key_of: 2,
             max_depth_mt_containers: 4,
             max_depth_mt_vds: 6,
         };
@@ -927,6 +954,7 @@ pub mod tests {
             max_custom_batch_size: 3,
             max_custom_predicate_wildcards: 4,
             max_custom_predicate_verifications: 2,
+            max_merkle_proofs_containers: 0,
             ..Default::default()
         };
         println!("{:#?}", params);
@@ -980,7 +1008,7 @@ pub mod tests {
         let st = builder
             .pub_op(frontend::Operation::new_entry(
                 "entry",
-                Set::new(params.max_merkle_proofs_containers, set).unwrap(),
+                Set::new(params.max_depth_mt_containers, set).unwrap(),
             ))
             .unwrap();
 

src/backends/plonky2/mainpod/operation.rs

@@ -8,7 +8,7 @@ use crate::{
         mainpod::Statement,
         primitives::merkletree::MerkleClaimAndProof,
     },
-    middleware::{self, OperationType},
+    middleware::{self, OperationType, Params},
 };
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
@@ -30,19 +30,36 @@ impl OperationArg {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
 pub enum OperationAux {
     None,
     MerkleProofIndex(usize),
+    PublicKeyOfIndex(usize),
     CustomPredVerifyIndex(usize),
 }
 
 impl OperationAux {
-    pub fn as_usizes(&self) -> [usize; 2] {
+    fn table_offset_merkle_proof(_params: &Params) -> usize {
+        // At index 0 we store a zero entry
+        1
+    }
+    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 {
+        Self::table_offset_public_key_of(params) + params.max_public_key_of
+    }
+    pub(crate) fn table_size(params: &Params) -> usize {
+        1 + params.max_merkle_proofs_containers
+            + params.max_public_key_of
+            + params.max_custom_predicate_verifications
+    }
+    pub fn table_index(&self, params: &Params) -> usize {
         match self {
-            Self::None => [0, 0],
+            Self::None => 0,
-            Self::MerkleProofIndex(i) => [*i, 0],
+            Self::MerkleProofIndex(i) => Self::table_offset_merkle_proof(params) + *i,
-            Self::CustomPredVerifyIndex(i) => [0, *i],
+            Self::PublicKeyOfIndex(i) => Self::table_offset_public_key_of(params) + *i,
+            Self::CustomPredVerifyIndex(i) => Self::table_offset_custom_pred_verify(params) + *i,
         }
     }
 }
@@ -87,6 +104,7 @@ impl Operation {
                     .proof
                     .clone(),
             ),
+            OperationAux::PublicKeyOfIndex(_) => crate::middleware::OperationAux::None,
         };
         Ok(middleware::Operation::op(
             self.0.clone(),
@@ -114,6 +132,7 @@ impl fmt::Display for Operation {
             OperationAux::None => (),
             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)?,
         }
         Ok(())
     }

src/backends/plonky2/mock/mainpod.rs

@@ -14,7 +14,7 @@ use crate::{
         mainpod::{
             calculate_id, extract_merkle_proofs, layout_statements,
             process_private_statements_operations, process_public_statements_operations, Operation,
-            Statement,
+            OperationAux, Statement,
         },
         mock::emptypod::MockEmptyPod,
         primitives::merkletree::MerkleClaimAndProof,
@@ -172,14 +172,15 @@ impl MockMainPod {
 
     pub fn new(params: &Params, inputs: MainPodInputs) -> Result<Self> {
         let (statements, public_statements) = layout_statements(params, true, &inputs)?;
+        let mut aux_list = vec![OperationAux::None; params.max_priv_statements()];
         // Extract Merkle proofs and pad.
-        let merkle_proofs = extract_merkle_proofs(params, inputs.operations, inputs.statements)?;
+        let merkle_proofs =
+            extract_merkle_proofs(params, &mut aux_list, inputs.operations, inputs.statements)?;
 
         let operations = process_private_statements_operations(
             params,
             &statements,
-            &merkle_proofs,
+            &aux_list,
-            None,
             inputs.operations,
         )?;
         let operations = process_public_statements_operations(params, &statements, operations)?;

src/backends/plonky2/serialization.rs

@@ -19,7 +19,7 @@ use serde::{de, ser, Deserialize, Serialize};
 
 use crate::backends::plonky2::{
     basetypes::{CircuitData, CommonCircuitData, VerifierCircuitData, C, D, F},
-    circuits::{common::LtMaskGenerator, utils::DebugGenerator},
+    circuits::{common::LtMaskGenerator, mux_table::TableGetGenerator, utils::DebugGenerator},
     primitives::ec::{
         bits::ConditionalZeroGenerator,
         curve::PointSquareRootGenerator,
@@ -92,7 +92,6 @@ use plonky2::{
 #[derive(Debug)]
 pub(crate) struct Pod2GeneratorSerializer {}
 
-// TODO: Add pod2 custom generators
 impl WitnessGeneratorSerializer<F, D> for Pod2GeneratorSerializer {
     impl_generator_serializer! {
         Pod2GeneratorSerializer,
@@ -130,7 +129,8 @@ impl WitnessGeneratorSerializer<F, D> for Pod2GeneratorSerializer {
         RecursiveGenerator<1, NNFMulSimple<5, QuinticExtension<F>>>,
         RecursiveGenerator<D, ECAddHomogOffset>,
         RecursiveGenerator<1, ECAddHomogOffset>,
-        ComparisonGenerator<F, D>
+        ComparisonGenerator<F, D>,
+        TableGetGenerator
     }
 }
 

src/middleware/mod.rs

@@ -764,8 +764,11 @@ pub struct Params {
     pub max_depth_mt_containers: usize,
     // maximum depth of the merkle tree gadget used for verifier_data membership
     // check.  This allows creating verifying sets of pod circuits of size
-    // 2^max_depth_mt_vds.
+    // 2^max_depth_mt_vds.  Limits the number of container operations of the type Contains,
+    // NotContains.
     pub max_depth_mt_vds: usize,
+    // maximum number of public key derivations used for PublicKeyOf operation
+    pub max_public_key_of: usize,
     //
     // The following parameters define how a pod id is calculated.  They need to be the same among
     // different circuits to be compatible in their verification.
@@ -803,6 +806,7 @@ impl Default for Params {
             max_merkle_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,
         }
     }
 }
@@ -828,10 +832,6 @@ impl Params {
         Self::predicate_size() + STATEMENT_ARG_F_LEN * self.max_statement_args
     }
 
-    pub fn operation_size(&self, operation_arg_f_len: usize) -> usize {
-        Self::operation_type_size() + operation_arg_f_len * self.max_operation_args
-    }
-
     pub const fn statement_tmpl_size(&self) -> usize {
         Self::predicate_size() + self.max_statement_args * Self::statement_tmpl_arg_size()
     }