From 111bdf8e4f96c423a0e234d051dd332307d137ad Mon Sep 17 00:00:00 2001
From: yoavGrs <97383386+yoavGrs@users.noreply.github.com>
Date: Sun, 19 May 2024 12:46:51 +0300
Subject: [PATCH] feat: filter out a contract with empty storage maps (#256)

---
 src/block_hash/state_diff_hash.rs      | 18 ++++++++++++------
 src/block_hash/state_diff_hash_test.rs | 20 ++++++++++++++++++++
 src/crypto/utils.rs                    |  6 ++++++
 3 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/src/block_hash/state_diff_hash.rs b/src/block_hash/state_diff_hash.rs
index b868966..8fab0c2 100644
--- a/src/block_hash/state_diff_hash.rs
+++ b/src/block_hash/state_diff_hash.rs
@@ -76,17 +76,23 @@ fn chain_deprecated_declared_classes(
 // ]
 fn chain_storage_diffs(
     storage_diffs: &IndexMap<ContractAddress, IndexMap<StorageKey, StarkFelt>>,
-    mut hash_chain: HashChain,
+    hash_chain: HashChain,
 ) -> HashChain {
-    hash_chain = hash_chain.chain(&storage_diffs.len().into());
+    let mut n_updated_contracts = 0_u64;
+    let mut storage_diffs_chain = HashChain::new();
     for (contract_address, key_value_map) in sorted_index_map(storage_diffs) {
-        hash_chain = hash_chain.chain(&contract_address);
-        hash_chain = hash_chain.chain(&key_value_map.len().into());
+        if key_value_map.is_empty() {
+            // Filter out a contract with empty storage maps.
+            continue;
+        }
+        n_updated_contracts += 1;
+        storage_diffs_chain = storage_diffs_chain.chain(&contract_address);
+        storage_diffs_chain = storage_diffs_chain.chain(&key_value_map.len().into());
         for (key, value) in sorted_index_map(&key_value_map) {
-            hash_chain = hash_chain.chain(&key).chain(&value);
+            storage_diffs_chain = storage_diffs_chain.chain(&key).chain(&value);
         }
     }
-    hash_chain
+    hash_chain.chain(&n_updated_contracts.into()).extend(storage_diffs_chain)
 }
 
 // Chains: [number_of_updated_contracts nonces,
diff --git a/src/block_hash/state_diff_hash_test.rs b/src/block_hash/state_diff_hash_test.rs
index 9bf795f..4f8b0f2 100644
--- a/src/block_hash/state_diff_hash_test.rs
+++ b/src/block_hash/state_diff_hash_test.rs
@@ -116,6 +116,26 @@ fn test_sorting_storage_diffs() {
     );
 }
 
+#[test]
+fn test_empty_storage_diffs() {
+    let storage_diffs_0 = indexmap! {
+        0u64.into() => indexmap! {
+            1u64.into() => 2u64.into(),
+        },
+        3u64.into() => indexmap! {
+        },
+    };
+    let storage_diffs_1 = indexmap! {
+        0u64.into() => indexmap! {
+            1u64.into() => 2u64.into(),
+        },
+    };
+    assert_eq!(
+        chain_storage_diffs(&storage_diffs_0, HashChain::new()).get_poseidon_hash(),
+        chain_storage_diffs(&storage_diffs_1, HashChain::new()).get_poseidon_hash(),
+    );
+}
+
 #[test]
 fn test_sorting_nonces() {
     let nonces_0 = indexmap! {
diff --git a/src/crypto/utils.rs b/src/crypto/utils.rs
index 83503dc..cc0bd1f 100644
--- a/src/crypto/utils.rs
+++ b/src/crypto/utils.rs
@@ -95,6 +95,12 @@ impl HashChain {
         self.chain(&felts.len().into()).chain_iter(felts.iter())
     }
 
+    // Chains a chain of felts to the hash chain.
+    pub fn extend(mut self, chain: HashChain) -> Self {
+        self.elements.extend(chain.elements);
+        self
+    }
+
     // Returns the pedersen hash of the chained felts, hashed with the length of the chain.
     pub fn get_pedersen_hash(&self) -> StarkHash {
         let current_hash = self