Add an associated type DigestGadget to MerkleTreeGadget

primitives/src/circuit/merkle_tree/mod.rs

@@ -7,7 +7,8 @@
 //! Trait definitions for a Merkle tree gadget.
 
 use crate::merkle_tree::MerkleTreeScheme;
-use jf_relation::{errors::CircuitError, BoolVar, Variable};
+use ark_ff::PrimeField;
+use jf_relation::{errors::CircuitError, BoolVar, PlonkCircuit, Variable};
 
 mod rescue_merkle_tree;
 
@@ -36,9 +37,10 @@ mod rescue_merkle_tree;
 /// circuit.enforce_merkle_proof(leaf_var, path_vars, root_var).unwrap();
 /// assert!(circuit.check_circuit_satisfiability(&[]).is_ok());
 /// ```
-pub trait MerkleTreeGadget<M>
+pub trait MerkleTreeGadget<M, F>
 where
     M: MerkleTreeScheme,
+    F: PrimeField,
 {
     /// Type to represent the leaf element of the concrete MT instantiation.
     type LeafVar;
@@ -48,6 +50,9 @@ where
     /// Merkle path.
     type MerklePathVar;
 
+    /// Gadget for the digest algorithm.
+    type DigestGadget: DigestAlgorithmGadget<F>;
+
     /// Allocate a variable for the leaf element.
     fn create_leaf_variable(
         &mut self,
@@ -82,3 +87,19 @@ where
         expected_merkle_root: Variable,
     ) -> Result<(), CircuitError>;
 }
+
+/// 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>;
+}

primitives/src/circuit/merkle_tree/rescue_merkle_tree.rs

@@ -23,7 +23,7 @@ type Element<F> = <RescueMerkleTree<F> as MerkleTreeScheme>::Element;
 type Index<F> = <RescueMerkleTree<F> as MerkleTreeScheme>::Index;
 use typenum::U3;
 
-use super::MerkleTreeGadget;
+use super::{DigestAlgorithmGadget, MerkleTreeGadget};
 
 #[derive(Debug, Clone)]
 /// Circuit variable for a Merkle authentication path for a Rescue-based, 3-ary
@@ -41,12 +41,30 @@ pub struct StandardLeafVar {
     pub elem: Variable,
 }
 
-impl<F> MerkleTreeGadget<RescueMerkleTree<F>> for PlonkCircuit<F>
+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: usize,
+        elem: Variable,
+    ) -> Result<Variable, CircuitError> {
+        let zero = circuit.zero();
+        Ok(RescueNativeGadget::<F>::rescue_sponge_no_padding(circuit, &[zero, pos, elem], 1)?[0])
+    }
+}
+
+impl<F> MerkleTreeGadget<RescueMerkleTree<F>, F> for PlonkCircuit<F>
 where
     F: RescueParameter,
 {
     type LeafVar = StandardLeafVar;
     type MerklePathVar = Rescue3AryMerklePathVar;
+    type DigestGadget = RescueDigestGadget;
 
     fn create_leaf_variable(
         &mut self,
@@ -166,7 +184,7 @@ pub struct Rescue3AryNodeVar {
     pub is_right_child: BoolVar,
 }
 
-trait MerkleTreeHelperGadget<F: RescueParameter> {
+trait MerkleTreeHelperGadget<F: RescueParameter>: MerkleTreeGadget<RescueMerkleTree<F>, F> {
     /// Produces an ordered list of variables based on the relative position of
     /// a node and its siblings.
     /// * `node` - node to be inserted in the final list.
@@ -265,14 +283,7 @@ impl<F: RescueParameter> MerkleTreeHelperGadget<F> for PlonkCircuit<F> {
         elem: StandardLeafVar,
         path_vars: &Rescue3AryMerklePathVar,
     ) -> Result<Variable, CircuitError> {
-        let zero_var = self.zero();
-
-        // leaf label = H(0, uid, arc)
-        let mut cur_label = RescueNativeGadget::<F>::rescue_sponge_no_padding(
-            self,
-            &[zero_var, elem.uid, elem.elem],
-            1,
-        )?[0];
+        let mut cur_label = Self::DigestGadget::digest_leaf(self, elem.uid, elem.elem)?;
         for cur_node in path_vars.nodes.iter() {
             let input_labels = self.permute(
                 cur_label,
@@ -283,8 +294,7 @@ impl<F: RescueParameter> MerkleTreeHelperGadget<F> for PlonkCircuit<F> {
             )?;
             // 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)
     }