Digest-generic MerkleTreeGadget (#167)

* Add an associated type DigestGadget to MerkleTreeGadget

* make struct fields private

* replace usages of RescueGadget with generic DigestGadget

* constrain M::NodeValue to be of type PrimeField

* position should be of type `Variable`, not `usize`

primitives/src/circuit/merkle_tree/mod.rs

@@ -10,12 +10,15 @@ use crate::{
     merkle_tree::{MerkleTreeScheme, UniversalMerkleTreeScheme},
     rescue::RescueParameter,
 };
+use ark_ff::PrimeField;
 use jf_relation::{errors::CircuitError, BoolVar, Circuit, PlonkCircuit, Variable};
 
 mod rescue_merkle_tree;
 mod sparse_merkle_tree;
 use ark_std::vec::Vec;
 
+use super::rescue::RescueNativeGadget;
+
 /// Gadget for a Merkle tree
 ///
 /// # Examples
@@ -60,12 +63,16 @@ use ark_std::vec::Vec;
 pub trait MerkleTreeGadget<M>
 where
     M: MerkleTreeScheme,
+    M::NodeValue: PrimeField,
 {
     /// Type to represent the merkle proof of the concrete MT instantiation.
     /// It is MT-specific, e.g arity will affect the exact definition of the
     /// underlying Merkle path.
     type MembershipProofVar;
 
+    /// Gadget for the digest algorithm.
+    type DigestGadget: DigestAlgorithmGadget<M::NodeValue>;
+
     /// Allocate a variable for the membership proof.
     fn create_membership_proof_variable(
         &mut self,
@@ -149,6 +156,7 @@ where
 pub trait UniversalMerkleTreeGadget<M>: MerkleTreeGadget<M>
 where
     M: UniversalMerkleTreeScheme,
+    M::NodeValue: PrimeField,
 {
     /// Type to represent the merkle non-membership proof of the concrete MT
     /// instantiation. It is MT-specific, e.g arity will affect the exact
@@ -212,13 +220,13 @@ fn constrain_sibling_order<F: RescueParameter>(
 /// Circuit variable for a node in the Merkle path.
 pub struct Merkle3AryNodeVar {
     /// First sibling of the node.
-    pub sibling1: Variable,
+    sibling1: Variable,
     /// Second sibling of the node.
-    pub sibling2: Variable,
+    sibling2: Variable,
     /// Boolean variable indicating whether the node is a left child.
-    pub is_left_child: BoolVar,
+    is_left_child: BoolVar,
     /// Boolean variable indicating whether the node is a right child.
-    pub is_right_child: BoolVar,
+    is_right_child: BoolVar,
 }
 
 /// Circuit variable for a Merkle non-membership proof of a 3-ary Merkle tree.
@@ -240,3 +248,36 @@ pub struct Merkle3AryMembershipProofVar {
     node_vars: Vec<Merkle3AryNodeVar>,
     elem_var: Variable,
 }
+/// Circuit counterpart to DigestAlgorithm
+pub trait DigestAlgorithmGadget<F>
+where
+    F: PrimeField,
+{
+    /// Digest a list of variables
+    fn digest(circuit: &mut PlonkCircuit<F>, data: &[Variable]) -> Result<Variable, CircuitError>;
+
+    /// Digest an indexed element
+    fn digest_leaf(
+        circuit: &mut PlonkCircuit<F>,
+        pos: usize,
+        elem: Variable,
+    ) -> Result<Variable, CircuitError>;
+}
+
+/// Digest gadget using for the Rescue hash function.
+pub struct RescueDigestGadget {}
+
+impl<F: RescueParameter> DigestAlgorithmGadget<F> for RescueDigestGadget {
+    fn digest(circuit: &mut PlonkCircuit<F>, data: &[Variable]) -> Result<Variable, CircuitError> {
+        Ok(RescueNativeGadget::<F>::rescue_sponge_no_padding(circuit, data, 1)?[0])
+    }
+
+    fn digest_leaf(
+        circuit: &mut PlonkCircuit<F>,
+        pos: Variable,
+        elem: Variable,
+    ) -> Result<Variable, CircuitError> {
+        let zero = circuit.zero();
+        Ok(RescueNativeGadget::<F>::rescue_sponge_no_padding(circuit, &[zero, pos, elem], 1)?[0])
+    }
+}

primitives/src/circuit/merkle_tree/rescue_merkle_tree.rs

@@ -8,21 +8,21 @@
 //! with a Rescue hash function.
 
 use crate::{
-    circuit::rescue::RescueNativeGadget,
+    circuit::merkle_tree::DigestAlgorithmGadget,
     merkle_tree::{
         internal::MerkleNode, prelude::RescueMerkleTree, MerkleTreeScheme, ToTraversalPath,
     },
     rescue::RescueParameter,
 };
 use ark_std::{string::ToString, vec::Vec};
 use jf_relation::{errors::CircuitError, BoolVar, Circuit, PlonkCircuit, Variable};
-
 type NodeVal<F> = <RescueMerkleTree<F> as MerkleTreeScheme>::NodeValue;
 type MembershipProof<F> = <RescueMerkleTree<F> as MerkleTreeScheme>::MembershipProof;
 use typenum::U3;
 
 use super::{
     constrain_sibling_order, Merkle3AryMembershipProofVar, Merkle3AryNodeVar, MerkleTreeGadget,
+    RescueDigestGadget,
 };
 
 impl<F> MerkleTreeGadget<RescueMerkleTree<F>> for PlonkCircuit<F>
@@ -31,6 +31,8 @@ where
 {
     type MembershipProofVar = Merkle3AryMembershipProofVar;
 
+    type DigestGadget = RescueDigestGadget;
+
     fn create_membership_proof_variable(
         &mut self,
         merkle_proof: &MembershipProof<F>,
@@ -96,14 +98,10 @@ where
     ) -> Result<BoolVar, CircuitError> {
         let computed_root_var = {
             let proof_var = &proof_var;
-            let zero_var = self.zero();
 
             // elem label = H(0, uid, elem)
-            let mut cur_label = RescueNativeGadget::<F>::rescue_sponge_no_padding(
-                self,
-                &[zero_var, elem_idx_var, proof_var.elem_var],
-                1,
-            )?[0];
+            let mut cur_label =
+                Self::DigestGadget::digest_leaf(self, elem_idx_var, proof_var.elem_var)?;
             for cur_node in proof_var.node_vars.iter() {
                 let input_labels = constrain_sibling_order(
                     self,
@@ -115,8 +113,7 @@ where
                 )?;
                 // check that the left child's label is non-zero
                 self.non_zero_gate(input_labels[0])?;
-                cur_label =
-                    RescueNativeGadget::<F>::rescue_sponge_no_padding(self, &input_labels, 1)?[0];
+                cur_label = Self::DigestGadget::digest(self, &input_labels)?;
             }
             Ok(cur_label)
         }?;

primitives/src/circuit/merkle_tree/sparse_merkle_tree.rs

@@ -8,7 +8,7 @@
 //! with a Rescue hash function.
 
 use crate::{
-    circuit::rescue::RescueNativeGadget,
+    circuit::merkle_tree::DigestAlgorithmGadget,
     merkle_tree::{
         internal::MerkleNode, prelude::RescueSparseMerkleTree, MerkleTreeScheme, ToTraversalPath,
     },
@@ -18,15 +18,15 @@ use ark_std::{string::ToString, vec::Vec};
 use jf_relation::{errors::CircuitError, BoolVar, Circuit, PlonkCircuit, Variable};
 
 type SparseMerkleTree<F> = RescueSparseMerkleTree<BigUint, F>;
-
 type NodeVal<F> = <SparseMerkleTree<F> as MerkleTreeScheme>::NodeValue;
 type MembershipProof<F> = <SparseMerkleTree<F> as MerkleTreeScheme>::MembershipProof;
 use num_bigint::BigUint;
 use typenum::U3;
 
 use super::{
     constrain_sibling_order, Merkle3AryMembershipProofVar, Merkle3AryNodeVar,
-    Merkle3AryNonMembershipProofVar, MerkleTreeGadget, UniversalMerkleTreeGadget,
+    Merkle3AryNonMembershipProofVar, MerkleTreeGadget, RescueDigestGadget,
+    UniversalMerkleTreeGadget,
 };
 
 impl<F> UniversalMerkleTreeGadget<SparseMerkleTree<F>> for PlonkCircuit<F>
@@ -57,8 +57,7 @@ where
                 )?;
                 // check that the left child's label is non-zero
                 self.non_zero_gate(input_labels[0])?;
-                cur_label =
-                    RescueNativeGadget::<F>::rescue_sponge_no_padding(self, &input_labels, 1)?[0];
+                cur_label = Self::DigestGadget::digest(self, &input_labels)?;
             }
             Ok(cur_label)
         }?;
@@ -126,6 +125,8 @@ where
 {
     type MembershipProofVar = Merkle3AryMembershipProofVar;
 
+    type DigestGadget = RescueDigestGadget;
+
     fn create_membership_proof_variable(
         &mut self,
         merkle_proof: &MembershipProof<F>,
@@ -191,14 +192,10 @@ where
     ) -> Result<BoolVar, CircuitError> {
         let computed_root_var = {
             let proof_var = &proof_var;
-            let zero_var = self.zero();
 
             // elem label = H(0, uid, elem)
-            let mut cur_label = RescueNativeGadget::<F>::rescue_sponge_no_padding(
-                self,
-                &[zero_var, elem_idx_var, proof_var.elem_var],
-                1,
-            )?[0];
+            let mut cur_label =
+                Self::DigestGadget::digest_leaf(self, elem_idx_var, proof_var.elem_var)?;
             for cur_node in proof_var.node_vars.iter() {
                 let input_labels = constrain_sibling_order(
                     self,
@@ -210,8 +207,7 @@ where
                 )?;
                 // check that the left child's label is non-zero
                 self.non_zero_gate(input_labels[0])?;
-                cur_label =
-                    RescueNativeGadget::<F>::rescue_sponge_no_padding(self, &input_labels, 1)?[0];
+                cur_label = Self::DigestGadget::digest(self, &input_labels)?;
             }
             Ok(cur_label)
         }?;