From 7669068451c079a789ebe261f29ba5f0df0bad22 Mon Sep 17 00:00:00 2001
From: therealyingtong <yingtong@z.cash>
Date: Fri, 4 Jun 2021 20:05:55 +0800
Subject: [PATCH] Take in root as a Cell instead of a value.

In the Orchard circuit, the root (i.e. anchor) is a public input.
We will not need direct access to its value, but will instead
constrain the output of the Merkle hash to be equal to the
corresponding cell in the public input column.
---
 src/circuit/gadget/merkle.rs | 71 ++++++++++++++----------------------
 1 file changed, 28 insertions(+), 43 deletions(-)

diff --git a/src/circuit/gadget/merkle.rs b/src/circuit/gadget/merkle.rs
index 5af5de3a2..a8031002c 100644
--- a/src/circuit/gadget/merkle.rs
+++ b/src/circuit/gadget/merkle.rs
@@ -1,6 +1,6 @@
 use halo2::{
-    circuit::{Chip, Layouter},
-    plonk::{Advice, Column, ConstraintSystem, Error, Expression, Fixed, Permutation, Selector},
+    circuit::{Cell, Chip, Layouter},
+    plonk::{Advice, Column, ConstraintSystem, Error, Expression, Fixed, Permutation},
     poly::Rotation,
 };
 use pasta_curves::arithmetic::{CurveAffine, FieldExt};
@@ -37,7 +37,7 @@ pub trait MerkleInstructions<
         layouter: impl Layouter<C::Base>,
         domain: &<Self as SinsemillaInstructions<C, K, MAX_WORDS>>::HashDomains,
         start_height: usize,
-        root: Option<C::Base>,
+        root: Cell,
         node: (<Self as UtilitiesInstructions<C::Base>>::Var, Option<u32>),
         merkle_path: Vec<Option<C::Base>>,
     ) -> Result<<Self as UtilitiesInstructions<C::Base>>::Var, Error>;
@@ -57,7 +57,6 @@ pub trait MerkleInstructions<
 
 #[derive(Clone, Debug)]
 pub struct MerkleConfig<const PATH_LENGTH: usize> {
-    q_merkle: Selector,
     a: Column<Advice>,
     b: Column<Advice>,
     c: Column<Advice>,
@@ -103,7 +102,6 @@ impl<C: CurveAffine, const PATH_LENGTH: usize, const K: usize, const MAX_WORDS:
         let left = advices[3];
         let right = advices[4];
 
-        let q_merkle = meta.selector();
         let l_star = meta.fixed_column();
 
         // Check that pieces have been decomposed correctly for Sinsemilla hash.
@@ -151,17 +149,7 @@ impl<C: CurveAffine, const PATH_LENGTH: usize, const K: usize, const MAX_WORDS:
                 .collect()
         });
 
-        // Check that root is recovered at the end of Merkle hash.
-        meta.create_gate("Root check", |meta| {
-            let q_merkle = meta.query_selector(q_merkle, Rotation::cur());
-            let root = meta.query_advice(a, Rotation::cur());
-            let computed_root = meta.query_advice(b, Rotation::cur());
-
-            vec![q_merkle * (root - computed_root)]
-        });
-
         MerkleConfig {
-            q_merkle,
             a,
             b,
             c,
@@ -235,7 +223,7 @@ impl<C: CurveAffine, const PATH_LENGTH: usize, const K: usize, const MAX_WORDS:
         mut layouter: impl Layouter<C::Base>,
         domain: &<Self as SinsemillaInstructions<C, K, MAX_WORDS>>::HashDomains,
         start_height: usize,
-        root: Option<C::Base>,
+        root: Cell,
         node: (<Self as UtilitiesInstructions<C::Base>>::Var, Option<u32>),
         merkle_path: Vec<Option<C::Base>>,
     ) -> Result<<Self as UtilitiesInstructions<C::Base>>::Var, Error> {
@@ -300,29 +288,10 @@ impl<C: CurveAffine, const PATH_LENGTH: usize, const K: usize, const MAX_WORDS:
             )?;
         }
 
