feat: Poseidon2-based CRHF and available as Merkle hash (#716)

* add Poseidon2Sponge impl

* fix CI attempt

* fix ci

* rename Sponge->SpongeState, Sponge=DuplexSponge<SpongeState>

* add Fixed/VariableLenCRHF from P2-sponge

* make P2Hash available in MT hash

Cargo.toml

@@ -6,7 +6,7 @@ resolver = "2"
 authors = ["Espresso Systems <[email protected]>"]
 edition = "2021"
 license = "MIT"
-rust-version = "1.64.0"
+rust-version = "1.73.0"
 homepage = "https://github.com/EspressoSystems/jellyfish"
 documentation = "https://jellyfish.docs.espressosys.com"
 repository = "https://github.com/EspressoSystems/jellyfish"
@@ -31,6 +31,6 @@ rand_chacha = { version = "0.3.1", default-features = false }
 serde = { version = "1.0", default-features = false, features = [ "derive", "rc" ] }
 sha2 = { version = "0.10", default-features = false }
 sha3 = { version = "0.10", default-features = false }
-itertools = { version = "0.12", default-features = false } 
+itertools = { version = "0.12", default-features = false }
 tagged-base64 = "0.4"
 zeroize = { version = "^1.8" }

crhf/src/lib.rs

@@ -14,8 +14,8 @@ use ark_std::{borrow::Borrow, fmt::Debug, hash::Hash};
 /// (based on ark-primitives' definition, but self-declared for minimal
 /// dependency and easier future upgradability.)
 pub trait CRHF {
-    /// Input to the CRHF
-    type Input;
+    /// Input to the CRHF, allowed to be dynamically sized
+    type Input: ?Sized;
     /// Output of the CRHF
     type Output: Clone + PartialEq + Eq + Hash + Debug + CanonicalSerialize + CanonicalDeserialize;
     /// Error type

merkle_tree/Cargo.toml

@@ -24,10 +24,15 @@ displaydoc = { workspace = true }
 hashbrown = { workspace = true }
 hex = "0.4.3"
 itertools = { workspace = true, features = ["use_alloc"] }
-jf-poseidon2 = { version = "0.1.0", git = "https://github.com/EspressoSystems/jellyfish" }
+jf-crhf = { path = "../crhf" }
+# TODO: switch to this once jf-poseidon2 cut a release
+# jf-poseidon2 = { version = "0.1.0", git = "https://github.com/EspressoSystems/jellyfish" }
+jf-poseidon2 = { path = "../poseidon2" }
 jf-relation = { version = "0.4.4", git = "https://github.com/EspressoSystems/jellyfish", tag = "0.4.5", optional = true, default-features = false }
 jf-rescue = { version = "0.1.0", git = "https://github.com/EspressoSystems/jellyfish", tag = "0.4.5", default-features = false }
 jf-utils = { version = "0.4.4", git = "https://github.com/EspressoSystems/jellyfish", tag = "0.4.5", default-features = false }
+# TODO: fix this once upstream release a new tag
+nimue = { git = "https://github.com/alxiong/nimue", branch = "ax/0.1.1", features = ["ark"] }
 num-bigint = { workspace = true }
 num-traits = { version = "0.2.15", default-features = false }
 serde = { workspace = true }

merkle_tree/src/errors.rs

@@ -7,6 +7,7 @@
 
 use ark_std::string::String;
 use displaydoc::Display;
+use jf_poseidon2::Poseidon2Error;
 use jf_rescue::RescueError;
 
 /// Error type for Merkle tree
@@ -35,3 +36,9 @@ impl From<RescueError> for MerkleTreeError {
         MerkleTreeError::DigestError(ark_std::format!("{}", err))
     }
 }
+
+impl From<Poseidon2Error> for MerkleTreeError {
+    fn from(err: Poseidon2Error) -> Self {
+        MerkleTreeError::DigestError(ark_std::format!("{}", err))
+    }
+}

merkle_tree/src/prelude.rs

@@ -24,8 +24,10 @@ use ark_serialize::{
     Write,
 };
 use ark_std::{fmt, marker::PhantomData, vec::Vec};
-use jf_poseidon2::{Poseidon2, Poseidon2Params};
+use jf_crhf::CRHF;
+use jf_poseidon2::{crhf::FixedLenPoseidon2Hash, Poseidon2, Poseidon2Params};
 use jf_rescue::{crhf::RescueCRHF, RescueParameter};
+use nimue::hash::sponge::Sponge;
 use sha3::{Digest, Keccak256, Sha3_256};
 
 /// Wrapper for rescue hash function
@@ -54,33 +56,25 @@ pub type RescueLightWeightMerkleTree<F> = LightWeightMerkleTree<F, RescueHash<F>
 /// Example instantiation of a SparseMerkleTree indexed by I
 pub type RescueSparseMerkleTree<I, F> = UniversalMerkleTree<F, RescueHash<F>, I, 3, F>;
 
-// TODO: (alex) move this compression to CRHF and wrap with better API?
-/// Wrapper for Poseidon2 compression function
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub struct Poseidon2Compression<F, P, const N: usize>(
-    (PhantomData<F>, PhantomData<P>, PhantomData<[(); N]>),
-)
-where
-    F: PrimeField,
-    P: Poseidon2Params<F, N>;
-
-impl<I, F, P, const N: usize> DigestAlgorithm<F, I, F> for Poseidon2Compression<F, P, N>
+// Make `FixedLenPoseidon2Hash<F, S, INPUT_SIZE, 1>` usable as Merkle tree hash
+impl<I, F, S, const INPUT_SIZE: usize> DigestAlgorithm<F, I, F>
+    for FixedLenPoseidon2Hash<F, S, INPUT_SIZE, 1>
 where
     I: Index,
-    F: PrimeField + From<I>,
-    P: Poseidon2Params<F, N>,
+    F: PrimeField + From<I> + nimue::Unit,
+    S: Sponge<U = F>,
 {
     fn digest(data: &[F]) -> Result<F, MerkleTreeError> {
-        let mut input = [F::default(); N];
+        let mut input = [F::default(); INPUT_SIZE];
         input.copy_from_slice(&data[..]);
-        Ok(Poseidon2::permute::<P, N>(&input)[0])
+        Ok(FixedLenPoseidon2Hash::<F, S, INPUT_SIZE, 1>::evaluate(input)?[0])
     }
 
     fn digest_leaf(pos: &I, elem: &F) -> Result<F, MerkleTreeError> {
-        let mut input = [F::default(); N];
-        input[N - 1] = F::from(pos.clone());
-        input[N - 2] = *elem;
-        Ok(Poseidon2::permute::<P, N>(&input)[0])
+        let mut input = [F::default(); INPUT_SIZE];
+        input[INPUT_SIZE - 1] = F::from(pos.clone());
+        input[INPUT_SIZE - 2] = *elem;
+        Ok(FixedLenPoseidon2Hash::<F, S, INPUT_SIZE, 1>::evaluate(input)?[0])
     }
 }
 

poseidon2/Cargo.toml

@@ -15,8 +15,11 @@ ark-bls12-381 = { workspace = true, optional = true }
 ark-bn254 = { workspace = true, optional = true }
 ark-ff = { workspace = true }
 ark-std = { workspace = true }
+displaydoc = { workspace = true }
 hex = "0.4.3"
-jf-crhf = { version = "0.1.0", git = "https://github.com/EspressoSystems/jellyfish", tag = "0.4.5", default-features = false }
+# TODO: need to update after jf-crhf have a new tag
+# jf-crhf = { version = "0.1.0", git = "https://github.com/EspressoSystems/jellyfish", tag = "0.4.5", default-features = false }
+jf-crhf = { path = "../crhf" }
 lazy_static = "1.5.0"
 # TODO: fix this once upstream release a new tag
 nimue = { git = "https://github.com/alxiong/nimue", branch = "ax/0.1.1", features = ["ark"] }

poseidon2/src/crhf.rs

@@ -0,0 +1,102 @@
+//! Collision-resistant Hash Function (CRHF) based on Poseidon2 permutation
+
+use core::marker::PhantomData;
+
+use ark_ff::{Field, PrimeField};
+use ark_std::{borrow::Borrow, string::ToString, vec::Vec};
+use jf_crhf::CRHF;
+use nimue::{
+    hash::sponge::{DuplexSponge, Sponge},
+    DuplexHash, Unit,
+};
+
+use crate::Poseidon2Error;
+
+/// Sponge-based CRHF where the Sponge uses Poseidon2 permutation
+/// Input length is fixed: the actual input can be shorter, but will internally
+/// be zero-padded to `INPUT_SIZE`
+///
+/// Example:
+/// `FixedLenPoseidon2Hash<ark_bn254::Fr, Poseidon2SpongeStateBnN3R1, 6, 2>`
+#[derive(Clone)]
+pub struct FixedLenPoseidon2Hash<F, S, const INPUT_SIZE: usize, const OUTPUT_SIZE: usize>
+where
+    F: PrimeField + Unit,
+    S: Sponge<U = F>,
+{
+    _field: PhantomData<F>,
+    _sponge: PhantomData<S>,
+}
+
+impl<F, S, const IN: usize, const OUT: usize> CRHF for FixedLenPoseidon2Hash<F, S, IN, OUT>
+where
+    F: PrimeField + Unit,
+    S: Sponge<U = F>,
+{
+    type Input = [F]; // length should be <= IN
+    type Output = [F; OUT];
+    type Error = Poseidon2Error;
+
+    fn evaluate<T: Borrow<Self::Input>>(input: T) -> Result<Self::Output, Self::Error> {
+        let input = input.borrow();
+        if input.len() > IN {
+            return Err(Poseidon2Error::ParamErr("hash input too long".to_string()));
+        }
+
+        let mut padded = Vec::from(input);
+        zero_padding(&mut padded, IN);
+
+        let mut sponge = DuplexSponge::<S>::default();
+        sponge.absorb_unchecked(&padded);
+        let mut output = [F::default(); OUT];
+        sponge.squeeze_unchecked(&mut output);
+        Ok(output)
+    }
+}
+
+/// Sponge-based CRHF where the Sponge uses Poseidon2 permutation, with
+/// variable-length input
+#[derive(Debug, Clone)]
+pub struct VariableLenPoseidon2Hash<F, S, const OUTPUT_SIZE: usize>
+where
+    F: PrimeField + Unit,
+    S: Sponge<U = F>,
+{
+    _field: PhantomData<F>,
+    _sponge: PhantomData<S>,
+}
+
+impl<F, S, const OUT: usize> CRHF for VariableLenPoseidon2Hash<F, S, OUT>
+where
+    F: PrimeField + Unit,
+    S: Sponge<U = F>,
+{
+    type Input = [F];
+    type Output = [F; OUT];
+    type Error = Poseidon2Error;
+
+    fn evaluate<T: Borrow<Self::Input>>(input: T) -> Result<Self::Output, Self::Error> {
+        let mut padded = Vec::from(input.borrow());
+        bit_padding(&mut padded, S::R);
+
+        let mut sponge = DuplexSponge::<S>::default();
+        sponge.absorb_unchecked(&padded);
+        let mut output = [F::default(); OUT];
+        sponge.squeeze_unchecked(&mut output);
+        Ok(output)
+    }
+}
+
+// pad `data` with zeros until the length is the next multiple of `multiple`
+#[inline(always)]
+fn zero_padding<F: Field>(data: &mut Vec<F>, multiple: usize) {
+    data.resize(data.len().next_multiple_of(multiple), F::zero());
+}
+
+// pad `data` with "10..0" (always pad "1"), until the length is the next
+// multiple of `multiple`
+#[inline(always)]
+fn bit_padding<F: Field>(data: &mut Vec<F>, multiple: usize) {
+    data.push(F::one());
+    zero_padding(data, multiple);
+}

poseidon2/src/lib.rs

@@ -20,9 +20,11 @@
 #![deny(missing_docs)]
 
 use ark_ff::PrimeField;
-use ark_std::{borrow::ToOwned, marker::PhantomData};
+use ark_std::{borrow::ToOwned, marker::PhantomData, string::String};
+use displaydoc::Display;
 
 pub mod constants;
+pub mod crhf;
 mod external;
 mod internal;
 pub mod sponge;
@@ -130,3 +132,11 @@ pub(crate) fn add_rc_and_sbox<F: PrimeField>(val: &mut F, rc: F, d: usize) {
     *val += rc;
     *val = val.pow([d as u64]);
 }
+
+/// Poseidon2 Error type
+#[derive(Debug, Clone, Display, Eq, PartialEq)]
+pub enum Poseidon2Error {
+    /// Bad parameter: {0}
+    ParamErr(String),
+}
+impl ark_std::error::Error for Poseidon2Error {}

poseidon2/src/sponge.rs

@@ -6,29 +6,33 @@ use ark_std::marker::PhantomData;
 use nimue::{hash::sponge::Sponge, Unit};
 use zeroize::Zeroize;
 
-/// Poseidon2-based Cryptographic Sponge
+/// the state of Poseidon2-based Cryptographic Sponge
 ///
 /// # Generic parameters:
 /// - N: state size = rate (R) + capacity (C)
 /// - R: rate (number of field abosrbed/squeezed)
 ///
 /// For security, for b=128-bit security, field size |F|, C*|F|>=2b:
 /// i.e. 128-bit for 256-bit fields, C>=1.
-/// This check is being down during `Poseidon2Sponge::new(&iv)`
+/// This check is being down during `Poseidon2SpongeState::new(&iv)`
 /// (See Poseidon2 paper Page 7 Footnote 5)
 ///
 /// For BLS12-381, we choose C=1 for 128 security
 /// For BN254, we choose C=1 for (100<*<128)-security
 #[derive(Clone, Debug)]
-pub struct Poseidon2Sponge<F: PrimeField, const N: usize, const R: usize, P: Poseidon2Params<F, N>>
-{
+pub struct Poseidon2SpongeState<
+    F: PrimeField,
+    const N: usize,
+    const R: usize,
+    P: Poseidon2Params<F, N>,
+> {
     /// state of sponge
     pub(crate) state: [F; N],
     _rate: PhantomData<[(); R]>,
     _p: PhantomData<P>,
 }
 
-impl<F, const N: usize, const R: usize, P> Sponge for Poseidon2Sponge<F, N, R, P>
+impl<F, const N: usize, const R: usize, P> Sponge for Poseidon2SpongeState<F, N, R, P>
 where
     F: PrimeField + Unit,
     P: Poseidon2Params<F, N>,
@@ -57,7 +61,7 @@ where
         Poseidon2::permute_mut::<P, N>(&mut self.state);
     }
 }
