feat: integrate mt in mock_sign (#23)

* feat: integrate mt in mock_sign

* fix: handle compile error

* fix: use PodClass::Main in MainPod origin

src/backend.rs

@@ -1,4 +1,3 @@
-// TODO: Move the SignedPod.id calculation to mock_signed
 // TODO: Move the MainPod logic to mock_main and implement the MainPod trait
 /*
 use anyhow::Result;
@@ -11,29 +10,6 @@ use std::iter;
 use crate::merkletree::MerkleTree;
 use crate::middleware::{Hash, Params, PodId, Value, NULL};
 
-#[derive(Clone, Debug)]
-pub struct SignedPod {
-    pub params: Params,
-    pub id: PodId,
-    pub kvs: MerkleTree,
-}
-
-impl SignedPod {
-    pub fn new(params: &Params, kvs: HashMap<Hash, Value>) -> Result<Self> {
-        let mt = MerkleTree::new(kvs);
-        let root = mt.root()?;
-        Ok(Self {
-            params: *params,
-            id: PodId(root),
-            kvs: mt,
-        })
-    }
-
-    pub fn is_null(&self) -> bool {
-        self.id.0 == NULL
-    }
-}
-
 #[derive(Clone, Debug)]
 pub struct MainPod {
     pub params: Params,

src/backends/mock_signed.rs

@@ -1,7 +1,8 @@
+use crate::merkletree::MerkleTree;
 use crate::middleware::{
     hash_str, Hash, Params, PodId, PodSigner, PodType, SignedPod, Value, KEY_SIGNER, KEY_TYPE,
 };
-use itertools::Itertools;
+use anyhow::Result;
 use std::any::Any;
 use std::collections::HashMap;
 
@@ -9,28 +10,17 @@ pub struct MockSigner {
     pub pk: String,
 }
 
-fn calculate_pod_id(kvs: &HashMap<Hash, Value>) -> PodId {
-    let mut s = String::new();
-    for (k, v) in kvs.iter().sorted_by_key(|kv| kv.0) {
-        s += &format!("{}:{},", k, v);
-    }
-    PodId(hash_str(&s))
-}
-
 impl PodSigner for MockSigner {
-    fn sign(&mut self, _params: &Params, kvs: &HashMap<Hash, Value>) -> Box<dyn SignedPod> {
+    fn sign(&mut self, _params: &Params, kvs: &HashMap<Hash, Value>) -> Result<Box<dyn SignedPod>> {
         let mut kvs = kvs.clone();
         let pk_hash = hash_str(&self.pk);
         kvs.insert(hash_str(&KEY_SIGNER), Value(pk_hash.0));
         kvs.insert(hash_str(&KEY_TYPE), Value::from(PodType::MockSigned));
 
-        let id = calculate_pod_id(&kvs);
+        let mt = MerkleTree::new(&kvs);
+        let id = PodId(mt.root()?);
         let signature = format!("{}_signed_by_{}", id, pk_hash);
-        Box::new(MockSignedPod {
+        Ok(Box::new(MockSignedPod { mt, id, signature }))
-            kvs: kvs.clone(),
-            id,
-            signature,
-        })
     }
 }
 
@@ -38,24 +28,28 @@ impl PodSigner for MockSigner {
 pub struct MockSignedPod {
     pub id: PodId,
     pub signature: String,
-    pub kvs: HashMap<Hash, Value>,
+    pub mt: MerkleTree,
 }
 
 impl SignedPod for MockSignedPod {
     fn verify(&self) -> bool {
         // Verify type
-        if Some(&Value::from(PodType::MockSigned)) != self.kvs.get(&hash_str(&KEY_TYPE)) {
+        if Some(&Value::from(PodType::MockSigned)) != self.mt.kvs().get(&hash_str(&KEY_TYPE)) {
             return false;
         }
 
         // Verify id
-        let id = calculate_pod_id(&self.kvs);
+        let mt = MerkleTree::new(&self.mt.kvs());
+        let id = match mt.root() {
+            Ok(id) => PodId(id),
+            Err(_) => return false,
+        };
         if id != self.id {
             return false;
         }
 
         // Verify signature
-        let pk_hash = match self.kvs.get(&hash_str(&KEY_SIGNER)) {
+        let pk_hash = match self.mt.kvs().get(&hash_str(&KEY_SIGNER)) {
             Some(v) => v,
             None => return false,
         };
@@ -72,7 +66,7 @@ impl SignedPod for MockSignedPod {
     }
 
     fn kvs(&self) -> HashMap<Hash, Value> {
-        self.kvs.clone()
+        self.mt.kvs().clone()
     }
 
     fn into_any(self: Box<Self>) -> Box<dyn Any> {
@@ -96,7 +90,7 @@ pub mod tests {
         pod.insert("socialSecurityNumber", "G2121210");
 
         let mut signer = MockSigner { pk: "Molly".into() };
-        let pod = pod.sign(&mut signer);
+        let pod = pod.sign(&mut signer).unwrap();
         let pod = pod.pod.into_any().downcast::<MockSignedPod>().unwrap();
 
         assert_eq!(pod.verify(), true);
@@ -112,13 +106,17 @@ pub mod tests {
         assert_eq!(bad_pod.verify(), false);
 
         let mut bad_pod = pod.clone();
-        bad_pod
+        let mut bad_kvs = bad_pod.kvs();
-            .kvs
+        bad_kvs.insert(hash_str(KEY_SIGNER), Value(PodId(NULL).0 .0));
-            .insert(hash_str(KEY_SIGNER), Value(PodId(NULL).0 .0));
+        let bad_mt = MerkleTree::new(&bad_kvs);
+        bad_pod.mt = bad_mt;
         assert_eq!(bad_pod.verify(), false);
 
         let mut bad_pod = pod.clone();
-        bad_pod.kvs.insert(hash_str(KEY_TYPE), Value::from(0));
+        let mut bad_kvs = bad_pod.kvs();
+        bad_kvs.insert(hash_str(KEY_TYPE), Value::from(0));
+        let bad_mt = MerkleTree::new(&bad_kvs);
+        bad_pod.mt = bad_mt;
         assert_eq!(bad_pod.verify(), false);
     }
 }

src/frontend.rs

@@ -94,7 +94,7 @@ impl SignedPodBuilder {
         self.kvs.insert(key.into(), value.into());
     }
 
-    pub fn sign<S: PodSigner>(&self, signer: &mut S) -> SignedPod {
+    pub fn sign<S: PodSigner>(&self, signer: &mut S) -> Result<SignedPod> {
         let mut kvs = HashMap::new();
         let mut key_string_map = HashMap::new();
         for (k, v) in self.kvs.iter() {
@@ -102,11 +102,11 @@ impl SignedPodBuilder {
             kvs.insert(k_hash, middleware::Value::from(v));
             key_string_map.insert(k_hash, k.clone());
         }
-        let pod = signer.sign(&self.params, &kvs);
+        let pod = signer.sign(&self.params, &kvs)?;
-        SignedPod {
+        Ok(SignedPod {
             pod,
             key_string_map,
-        }
+        })
     }
 }
 
@@ -254,7 +254,7 @@ impl MainPodBuilder {
         let mut st_args = Vec::new();
         for arg in args.iter_mut() {
             match arg {
-                OperationArg::Statement(s) => panic!("can't convert Statement to StatementArg"),
+                OperationArg::Statement(_s) => panic!("can't convert Statement to StatementArg"),
                 OperationArg::Key(k) => st_args.push(StatementArg::Key(k.clone())),
                 OperationArg::Literal(v) => {
                     let k = format!("c{}", self.const_cnt);
@@ -307,7 +307,7 @@ impl MainPodBuilder {
         &self.statements[self.statements.len() - 1]
     }
 
-    pub fn prove<P: PodProver>(&self, prover: &mut P) -> MainPod {
+    pub fn prove<P: PodProver>(&self, prover: &mut P) -> Result<MainPod> {
         let compiler = MainPodCompiler::new(&self.params);
         let inputs = MainPodCompilerInputs {
             // signed_pods: &self.input_signed_pods,
@@ -315,9 +315,9 @@ impl MainPodBuilder {
             statements: &self.statements,
             operations: &self.operations,
         };
-        let (statements, operations) = compiler.compile(inputs).expect("TODO");
+        let (statements, operations) = compiler.compile(inputs)?;
 
-        let inputs = middleware::MainPodInputs {
+        let inputs = MainPodInputs {
             signed_pods: &self
                 .input_signed_pods
                 .iter()
@@ -331,8 +331,8 @@ impl MainPodBuilder {
             statements: &statements,
             operations: &operations,
         };
-        let pod = prover.prove(&self.params, inputs);
+        let pod = prover.prove(&self.params, inputs)?;
-        MainPod { pod }
+        Ok(MainPod { pod })
     }
 }
 
@@ -347,7 +347,7 @@ impl MainPod {
         self.pod.id()
     }
     pub fn origin(&self) -> Origin {
-        Origin(PodClass::Signed, self.id())
+        Origin(PodClass::Main, self.id())
     }
 }
 
@@ -374,10 +374,6 @@ impl MainPodCompiler {
         }
     }
 
-    fn max_priv_statements(&self) -> usize {
-        self.params.max_statements - self.params.max_public_statements
-    }
-
     fn push_st_op(&mut self, st: middleware::Statement, op: middleware::Operation) {
         self.statements.push(st);
         self.operations.push(op);
@@ -394,7 +390,8 @@ impl MainPodCompiler {
             }
             OperationArg::Entry(_k, _v) => {
                 // OperationArg::Entry is only used in the frontend.  The (key, value) will only
-                // appear in the ValueOf statement in the backend.
+                // appear in the ValueOf statement in the backend.  This is because a new ValueOf
+                // statement doesn't have any requirement on the key and value.
                 middleware::OperationArg::None
             }
         }
@@ -521,14 +518,8 @@ impl Printer {
 pub mod tests {
     use super::*;
     use crate::backends::mock_signed::MockSigner;
-    use crate::middleware::Hash;
-    use hex::FromHex;
     use std::io;
 
-    fn pod_id(hex: &str) -> PodId {
-        PodId(Hash::from_hex(hex).unwrap())
-    }
-
     macro_rules! args {
         ($($arg:expr),+) => {vec![$(OperationArg::from($arg)),*]}
     }
@@ -592,13 +583,13 @@ pub mod tests {
         let mut signer = MockSigner {
             pk: "ZooGov".into(),
         };
-        let gov_id = gov_id.sign(&mut signer);
+        let gov_id = gov_id.sign(&mut signer).unwrap();
         printer.fmt_signed_pod(&mut w, &gov_id).unwrap();
 
         let mut signer = MockSigner {
             pk: "ZooDeel".into(),
         };
-        let pay_stub = pay_stub.sign(&mut signer);
+        let pay_stub = pay_stub.sign(&mut signer).unwrap();
         printer.fmt_signed_pod(&mut w, &pay_stub).unwrap();
 
         let kyc = zu_kyc_pod_builder(&params, &gov_id, &pay_stub);

src/merkletree.rs

@@ -42,18 +42,18 @@ pub struct MerkleProof {
 }
 
 impl MerkleTree {
-    pub fn new(kvs: HashMap<Hash, Value>) -> Self {
+    pub fn new(kvs: &HashMap<Hash, Value>) -> Self {
         let mut keyindex: HashMap<Hash, usize> = HashMap::new();
         let mut leaves: Vec<Vec<F>> = Vec::new();
         // Note: current version iterates sorting by keys of the kvs, but the merkletree defined at
         // 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 (i, (k, v)) in kvs.clone().into_iter().sorted_by_key(|kv| kv.0).enumerate() {
+        for (i, (k, v)) in kvs.iter().sorted_by_key(|kv| kv.0).enumerate() {
             let input: Vec<F> = [k.0, v.0].concat();
             let leaf = PoseidonHash::hash_no_pad(&input).elements;
             leaves.push(leaf.into());
-            keyindex.insert(k, i);
+            keyindex.insert(*k, i);
         }
 
         // pad to a power of two if needed
@@ -66,7 +66,7 @@ impl MerkleTree {
         Self {
             tree,
             keyindex,
-            kvs,
+            kvs: kvs.clone(),
         }
     }
 
@@ -127,6 +127,10 @@ impl MerkleTree {
     pub fn iter(&self) -> std::collections::hash_map::Iter<Hash, Value> {
         self.kvs.iter()
     }
+
+    pub fn kvs(&self) -> &HashMap<Hash, Value> {
+        &self.kvs
+    }
 }
 
 impl<'a> IntoIterator for &'a MerkleTree {
@@ -164,7 +168,7 @@ pub mod tests {
         kvs.insert(k1, v1);
         kvs.insert(k2, v2);
 
-        let tree = MerkleTree::new(kvs);
+        let tree = MerkleTree::new(&kvs);
 
         let proof = tree.prove(&k2)?;
         MerkleTree::verify(tree.root()?, &proof, &k2, &v2)?;

src/middleware.rs

@@ -1,6 +1,7 @@
 //! The middleware includes the type definitions and the traits used to connect the frontend and
 //! the backend.
 
+use anyhow::Result;
 use dyn_clone::DynClone;
 use hex::{FromHex, FromHexError};
 use itertools::Itertools;
@@ -220,7 +221,7 @@ pub trait SignedPod: fmt::Debug + DynClone {
 dyn_clone::clone_trait_object!(SignedPod);
 
 pub trait PodSigner {
-    fn sign(&mut self, params: &Params, kvs: &HashMap<Hash, Value>) -> Box<dyn SignedPod>;
+    fn sign(&mut self, params: &Params, kvs: &HashMap<Hash, Value>) -> Result<Box<dyn SignedPod>>;
 }
 
 #[derive(Clone, Copy, Debug, FromRepr, PartialEq, Eq)]
@@ -313,5 +314,5 @@ pub struct MainPodInputs<'a> {
 }
 
 pub trait PodProver {
-    fn prove(&mut self, params: &Params, inputs: MainPodInputs) -> Box<dyn MainPod>;
+    fn prove(&mut self, params: &Params, inputs: MainPodInputs) -> Result<Box<dyn MainPod>>;
 }