Merge pull request #708 from EspressoSystems/ax/poseidon2

feat: Poseidon2 implementation and availabe as Merkle Hash

.gitignore

@@ -1,7 +1,7 @@
 .*.sw*
 .DS_Store
 .idea
-/target
+**/target
 cargo-system-config.toml
 Cargo.lock
 *.org

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-members = ["aead", "commitment", "crhf", "elgamal", "merkle_tree", "pcs", "plonk", "prf", "relation", "rescue", "signature", "utilities", "vid", "vrf"]
+members = ["aead", "commitment", "crhf", "elgamal", "merkle_tree", "pcs", "plonk", "poseidon2", "prf", "relation", "rescue", "signature", "utilities", "vid", "vrf"]
 resolver = "2"
 
 [workspace.package]
@@ -22,7 +22,7 @@ ark-poly = { version = "0.4.0", default-features = false }
 ark-serialize = { version = "0.4.0", default-features = false, features = [ "derive" ] }
 ark-std = { version = "0.4.0", default-features = false }
 derivative = { version = "2", features = ["use_core"] }
-digest = { version = "0.10.1", default-features = false, features = [ "alloc" ] }
+digest = { version = "0.10.7", default-features = false, features = [ "alloc" ] }
 displaydoc = { version = "0.2", default-features = false }
 hashbrown = "0.14.3"
 merlin = { version = "3.0.0", default-features = false }

merkle_tree/Cargo.toml

@@ -24,6 +24,7 @@ displaydoc = { workspace = true }
 hashbrown = { workspace = true }
 hex = "0.4.3"
 itertools = { workspace = true, features = ["use_alloc"] }
+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 }

merkle_tree/src/hasher.rs

@@ -73,7 +73,7 @@ pub type GenericHasherMerkleTree<H, E, I, const ARITY: usize> =
 
 /// Convenience trait and blanket impl for downstream trait bounds.
 ///
-/// Useful for downstream code that's generic offer [`Digest`] hasher `H`.
+/// Useful for downstream code that's generic over [`Digest`] hasher `H`.
 ///
 /// # Example
 ///
@@ -120,19 +120,19 @@ pub type GenericHasherMerkleTree<H, E, I, const ARITY: usize> =
 ///     let mt = HasherMerkleTree::<H, usize>::from_elems(None, &my_data).unwrap();
 /// }
 /// ```
