starknet_committer: fetch patricia paths concurrently

Author: ArielElp

crates/starknet_committer/src/db/trie_traversal.rs

@@ -927,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<'_>>,
@@ -961,6 +960,7 @@ where
 /// The contract address might not exist in the contracts trie in the following cases:
 /// 1. We are looking at the previous tree and the contract is new.
 /// 2. We are looking at the new tree and the contract is deleted (revert).
+///
 /// In either case, the storage trie of this contract is empty, so there is nothing to
 /// prove regarding the contract storage.
 pub(crate) fn get_address_and_storage_root<ContractLeaf: AsRef<ContractState>>(
@@ -978,7 +978,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<'_>>,
@@ -1016,7 +1015,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

@@ -1,27 +1,22 @@
 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;
 pub use starknet_patricia::patricia_merkle_tree::types::SortedLeafIndices;
 use starknet_patricia_storage::db_object::EmptyKeyContext;
 use starknet_patricia_storage::storage_trait::ReadOnlyStorage;
 
-use crate::block_committer::input::{
-    contract_address_into_node_index,
-    try_node_index_into_contract_address,
-    StarknetStorageKey,
-};
+use crate::block_committer::input::{contract_address_into_node_index, StarknetStorageKey};
 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,
-    ContractsTrieProof,
     RootHashes,
     StarknetForestProofs,
 };
@@ -132,6 +127,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(StarknetForestProofs::build::<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,53 +209,7 @@ 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);
-    }
-
-    // Convert contract_leaves_data keys from NodeIndex to ContractAddress.
-    let contract_leaves_data: HashMap<ContractAddress, ContractState> = leaves
-        .into_iter()
-        .map(|(idx, contract_state_leaf)| {
-            (
-                try_node_index_into_contract_address(&idx).unwrap_or_else(|_| {
-                    panic!(
-                        "Converting leaf NodeIndex to ContractAddress should succeed; failed to \
-                         convert {idx:?}."
-                    )
-                }),
-                contract_state_leaf.into(),
-            )
-        })
-        .collect();
-
-    Ok(StarknetForestProofs {
-        classes_trie_proof,
-        contracts_trie_proof: ContractsTrieProof {
-            nodes: contracts_proof_nodes,
-            leaves: contract_leaves_data,
-        },
-        contracts_trie_storage_proofs,
-    })
+    Ok((classes_trie_proof, contracts_proof_nodes, leaves))
 }
 
 /// Fetch the Patricia paths (inner nodes) in the classes trie, contracts trie,

crates/starknet_committer/src/patricia_merkle_tree/types.rs

@@ -10,7 +10,8 @@ use starknet_patricia::patricia_merkle_tree::node_data::leaf::SkeletonLeaf;
 use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
 use starknet_types_core::felt::{Felt, FromStrError};
 
-use crate::block_committer::input::StarknetStorageValue;
+use crate::block_committer::input::{try_node_index_into_contract_address, StarknetStorageValue};
+use crate::db::db_layout::DbLayout;
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 
 pub fn fixed_hex_string_no_prefix(felt: &Felt) -> String {
@@ -57,6 +58,42 @@ pub struct StarknetForestProofs {
 }
 
 impl StarknetForestProofs {
+    pub fn build<Layout>(
+        classes_trie_proof: PreimageMap,
+        contracts_proof_nodes: PreimageMap,
+        contract_leaves: HashMap<NodeIndex, Layout::ContractStateDbLeaf>,
+        contracts_trie_storage_proofs: HashMap<ContractAddress, PreimageMap>,
+    ) -> Self
+    where
+        Layout: DbLayout,
+        Layout::ContractStateDbLeaf: Into<ContractState>,
+    {
+        // Convert contract_leaves_data keys from NodeIndex to ContractAddress.
+        let contract_leaves_data: HashMap<ContractAddress, ContractState> = contract_leaves
+            .into_iter()
+            .map(|(idx, contract_state_leaf)| {
+                (
+                    try_node_index_into_contract_address(&idx).unwrap_or_else(|_| {
+                        panic!(
+                            "Converting leaf NodeIndex to ContractAddress should succeed; failed \
+                             to convert {idx:?}."
+                        )
+                    }),
+                    contract_state_leaf.into(),
+                )
+            })
+            .collect();
+
+        Self {
+            classes_trie_proof,
+            contracts_trie_proof: ContractsTrieProof {
+                nodes: contracts_proof_nodes,
+                leaves: contract_leaves_data,
+            },
+            contracts_trie_storage_proofs,
+        }
+    }
+
     pub fn extend(&mut self, other: Self) {
         self.classes_trie_proof.extend(other.classes_trie_proof);
         self.contracts_trie_proof.nodes.extend(other.contracts_trie_proof.nodes);