Allow literals in statements (#276)

Implements #229 and #261.

src/backends/plonky2/circuits/common.rs

@@ -864,6 +864,12 @@ pub trait CircuitBuilderPod<F: RichField + Extendable<D>, const D: usize> {
     fn add_virtual_custom_predicate_entry(&mut self, params: &Params)
         -> CustomPredicateEntryTarget;
     fn select_value(&mut self, b: BoolTarget, x: ValueTarget, y: ValueTarget) -> ValueTarget;
+    fn select_statement_arg(
+        &mut self,
+        b: BoolTarget,
+        x: &StatementArgTarget,
+        y: &StatementArgTarget,
+    ) -> StatementArgTarget;
     fn select_bool(&mut self, b: BoolTarget, x: BoolTarget, y: BoolTarget) -> BoolTarget;
     fn constant_value(&mut self, v: RawValue) -> ValueTarget;
     fn is_equal_slice(&mut self, xs: &[Target], ys: &[Target]) -> BoolTarget;
@@ -1038,6 +1044,17 @@ impl CircuitBuilderPod<F, D> for CircuitBuilder {
         }
     }
 
+    fn select_statement_arg(
+        &mut self,
+        b: BoolTarget,
+        x: &StatementArgTarget,
+        y: &StatementArgTarget,
+    ) -> StatementArgTarget {
+        StatementArgTarget {
+            elements: std::array::from_fn(|i| self.select(b, x.elements[i], y.elements[i])),
+        }
+    }
+
     fn select_bool(&mut self, b: BoolTarget, x: BoolTarget, y: BoolTarget) -> BoolTarget {
         BoolTarget::new_unsafe(self.select(b, x.target, y.target))
     }

src/backends/plonky2/circuits/signedpod.rs

@@ -94,8 +94,7 @@ impl SignedPodVerifyTarget {
         self_id: bool,
     ) -> Vec<StatementTarget> {
         let mut statements = Vec::new();
-        let predicate =
-            PredicateTarget::new_native(builder, &self.params, NativePredicate::ValueOf);
+        let predicate = PredicateTarget::new_native(builder, &self.params, NativePredicate::Equal);
         let pod_id = if self_id {
             builder.constant_value(SELF.0.into())
         } else {

src/backends/plonky2/emptypod.rs

@@ -40,7 +40,7 @@ struct EmptyPodVerifyCircuit {
 }
 
 fn type_statement() -> Statement {
-    Statement::ValueOf(
+    Statement::equal(
         AnchoredKey::from((SELF, KEY_TYPE)),
         Value::from(PodType::Empty),
     )

src/backends/plonky2/mainpod/mod.rs

@@ -26,9 +26,9 @@ use crate::{
         STANDARD_REC_MAIN_POD_CIRCUIT_DATA,
     },
     middleware::{
-        self, resolve_wildcard_values, AnchoredKey, CustomPredicateBatch, DynError, Hash,
-        MainPodInputs, NativeOperation, OperationType, Params, Pod, PodId, PodProver, PodType,
-        RecursivePod, StatementArg, ToFields, VDSet, F, KEY_TYPE, SELF,
+        self, resolve_wildcard_values, value_from_op, AnchoredKey, CustomPredicateBatch, DynError,
+        Hash, MainPodInputs, NativeOperation, OperationType, Params, Pod, PodId, PodProver,
+        PodType, RecursivePod, StatementArg, ToFields, VDSet, F, KEY_TYPE, SELF,
     },
 };
 
@@ -125,31 +125,40 @@ pub(crate) fn extract_custom_predicate_verifications(
 pub(crate) fn extract_merkle_proofs(
     params: &Params,
     operations: &[middleware::Operation],
+    statements: &[middleware::Statement],
 ) -> Result<Vec<MerkleClaimAndProof>> {
+    assert_eq!(operations.len(), statements.len());
     let merkle_proofs: Vec<_> = operations
         .iter()
-        .flat_map(|op| match op {
-            middleware::Operation::ContainsFromEntries(
-                middleware::Statement::ValueOf(_, root),
-                middleware::Statement::ValueOf(_, key),
-                middleware::Statement::ValueOf(_, value),
-                pf,
-            ) => Some(MerkleClaimAndProof::new(
-                Hash::from(root.raw()),
-                key.raw(),
-                Some(value.raw()),
-                pf.clone(),
-            )),
-            middleware::Operation::NotContainsFromEntries(
-                middleware::Statement::ValueOf(_, root),
-                middleware::Statement::ValueOf(_, key),
-                pf,
-            ) => Some(MerkleClaimAndProof::new(
-                Hash::from(root.raw()),
-                key.raw(),
-                None,
-                pf.clone(),
-            )),
+        .zip(statements.iter())
+        .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 key = value_from_op(key_s, key_ref)?;
+                let value = value_from_op(value_s, value_ref)?;
+                Some(MerkleClaimAndProof::new(
+                    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 key = value_from_op(key_s, key_ref)?;
+                Some(MerkleClaimAndProof::new(
+                    Hash::from(root.raw()),
+                    key.raw(),
+                    None,
+                    pf.clone(),
+                ))
+            }
             _ => None,
         })
         .collect();
@@ -320,9 +329,9 @@ pub(crate) fn layout_statements(
 
     // Public statements
     assert!(inputs.public_statements.len() < params.max_public_statements);
-    let mut type_st = middleware::Statement::ValueOf(
-        AnchoredKey::from((SELF, KEY_TYPE)),
-        middleware::Value::from(PodType::MockMain),
+    let mut type_st = middleware::Statement::Equal(
+        AnchoredKey::from((SELF, KEY_TYPE)).into(),
+        middleware::Value::from(PodType::MockMain).into(),
     )
     .into();
     pad_statement(params, &mut type_st);
@@ -470,7 +479,7 @@ impl Prover {
             })
             .collect_vec();
 
-        let merkle_proofs = extract_merkle_proofs(params, inputs.operations)?;
+        let merkle_proofs = extract_merkle_proofs(params, inputs.operations, inputs.statements)?;
         let custom_predicate_batches = extract_custom_predicate_batches(params, inputs.operations)?;
         let custom_predicate_verifications = extract_custom_predicate_verifications(
             params,
@@ -805,12 +814,12 @@ pub mod tests {
             max_signed_pod_values: 2,
             max_public_statements: 2,
             num_public_statements_id: 4,
-            max_statement_args: 2,
+            max_statement_args: 3,
             max_operation_args: 3,
             max_custom_predicate_batches: 2,
             max_custom_predicate_verifications: 2,
             max_custom_predicate_arity: 2,
-            max_custom_predicate_wildcards: 2,
+            max_custom_predicate_wildcards: 3,
             max_custom_batch_size: 2,
             max_merkle_proofs_containers: 2,
             max_depth_mt_containers: 4,
@@ -905,7 +914,7 @@ pub mod tests {
         let vd_set = &*DEFAULT_VD_SET;
 
         let mut cpb_builder = CustomPredicateBatchBuilder::new(params.clone(), "cpb".into());
-        let stb0 = STB::new(NP::ValueOf)
+        let stb0 = STB::new(NP::Equal)
             .arg(("id", key("score")))
             .arg(literal(42));
         let stb1 = STB::new(NP::Equal)

src/backends/plonky2/mainpod/operation.rs

@@ -75,7 +75,11 @@ impl Operation {
             .iter()
             .flat_map(|arg| match arg {
                 OperationArg::None => None,
-                OperationArg::Index(i) => Some(statements[*i].clone().try_into()),
+                OperationArg::Index(i) => {
+                    let st: Result<crate::middleware::Statement> =
+                        statements[*i].clone().try_into();
+                    Some(st)
+                }
             })
             .collect::<Result<Vec<_>>>()?;
         let deref_aux = match self.2 {

src/backends/plonky2/mainpod/statement.rs

@@ -48,41 +48,25 @@ impl TryFrom<Statement> for middleware::Statement {
         type NP = NativePredicate;
         type SA = StatementArg;
         let proper_args = s.args();
-        let args = (
-            proper_args.first().cloned(),
-            proper_args.get(1).cloned(),
-            proper_args.get(2).cloned(),
-        );
         Ok(match s.0 {
-            Predicate::Native(np) => match (np, args, proper_args.len()) {
-                (NP::None, _, 0) => S::None,
-                (NP::ValueOf, (Some(SA::Key(ak)), Some(SA::Literal(v)), None), 2) => {
-                    S::ValueOf(ak, v)
+            Predicate::Native(np) => match (np, &proper_args.as_slice()) {
+                (NP::None, &[]) => S::None,
+                (NP::Equal, &[a1, a2]) => S::Equal(a1.try_into()?, a2.try_into()?),
+                (NP::NotEqual, &[a1, a2]) => S::NotEqual(a1.try_into()?, a2.try_into()?),
+                (NP::LtEq, &[a1, a2]) => S::LtEq(a1.try_into()?, a2.try_into()?),
+                (NP::Lt, &[a1, a2]) => S::Lt(a1.try_into()?, a2.try_into()?),
+                (NP::Contains, &[a1, a2, a3]) => {
+                    S::Contains(a1.try_into()?, a2.try_into()?, a3.try_into()?)
                 }
-                (NP::Equal, (Some(SA::Key(ak1)), Some(SA::Key(ak2)), None), 2) => {
-                    S::Equal(ak1, ak2)
+                (NP::NotContains, &[a1, a2]) => S::NotContains(a1.try_into()?, a2.try_into()?),
+                (NP::SumOf, &[a1, a2, a3]) => {
+                    S::SumOf(a1.try_into()?, a2.try_into()?, a3.try_into()?)
                 }
-                (NP::NotEqual, (Some(SA::Key(ak1)), Some(SA::Key(ak2)), None), 2) => {
-                    S::NotEqual(ak1, ak2)
+                (NP::ProductOf, &[a1, a2, a3]) => {
+                    S::ProductOf(a1.try_into()?, a2.try_into()?, a3.try_into()?)
                 }
-                (NP::LtEq, (Some(SA::Key(ak1)), Some(SA::Key(ak2)), None), 2) => S::LtEq(ak1, ak2),
-                (NP::Lt, (Some(SA::Key(ak1)), Some(SA::Key(ak2)), None), 2) => S::Lt(ak1, ak2),
-                (NP::Contains, (Some(SA::Key(ak1)), Some(SA::Key(ak2)), Some(SA::Key(ak3))), 3) => {
-                    S::Contains(ak1, ak2, ak3)
-                }
-                (NP::NotContains, (Some(SA::Key(ak1)), Some(SA::Key(ak2)), None), 2) => {
-                    S::NotContains(ak1, ak2)
-                }
-                (NP::SumOf, (Some(SA::Key(ak1)), Some(SA::Key(ak2)), Some(SA::Key(ak3))), 3) => {
-                    S::SumOf(ak1, ak2, ak3)
-                }
-                (
-                    NP::ProductOf,
-                    (Some(SA::Key(ak1)), Some(SA::Key(ak2)), Some(SA::Key(ak3))),
-                    3,
-                ) => S::ProductOf(ak1, ak2, ak3),
-                (NP::MaxOf, (Some(SA::Key(ak1)), Some(SA::Key(ak2)), Some(SA::Key(ak3))), 3) => {
-                    S::MaxOf(ak1, ak2, ak3)
+                (NP::MaxOf, &[a1, a2, a3]) => {
+                    S::MaxOf(a1.try_into()?, a2.try_into()?, a3.try_into()?)
                 }
                 _ => Err(Error::custom(format!(
                     "Ill-formed statement expression {:?}",

src/backends/plonky2/mock/emptypod.rs

@@ -19,7 +19,7 @@ pub struct MockEmptyPod {
 }
 
 fn type_statement() -> Statement {
-    Statement::ValueOf(
+    Statement::equal(
         AnchoredKey::from((SELF, KEY_TYPE)),
         Value::from(PodType::Empty),
     )

src/backends/plonky2/mock/mainpod.rs

@@ -4,6 +4,7 @@
 
 use std::fmt;
 
+use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 
 use crate::{
@@ -18,8 +19,9 @@ use crate::{
         primitives::merkletree::MerkleClaimAndProof,
     },
     middleware::{
-        self, hash_str, AnchoredKey, DynError, Hash, MainPodInputs, NativePredicate, Params, Pod,
-        PodId, PodProver, PodType, Predicate, RecursivePod, StatementArg, VDSet, KEY_TYPE, SELF,
+        self, hash_str, AnchoredKey, DynError, Hash, MainPodInputs, NativeOperation,
+        NativePredicate, OperationType, Params, Pod, PodId, PodProver, PodType, Predicate,
+        RecursivePod, StatementArg, VDSet, KEY_TYPE, SELF,
     },
 };
 
@@ -160,7 +162,7 @@ impl MockMainPod {
         // value=PodType::MockMainPod`
         let (statements, public_statements) = layout_statements(params, true, &inputs)?;
         // Extract Merkle proofs and pad.
-        let merkle_proofs = extract_merkle_proofs(params, inputs.operations)?;
+        let merkle_proofs = extract_merkle_proofs(params, inputs.operations, inputs.statements)?;
 
         let operations = process_private_statements_operations(
             params,
@@ -225,44 +227,30 @@ impl MockMainPod {
         // find a ValueOf statement from the public statements with key=KEY_TYPE and check that the
         // value is PodType::MockMainPod
         let has_type_statement = self.public_statements.iter().any(|s| {
-            s.0 == Predicate::Native(NativePredicate::ValueOf)
-                && !s.1.is_empty()
-                && if let StatementArg::Key(AnchoredKey { pod_id, ref key }) = s.1[0] {
-                    pod_id == SELF && key.hash() == hash_str(KEY_TYPE)
+            s.0 == Predicate::Native(NativePredicate::Equal) && {
+                if let [StatementArg::Key(AnchoredKey { pod_id, ref key }), StatementArg::Literal(_)] = &s.1[..2] {
+                    pod_id == &SELF && key.hash() == hash_str(KEY_TYPE)
                 } else {
                     false
                 }
-        });
-        // 3. check that all `input_statements` of type `ValueOf` with origin=SELF have unique keys
+        }});
+        // 3. check that all `NewEntry` operations have unique keys
         // (no duplicates)
-        // TODO: Instead of doing this, do a uniqueness check when verifying the output of a
-        // `NewValue` operation.
-        let value_ofs_unique = {
-            let key_id_pairs = input_statements
-                .iter()
-                .enumerate()
-                .map(|(i, s)| {
-                    (
-                        // Separate private from public statements.
-                        if i < self.params.max_priv_statements() {
-                            0
-                        } else {
-                            1
-                        },
-                        s,
-                    )
-                })
-                .filter(|(_, s)| s.0 == Predicate::Native(NativePredicate::ValueOf))
-                .flat_map(|(i, s)| {
-                    if let StatementArg::Key(ak) = &s.1[0] {
-                        vec![(i, ak.pod_id, ak.key.hash())]
-                    } else {
-                        vec![]
+        let value_ofs_unique = input_statements
+            .iter()
+            .zip(self.operations.iter())
+            .filter_map(|(s, o)| {
+                if matches!(o.0, OperationType::Native(NativeOperation::NewEntry)) {
+                    match s.1.get(0) {
+                        Some(StatementArg::Key(k)) => Some(k),
+                        // malformed NewEntry operations are caught in step 5
+                        _ => None,
                     }
-                })
-                .collect::<Vec<_>>();
-            !(0..key_id_pairs.len() - 1).any(|i| key_id_pairs[i + 1..].contains(&key_id_pairs[i]))
-        };
+                } else {
+                    None
+                }
+            })
+            .all_unique();
         // 4. TODO: Verify type
 
         // 5. verify that all `input_statements` are correctly generated

src/backends/plonky2/mock/signedpod.rs

@@ -156,7 +156,7 @@ impl Pod for MockSignedPod {
         [(key_type, value_type), (key_signer, value_signer)]
             .into_iter()
             .chain(kvs.into_iter().sorted_by_key(|kv| kv.0.hash()))
-            .map(|(k, v)| Statement::ValueOf(AnchoredKey::from((SELF, k)), v))
+            .map(|(k, v)| Statement::equal(AnchoredKey::from((SELF, k)), v))
             .collect()
     }
 

src/backends/plonky2/signedpod.rs

@@ -192,7 +192,7 @@ impl Pod for SignedPod {
         [(key_type, value_type), (key_signer, value_signer)]
             .into_iter()
             .chain(kvs.into_iter().sorted_by_key(|kv| kv.0.hash()))
-            .map(|(k, v)| Statement::ValueOf(AnchoredKey::from((SELF, k)), v))
+            .map(|(k, v)| Statement::equal(AnchoredKey::from((SELF, k)), v))
             .collect()
     }