Add a generic DigestAlgorithmGadget to MerkleTreeGadget instead

EspressoSystems · Dec 19, 2022 · a8c56d3 · a8c56d3
1 parent b6e7bf0
commit a8c56d3
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 34 deletions.
diff --git a/primitives/src/circuit/merkle_tree/mod.rs b/primitives/src/circuit/merkle_tree/mod.rs
@@ -7,8 +7,7 @@
 //! Trait definitions for a Merkle tree gadget.
 
 use crate::merkle_tree::MerkleTreeScheme;
-use ark_ff::PrimeField;
-use jf_relation::{errors::CircuitError, BoolVar, PlonkCircuit, Variable};
+use jf_relation::{errors::CircuitError, BoolVar, Variable};
 
 mod rescue_merkle_tree;
 
@@ -37,10 +36,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, F>
+pub trait MerkleTreeGadget<M, D>
 where
     M: MerkleTreeScheme,
-    F: PrimeField,
+    D: DigestAlgorithmGadget<M> + ?Sized,
 {
     /// Type to represent the leaf element of the concrete MT instantiation.
     type LeafVar;
@@ -50,9 +49,6 @@ 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,
@@ -89,17 +85,10 @@ where
 }
 
 /// Circuit counterpart to DigestAlgorithm
-pub trait DigestAlgorithmGadget<F>
-where
-    F: PrimeField,
-{
+pub trait DigestAlgorithmGadget<M: MerkleTreeScheme> {
     /// Digest a list of variables
-    fn digest(circuit: &mut PlonkCircuit<F>, data: &[Variable]) -> Result<Variable, CircuitError>;
+    fn digest(&mut self, data: &[Variable]) -> Result<Variable, CircuitError>;
 
     /// Digest an indexed element
-    fn digest_leaf(
-        circuit: &mut PlonkCircuit<F>,
-        pos: usize,
-        elem: Variable,
-    ) -> Result<Variable, CircuitError>;
+    fn digest_leaf(&mut self, pos: usize, elem: Variable) -> Result<Variable, CircuitError>;
 }
diff --git a/primitives/src/circuit/merkle_tree/rescue_merkle_tree.rs b/primitives/src/circuit/merkle_tree/rescue_merkle_tree.rs
@@ -41,30 +41,24 @@ pub struct StandardLeafVar {
     pub elem: Variable,
 }
 
-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])
+impl<F: RescueParameter> DigestAlgorithmGadget<RescueMerkleTree<F>> for PlonkCircuit<F> {
+    fn digest(&mut self, data: &[Variable]) -> Result<Variable, CircuitError> {
+        Ok(RescueNativeGadget::<F>::rescue_sponge_no_padding(self, &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])
+    fn digest_leaf(&mut self, pos: usize, elem: Variable) -> Result<Variable, CircuitError> {
+        let zero = self.zero();
+        Ok(RescueNativeGadget::<F>::rescue_sponge_no_padding(self, &[zero, pos, elem], 1)?[0])
     }
 }
 
-impl<F> MerkleTreeGadget<RescueMerkleTree<F>, F> for PlonkCircuit<F>
+impl<F> MerkleTreeGadget<RescueMerkleTree<F>, dyn DigestAlgorithmGadget<RescueMerkleTree<F>>>
+    for PlonkCircuit<F>
 where
     F: RescueParameter,
 {
     type LeafVar = StandardLeafVar;
     type MerklePathVar = Rescue3AryMerklePathVar;
-    type DigestGadget = RescueDigestGadget;
 
     fn create_leaf_variable(
         &mut self,
@@ -184,7 +178,9 @@ pub struct Rescue3AryNodeVar {
     pub is_right_child: BoolVar,
 }
 
-trait MerkleTreeHelperGadget<F: RescueParameter>: MerkleTreeGadget<RescueMerkleTree<F>, F> {
+trait MerkleTreeHelperGadget<F: RescueParameter>:
+    MerkleTreeGadget<RescueMerkleTree<F>, dyn DigestAlgorithmGadget<RescueMerkleTree<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.
@@ -283,7 +279,7 @@ impl<F: RescueParameter> MerkleTreeHelperGadget<F> for PlonkCircuit<F> {
         elem: StandardLeafVar,
         path_vars: &Rescue3AryMerklePathVar,
     ) -> Result<Variable, CircuitError> {
-        let mut cur_label = Self::DigestGadget::digest_leaf(self, elem.uid, elem.elem)?;
+        let mut cur_label = Self::digest_leaf(self, elem.uid, elem.elem)?;
         for cur_node in path_vars.nodes.iter() {
             let input_labels = self.permute(
                 cur_label,
@@ -294,7 +290,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 = Self::DigestGadget::digest(self, &input_labels)?;
+            cur_label = Self::digest(self, &input_labels)?;
         }
         Ok(cur_label)
     }