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

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::MerkleProofIndex(i) => [*i, 0],
-            Self::CustomPredVerifyIndex(i) => [0, *i],
+            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::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(())
     }