From 995a8b7cf321aec0bec848500e9cb60dbb533efa Mon Sep 17 00:00:00 2001
From: Nimrod Weiss <nimrod@starkware.co>
Date: Sun, 14 Jul 2024 15:29:20 +0300
Subject: [PATCH] refactor: no need to recalculate the actual updates

---
 crates/committer/src/block_committer/commit.rs           | 9 ++++++---
 .../original_skeleton_tree/skeleton_forest.rs            | 8 ++++----
 .../original_skeleton_tree/skeleton_forest_test.rs       | 3 ++-
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/crates/committer/src/block_committer/commit.rs b/crates/committer/src/block_committer/commit.rs
index f5faa509..a3d2cec8 100644
--- a/crates/committer/src/block_committer/commit.rs
+++ b/crates/committer/src/block_committer/commit.rs
@@ -31,11 +31,14 @@ pub async fn commit_block(input: Input<ConfigImpl>) -> BlockCommitmentResult<Fil
         contracts_trie_sorted_indices: SortedLeafIndices::new(&mut contracts_trie_indices),
         classes_trie_sorted_indices: SortedLeafIndices::new(&mut classes_trie_indices),
     };
+    let actual_storage_updates = input.state_diff.actual_storage_updates();
+    let actual_classes_updates = input.state_diff.actual_classes_updates();
     let (mut original_forest, original_contracts_trie_leaves) = OriginalSkeletonForest::create(
         MapStorage::from(input.storage),
         input.contracts_trie_root_hash,
         input.classes_trie_root_hash,
-        &input.state_diff,
+        &actual_storage_updates,
+        &actual_classes_updates,
         &forest_sorted_indices,
         &input.config,
     )?;
@@ -59,8 +62,8 @@ pub async fn commit_block(input: Input<ConfigImpl>) -> BlockCommitmentResult<Fil
 
     Ok(FilledForest::create::<TreeHashFunctionImpl>(
         updated_forest,
-        input.state_diff.actual_storage_updates(),
-        input.state_diff.actual_classes_updates(),
+        actual_storage_updates,
+        actual_classes_updates,
         &original_contracts_trie_leaves,
         &input.state_diff.address_to_class_hash,
         &input.state_diff.address_to_nonce,
diff --git a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest.rs b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest.rs
index 7298e0c8..f3704c32 100644
--- a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest.rs
+++ b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest.rs
@@ -1,7 +1,6 @@
 use crate::block_committer::input::Config;
 use crate::block_committer::input::ContractAddress;
 use crate::block_committer::input::StarknetStorageValue;
-use crate::block_committer::input::StateDiff;
 use crate::forest_errors::ForestError;
 use crate::forest_errors::ForestResult;
 use crate::hash::hash_trait::HashOutput;
@@ -37,7 +36,8 @@ impl<'a> OriginalSkeletonForest<'a> {
         storage: impl Storage,
         contracts_trie_root_hash: HashOutput,
         classes_trie_root_hash: HashOutput,
-        state_diff: &StateDiff,
+        storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
+        classes_updates: &LeafModifications<CompiledClassHash>,
         forest_sorted_indices: &ForestSortedIndices<'a>,
         config: &impl Config,
     ) -> ForestResult<(Self, HashMap<NodeIndex, ContractState>)>
@@ -50,14 +50,14 @@ impl<'a> OriginalSkeletonForest<'a> {
             forest_sorted_indices.contracts_trie_sorted_indices,
         )?;
         let storage_tries = Self::create_storage_tries(
-            &state_diff.actual_storage_updates(),
+            storage_updates,
             &original_contracts_trie_leaves,
             &storage,
             config,
             &forest_sorted_indices.storage_tries_sorted_indices,
         )?;
         let classes_trie = Self::create_classes_trie(
-            &state_diff.actual_classes_updates(),
+            classes_updates,
             classes_trie_root_hash,
             &storage,
             config,
diff --git a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest_test.rs b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest_test.rs
index 193137c8..0392dcfc 100644
--- a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest_test.rs
+++ b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest_test.rs
@@ -284,7 +284,8 @@ fn test_create_original_skeleton_forest(
         MapStorage::from(input.storage),
         input.contracts_trie_root_hash,
         input.classes_trie_root_hash,
-        &input.state_diff,
+        &input.state_diff.actual_storage_updates(),
+        &input.state_diff.actual_classes_updates(),
         &forest_sorted_indices,
         &ConfigImpl::new(false),
     )