feat: unify SignedPod and MainPod traits (#64)

* feat: unify SignedPod and MainPod traits

* fix: test

* feat: enable workflows after draft

.github/workflows/mdbook-check.yml

@@ -2,6 +2,8 @@ name: Check mdbook compilation
 
 on:
   pull_request:
+    branches: [ main ]
+    types: [ready_for_review, opened, synchronize, reopened]
   push:
     branches:
       - main

.github/workflows/tests.yml

@@ -3,6 +3,7 @@ name: Rust Tests
 on:
   pull_request:
     branches: [ main ]
+    types: [ready_for_review, opened, synchronize, reopened]
   push:
     branches: [ main ]
 

.github/workflows/typos.yml

@@ -1,6 +1,8 @@
 name: typos
 on:
   pull_request:
+    branches: [ main ]
+    types: [ready_for_review, opened, synchronize, reopened]
   push:
     branches:
       - main

src/backends/mock_main.rs

@@ -1,7 +1,6 @@
 use crate::middleware::{
-    self, hash_str, AnchoredKey, Hash, MainPod, MainPodInputs, NativeOperation, NativeStatement,
+    self, hash_str, AnchoredKey, Hash, MainPodInputs, NativeOperation, NativeStatement, NonePod,
-    NoneMainPod, NoneSignedPod, Params, PodId, PodProver, SignedPod, Statement, StatementArg,
+    Params, Pod, PodId, PodProver, Statement, StatementArg, ToFields, KEY_TYPE, SELF,
-    ToFields, KEY_TYPE, SELF,
 };
 use anyhow::Result;
 use itertools::Itertools;
@@ -14,7 +13,7 @@ use std::fmt;
 pub struct MockProver {}
 
 impl PodProver for MockProver {
-    fn prove(&mut self, params: &Params, inputs: MainPodInputs) -> Result<Box<dyn MainPod>> {
+    fn prove(&mut self, params: &Params, inputs: MainPodInputs) -> Result<Box<dyn Pod>> {
         Ok(Box::new(MockMainPod::new(params, inputs)?))
     }
 }
@@ -89,8 +88,8 @@ impl fmt::Display for Operation {
 pub struct MockMainPod {
     params: Params,
     id: PodId,
-    input_signed_pods: Vec<Box<dyn SignedPod>>,
+    input_signed_pods: Vec<Box<dyn Pod>>,
-    input_main_pods: Vec<Box<dyn MainPod>>,
+    input_main_pods: Vec<Box<dyn Pod>>,
     // New statements introduced by this pod
     input_statements: Vec<Statement>,
     public_statements: Vec<Statement>,
@@ -196,7 +195,7 @@ impl MockMainPod {
         let st_none = Self::statement_none(params);
 
         // Input signed pods region
-        let none_sig_pod: Box<dyn SignedPod> = Box::new(NoneSignedPod {});
+        let none_sig_pod: Box<dyn Pod> = Box::new(NonePod {});
         assert!(inputs.signed_pods.len() <= params.max_input_signed_pods);
         for i in 0..params.max_input_signed_pods {
             let pod = inputs
@@ -214,7 +213,7 @@ impl MockMainPod {
         }
 
         // Input main pods region
-        let none_main_pod: Box<dyn MainPod> = Box::new(NoneMainPod {});
+        let none_main_pod: Box<dyn Pod> = Box::new(NonePod {});
         assert!(inputs.main_pods.len() <= params.max_input_main_pods);
         for i in 0..params.max_input_main_pods {
             let pod = inputs
@@ -404,7 +403,7 @@ pub fn hash_statements(statements: &[middleware::Statement]) -> Result<middlewar
     Ok(Hash(PoseidonHash::hash_no_pad(&field_elems).elements))
 }
 
-impl MainPod for MockMainPod {
+impl Pod for MockMainPod {
     fn verify(&self) -> bool {
         let input_statement_offset = self.offset_input_statements();
         // get the input_statements from the self.statements

src/backends/mock_signed.rs

@@ -1,6 +1,6 @@
 use crate::middleware::{
-    containers::Dictionary, hash_str, Hash, Params, PodId, PodSigner, PodType, SignedPod, Value,
+    containers::Dictionary, hash_str, AnchoredKey, Hash, NativeStatement, Params, Pod, PodId,
-    KEY_SIGNER, KEY_TYPE,
+    PodSigner, PodType, Statement, StatementArg, Value, KEY_SIGNER, KEY_TYPE,
 };
 use crate::primitives::merkletree::MerkleTree;
 use anyhow::Result;
@@ -12,7 +12,7 @@ pub struct MockSigner {
 }
 
 impl PodSigner for MockSigner {
-    fn sign(&mut self, _params: &Params, kvs: &HashMap<Hash, Value>) -> Result<Box<dyn SignedPod>> {
+    fn sign(&mut self, _params: &Params, kvs: &HashMap<Hash, Value>) -> Result<Box<dyn Pod>> {
         let mut kvs = kvs.clone();
         let pk_hash = hash_str(&self.pk);
         kvs.insert(hash_str(&KEY_SIGNER), Value(pk_hash.0));
@@ -36,7 +36,7 @@ pub struct MockSignedPod {
     dict: Dictionary,
 }
 
-impl SignedPod for MockSignedPod {
+impl Pod for MockSignedPod {
     fn verify(&self) -> bool {
         // Verify type
         let value_at_type = match self.dict.get(&hash_str(&KEY_TYPE).into()) {
@@ -77,10 +77,19 @@ impl SignedPod for MockSignedPod {
         self.id
     }
 
-    fn kvs(&self) -> HashMap<Hash, Value> {
+    fn pub_statements(&self) -> Vec<Statement> {
+        let id = self.id();
         self.dict
-            .into_iter()
+            .iter()
-            .map(|(&k, &v)| (Hash(k.0), v))
+            .map(|(k, v)| {
+                Statement(
+                    NativeStatement::ValueOf,
+                    vec![
+                        StatementArg::Key(AnchoredKey(id, Hash(k.0))),
+                        StatementArg::Literal(*v),
+                    ],
+                )
+            })
             .collect()
     }
 
@@ -95,6 +104,7 @@ pub mod tests {
     use crate::frontend;
     use crate::middleware::{self, F, NULL};
     use plonky2::field::types::Field;
+    use std::iter;
 
     #[test]
     fn test_mock_signed_0() {
@@ -121,22 +131,24 @@ pub mod tests {
         assert_eq!(bad_pod.verify(), false);
 
         let mut bad_pod = pod.clone();
-        let mut bad_kvs = bad_pod.kvs();
+        let bad_kv = (hash_str(KEY_SIGNER).into(), Value(PodId(NULL).0 .0));
-        bad_kvs.insert(hash_str(KEY_SIGNER), Value(PodId(NULL).0 .0));
+        let bad_kvs_mt = &bad_pod
-        let bad_kvs_mt = &bad_kvs
+            .kvs()
             .into_iter()
-            .map(|(k, v)| (Value(k.0), v))
+            .map(|(AnchoredKey(_, k), v)| (Value(k.0), v))
+            .chain(iter::once(bad_kv))
             .collect::<HashMap<Value, Value>>();
         let bad_mt = MerkleTree::new(&bad_kvs_mt);
         bad_pod.dict.mt = bad_mt;
         assert_eq!(bad_pod.verify(), false);
 
         let mut bad_pod = pod.clone();
-        let mut bad_kvs = bad_pod.kvs();
+        let bad_kv = (hash_str(KEY_TYPE).into(), Value::from(0));
-        bad_kvs.insert(hash_str(KEY_TYPE), Value::from(0));
+        let bad_kvs_mt = &bad_pod
-        let bad_kvs_mt = &bad_kvs
+            .kvs()
             .into_iter()
-            .map(|(k, v)| (Value(k.0), v))
+            .map(|(AnchoredKey(_, k), v)| (Value(k.0), v))
+            .chain(iter::once(bad_kv))
             .collect::<HashMap<Value, Value>>();
         let bad_mt = MerkleTree::new(&bad_kvs_mt);
         bad_pod.dict.mt = bad_mt;

src/frontend.rs

@@ -109,7 +109,7 @@ impl SignedPodBuilder {
 /// string<-->hash relation of the keys.
 #[derive(Debug, Clone)]
 pub struct SignedPod {
-    pub pod: Box<dyn middleware::SignedPod>,
+    pub pod: Box<dyn middleware::Pod>,
     /// HashMap to store the reverse relation between key strings and key hashes
     pub key_string_map: HashMap<Hash, String>,
 }
@@ -121,7 +121,7 @@ impl fmt::Display for SignedPod {
         // https://0xparc.github.io/pod2/merkletree.html will not need it since it will be
         // deterministic based on the keys values not on the order of the keys when added into the
         // tree.
-        for (k, v) in self.pod.kvs().iter().sorted_by_key(|kv| kv.0) {
+        for (k, v) in self.kvs().iter().sorted_by_key(|kv| kv.0) {
             writeln!(f, "  - {}: {}", k, v)?;
         }
         Ok(())
@@ -139,7 +139,11 @@ impl SignedPod {
         self.pod.verify()
     }
     pub fn kvs(&self) -> HashMap<Hash, middleware::Value> {
-        self.pod.kvs()
+        self.pod
+            .kvs()
+            .into_iter()
+            .map(|(middleware::AnchoredKey(_, k), v)| (k, v))
+            .collect()
     }
 }
 
@@ -411,7 +415,7 @@ impl MainPodBuilder {
 
 #[derive(Debug)]
 pub struct MainPod {
-    pub pod: Box<dyn middleware::MainPod>,
+    pub pod: Box<dyn middleware::Pod>,
     // TODO: metadata
 }
 

src/middleware/mod.rs

@@ -4,7 +4,6 @@
 use anyhow::{anyhow, Error, Result};
 use dyn_clone::DynClone;
 use hex::{FromHex, FromHexError};
-use itertools::Itertools;
 use plonky2::field::goldilocks_field::GoldilocksField;
 use plonky2::field::types::{Field, PrimeField64};
 use plonky2::hash::poseidon::PoseidonHash;
@@ -232,61 +231,6 @@ impl Default for Params {
     }
 }
 
-pub trait SignedPod: fmt::Debug + DynClone {
-    fn verify(&self) -> bool;
-    fn id(&self) -> PodId;
-    // NOTE: Maybe replace this by
-    // - `get(key: Hash) -> Option<Value>`
-    // - `iter() -> impl Iter<(Hash, Value)>`
-    fn kvs(&self) -> HashMap<Hash, Value>;
-    fn pub_statements(&self) -> Vec<Statement> {
-        let id = self.id();
-        let mut statements = Vec::new();
-        for (k, v) in self.kvs().iter().sorted_by_key(|kv| kv.0) {
-            statements.push(Statement(
-                NativeStatement::ValueOf,
-                vec![
-                    StatementArg::Key(AnchoredKey(id, *k)),
-                    StatementArg::Literal(*v),
-                ],
-            ));
-        }
-        statements
-    }
-    // Used for downcasting
-    fn into_any(self: Box<Self>) -> Box<dyn Any>;
-}
-
-// impl Clone for Box<dyn SignedPod>
-dyn_clone::clone_trait_object!(SignedPod);
-
-/// This is a filler type that fulfills the SignedPod trait and always verifies.  It's empty.  This
-/// can be used to simulate padding in a circuit.
-#[derive(Debug, Clone)]
-pub struct NoneSignedPod {}
-
-impl SignedPod for NoneSignedPod {
-    fn verify(&self) -> bool {
-        true
-    }
-    fn id(&self) -> PodId {
-        PodId(NULL)
-    }
-    fn kvs(&self) -> HashMap<Hash, Value> {
-        HashMap::new()
-    }
-    fn pub_statements(&self) -> Vec<Statement> {
-        Vec::new()
-    }
-    fn into_any(self: Box<Self>) -> Box<dyn Any> {
-        self
-    }
-}
-
-pub trait PodSigner {
-    fn sign(&mut self, params: &Params, kvs: &HashMap<Hash, Value>) -> Result<Box<dyn SignedPod>>;
-}
-
 #[derive(Clone, Copy, Debug, FromRepr, PartialEq, Eq)]
 pub enum NativeStatement {
     None = 0,
@@ -587,23 +531,40 @@ impl Operation {
     }
 }
 
-pub trait MainPod: fmt::Debug + DynClone {
+pub trait Pod: fmt::Debug + DynClone {
     fn verify(&self) -> bool;
     fn id(&self) -> PodId;
     fn pub_statements(&self) -> Vec<Statement>;
+    /// Extract key-values from ValueOf public statements
+    fn kvs(&self) -> HashMap<AnchoredKey, Value> {
+        self.pub_statements()
+            .into_iter()
+            .filter_map(|st| match st.0 {
+                NativeStatement::ValueOf => Some((
+                    st.1[0].key().expect("key"),
+                    st.1[1].literal().expect("literal"),
+                )),
+                _ => None,
+            })
+            .collect()
+    }
     // Used for downcasting
     fn into_any(self: Box<Self>) -> Box<dyn Any>;
 }
 
 // impl Clone for Box<dyn SignedPod>
-dyn_clone::clone_trait_object!(MainPod);
+dyn_clone::clone_trait_object!(Pod);
 
-/// This is a filler type that fulfills the MainPod trait and always verifies.  It's empty.  This
+pub trait PodSigner {
+    fn sign(&mut self, params: &Params, kvs: &HashMap<Hash, Value>) -> Result<Box<dyn Pod>>;
+}
+
+/// This is a filler type that fulfills the Pod trait and always verifies.  It's empty.  This
 /// can be used to simulate padding in a circuit.
 #[derive(Debug, Clone)]
-pub struct NoneMainPod {}
+pub struct NonePod {}
 
-impl MainPod for NoneMainPod {
+impl Pod for NonePod {
     fn verify(&self) -> bool {
         true
     }
@@ -620,8 +581,8 @@ impl MainPod for NoneMainPod {
 
 #[derive(Debug)]
 pub struct MainPodInputs<'a> {
-    pub signed_pods: &'a [&'a Box<dyn SignedPod>],
+    pub signed_pods: &'a [&'a Box<dyn Pod>],
-    pub main_pods: &'a [&'a Box<dyn MainPod>],
+    pub main_pods: &'a [&'a Box<dyn Pod>],
     pub statements: &'a [Statement],
     pub operations: &'a [Operation],
     /// Statements that need to be made public (they can come from input pods or input
@@ -630,7 +591,7 @@ pub struct MainPodInputs<'a> {
 }
 
 pub trait PodProver {
-    fn prove(&mut self, params: &Params, inputs: MainPodInputs) -> Result<Box<dyn MainPod>>;
+    fn prove(&mut self, params: &Params, inputs: MainPodInputs) -> Result<Box<dyn Pod>>;
 }
 
 pub trait ToFields {