From 96e8021ee7d0078cc0198779df00999bd4bb2279 Mon Sep 17 00:00:00 2001
From: Alex Xiong <alex.xiong.tech@gmail.com>
Date: Wed, 18 Sep 2024 14:47:01 +0800
Subject: [PATCH] add keccak256 merkle tree

---
 merkle_tree/src/prelude.rs | 141 +++++++++++++++++++++----------------
 1 file changed, 80 insertions(+), 61 deletions(-)

diff --git a/merkle_tree/src/prelude.rs b/merkle_tree/src/prelude.rs
index 8a9c26e3a..f65b4bd20 100644
--- a/merkle_tree/src/prelude.rs
+++ b/merkle_tree/src/prelude.rs
@@ -24,7 +24,7 @@ use ark_serialize::{
 };
 use ark_std::{marker::PhantomData, vec::Vec};
 use jf_rescue::{crhf::RescueCRHF, RescueParameter};
-use sha3::{Digest, Sha3_256};
+use sha3::{Digest, Keccak256, Sha3_256};
 
 /// Wrapper for rescue hash function
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -52,71 +52,90 @@ 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>;
 
-/// Update the array length here
-#[derive(Default, Eq, PartialEq, Clone, Copy, Debug, Ord, PartialOrd, Hash)]
-pub struct Sha3Node(pub(crate) [u8; 32]);
-
-impl AsRef<[u8]> for Sha3Node {
-    fn as_ref(&self) -> &[u8] {
-        &self.0
-    }
-}
-
-impl CanonicalSerialize for Sha3Node {
-    fn serialize_with_mode<W: Write>(
-        &self,
-        mut writer: W,
-        _compress: Compress,
-    ) -> Result<(), SerializationError> {
-        writer.write_all(&self.0)?;
-        Ok(())
-    }
-
-    fn serialized_size(&self, _compress: Compress) -> usize {
-        32
-    }
-}
-impl CanonicalDeserialize for Sha3Node {
-    fn deserialize_with_mode<R: Read>(
-        mut reader: R,
-        _compress: Compress,
-        _validate: Validate,
-    ) -> Result<Self, SerializationError> {
-        let mut ret = [0u8; 32];
-        reader.read_exact(&mut ret)?;
-        Ok(Sha3Node(ret))
-    }
-}
-
-impl Valid for Sha3Node {
-    fn check(&self) -> Result<(), SerializationError> {
-        Ok(())
-    }
-}
-
-/// Wrapper for SHA3_512 hash function
-#[derive(Clone, Debug, Hash, Eq, PartialEq)]
-pub struct Sha3Digest();
-
-impl<E: Element + CanonicalSerialize, I: Index> DigestAlgorithm<E, I, Sha3Node> for Sha3Digest {
-    fn digest(data: &[Sha3Node]) -> Result<Sha3Node, MerkleTreeError> {
-        let mut hasher = Sha3_256::new();
-        for value in data {
-            hasher.update(value);
+/// Implement Internal node type and implement DigestAlgorithm for a hash
+/// function with 32 bytes output size
+///
+/// # Usage
+/// `impl_mt_hash_256!(Sha3_256, Sha3Node, Sha3Digest)` will
+/// - introduce `struct Sha3Node` for internal node in a merkle tree that use
+///   Sha3_256 as hash
+/// - introduce `struct Sha3Digest` which implements `DigestAlgorithm`
+macro_rules! impl_mt_hash_256 {
+    ($hasher:ident, $node_name:ident, $digest_name:ident) => {
+        /// Internal node for merkle tree
+        #[derive(Default, Eq, PartialEq, Clone, Copy, Debug, Ord, PartialOrd, Hash)]
+        pub struct $node_name(pub(crate) [u8; 32]);
+
+        impl AsRef<[u8]> for $node_name {
+            fn as_ref(&self) -> &[u8] {
+                &self.0
+            }
+        }
+        impl CanonicalSerialize for $node_name {
+            fn serialize_with_mode<W: Write>(
+                &self,
+                mut writer: W,
+                _compress: Compress,
+            ) -> Result<(), SerializationError> {
+                writer.write_all(&self.0)?;
+                Ok(())
+            }
+            fn serialized_size(&self, _compress: Compress) -> usize {
+                32
+            }
+        }
+        impl CanonicalDeserialize for $node_name {
+            fn deserialize_with_mode<R: Read>(
+                mut reader: R,
+                _compress: Compress,
+                _validate: Validate,
+            ) -> Result<Self, SerializationError> {
+                let mut ret = [0u8; 32];
+                reader.read_exact(&mut ret)?;
+                Ok(Self(ret))
+            }
+        }
+        impl Valid for $node_name {
+            fn check(&self) -> Result<(), SerializationError> {
+                Ok(())
+            }
         }
-        Ok(Sha3Node(hasher.finalize().into()))
-    }
 
-    fn digest_leaf(_pos: &I, elem: &E) -> Result<Sha3Node, MerkleTreeError> {
-        let mut writer = Vec::new();
-        elem.serialize_compressed(&mut writer).unwrap();
-        let mut hasher = Sha3_256::new();
-        hasher.update(writer);
-        Ok(Sha3Node(hasher.finalize().into()))
-    }
+        /// Wrapper for the actual hash function
+        #[derive(Clone, Debug, Hash, Eq, PartialEq)]
+        pub struct $digest_name;
+        impl<E: Element + CanonicalSerialize, I: Index> DigestAlgorithm<E, I, $node_name>
+            for $digest_name
+        {
+            fn digest(data: &[$node_name]) -> Result<$node_name, MerkleTreeError> {
+                let mut h = $hasher::new();
+                for value in data {
+                    h.update(value);
+                }
+                Ok($node_name(h.finalize().into()))
+            }
+
+            fn digest_leaf(_pos: &I, elem: &E) -> Result<$node_name, MerkleTreeError> {
+                let mut writer = Vec::new();
+                elem.serialize_compressed(&mut writer).unwrap();
+                let mut h = $hasher::new();
+                h.update(writer);
+                Ok($node_name(h.finalize().into()))
+            }
+        }
+    };
 }
 
+impl_mt_hash_256!(Sha3_256, Sha3Node, Sha3Digest);
+impl_mt_hash_256!(Keccak256, Keccak256Node, Keccak256Digest);
+
 /// Merkle tree using SHA3 hash
 pub type SHA3MerkleTree<E> = MerkleTree<E, Sha3Digest, u64, 3, Sha3Node>;
 /// Light weight merkle tree using SHA3 hash
 pub type LightWeightSHA3MerkleTree<E> = LightWeightMerkleTree<E, Sha3Digest, u64, 3, Sha3Node>;
+
+/// Merkle tree using keccak256 hash
+pub type Keccak256MerkleTree<E> = MerkleTree<E, Keccak256Node, u64, 3, Keccak256Digest>;
+/// Light weight merkle tree using Keccak256 hash
+pub type LightWeightKeccak256MerkleTree<E> =
+    LightWeightMerkleTree<E, Keccak256Digest, u64, 3, Keccak256Node>;