-        // Check that the final node is equal to the claimed root.
+        // Constrain the final hash output to equal the claimed root.
         layouter.assign_region(
             || "Root check",
-            |mut region| {
-                region.assign_advice(
-                    || "Witness root",
-                    config.a,
-                    0,
-                    || root.ok_or(Error::SynthesisError),
-                )?;
-
-                // Copy final output of Sinsemilla hash
-                copy(
-                    &mut region,
-                    || "Copy final hash output",
-                    config.b,
-                    0,
-                    &node,
-                    &config.perm,
-                )?;
-
-                config.q_merkle.enable(&mut region, 0)
-            },
+            |mut region| region.constrain_equal(&config.perm, node.cell(), root),
         )?;
 
         Ok(node)
@@ -685,7 +654,7 @@ pub mod tests {
 
     use crate::circuit::gadget::{
         sinsemilla::chip::{SinsemillaChip, SinsemillaHashDomains},
-        utilities::UtilitiesInstructions,
+        utilities::{UtilitiesInstructions, Var},
     };
     use crate::constants::{
         util::i2lebsp, L_ORCHARD_BASE, MERKLE_CRH_PERSONALIZATION, MERKLE_DEPTH_ORCHARD,
@@ -710,7 +679,7 @@ pub mod tests {
         const K: usize,
         const MAX_WORDS: usize,
     > {
-        root: Option<C::Base>,
+        root: C::Base,
         leaf: (Option<C::Base>, Option<u32>),
         merkle_path: Vec<Option<C::Base>>,
         _marker: PhantomData<C>,
@@ -784,7 +753,7 @@ pub mod tests {
                 config.0.clone(),
             );
             let merkle_chip_2 =
-                MerkleChip::<pallas::Affine, PATH_LENGTH, K, MAX_WORDS>::construct(config.1);
+                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(
@@ -800,6 +769,7 @@ pub mod tests {
                 .into_iter()
                 .collect();
 
+            // Compute the intermediate root
             let intermediate_root = self
                 .leaf
                 .0
@@ -807,15 +777,30 @@ pub mod tests {
                 .zip(path_lo)
                 .map(|((leaf, pos), path)| hash_path(&merkle_crh, 0, leaf, &pos, &path));
 
+            // Witness the intermediate root
+            let intermediate_root = merkle_chip_1.load_private(
+                layouter.namespace(|| "intermediate root"),
+                config.0.a,
+                intermediate_root,
+            )?;
+
             let intermediate_root = merkle_chip_1.merkle_path_check(
                 layouter.namespace(|| ""),
                 &SinsemillaHashDomains::MerkleCrh,
                 0,
-                intermediate_root,
+                intermediate_root.cell(),
                 (leaf, pos_lo),
                 self.merkle_path[0..PATH_LENGTH].to_vec(),
             )?;
 
+            // Witness the final root. In the Orchard circuit, this is a public
+            // input that will be constrained to equal some cell in an advice column.
+            let final_root = merkle_chip_2.load_private(
+                layouter.namespace(|| "final root"),
+                config.1.a,
+                Some(self.root),
+            )?;
+
             // Process hi half of the Merkle path from intermediate root to root.
             let pos_hi = self.leaf.1.map(|pos| pos >> (PATH_LENGTH));
 
@@ -823,7 +808,7 @@ pub mod tests {
                 layouter.namespace(|| ""),
                 &SinsemillaHashDomains::MerkleCrh,
                 PATH_LENGTH,
-                self.root,
+                final_root.cell(),
                 (intermediate_root, pos_hi),
                 self.merkle_path[PATH_LENGTH..].to_vec(),
             )?;
@@ -891,7 +876,7 @@ pub mod tests {
         let root = hash_path(&merkle_crh, 0, leaf, &pos_bool, &path);
 
         let circuit = MyCircuit::<pallas::Affine, { MERKLE_DEPTH_ORCHARD / 2 }, K, C> {
-            root: Some(root),
+            root,
             leaf: (Some(leaf), Some(pos)),
             merkle_path: path.into_iter().map(Some).collect(),
             _marker: PhantomData,