starknet_committer: define fetch_contract_storage_paths_concurrently

Author: ArielElp

crates/starknet_committer/src/db/trie_traversal.rs

@@ -45,6 +45,7 @@ use tracing::warn;
 
 use crate::block_committer::input::{
     contract_address_into_node_index,
+    try_node_index_into_contract_address,
     ReaderConfig,
     StarknetStorageValue,
 };
@@ -851,7 +852,6 @@ where
 }
 
 /// Storage task for fetching Patricia paths in a single storage trie.
-#[expect(dead_code)]
 struct StoragePathsReadTask<'indices, Layout: NodeLayoutFor<StarknetStorageValue>> {
     address: ContractAddress,
     storage_root_hash: HashOutput,
@@ -925,6 +925,127 @@ where
     storage.gather(tasks).await.into_iter().collect()
 }
 
+/// 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<'_>>,
+    contract_leaves: &HashMap<NodeIndex, ContractLeaf>,
+) -> TraversalResult<HashMap<ContractAddress, PreimageMap>>
+where
+    StorageLayout: NodeLayoutFor<StarknetStorageValue> + Send + 'static,
+    <StorageLayout as NodeLayoutFor<StarknetStorageValue>>::DbLeaf:
+        HasStaticPrefix<KeyContext = ContractAddress>,
+    ContractLeaf: AsRef<ContractState>,
+{
+    if let Some(gatherable_storage) = storage.as_gatherable_storage() {
+        return fetch_contract_storage_paths_concurrently::<_, StorageLayout, ContractLeaf>(
+            gatherable_storage,
+            contract_storage_sorted_leaf_indices,
+            contract_leaves,
+        )
+        .await;
+    }
+    fetch_contract_storage_paths_sequentially::<StorageLayout, ContractLeaf>(
+        storage,
+        contract_storage_sorted_leaf_indices,
+        contract_leaves,
+    )
+    .await
+}
+
+/// Returns the contract address and storage root hash for the given leaf index, if the contract
+/// exists in the contracts trie.
+///
+/// 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>>(
+    idx: &NodeIndex,
+    contract_leaves: &HashMap<NodeIndex, ContractLeaf>,
+) -> Option<(ContractAddress, HashOutput)> {
+    let contract_address = try_node_index_into_contract_address(idx).unwrap_or_else(|_| {
+        panic!(
+            "Converting leaf NodeIndex to ContractAddress should succeed; failed to convert \
+             {idx:?}."
+        )
+    });
+    let storage_root_hash = contract_leaves.get(idx).map(|leaf| leaf.as_ref().storage_root_hash)?;
+    Some((contract_address, storage_root_hash))
+}
+
+/// 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<'_>>,
+    contract_leaves: &HashMap<NodeIndex, ContractLeaf>,
+) -> TraversalResult<HashMap<ContractAddress, PreimageMap>>
+where
+    StorageLayout: NodeLayoutFor<StarknetStorageValue> + Send + 'static,
+    <StorageLayout as NodeLayoutFor<StarknetStorageValue>>::DbLeaf:
+        HasStaticPrefix<KeyContext = ContractAddress>,
+    ContractLeaf: AsRef<ContractState>,
+{
+    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, contract_leaves)
+        else {
+            continue;
+        };
+
+        let leaves = None;
+        let proof = fetch_patricia_paths::<StorageLayout::DbLeaf, StorageLayout>(
+            storage,
+            storage_root_hash,
+            *sorted_leaf_indices,
+            leaves,
+            &contract_address,
+        )
+        .await?;
+        contracts_trie_storage_proofs.insert(contract_address, proof);
+    }
+
+    Ok(contracts_trie_storage_proofs)
+}
+
+/// 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<'_>>,
+    contract_leaves: &HashMap<NodeIndex, ContractLeaf>,
+) -> TraversalResult<HashMap<ContractAddress, PreimageMap>>
+where
+    S: GatherableStorage,
+    StorageLayout: NodeLayoutFor<StarknetStorageValue> + Send + 'static,
+    <StorageLayout as NodeLayoutFor<StarknetStorageValue>>::DbLeaf:
+        HasStaticPrefix<KeyContext = ContractAddress>,
+    ContractLeaf: AsRef<ContractState>,
+{
+    let mut tasks = Vec::new();
+    for (idx, sorted_leaf_indices) in contract_storage_sorted_leaf_indices {
+        let Some((contract_address, storage_root_hash)) =
+            get_address_and_storage_root(idx, contract_leaves)
+        else {
+            continue;
+        };
+        tasks.push(StoragePathsReadTask::<StorageLayout> {
+            address: contract_address,
+            storage_root_hash,
+            sorted_leaf_indices: *sorted_leaf_indices,
+            _layout: PhantomData,
+        });
+    }
+    storage.gather(tasks).await.into_iter().collect()
+}
+
 /// Helper function to create a storage trie for a single contract.
 async fn create_storage_trie<'a, Layout: NodeLayoutFor<StarknetStorageValue>>(
     storage: &mut impl ReadOnlyStorage,

crates/starknet_committer/src/patricia_merkle_tree/tree.rs

@@ -17,7 +17,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;
+use crate::db::trie_traversal::{fetch_patricia_paths, get_address_and_storage_root};
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 use crate::patricia_merkle_tree::types::{
     class_hash_into_node_index,
@@ -181,19 +181,8 @@ where
         HashMap::with_capacity(contract_storage_sorted_leaf_indices.len());
 
     for (idx, sorted_leaf_indices) in contract_storage_sorted_leaf_indices {
-        let contract_address = try_node_index_into_contract_address(idx).unwrap_or_else(|_| {
-            panic!(
-                "Converting leaf NodeIndex to ContractAddress should succeed; failed to convert \
-                 {idx:?}."
-            )
-        });
-
-        // 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.
-        let Some(storage_root_hash) = leaves.get(idx).map(|leaf| leaf.as_ref().storage_root_hash)
+        let Some((contract_address, storage_root_hash)) =
+            get_address_and_storage_root(idx, &leaves)
         else {
             continue;
         };