sinsemilla::merkle.rs: Add test for MerkleChip.

Author: therealyingtong

src/circuit/gadget/sinsemilla/merkle.rs

@@ -545,3 +545,226 @@ impl<C: CurveAffine, const PATH_LENGTH: usize, const K: usize, const MAX_WORDS:
         SinsemillaChip::<C, K, MAX_WORDS>::extract(point)
     }
 }
+
+#[cfg(test)]
+pub mod tests {
+    use super::{MerkleChip, MerkleConfig, MerkleInstructions};
+
+    use crate::circuit::gadget::{
+        sinsemilla::chip::SinsemillaChip,
+        utilities::{UtilitiesInstructions, Var},
+    };
+    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::{convert::TryInto, marker::PhantomData};
+
+    struct MyCircuit<
+        C: CurveAffine,
+        const PATH_LENGTH: usize,
+        const K: usize,
+        const MAX_WORDS: usize,
+    > {
+        leaf: (Option<C::Base>, Option<u32>),
+        merkle_path: Vec<Option<C::Base>>,
+        _marker: PhantomData<C>,
+    }
+
+    impl<const PATH_LENGTH: usize, const K: usize, const MAX_WORDS: usize> Circuit<pallas::Base>
+        for MyCircuit<pallas::Affine, PATH_LENGTH, K, MAX_WORDS>
+    {
+        type Config = (
+            MerkleConfig<pallas::Affine, PATH_LENGTH, K, MAX_WORDS>,
+            MerkleConfig<pallas::Affine, PATH_LENGTH, K, MAX_WORDS>,
+        );
+
+        fn configure(meta: &mut ConstraintSystem<pallas::Base>) -> Self::Config {
+            let advices = [
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                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(),
+            );
+
+            let sinsemilla_config_1 = SinsemillaChip::<pallas::Affine, K, MAX_WORDS>::configure(
+                meta,
+                advices[..5].try_into().unwrap(),
+                lookup,
+                perm.clone(),
+            );
+            let config1 = MerkleChip::<pallas::Affine, PATH_LENGTH, K, MAX_WORDS>::configure(
+                meta,
+                sinsemilla_config_1,
+            );
+
+            let sinsemilla_config_2 = SinsemillaChip::<pallas::Affine, K, MAX_WORDS>::configure(
+                meta,
+                advices[..5].try_into().unwrap(),
+                lookup,
+                perm,
+            );
+            let config2 = MerkleChip::<pallas::Affine, PATH_LENGTH, K, MAX_WORDS>::configure(
+                meta,
+                sinsemilla_config_2,
+            );
+
+            (config1, config2)
+        }
+
+        fn synthesize(
+            &self,
+            cs: &mut impl Assignment<pallas::Base>,
+            config: Self::Config,
+        ) -> Result<(), Error> {
+            let mut layouter = SingleChipLayouter::new(cs)?;
+
+            // Load generator table (shared across both configs)
+            SinsemillaChip::<pallas::Affine, K, MAX_WORDS>::load(
+                config.0.sinsemilla_config.clone(),
+                &mut layouter,
+            )?;
+
+            // Construct Merkle chips which will be placed side-by-side in the circuit.
+            let merkle_chip_1 = MerkleChip::<pallas::Affine, PATH_LENGTH, K, MAX_WORDS>::construct(
+                config.0.clone(),
+            );
+            let merkle_chip_2 = MerkleChip::<pallas::Affine, PATH_LENGTH, K, MAX_WORDS>::construct(
+                config.1.clone(),
+            );
+
+            // Process lo half of the Merkle path from leaf to intermediate root.
+            let leaf = merkle_chip_1.load_private(
+                layouter.namespace(|| ""),
+                config.0.cond_swap_config.x,
+                self.leaf.0,
+            )?;
+            let pos_lo = self.leaf.1.map(|pos| pos & ((1 << PATH_LENGTH) - 1));
+
+            let intermediate_root = merkle_chip_1.merkle_path_check(
+                layouter.namespace(|| ""),
+                0,
+                (leaf, pos_lo),
+                self.merkle_path[0..PATH_LENGTH].to_vec(),
+            )?;
+
+            // Process hi half of the Merkle path from intermediate root to root.
+            let pos_hi = self.leaf.1.map(|pos| pos >> (PATH_LENGTH));
+
+            let computed_final_root = merkle_chip_2.merkle_path_check(
+                layouter.namespace(|| ""),
+                PATH_LENGTH,
+                (intermediate_root, pos_hi),
+                self.merkle_path[PATH_LENGTH..].to_vec(),
+            )?;
+
+            // The expected final root
+            let pos_bool = i2lebsp::<32>(self.leaf.1.unwrap());
+            let path: Option<Vec<pallas::Base>> = self.merkle_path.to_vec().into_iter().collect();
+            let final_root = hash_path(0, self.leaf.0.unwrap(), &pos_bool, &path.unwrap());
+
+            // Check the computed final root against the expected final root.
+            assert_eq!(computed_final_root.value().unwrap(), final_root);
+
+            Ok(())
+        }
+    }
+
+    fn hash_path(
+        offset: usize,
+        leaf: pallas::Base,
+        pos_bool: &[bool],
+        path: &[pallas::Base],
+    ) -> pallas::Base {
+        let domain = HashDomain::new(MERKLE_CRH_PERSONALIZATION);
+
+        // Compute the root
+        let mut node = leaf;
+        for (l_star, (sibling, pos)) in path.iter().zip(pos_bool.iter()).enumerate() {
+            let l_star = l_star + offset;
+
+            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 = domain.hash(message.into_iter()).unwrap();
+        }
+        node
+    }
+
+    #[test]
+    fn merkle_chip() {
+        // 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();
+
+        // This root is provided as a public input in the Orchard circuit.
+        let _root = hash_path(0, leaf, &pos_bool, &path);
+
+        let circuit = MyCircuit::<pallas::Affine, { MERKLE_DEPTH_ORCHARD / 2 }, K, C> {
+            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(()))
+    }
+}