Serialize and hash custom predicates (#90)

* Print pods from SignedPodBuilder

* Add additional print to test printing SignedPodBuilder

* Mock-prove and print MainPod

* Implement ToFields for custom predicates and dependencies

* Test: print serialization of a recursive batch

* Rearrange serialization of CustomPredicate so args_len is always in the same position

* Serialize predicates with first entry nonzero to avoid collision with padding

* Off by one error in ethdos test BatchSelf(2)

* cargo fmt

* not a typo

* Typos, trying again

Author: tideofwords

.github/workflows/typos.toml

@@ -1,3 +1,4 @@
 [default.extend-words]
 groth = "groth" # to avoid it dectecting it as 'growth'
 BA = "BA"
+Ded = "Ded" # "ANDed", it thought "Ded" should be "Dead"

src/backends/mock_main/mod.rs

@@ -327,7 +327,7 @@ impl MockMainPod {
             statements[statements.len() - params.max_public_statements..].to_vec();
 
         // get the id out of the public statements
-        let id: PodId = PodId(hash_statements(&public_statements));
+        let id: PodId = PodId(hash_statements(&public_statements, *params));
 
         Ok(Self {
             params: params.clone(),
@@ -362,10 +362,10 @@ impl MockMainPod {
     }
 }
 
-pub fn hash_statements(statements: &[Statement]) -> middleware::Hash {
+pub fn hash_statements(statements: &[Statement], params: Params) -> middleware::Hash {
     let field_elems = statements
         .into_iter()
-        .flat_map(|statement| statement.clone().to_fields().0)
+        .flat_map(|statement| statement.clone().to_fields(params).0)
         .collect::<Vec<_>>();
     Hash(PoseidonHash::hash_no_pad(&field_elems).elements)
 }
@@ -376,7 +376,7 @@ impl Pod for MockMainPod {
         // get the input_statements from the self.statements
         let input_statements = &self.statements[input_statement_offset..];
         // get the id out of the public statements, and ensure it is equal to self.id
-        let ids_match = self.id == PodId(hash_statements(&self.public_statements));
+        let ids_match = self.id == PodId(hash_statements(&self.public_statements, self.params));
         // 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

src/backends/mock_main/statement.rs

@@ -1,7 +1,7 @@
 use anyhow::{anyhow, Result};
 use std::fmt;
 
-use crate::middleware::{self, NativePredicate, StatementArg, ToFields};
+use crate::middleware::{self, NativePredicate, Params, StatementArg, ToFields};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct Statement(pub NativePredicate, pub Vec<StatementArg>);
@@ -21,12 +21,13 @@ impl Statement {
 }
 
 impl ToFields for Statement {
-    fn to_fields(self) -> (Vec<middleware::F>, usize) {
-        let (native_statement_f, native_statement_f_len) = self.0.to_fields();
+    fn to_fields(&self, params: Params) -> (Vec<middleware::F>, usize) {
+        let (native_statement_f, native_statement_f_len) = self.0.to_fields(params);
         let (vec_statementarg_f, vec_statementarg_f_len) = self
             .1
+            .clone()
             .into_iter()
-            .map(|statement_arg| statement_arg.to_fields())
+            .map(|statement_arg| statement_arg.to_fields(params))
             .fold((Vec::new(), 0), |mut acc, (f, l)| {
                 acc.0.extend(f);
                 acc.1 += l;

src/frontend/custom.rs

@@ -171,13 +171,16 @@ fn resolve_wildcard(args: &[&str], priv_args: &[&str], v: &HashOrWildcardStr) ->
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::middleware::{CustomPredicateRef, PodType};
+    use crate::middleware::{CustomPredicateRef, Params, PodType};
 
     #[test]
     fn test_custom_pred() {
         use NativePredicate as NP;
         use StatementTmplBuilder as STB;
 
+        let params = Params::default();
+        params.print_serialized_sizes();
+
         let mut builder = CustomPredicateBatchBuilder::new("eth_friend".into());
         let _eth_friend = builder.predicate_and(
             // arguments:
@@ -239,7 +242,7 @@ mod tests {
             builder.predicates.last().unwrap()
         );
 
-        let eth_dos_distance = Predicate::BatchSelf(3);
+        let eth_dos_distance = Predicate::BatchSelf(2);
 
         // next chunk builds:
         let eth_dos_distance_ind = builder.predicate_and(
@@ -311,5 +314,9 @@ mod tests {
             "b.2. eth_dos_distance = {}",
             builder.predicates.last().unwrap()
         );
+
+        let eth_dos_batch_b = builder.finish();
+        let fields = eth_dos_batch_b.to_fields(params);
+        println!("Batch b, serialized: {:?}", fields);
     }
 }

src/frontend/mod.rs

@@ -96,6 +96,16 @@ pub struct SignedPodBuilder {
     pub kvs: HashMap<String, Value>,
 }
 
+impl fmt::Display for SignedPodBuilder {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        writeln!(f, "SignedPodBuilder:")?;
+        for (k, v) in self.kvs.iter().sorted_by_key(|kv| kv.0) {
+            writeln!(f, "  - {}: {}", k, v)?;
+        }
+        Ok(())
+    }
+}
+
 impl SignedPodBuilder {
     pub fn new(params: &Params) -> Self {
         Self {
@@ -347,6 +357,22 @@ pub struct MainPod {
     // TODO: metadata
 }
 
+impl fmt::Display for MainPod {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        writeln!(f, "MainPod: {}", self.pod.id())?;
+        writeln!(f, "  valid?  {}", self.pod.verify())?;
+        writeln!(f, "  statements:")?;
+        for st in &self.pod.pub_statements() {
+            writeln!(f, "    - {}", st)?;
+        }
+        writeln!(f, "  kvs:")?;
+        for (k, v) in &self.pod.kvs() {
+            writeln!(f, "    - {}: {}", k, v)?;
+        }
+        Ok(())
+    }
+}
+
 impl MainPod {
     pub fn id(&self) -> PodId {
         self.pod.id()
@@ -487,6 +513,7 @@ pub mod build_utils {
 #[cfg(test)]
 pub mod tests {
     use super::*;
+    use crate::backends::mock_main::MockProver;
     use crate::backends::mock_signed::MockSigner;
     use crate::examples::{
         great_boy_pod_full_flow, tickets_pod_full_flow, zu_kyc_pod_builder,
@@ -498,7 +525,8 @@ pub mod tests {
         let params = Params::default();
         let (gov_id, pay_stub) = zu_kyc_sign_pod_builders(&params);
 
-        // TODO: print pods from the builder
+        println!("{}", gov_id);
+        println!("{}", pay_stub);
 
         let mut signer = MockSigner {
             pk: "ZooGov".into(),
@@ -515,7 +543,11 @@ pub mod tests {
         let kyc = zu_kyc_pod_builder(&params, &gov_id, &pay_stub)?;
         println!("{}", kyc);
 
+        let mut prover = MockProver {};
+        let kyc = kyc.prove(&mut prover)?;
+
         // TODO: prove kyc with MockProver and print it
+        println!("{}", kyc);
 
         Ok(())
     }

src/middleware/custom.rs

@@ -2,10 +2,14 @@ use std::sync::Arc;
 use std::{fmt, hash as h, iter::zip};
 
 use anyhow::{anyhow, Result};
+use plonky2::field::goldilocks_field::GoldilocksField;
+use plonky2::field::types::Field;
+use plonky2::hash::poseidon::PoseidonHash;
+use plonky2::plonk::config::Hasher;
 
 use super::{
-    hash_str, AnchoredKey, Hash, NativePredicate, PodId, Statement, StatementArg, ToFields, Value,
-    F,
+    hash_str, AnchoredKey, Hash, NativePredicate, Params, PodId, Statement, StatementArg, ToFields,
+    Value, F,
 };
 
 // BEGIN Custom 1b
@@ -38,6 +42,22 @@ impl fmt::Display for HashOrWildcard {
     }
 }
 
+impl ToFields for HashOrWildcard {
+    fn to_fields(&self, params: Params) -> (Vec<F>, usize) {
+        match self {
+            HashOrWildcard::Hash(h) => h.to_fields(params),
+            HashOrWildcard::Wildcard(w) => {
+                let usizes: Vec<usize> = vec![0, 0, 0, *w];
+                let fields: Vec<F> = usizes
+                    .iter()
+                    .map(|x| F::from_canonical_u64(*x as u64))
+                    .collect();
+                (fields, 4)
+            }
+        }
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, h::Hash)]
 pub enum StatementTmplArg {
     None,
@@ -64,6 +84,40 @@ impl StatementTmplArg {
     }
 }
 
+impl ToFields for StatementTmplArg {
+    fn to_fields(&self, params: Params) -> (Vec<F>, usize) {
+        // None => (0, ...)
+        // Literal(value) => (1, [value], 0, 0, 0, 0)
+        // Key(hash_or_wildcard1, hash_or_wildcard2)
+        //    => (2, [hash_or_wildcard1], [hash_or_wildcard2])
+        // In all three cases, we pad to 2 * hash_size + 1 = 9 field elements
+        let hash_size = 4;
+        let statement_tmpl_arg_size = 2 * hash_size + 1;
+        match self {
+            StatementTmplArg::None => {
+                let fields: Vec<F> = std::iter::repeat_with(|| F::from_canonical_u64(0))
+                    .take(statement_tmpl_arg_size)
+                    .collect();
+                (fields, statement_tmpl_arg_size)
+            }
+            StatementTmplArg::Literal(v) => {
+                let fields: Vec<F> = std::iter::once(F::from_canonical_u64(1))
+                    .chain(v.to_fields(params).0.into_iter())
+                    .chain(std::iter::repeat_with(|| F::from_canonical_u64(0)).take(hash_size))
+                    .collect();
+                (fields, statement_tmpl_arg_size)
+            }
+            StatementTmplArg::Key(hw1, hw2) => {
+                let fields: Vec<F> = std::iter::once(F::from_canonical_u64(2))
+                    .chain(hw1.to_fields(params).0.into_iter())
+                    .chain(hw2.to_fields(params).0.into_iter())
+                    .collect();
+                (fields, statement_tmpl_arg_size)
+            }
+        }
+    }
+}
+
 impl fmt::Display for StatementTmplArg {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
@@ -117,6 +171,26 @@ impl StatementTmpl {
     }
 }
 
+impl ToFields for StatementTmpl {
+    fn to_fields(&self, params: Params) -> (Vec<F>, usize) {
+        // serialize as:
+        // predicate (6 field elements)
+        // then the StatementTmplArgs
+        if self.1.len() > params.max_statement_args {
+            panic!("Statement template has too many arguments");
+        }
+        let mut fields: Vec<F> = self
+            .0
+            .to_fields(params)
+            .0
+            .into_iter()
+            .chain(self.1.iter().flat_map(|sta| sta.to_fields(params).0))
+            .collect();
+        fields.resize_with(params.statement_tmpl_size(), || F::from_canonical_u64(0));
+        (fields, params.statement_tmpl_size())
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct CustomPredicate {
     /// true for "and", false for "or"
@@ -128,10 +202,22 @@ pub struct CustomPredicate {
 }
 
 impl ToFields for CustomPredicate {
-    fn to_fields(self) -> (Vec<F>, usize) {
-        todo!()
-        // let f: Vec<F> = Vec::new();
-        // (self.conjunction.to_f(), 1)
+    fn to_fields(&self, params: Params) -> (Vec<F>, usize) {
+        // serialize as:
+        // conjunction (one field element)
+        // args_len (one field element)
+        // statements
+        //   (params.max_custom_predicate_arity * params.statement_tmpl_size())
+        //   field elements
+        if self.statements.len() > params.max_custom_predicate_arity {
+            panic!("Custom predicate depends on too many statements");
+        }
+        let mut fields: Vec<F> = std::iter::once(F::from_bool(self.conjunction))
+            .chain(std::iter::once(F::from_canonical_usize(self.args_len)))
+            .chain(self.statements.iter().flat_map(|st| st.to_fields(params).0))
+            .collect();
+        fields.resize_with(params.custom_predicate_size(), || F::from_canonical_u64(0));
+        (fields, params.custom_predicate_size())
     }
 }
 
@@ -166,10 +252,30 @@ pub struct CustomPredicateBatch {
     pub predicates: Vec<CustomPredicate>,
 }
 
+impl ToFields for CustomPredicateBatch {
+    fn to_fields(&self, params: Params) -> (Vec<F>, usize) {
+        // all the custom predicates in order
+        if self.predicates.len() > params.max_custom_batch_size {
+            panic!("Predicate batch exceeds maximum size");
+        }
+        let mut fields: Vec<F> = self
+            .predicates
+            .iter()
+            .flat_map(|p| p.to_fields(params).0)
+            .collect();
+        fields.resize_with(params.custom_predicate_batch_size_field_elts(), || {
+            F::from_canonical_u64(0)
+        });
+
+        (fields, params.custom_predicate_batch_size_field_elts())
+    }
+}
+
 impl CustomPredicateBatch {
-    pub fn hash(&self) -> Hash {
-        // TODO
-        hash_str(&format!("{:?}", self))
+    pub fn hash(&self, params: Params) -> Hash {
+        let input = self.to_fields(params).0;
+        let h = Hash(PoseidonHash::hash_no_pad(&input).elements);
+        h
     }
 }
 
@@ -190,12 +296,36 @@ impl From<NativePredicate> for Predicate {
 }
 
 impl ToFields for Predicate {
-    fn to_fields(self) -> (Vec<F>, usize) {
+    fn to_fields(&self, params: Params) -> (Vec<F>, usize) {
+        // serialize:
+        // NativePredicate(id) as (0, id, 0, 0, 0, 0) -- id: usize
+        // BatchSelf(i) as (1, i, 0, 0, 0, 0) -- i: usize
+        // CustomPredicateRef(pb, i) as
+        // (2, [hash of pb], i) -- pb hashes to 4 field elements
+        //                      -- i: usize
+
+        // in every case: pad to (hash_size + 2) field elements
+        let mut fields: Vec<F> = Vec::new();
         match self {
-            Self::Native(p) => p.to_fields(),
-            Self::BatchSelf(i) => Value::from(i as i64).to_fields(),
-            Self::Custom(_) => todo!(), // TODO
+            Self::Native(p) => {
+                fields = std::iter::once(F::from_canonical_u64(1))
+                    .chain(p.to_fields(params).0.into_iter())
+                    .collect();
+            }
+            Self::BatchSelf(i) => {
+                fields = std::iter::once(F::from_canonical_u64(2))
+                    .chain(std::iter::once(F::from_canonical_usize(*i)))
+                    .collect();
+            }
+            Self::Custom(CustomPredicateRef(pb, i)) => {
+                fields = std::iter::once(F::from_canonical_u64(3))
+                    .chain(pb.hash(params).0)
+                    .chain(std::iter::once(F::from_canonical_usize(*i)))
+                    .collect();
+            }
         }
+        fields.resize_with(params.predicate_size(), || F::from_canonical_u64(0));
+        (fields, params.predicate_size())
     }
 }