Unify sponge gadget

CHANGELOG.md

@@ -45,6 +45,9 @@ and follow [semantic versioning](https://semver.org/) for our releases.
 - [#148](https://github.com/EspressoSystems/jellyfish/pull/148), [#156](https://github.com/EspressoSystems/jellyfish/pull/156) (`jf-primitives`) Refactored BLS Signature implementation
   - #148 Added trait bounds on associated types of `trait SignatureScheme`
   - #156 Improved BLS correctness and API compliance with IRTF standard with better doc
+- [#150](https://github.com/EspressoSystems/jellyfish/pull/150) (`jf-primitives`) Refactor `RescueGadget`
+    - Introduce `SpongeStateVar` to abstract over `RescueStateVar` and `RescueNonNativeStateVar` structs.
+    - Unify `RescueGadget` and `RescueNonNativeGadget` traits into `RescueGadget`.
 
 ### Fixed
 

plonk/src/circuit/transcript.rs

@@ -12,7 +12,7 @@ use ark_ff::PrimeField;
 use ark_std::{string::ToString, vec::Vec};
 use core::marker::PhantomData;
 use jf_primitives::{
-    circuit::rescue::RescueGadget,
+    circuit::rescue::{RescueGadget, RescueStateVar},
     rescue::{RescueParameter, STATE_SIZE},
 };
 use jf_relation::{
@@ -208,9 +208,10 @@ where
 
         // step 1. state: [F: STATE_SIZE] = hash(state|transcript)
         let input_var = [self.state_var.as_ref(), self.transcript_var.as_ref()].concat();
-        let res_var = circuit
-            .rescue_sponge_with_padding(&input_var, STATE_SIZE)
-            .unwrap();
+        let res_var = RescueGadget::<RescueStateVar, F, F>::rescue_sponge_with_padding(
+            circuit, &input_var, STATE_SIZE,
+        )
+        .unwrap();
         let out_var = res_var[0];
 
         // step 2. challenge = state[0] in Fr

primitives/src/circuit/commitment.rs

@@ -14,6 +14,8 @@ use crate::{
 use ark_std::vec;
 use jf_relation::{errors::CircuitError, Circuit, PlonkCircuit, Variable};
 
+use super::rescue::RescueStateVar;
+
 /// Circuit implementation of the commitment scheme.
 pub trait CommitmentGadget {
     // Commitment scheme
@@ -34,7 +36,7 @@ where
         let mut msg = vec![blinding];
         msg.extend_from_slice(input);
         pad_with(&mut msg, CRHF_RATE, self.zero());
-        Ok(self.rescue_sponge_no_padding(&msg, 1)?[0])
+        Ok(RescueGadget::<RescueStateVar, F, F>::rescue_sponge_no_padding(self, &msg, 1)?[0])
     }
 }
 

primitives/src/circuit/merkle_tree/rescue_merkle_tree.rs

@@ -8,7 +8,7 @@
 //! with a Rescue hash function.
 
 use crate::{
-    circuit::rescue::RescueGadget,
+    circuit::rescue::{RescueGadget, RescueStateVar},
     merkle_tree::{
         internal::MerkleNode, prelude::RescueMerkleTree, MerkleTreeScheme, ToTraversalPath,
     },
@@ -268,7 +268,11 @@ impl<F: RescueParameter> MerkleTreeHelperGadget<F> for PlonkCircuit<F> {
         let zero_var = self.zero();
 
         // leaf label = H(0, uid, arc)
-        let mut cur_label = self.rescue_sponge_no_padding(&[zero_var, elem.uid, elem.elem], 1)?[0];
+        let mut cur_label = RescueGadget::<RescueStateVar, F, F>::rescue_sponge_no_padding(
+            self,
+            &[zero_var, elem.uid, elem.elem],
+            1,
+        )?[0];
         for cur_node in path_vars.nodes.iter() {
             let input_labels = self.permute(
                 cur_label,
@@ -279,7 +283,11 @@ 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.rescue_sponge_no_padding(&input_labels, 1)?[0];
+            cur_label = RescueGadget::<RescueStateVar, F, F>::rescue_sponge_no_padding(
+                self,
+                &input_labels,
+                1,
+            )?[0];
         }
         Ok(cur_label)
     }

primitives/src/circuit/prf.rs

@@ -13,6 +13,8 @@ use crate::{
 };
 use jf_relation::{errors::CircuitError, Circuit, PlonkCircuit, Variable};
 
+use super::rescue::RescueStateVar;
+
 /// Circuit implementation of a PRF.
 pub trait PrfGadget {
     /// PRF many to one
@@ -32,7 +34,9 @@ where
 
         let mut input_vec = input.to_vec();
         pad_with(&mut input_vec, STATE_SIZE, self.zero());
-        self.rescue_full_state_keyed_sponge_no_padding(key, &input_vec)
+        RescueGadget::<RescueStateVar, F, F>::rescue_full_state_keyed_sponge_no_padding(
+            self, key, &input_vec,
+        )
     }
 }
 

primitives/src/circuit/rescue/mod.rs

@@ -10,5 +10,175 @@
 mod native;
 mod non_native;
 
-pub use native::{RescueGadget, RescueStateVar};
-pub use non_native::{RescueNonNativeGadget, RescueNonNativeStateVar};
+use ark_ff::PrimeField;
+use ark_std::vec::Vec;
+use jf_relation::{errors::CircuitError, Circuit};
+pub use native::RescueStateVar;
+pub use non_native::RescueNonNativeStateVar;
+
+use crate::rescue::{RescueMatrix, RescueVector, PRP};
+
+/// Variable to represent the state of the sponge.
+pub trait SpongeStateVar<T, F> {
+    /// The native field.
+    type Native;
+    /// Non-native field.
+    type NonNative;
+    /// How variable is represented in this Rescue(NonNative)StateVar.
+    type Var;
+}
+
+/// Trait for rescue circuit over native field.
+pub trait RescueGadget<R, T, F>
+where
+    R: SpongeStateVar<T, F>,
+{
+    /// Given an input state st_0 and an output state st_1, ensure that st_1 =
+    /// rescue_permutation(st_0)  where rescue_permutation is the instance
+    /// of the Rescue permutation defined by its respective constants
+    /// * `input_var` - variables corresponding to the input state
+    /// * `returns` - variables corresponding to the output state
+    fn rescue_permutation(&mut self, input_var: R) -> Result<R, CircuitError>;
+
+    /// Rescue based Pseudo Random Permutation (PRP)
+    /// * `key_var` - rescue state variable corresponding to the cipher key
+    /// * `input_var` - rescue state variable corresponding to the plaintext
+    /// * `returns` - state variable corresponding to the cipher text
+    fn prp(&mut self, key_var: &R, input_var: &R) -> Result<R, CircuitError>;
+
+    /// Sponge-based hashes from Rescue permutations
+    /// * `data_vars` - sponge input variables, `data_vars.len()` should be a
+    ///   positive integer that is a multiple of the sponge rate (i.e. 3)
+    /// * `num_output` - number of output variables
+    /// * `returns` - a vector of variables that refers to the sponge hash
+    ///   output
+    fn rescue_sponge_no_padding(
+        &mut self,
+        data_vars: &[R::Var],
+        num_output: usize,
+    ) -> Result<Vec<R::Var>, CircuitError>;
+
+    /// Sponge-based hashes from Rescue permutations
+    /// * `data_vars` - sponge input variables,
+    /// * `num_output` - number of output variables
+    /// * `returns` - a vector of variables that refers to the sponge hash
+    ///   output
+    fn rescue_sponge_with_padding(
+        &mut self,
+        data_vars: &[R::Var],
+        num_output: usize,
+    ) -> Result<Vec<R::Var>, CircuitError>;
+
+    /// Full-State-Keyed-Sponge with a single output
+    /// * `key` - key variable
+    /// * `input` - input variables,
+    /// * `returns` a variable that refers to the output
+    fn rescue_full_state_keyed_sponge_no_padding(
+        &mut self,
+        key: R::Var,
+        data_vars: &[R::Var],
+    ) -> Result<R::Var, CircuitError>;
+
+    /// Return the round keys variables for the Rescue block cipher
+    /// * `mds_states` - Rescue MDS matrix
+    /// * `key_var` - state variable representing the cipher key
+    /// * `returns` - state variables corresponding to the scheduled keys
+    fn key_schedule(
+        &mut self,
+        mds_states: &RescueMatrix<R::Native>,
+        key_var: &R,
+        prp_instance: &PRP<R::Native>,
+    ) -> Result<Vec<R>, CircuitError>;
+
+    /// Create a variable representing a rescue state
+    /// * `state` - Rescue state
+    /// * `returns` - state variables corresponding to the state
+    fn create_rescue_state_variable(
+        &mut self,
+        state: &RescueVector<R::Native>,
+    ) -> Result<R, CircuitError>;
+
+    /// Return the variable corresponding to the output of the of the Rescue
+    /// PRP where the rounds keys have already been computed "dynamically"
+    /// * `input_var` - variable corresponding to the plain text
+    /// * `mds_states` - Rescue MDS matrix
+    /// * `key_vars` - variables corresponding to the scheduled keys
+    /// * `returns` -
+    fn prp_with_round_keys(
+        &mut self,
+        input_var: &R,
+        mds: &RescueMatrix<R::Native>,
+        keys_vars: &[R],
+    ) -> Result<R, CircuitError>;
+}
+
+pub(crate) trait PermutationGadget<R, T, F>: Circuit<F>
+where
+    R: SpongeStateVar<T, F>,
+    F: PrimeField,
+{
+    fn check_var_bound_rescue_state(&self, rescue_state: &R) -> Result<(), CircuitError>;
+
+    fn add_constant_state(
+        &mut self,
+        input_var: &R,
+        constant: &RescueVector<R::Native>,
+    ) -> Result<R, CircuitError>;
+
+    fn add_state(&mut self, left_state_var: &R, right_state_var: &R) -> Result<R, CircuitError>;
+
+    /// Given a state st_0=(x_1,...,x_w) and st_1=(y_1,...,y_w),
+    /// add the constraints that ensure we have y_i=x_i ^{1/5} for i in
+    /// [1,...,w]
+    /// * `input_var` - rescue state variables st_0
+    /// * `returns` - rescue state variables st_1
+    fn pow_alpha_inv_state(&mut self, input_var: &R) -> Result<R, CircuitError>;
+
+    /// Given an input state st_0 and an output state st_1, ensure that st_1 = M
+    /// st_0 + C where M is a Rescue matrix and c is a constant vector
+    /// * `input_var` - variables corresponding to the input state
+    /// * `matrix` - matrix M in the description above
+    /// * `constant` - constant c in the description above
+    /// * `returns` - variables corresponding to the output state
+    fn affine_transform(
+        &mut self,
+        input_var: &R,
+        matrix: &RescueMatrix<R::Native>,
+        constant: &RescueVector<R::Native>,
+    ) -> Result<R, CircuitError>;
+
+    /// Given an input state st_0=(x_1,...,x_w) and an output state
+    /// st_1=(y_1,...,y_m) y_i = \sum_{j=1}^w M_{i,j}x_j^alpha+c_i for all i in
+    /// [1,..,w] where M is a Rescue matrix and c=(c_1,...,c_w) is a
+    /// constant vector
+    /// * `input_var` - variables corresponding to the input state
+    /// * `matrix` - matrix M in the description above
+    /// * `constant` - constant c in the description above
+    /// * `returns` - variables corresponding to the output state
+    fn non_linear_transform(
+        &mut self,
+        input_var: &R,
+        matrix: &RescueMatrix<R::Native>,
+        constant: &RescueVector<R::Native>,
+    ) -> Result<R, CircuitError>;
+
+    /// Define a constraint such that y = x^(1/alpha).
+    /// It is implemented by setting q_{H1} y^alpha = q_O x
+    /// * `input_var`  - variable id corresponding to x in the equation above
+    /// * `returns` - the variable id corresponding to y
+    fn pow_alpha_inv(&mut self, input_var: R::Var) -> Result<R::Var, CircuitError>;
+
+    /// Given an input state st_0 and an output state st_1, ensure that st_1 is
+    /// obtained by applying the rescue permutation with a specific  list of
+    /// round keys (i.e. the keys are constants) and a matrix
+    /// * `input_var` - variables corresponding to the input state
+    /// * `mds` - Rescue matrix
+    /// * `round_keys` - list of round keys
+    /// * `returns` - variables corresponding to the output state
+    fn permutation_with_const_round_keys(
+        &mut self,
+        input_var: R,
+        mds: &RescueMatrix<R::Native>,
+        round_keys: &[RescueVector<R::Native>],
+    ) -> Result<R, CircuitError>;
+}