-pub trait HasherDigest: Digest<OutputSize = Self::Foo> + Write + Send + Sync {
-    /// Associated type needed to express trait bounds.
-    type Foo: ArrayLength<u8, ArrayType = Self::Bar>;
-    /// Associated type needed to express trait bounds.
-    type Bar: Copy;
+pub trait HasherDigest: Digest<OutputSize = Self::OutSize> + Write + Send + Sync {
+    /// Type for the output size
+    type OutSize: ArrayLength<u8, ArrayType = Self::ArrayType>;
+    /// Type for the array
+    type ArrayType: Copy;
 }
 impl<T> HasherDigest for T
 where
     T: Digest + Write + Send + Sync,
     <T::OutputSize as ArrayLength<u8>>::ArrayType: Copy,
 {
-    type Foo = T::OutputSize;
-    type Bar = <<T as HasherDigest>::Foo as ArrayLength<u8>>::ArrayType;
+    type OutSize = T::OutputSize;
+    type ArrayType = <<T as HasherDigest>::OutSize as ArrayLength<u8>>::ArrayType;
 }
 
 /// A struct that impls [`DigestAlgorithm`] for use with [`MerkleTree`].

merkle_tree/src/prelude.rs

@@ -18,11 +18,13 @@ pub use crate::{
 
 use super::light_weight::LightWeightMerkleTree;
 use crate::errors::MerkleTreeError;
+use ark_ff::PrimeField;
 use ark_serialize::{
     CanonicalDeserialize, CanonicalSerialize, Compress, Read, SerializationError, Valid, Validate,
     Write,
 };
 use ark_std::{fmt, marker::PhantomData, vec::Vec};
+use jf_poseidon2::{Poseidon2, Poseidon2Params};
 use jf_rescue::{crhf::RescueCRHF, RescueParameter};
 use sha3::{Digest, Keccak256, Sha3_256};
 
@@ -52,6 +54,36 @@ 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>
+where
+    I: Index,
+    F: PrimeField + From<I>,
+    P: Poseidon2Params<F, N>,
+{
+    fn digest(data: &[F]) -> Result<F, MerkleTreeError> {
+        let mut input = [F::default(); N];
+        input.copy_from_slice(&data[..]);
+        Ok(Poseidon2::permute::<P, N>(&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])
+    }
+}
+
 /// Implement Internal node type and implement DigestAlgorithm for a hash
 /// function with 32 bytes output size
 ///

poseidon2/CHANGELOG.md

@@ -0,0 +1,9 @@
+# CHANGELOG
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+
+## 0.1.0
+
+- Initial release with Poseidon2 Permutation (w/ BLS12-381, BN254 instances)

poseidon2/Cargo.toml

@@ -0,0 +1,32 @@
+[package]
+name = "jf-poseidon2"
+version = "0.1.0"
+description = "Poseidon2 algebraic hash functions implementation."
+authors = { workspace = true }
+edition = { workspace = true }
+license = { workspace = true }
+rust-version = { workspace = true }
+homepage = { workspace = true }
+documentation = { workspace = true }
+repository = { workspace = true }
+
+[dependencies]
+ark-bls12-381 = { workspace = true, optional = true }
+ark-bn254 = { workspace = true, optional = true }
+ark-ff = { workspace = true }
+ark-std = { workspace = true }
+hex = "0.4.3"
+lazy_static = "1.5.0"
+
+[dev-dependencies]
+criterion = "0.5.1"
+
+[features]
+default = ["bls12-381", "bn254"]
+# curve-named features contains constants for scalar fields of that curve
+bls12-381 = ["dep:ark-bls12-381"]
+bn254 = ["dep:ark-bn254"]
+
+[[bench]]
+name = "p2_native"
+harness = false

poseidon2/benches/p2_native.rs

@@ -0,0 +1,98 @@
+//! Benchmark for native speed of Poseidon2
+//! `cargo bench --bench p2_native`
+#[macro_use]
+extern crate criterion;
+use std::time::Duration;
+
+use ark_std::{test_rng, UniformRand};
+use criterion::Criterion;
+use jf_poseidon2::{
+    constants::{
+        bls12_381::{Poseidon2ParamsBls2, Poseidon2ParamsBls3},
+        bn254::Poseidon2ParamsBn3,
+    },
+    Poseidon2,
+};
+
+// BLS12-381 scalar field, state size = 2
+fn bls2(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Poseidon2 over (Bls12_381::Fr, t=2)");
+    group.sample_size(10).measurement_time(Duration::new(20, 0));
+    type Fr = ark_bls12_381::Fr;
+    let rng = &mut test_rng();
+
+    group.bench_function("1k iter", |b| {
+        b.iter(|| {
+            let mut input = [Fr::rand(rng), Fr::rand(rng)];
+            for _ in 0..1000 {
+                Poseidon2::permute_mut::<Poseidon2ParamsBls2, 2>(&mut input);
+            }
+        })
+    });
+    group.bench_function("100k iter", |b| {
+        b.iter(|| {
+            let mut input = [Fr::rand(rng), Fr::rand(rng)];
+            for _ in 0..100_000 {
+                Poseidon2::permute_mut::<Poseidon2ParamsBls2, 2>(&mut input);
+            }
+        })
+    });
+    group.finish();
+}
+
+// BLS12-381 scalar field, state size = 3
+fn bls3(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Poseidon2 over (Bls12_381::Fr, t=3)");
+    group.sample_size(10).measurement_time(Duration::new(20, 0));
+    type Fr = ark_bls12_381::Fr;
+    let rng = &mut test_rng();
+
+    group.bench_function("1k iter", |b| {
+        b.iter(|| {
+            let mut input = [Fr::rand(rng), Fr::rand(rng), Fr::rand(rng)];
+            for _ in 0..1000 {
+                Poseidon2::permute_mut::<Poseidon2ParamsBls3, 3>(&mut input);
+            }
+        })
+    });
+    group.bench_function("100k iter", |b| {
+        b.iter(|| {
+            let mut input = [Fr::rand(rng), Fr::rand(rng), Fr::rand(rng)];
+            for _ in 0..100_000 {
+                Poseidon2::permute_mut::<Poseidon2ParamsBls3, 3>(&mut input);
+            }
+        })
+    });
+    group.finish();
+}
+
+// BN254 scalar field, state size = 3
+fn bn3(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Poseidon2 over (Bn254::Fr, t=3)");
+    group.sample_size(10).measurement_time(Duration::new(20, 0));
+    type Fr = ark_bn254::Fr;
+    let rng = &mut test_rng();
+
+    group.bench_function("1k iter", |b| {
+        b.iter(|| {
+            let mut input = [Fr::rand(rng), Fr::rand(rng), Fr::rand(rng)];
+            for _ in 0..1000 {
+                Poseidon2::permute_mut::<Poseidon2ParamsBn3, 3>(&mut input);
+            }
+        })
+    });
+    group.bench_function("100k iter", |b| {
+        b.iter(|| {
+            let mut input = [Fr::rand(rng), Fr::rand(rng), Fr::rand(rng)];
+            for _ in 0..100_000 {
+                Poseidon2::permute_mut::<Poseidon2ParamsBn3, 3>(&mut input);
+            }
+        })
+    });
+
+    group.finish();
+}
+
+criterion_group!(benches, bls2, bls3, bn3);
+
+criterion_main!(benches);

poseidon2/src/constants.rs

@@ -0,0 +1,55 @@
+//! Poseidon2 Constants copied from <https://github.com/HorizenLabs/poseidon2/blob/main/plain_implementations/src/poseidon2/>
+
+use ark_ff::PrimeField;
+use hex::FromHex;
+
+#[cfg(feature = "bls12-381")]
+pub mod bls12_381;
+#[cfg(feature = "bn254")]
+pub mod bn254;
+
+#[inline]
+pub(crate) fn from_hex<F: PrimeField>(s: &str) -> F {
+    F::from_be_bytes_mod_order(&<[u8; 32]>::from_hex(s).expect("Invalid HexStr"))
+}
+
+/// macros to derive instances that implements `trait Poseidon2Params`
+#[macro_export]
+macro_rules! define_poseidon2_params {
+    (
+        $struct_name:ident,
+        $state_size:expr,
+        $sbox_size:expr,
+        $ext_rounds:expr,
+        $int_rounds:expr,
+        $rc_ext:ident,
+        $rc_int:ident,
+        $mat_diag_m_1:ident
+    ) => {
+        /// Poseidon parameters for Bls12-381 scalar field, with
+        /// - state size = $state_size
+        /// - sbox size = $sbox_size
+        /// - external rounds = $ext_rounds
+        /// - internal rounds = $int_rounds
+        pub struct $struct_name;
+
+        impl Poseidon2Params<Fr, $state_size> for $struct_name {
+            const T: usize = $state_size;
+            const D: usize = $sbox_size;
+            const EXT_ROUNDS: usize = $ext_rounds;
+            const INT_ROUNDS: usize = $int_rounds;
+
+            fn external_rc() -> &'static [[Fr; $state_size]] {
+                &*$rc_ext
+            }
+
+            fn internal_rc() -> &'static [Fr] {
+                &*$rc_int
+            }
+
+            fn internal_mat_diag_m_1() -> &'static [Fr; $state_size] {
+                &$mat_diag_m_1
+            }
+        }
+    };
+}