Add test for merkle statements, and some debug prints to verify() (#140)

Author: tideofwords

src/backends/plonky2/mock_main/mod.rs

@@ -430,7 +430,7 @@ impl Pod for MockMainPod {
                 self.operations[i]
                     .deref(&self.statements[..input_statement_offset + i])
                     .unwrap()
-                    .check(&self.params, &s.clone().try_into().unwrap())
+                    .check_and_print(&self.params, &s.clone().try_into().unwrap())
             })
             .collect::<Result<Vec<_>>>()
             .unwrap();

src/backends/plonky2/primitives/merkletree.rs

@@ -13,8 +13,8 @@ use crate::middleware::{keypath, kv_hash};
 /// https://0xparc.github.io/pod2/merkletree.html
 #[derive(Clone, Debug)]
 pub struct MerkleTree {
-    max_depth: usize,
-    root: Node,
+    pub max_depth: usize,
+    pub root: Node,
 }
 
 impl MerkleTree {
@@ -248,7 +248,7 @@ impl MerkleProof {
 }
 
 #[derive(Clone, Debug)]
-enum Node {
+pub enum Node {
     None,
     Leaf(Leaf),
     Intermediate(Intermediate),
@@ -285,7 +285,7 @@ impl fmt::Display for Node {
 }
 
 impl Node {
-    fn is_empty(&self) -> bool {
+    pub fn is_empty(&self) -> bool {
         match self {
             Self::None => true,
             Self::Leaf(_l) => false,
@@ -299,14 +299,38 @@ impl Node {
             Self::Intermediate(n) => n.compute_hash(),
         }
     }
-    fn hash(&self) -> Hash {
+    pub fn hash(&self) -> Hash {
         match self {
             Self::None => NULL,
             Self::Leaf(l) => l.hash(),
             Self::Intermediate(n) => n.hash(),
         }
     }
 
+    pub fn left(&self) -> Option<&Box<Node>> {
+        match self {
+            Self::None => None,
+            Self::Leaf(_l) => None,
+            Self::Intermediate(Intermediate {
+                hash: _h,
+                left: l, 
+                right: _r
+            }) => Some(l),
+        }
+    }
+
+    pub fn right(&self) -> Option<&Box<Node>> {
+        match self {
+            Self::None => None,
+            Self::Leaf(_l) => None,
+            Self::Intermediate(Intermediate {
+                hash: _h,
+                left: _l, 
+                right: r
+            }) => Some(r),
+        }
+    }
+
     /// Goes down from the current node until it encounters a terminal node,
     /// viz. a leaf or empty node, or until it reaches the maximum depth. The
     /// `siblings` parameter is used to store the siblings while going down to

src/frontend/mod.rs

@@ -953,6 +953,7 @@ pub mod tests {
     use super::*;
     use crate::backends::plonky2::mock_main::MockProver;
     use crate::backends::plonky2::mock_signed::MockSigner;
+    use crate::backends::plonky2::primitives::merkletree::MerkleTree;
     use crate::examples::{
         eth_dos_pod_builder, eth_friend_signed_pod_builder, great_boy_pod_full_flow,
         tickets_pod_full_flow, zu_kyc_pod_builder, zu_kyc_sign_pod_builders,
@@ -1165,4 +1166,74 @@ pub mod tests {
         println!("{}", builder);
         println!("{}", false_pod);
     }
+
+    #[test]
+    fn test_merkle_proofs() -> Result<()> {
+        let mut kvs = HashMap::new();
+        for i in 0..8 {
+            if i == 1 {
+                continue;
+            }
+            kvs.insert(middleware::Value::from(i), middleware::Value::from(1000 + i));
+        }
+        let key = middleware::Value::from(13);
+        let value = middleware::Value::from(1013);
+        kvs.insert(key, value);
+
+        let tree = MerkleTree::new(32, &kvs)?;
+        // when printing the tree, it should print the same tree as in
+        // https://0xparc.github.io/pod2/merkletree.html#example-2
+        println!("{}", tree);
+
+        println!("{}", tree.root());
+
+        let root: Hash = tree.root();
+        let left: Hash = (*tree.root.left().unwrap().clone()).hash();
+        let right: Hash = (*tree.root.right().unwrap().clone()).hash();
+
+        let params = Params::default();
+        let mut signed_builder = SignedPodBuilder::new(&params);
+
+        let mut builder = MainPodBuilder::new(&params);
+        let introduce_root_op = Operation(
+            OperationType::Native(NativeOperation::NewEntry),
+            vec![
+                OperationArg::Entry("root".into(), Value::Raw(middleware::Value::from(root))),
+            ]
+        );
+        let introduce_left_op = Operation(
+            OperationType::Native(NativeOperation::NewEntry),
+            vec![
+                OperationArg::Entry("left".into(), Value::Raw(middleware::Value::from(left))),
+            ]
+        );
+        let introduce_right_op = Operation(
+            OperationType::Native(NativeOperation::NewEntry),
+            vec![
+                OperationArg::Entry("right".into(), Value::Raw(middleware::Value::from(right))),
+            ]
+        );
+        let st1 = builder.op(false, introduce_root_op).unwrap();
+        let st2 = builder.op(false, introduce_left_op).unwrap();
+        let st3 = builder.op(false, introduce_right_op).unwrap();
+        
+        // verify Branches statement
+        let branches_op = Operation(
+            OperationType::Native(NativeOperation::BranchesFromEntries),
+            vec![
+                OperationArg::Statement(st1),
+                OperationArg::Statement(st2),
+                OperationArg::Statement(st3),
+            ]
+        );
+
+        let _branches_st = builder.op(true, branches_op).unwrap();
+
+        let mut prover = MockProver {};
+        let pod = builder.prove(&mut prover, &params).unwrap();
+        print!("{}", pod);
+        assert_eq!(pod.pod.verify(), true);
+
+        Ok(())
+    }
 }

src/middleware/operation.rs

@@ -347,6 +347,15 @@ impl Operation {
             .map(|(pred, st_args)| Statement::from_args(pred, st_args));
         x.transpose()
     }
+    /// Checks the given operation against a statement, and prints information if the check does not pass
+    pub fn check_and_print(&self, params: &Params, output_statement: &Statement) -> Result<bool> {
+        let valid: bool = self.check(params, output_statement)?;
+        if !valid {
+            println!("Check failed on the following statement");
+            println!("{}", output_statement);
+        }
+        Ok(valid)
+    }
     /// Checks the given operation against a statement.
     pub fn check(&self, _params: &Params, output_statement: &Statement) -> Result<bool> {
         use Statement::*;