-impl<F, const N: usize, const R: usize, P> Default for Poseidon2Sponge<F, N, R, P>
+impl<F, const N: usize, const R: usize, P> Default for Poseidon2SpongeState<F, N, R, P>
 where
     F: PrimeField,
     P: Poseidon2Params<F, N>,
@@ -71,7 +75,7 @@ where
     }
 }
 
-impl<F, const N: usize, const R: usize, P> AsRef<[F]> for Poseidon2Sponge<F, N, R, P>
+impl<F, const N: usize, const R: usize, P> AsRef<[F]> for Poseidon2SpongeState<F, N, R, P>
 where
     F: PrimeField,
     P: Poseidon2Params<F, N>,
@@ -80,7 +84,7 @@ where
         &self.state
     }
 }
-impl<F, const N: usize, const R: usize, P> AsMut<[F]> for Poseidon2Sponge<F, N, R, P>
+impl<F, const N: usize, const R: usize, P> AsMut<[F]> for Poseidon2SpongeState<F, N, R, P>
 where
     F: PrimeField,
     P: Poseidon2Params<F, N>,
@@ -90,7 +94,7 @@ where
     }
 }
 
-impl<F, const N: usize, const R: usize, P> Zeroize for Poseidon2Sponge<F, N, R, P>
+impl<F, const N: usize, const R: usize, P> Zeroize for Poseidon2SpongeState<F, N, R, P>
 where
     F: PrimeField,
     P: Poseidon2Params<F, N>,
