From e103dbe321ed8b6af5ca3b035adb6befdff89695 Mon Sep 17 00:00:00 2001
From: Nimrod Weiss <nimrod@starkware.co>
Date: Wed, 17 Jul 2024 10:28:11 +0300
Subject: [PATCH] refactor: commit classes trie concurrently

---
 .../filled_tree/forest.rs                     | 21 ++++++++-----------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/crates/committer/src/patricia_merkle_tree/filled_tree/forest.rs b/crates/committer/src/patricia_merkle_tree/filled_tree/forest.rs
index 0c3c16e4..cadbe2a7 100644
--- a/crates/committer/src/patricia_merkle_tree/filled_tree/forest.rs
+++ b/crates/committer/src/patricia_merkle_tree/filled_tree/forest.rs
@@ -55,15 +55,13 @@ impl FilledForest {
         address_to_class_hash: &HashMap<ContractAddress, ClassHash>,
         address_to_nonce: &HashMap<ContractAddress, Nonce>,
     ) -> ForestResult<Self> {
-        let classes_trie = ClassesTrie::create::<TH>(
+        let classes_trie_task = tokio::spawn(ClassesTrie::create::<TH>(
             Arc::new(updated_forest.classes_trie),
             Arc::new(classes_updates),
-        )
-        .await?;
-
+        ));
         let mut contracts_trie_modifications = HashMap::new();
         let mut filled_storage_tries = HashMap::new();
-        let mut tasks = JoinSet::new();
+        let mut contracts_state_tasks = JoinSet::new();
 
         for (address, inner_updates) in storage_updates {
             let updated_storage_trie = updated_forest
@@ -74,7 +72,7 @@ impl FilledForest {
             let original_contract_state = original_contracts_trie_leaves
                 .get(&NodeIndex::from_contract_address(&address))
                 .ok_or(ForestError::MissingContractCurrentState(address))?;
-            tasks.spawn(Self::new_contract_state::<TH>(
+            contracts_state_tasks.spawn(Self::new_contract_state::<TH>(
                 address,
                 *(address_to_nonce
                     .get(&address)
@@ -87,7 +85,7 @@ impl FilledForest {
             ));
         }
 
-        while let Some(result) = tasks.join_next().await {
+        while let Some(result) = contracts_state_tasks.join_next().await {
             let (address, new_contract_state, filled_storage_trie) = result??;
             contracts_trie_modifications.insert(
                 NodeIndex::from_contract_address(&address),
@@ -96,16 +94,15 @@ impl FilledForest {
             filled_storage_tries.insert(address, filled_storage_trie);
         }
 
-        let contracts_trie = ContractsTrie::create::<TH>(
+        let contracts_trie_task = tokio::spawn(ContractsTrie::create::<TH>(
             Arc::new(updated_forest.contracts_trie),
             Arc::new(contracts_trie_modifications),
-        )
-        .await?;
+        ));
 
         Ok(Self {
             storage_tries: filled_storage_tries,
-            contracts_trie,
-            classes_trie,
+            contracts_trie: contracts_trie_task.await??,
+            classes_trie: classes_trie_task.await??,
         })
     }