starknet_committer: fetch patricia paths concurrently

Author: ArielElp

crates/starknet_committer/src/db/trie_traversal.rs

@@ -852,7 +852,6 @@ where
 }
 
 /// Storage task for fetching Patricia paths in a single storage trie.
-#[allow(dead_code)]
 struct StoragePathsReadTask<'indices, Layout: NodeLayoutFor<StarknetStorageValue>> {
     address: ContractAddress,
     storage_root_hash: HashOutput,
@@ -928,7 +927,6 @@ where
 
 /// Fetches Patricia proofs for the storage tries. If the storage has a [GatherableStorage] version,
 /// then the paths are fetched concurrently. Otherwise, they are fetched sequentially.
-#[allow(dead_code)]
 pub(crate) async fn fetch_contract_storage_paths<StorageLayout, ContractLeaf>(
     storage: &mut impl ReadOnlyStorage,
     contract_storage_sorted_leaf_indices: &HashMap<NodeIndex, SortedLeafIndices<'_>>,
@@ -979,7 +977,6 @@ pub(crate) fn get_address_and_storage_root<ContractLeaf: AsRef<ContractState>>(
 }
 
 /// Sequentially fetches Patricia proofs for the storage tries.
-#[allow(dead_code)]
 async fn fetch_contract_storage_paths_sequentially<StorageLayout, ContractLeaf>(
     storage: &mut impl ReadOnlyStorage,
     contract_storage_sorted_leaf_indices: &HashMap<NodeIndex, SortedLeafIndices<'_>>,
@@ -1017,7 +1014,6 @@ where
 }
 
 /// Concurrently fetches Patricia proofs for the storage tries.
-#[allow(dead_code)]
 async fn fetch_contract_storage_paths_concurrently<S, StorageLayout, ContractLeaf>(
     storage: &mut S,
     contract_storage_sorted_leaf_indices: &HashMap<NodeIndex, SortedLeafIndices<'_>>,

crates/starknet_committer/src/patricia_merkle_tree/tree.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 use blockifier::state::accessed_keys::AccessedKeys;
 use starknet_api::core::ContractAddress;
 use starknet_api::hash::HashOutput;
+use starknet_patricia::patricia_merkle_tree::node_data::inner_node::PreimageMap;
 use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonTreeConfig;
 use starknet_patricia::patricia_merkle_tree::traversal::TraversalResult;
 use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
@@ -17,7 +18,7 @@ use crate::block_committer::input::{
 };
 use crate::db::db_layout::DbLayout;
 use crate::db::facts_db::FactsNodeLayout;
-use crate::db::trie_traversal::{fetch_patricia_paths, get_address_and_storage_root};
+use crate::db::trie_traversal::{fetch_contract_storage_paths, fetch_patricia_paths};
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 use crate::patricia_merkle_tree::types::{
     class_hash_into_node_index,
@@ -132,6 +133,44 @@ pub async fn fetch_all_patricia_paths<Layout>(
     contract_sorted_leaf_indices: SortedLeafIndices<'_>,
     contract_storage_sorted_leaf_indices: &HashMap<NodeIndex, SortedLeafIndices<'_>>,
 ) -> TraversalResult<StarknetForestProofs>
+where
+    Layout: DbLayout,
+    Layout::ContractStateDbLeaf: AsRef<ContractState> + Into<ContractState>,
+    Layout::NodeLayout: Send + 'static,
+{
+    let (classes_trie_proof, contracts_proof_nodes, contract_leaves) =
+        fetch_classes_and_contracts_patricia_paths::<Layout>(
+            storage,
+            classes_trie_root_hash,
+            contracts_trie_root_hash,
+            class_sorted_leaf_indices,
+            contract_sorted_leaf_indices,
+            contract_storage_sorted_leaf_indices,
+        )
+        .await?;
+    let contracts_trie_storage_proofs =
+        fetch_contract_storage_paths::<Layout::NodeLayout, Layout::ContractStateDbLeaf>(
+            storage,
+            contract_storage_sorted_leaf_indices,
+            &contract_leaves,
+        )
+        .await?;
+    Ok(build_starknet_forest_proofs::<Layout>(
+        classes_trie_proof,
+        contracts_proof_nodes,
+        contract_leaves,
+        contracts_trie_storage_proofs,
+    ))
+}
+
+async fn fetch_classes_and_contracts_patricia_paths<Layout>(
+    storage: &mut impl ReadOnlyStorage,
+    classes_trie_root_hash: HashOutput,
+    contracts_trie_root_hash: HashOutput,
+    class_sorted_leaf_indices: SortedLeafIndices<'_>,
+    contract_sorted_leaf_indices: SortedLeafIndices<'_>,
+    contract_storage_sorted_leaf_indices: &HashMap<NodeIndex, SortedLeafIndices<'_>>,
+) -> TraversalResult<(PreimageMap, PreimageMap, HashMap<NodeIndex, Layout::ContractStateDbLeaf>)>
 where
     Layout: DbLayout,
     Layout::ContractStateDbLeaf: AsRef<ContractState> + Into<ContractState>,
@@ -176,31 +215,21 @@ where
         )
         .await?;
 
-    // Contracts storage tries.
-    let mut contracts_trie_storage_proofs =
-        HashMap::with_capacity(contract_storage_sorted_leaf_indices.len());
-
-    for (idx, sorted_leaf_indices) in contract_storage_sorted_leaf_indices {
-        let Some((contract_address, storage_root_hash)) =
-            get_address_and_storage_root(idx, &leaves)
-        else {
-            continue;
-        };
-        // No need to fetch the leaves.
-        let leaves = None;
-        let proof = fetch_patricia_paths::<Layout::StarknetStorageValueDbLeaf, Layout::NodeLayout>(
-            storage,
-            storage_root_hash,
-            *sorted_leaf_indices,
-            leaves,
-            &contract_address,
-        )
-        .await?;
-        contracts_trie_storage_proofs.insert(contract_address, proof);
-    }
+    Ok((classes_trie_proof, contracts_proof_nodes, leaves))
+}
 
+fn build_starknet_forest_proofs<Layout>(
+    classes_trie_proof: PreimageMap,
+    contracts_proof_nodes: PreimageMap,
+    contract_leaves: HashMap<NodeIndex, Layout::ContractStateDbLeaf>,
+    contracts_trie_storage_proofs: HashMap<ContractAddress, PreimageMap>,
+) -> StarknetForestProofs
+where
+    Layout: DbLayout,
+    Layout::ContractStateDbLeaf: Into<ContractState>,
+{
     // Convert contract_leaves_data keys from NodeIndex to ContractAddress.
-    let contract_leaves_data: HashMap<ContractAddress, ContractState> = leaves
+    let contract_leaves_data: HashMap<ContractAddress, ContractState> = contract_leaves
         .into_iter()
         .map(|(idx, contract_state_leaf)| {
             (
@@ -215,14 +244,14 @@ where
         })
         .collect();
 
-    Ok(StarknetForestProofs {
+    StarknetForestProofs {
         classes_trie_proof,
         contracts_trie_proof: ContractsTrieProof {
             nodes: contracts_proof_nodes,
             leaves: contract_leaves_data,
         },
         contracts_trie_storage_proofs,
-    })
+    }
 }
 
 /// Fetch the Patricia paths (inner nodes) in the classes trie, contracts trie,