impl HasherAlgorithm for Poseidon2

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
 ///