From 39ca3113d8817169ba1c4ed681ad0df3aadca633 Mon Sep 17 00:00:00 2001
From: peilun-conflux <48905552+peilun-conflux@users.noreply.github.com>
Date: Fri, 1 Mar 2024 10:08:10 +0800
Subject: [PATCH] Cache zero hash results. (#24)

* Cache zero hash results.

* Fix clippy.
---
 Cargo.lock                              |  5 ++--
 common/append_merkle/Cargo.toml         |  3 +-
 common/append_merkle/src/merkle_tree.rs | 25 ++++++++--------
 common/append_merkle/src/sha3.rs        | 39 +++++++++++++++++++++----
 4 files changed, 51 insertions(+), 21 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 072c7ba..35578a2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -150,6 +150,7 @@ dependencies = [
  "eth2_ssz_derive",
  "ethereum-types 0.14.1",
  "lazy_static",
+ "once_cell",
  "serde",
  "tiny-keccak",
  "tracing",
@@ -4615,9 +4616,9 @@ dependencies = [
 
 [[package]]
 name = "once_cell"
-version = "1.18.0"
+version = "1.19.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d"
+checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92"
 
 [[package]]
 name = "oorandom"
diff --git a/common/append_merkle/Cargo.toml b/common/append_merkle/Cargo.toml
index 37c0b3d..1ce0072 100644
--- a/common/append_merkle/Cargo.toml
+++ b/common/append_merkle/Cargo.toml
@@ -11,4 +11,5 @@ eth2_ssz = "0.4.0"
 eth2_ssz_derive = "0.3.0"
 serde = { version = "1.0.137", features = ["derive"] }
 lazy_static = "1.4.0"
-tracing = "0.1.36"
\ No newline at end of file
+tracing = "0.1.36"
+once_cell = "1.19.0"
\ No newline at end of file
diff --git a/common/append_merkle/src/merkle_tree.rs b/common/append_merkle/src/merkle_tree.rs
index 0dd928e..1747068 100644
--- a/common/append_merkle/src/merkle_tree.rs
+++ b/common/append_merkle/src/merkle_tree.rs
@@ -2,7 +2,7 @@ use crate::sha3::Sha3Algorithm;
 use crate::Proof;
 use anyhow::{bail, Result};
 use ethereum_types::H256;
-use lazy_static::lazy_static;
+use once_cell::sync::Lazy;
 use ssz::{Decode, Encode};
 use std::fmt::Debug;
 use std::hash::Hash;
@@ -28,22 +28,21 @@ impl HashElement for H256 {
     }
 }
 
-lazy_static! {
-    static ref ZERO_HASHES: [H256; 64] = {
-        let leaf_zero_hash: H256 = Sha3Algorithm::leaf(&[0u8; 256]);
-        let mut list = [H256::zero(); 64];
-        list[0] = leaf_zero_hash;
-        for i in 1..list.len() {
-            list[i] = Sha3Algorithm::parent(&list[i - 1], &list[i - 1]);
-        }
-        list
-    };
-}
+pub static ZERO_HASHES: Lazy<[H256; 64]> = Lazy::new(|| {
+    let leaf_zero_hash: H256 = Sha3Algorithm::leaf_raw(&[0u8; 256]);
+    let mut list = [H256::zero(); 64];
+    list[0] = leaf_zero_hash;
+    for i in 1..list.len() {
+        list[i] = Sha3Algorithm::parent_raw(&list[i - 1], &list[i - 1]);
+    }
+    list
+});
 
 pub trait Algorithm<E: HashElement> {
     fn parent(left: &E, right: &E) -> E;
     fn parent_single(r: &E, height: usize) -> E {
-        Self::parent(r, &E::end_pad(height))
+        let right = E::end_pad(height);
+        Self::parent(r, &right)
     }
     fn leaf(data: &[u8]) -> E;
 }
diff --git a/common/append_merkle/src/sha3.rs b/common/append_merkle/src/sha3.rs
index f7d983c..c1008ec 100644
--- a/common/append_merkle/src/sha3.rs
+++ b/common/append_merkle/src/sha3.rs
@@ -1,23 +1,52 @@
+use crate::merkle_tree::ZERO_HASHES;
 use crate::{Algorithm, HashElement};
+use ethereum_types::H256;
+use once_cell::sync::Lazy;
+use std::collections::BTreeMap;
 use tiny_keccak::{Hasher, Keccak};
 
 pub struct Sha3Algorithm {}
 
-impl<E: HashElement> Algorithm<E> for Sha3Algorithm {
-    fn parent(left: &E, right: &E) -> E {
+static ZERO_HASHES_MAP: Lazy<BTreeMap<H256, H256>> = Lazy::new(|| {
+    let mut map = BTreeMap::new();
+    for i in 1..ZERO_HASHES.len() {
+        map.insert(ZERO_HASHES[i - 1], ZERO_HASHES[i]);
+    }
+    map
+});
+
+impl Sha3Algorithm {
+    pub fn parent_raw(left: &H256, right: &H256) -> H256 {
         let mut h = Keccak::v256();
-        let mut e = E::null();
+        let mut e = H256::null();
         h.update(left.as_ref());
         h.update(right.as_ref());
         h.finalize(e.as_mut());
         e
     }
 
-    fn leaf(data: &[u8]) -> E {
+    pub fn leaf_raw(data: &[u8]) -> H256 {
         let mut h = Keccak::v256();
-        let mut e = E::null();
+        let mut e = H256::null();
         h.update(data.as_ref());
         h.finalize(e.as_mut());
         e
     }
 }
+impl Algorithm<H256> for Sha3Algorithm {
+    fn parent(left: &H256, right: &H256) -> H256 {
+        if left == right {
+            if let Some(v) = ZERO_HASHES_MAP.get(left) {
+                return *v;
+            }
+        }
+        Self::parent_raw(left, right)
+    }
+
+    fn leaf(data: &[u8]) -> H256 {
+        if *data == [0u8; 256] {
+            return ZERO_HASHES[0];
+        }
+        Self::leaf_raw(data)
+    }
+}