From 9882fe833809294107487b8323dc298cdf0ada1f Mon Sep 17 00:00:00 2001
From: therealyingtong <yingtong@z.cash>
Date: Fri, 4 Jun 2021 02:19:17 +0800
Subject: [PATCH] Add Merkle chip test

---
 src/circuit/gadget/orchard_action/merkle.rs | 150 ++++++++++++++++++++
 1 file changed, 150 insertions(+)

diff --git a/src/circuit/gadget/orchard_action/merkle.rs b/src/circuit/gadget/orchard_action/merkle.rs
index 84fa18beb..6f5026234 100644
--- a/src/circuit/gadget/orchard_action/merkle.rs
+++ b/src/circuit/gadget/orchard_action/merkle.rs
@@ -675,3 +675,153 @@ impl<C: CurveAffine, const PATH_LENGTH: usize, const K: usize, const MAX_WORDS:
         chip.swap(layouter, pair, swap)
     }
 }
+
+#[cfg(test)]
+pub mod tests {
+    use super::{MerkleChip, MerkleConfig, MerkleInstructions};
+
+    use crate::circuit::gadget::sinsemilla::chip::{SinsemillaChip, SinsemillaHashDomains};
+    use crate::constants::{
+        util::i2lebsp, L_ORCHARD_BASE, MERKLE_CRH_PERSONALIZATION, MERKLE_DEPTH_ORCHARD,
+    };
+    use crate::primitives::sinsemilla::{HashDomain, C, K};
+
+    use ff::PrimeField;
+    use halo2::{
+        arithmetic::{CurveAffine, FieldExt},
+        circuit::{layouter::SingleChipLayouter, Layouter},
+        dev::MockProver,
+        pasta::pallas,
+        plonk::{Assignment, Circuit, ConstraintSystem, Error},
+    };
+
+    use rand::random;
+    use std::marker::PhantomData;
+
+    struct MyCircuit<
+        C: CurveAffine,
+        const PATH_LENGTH: usize,
+        const K: usize,
+        const MAX_WORDS: usize,
+    > {
+        root: Option<C::Base>,
+        leaf: (Option<C::Base>, Option<u32>),
+        merkle_path: Vec<Option<C::Base>>,
+        _marker: PhantomData<C>,
+    }
+
+    impl<C: CurveAffine, const PATH_LENGTH: usize, const K: usize, const MAX_WORDS: usize>
+        Circuit<C::Base> for MyCircuit<C, PATH_LENGTH, K, MAX_WORDS>
+    {
+        type Config = MerkleConfig<PATH_LENGTH>;
+
+        fn configure(meta: &mut ConstraintSystem<C::Base>) -> Self::Config {
+            let advices = [
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+            ];
+
+            let perm = meta.permutation(
+                &advices
+                    .iter()
+                    .map(|advice| (*advice).into())
+                    .collect::<Vec<_>>(),
+            );
+
+            // Fixed columns for the Sinsemilla generator lookup table
+            let lookup = (
+                meta.fixed_column(),
+                meta.fixed_column(),
+                meta.fixed_column(),
+            );
+
+            MerkleChip::<C, PATH_LENGTH, K, MAX_WORDS>::configure(meta, advices, lookup, perm)
+        }
+
+        fn synthesize(
+            &self,
+            cs: &mut impl Assignment<C::Base>,
+            config: Self::Config,
+        ) -> Result<(), Error> {
+            let mut layouter = SingleChipLayouter::new(cs)?;
+
+            // Load generator table
+            SinsemillaChip::<C, K, MAX_WORDS>::load(
+                config.sinsemilla_config.clone(),
+                &mut layouter,
+            )?;
+
+            let merkle_chip = MerkleChip::<C, PATH_LENGTH, K, MAX_WORDS>::construct(config);
+
+            let domain = SinsemillaHashDomains::MerkleCrh;
+
+            merkle_chip.merkle_path_check(
+                layouter.namespace(|| ""),
+                &domain,
+                self.root,
+                self.leaf,
+                self.merkle_path.clone(),
+            )
+        }
+    }
+
+    #[test]
+    fn merkle_chip() {
+        // Initialize MerkleCRH HashDomain
+        let merkle_crh = HashDomain::new(MERKLE_CRH_PERSONALIZATION);
+
+        // Choose a random leaf and position
+        let leaf = pallas::Base::rand();
+        let pos = random::<u32>();
+        let pos_bool = i2lebsp::<32>(pos);
+
+        // Choose a path of random inner nodes
+        let path: Vec<_> = (0..(MERKLE_DEPTH_ORCHARD))
+            .map(|_| pallas::Base::rand())
+            .collect();
+
+        // Compute the root
+        let mut node = leaf;
+        for (l_star, (sibling, pos)) in path.iter().zip(pos_bool.iter()).enumerate() {
+            let (left, right) = if *pos {
+                (*sibling, node)
+            } else {
+                (node, *sibling)
+            };
+
+            let l_star = i2lebsp::<10>(l_star as u32);
+            let left: Vec<_> = left
+                .to_le_bits()
+                .iter()
+                .by_val()
+                .take(L_ORCHARD_BASE)
+                .collect();
+            let right: Vec<_> = right
+                .to_le_bits()
+                .iter()
+                .by_val()
+                .take(L_ORCHARD_BASE)
+                .collect();
+
+            let mut message = l_star.to_vec();
+            message.extend_from_slice(&left);
+            message.extend_from_slice(&right);
+
+            node = merkle_crh.hash(message.into_iter()).unwrap();
+        }
+        let root = node;
+
+        let circuit = MyCircuit::<pallas::Affine, MERKLE_DEPTH_ORCHARD, K, C> {
+            root: Some(root),
+            leaf: (Some(leaf), Some(pos)),
+            merkle_path: path.into_iter().map(Some).collect(),
+            _marker: PhantomData,
+        };
+
+        let prover = MockProver::run(11, &circuit, vec![]).unwrap();
+        assert_eq!(prover.verify(), Ok(()))
+    }
+}