@@ -107,12 +111,20 @@ mod bls12_381 {
     use crate::constants::bls12_381::*;
     use ark_bls12_381::Fr;
     use nimue::hash::sponge::DuplexSponge;
+    /// State of a sponge over BLS12-381 scalar field, state_size=2, rate=1.
+    pub type Poseidon2SpongeStateBlsN2R1 = Poseidon2SpongeState<Fr, 2, 1, Poseidon2ParamsBls2>;
     /// A sponge over BLS12-381 scalar field, state_size=2, rate=1.
-    pub type Poseidon2SpongeBlsN2R1 = DuplexSponge<Poseidon2Sponge<Fr, 2, 1, Poseidon2ParamsBls2>>;
+    pub type Poseidon2SpongeBlsN2R1 = DuplexSponge<Poseidon2SpongeStateBlsN2R1>;
+
+    /// State of a sponge over BLS12-381 scalar field, state_size=3, rate=1.
+    pub type Poseidon2SpongeStateBlsN3R1 = Poseidon2SpongeState<Fr, 3, 1, Poseidon2ParamsBls3>;
     /// A sponge over BLS12-381 scalar field, state_size=3, rate=1.
-    pub type Poseidon2SpongeBlsN3R1 = DuplexSponge<Poseidon2Sponge<Fr, 3, 1, Poseidon2ParamsBls3>>;
+    pub type Poseidon2SpongeBlsN3R1 = DuplexSponge<Poseidon2SpongeStateBlsN3R1>;
+
+    /// State of a sponge over BLS12-381 scalar field, state_size=3, rate=2.
+    pub type Poseidon2SpongeStateBlsN3R2 = Poseidon2SpongeState<Fr, 3, 2, Poseidon2ParamsBls3>;
     /// A sponge over BLS12-381 scalar field, state_size=3, rate=2.
-    pub type Poseidon2SpongeBlsN3R2 = DuplexSponge<Poseidon2Sponge<Fr, 3, 2, Poseidon2ParamsBls3>>;
+    pub type Poseidon2SpongeBlsN3R2 = DuplexSponge<Poseidon2SpongeStateBlsN3R2>;
 
     #[test]
     fn test_bls_sponge() {
@@ -130,10 +142,15 @@ mod bn254 {
     use crate::constants::bn254::*;
     use ark_bn254::Fr;
     use nimue::hash::sponge::DuplexSponge;
+    /// State of a sponge over BN254 scalar field, state_size=3, rate=1.
+    pub type Poseidon2SpongeStateBnN3R1 = Poseidon2SpongeState<Fr, 3, 1, Poseidon2ParamsBn3>;
     /// A sponge over BN254 scalar field, state_size=3, rate=1.
-    pub type Poseidon2SpongeBnN3R1 = DuplexSponge<Poseidon2Sponge<Fr, 3, 1, Poseidon2ParamsBn3>>;
+    pub type Poseidon2SpongeBnN3R1 = DuplexSponge<Poseidon2SpongeStateBnN3R1>;
+
+    /// State of a sponge over BN254 scalar field, state_size=3, rate=2.
+    pub type Poseidon2SpongeStateBnN3R2 = Poseidon2SpongeState<Fr, 3, 2, Poseidon2ParamsBn3>;
     /// A sponge over BN254 scalar field, state_size=3, rate=2.
-    pub type Poseidon2SpongeBnN3R2 = DuplexSponge<Poseidon2Sponge<Fr, 3, 2, Poseidon2ParamsBn3>>;
+    pub type Poseidon2SpongeBnN3R2 = DuplexSponge<Poseidon2SpongeStateBnN3R2>;
 
     #[test]
     fn test_bn_sponge() {