starknet_committer: remove code dup for single trie

Author: nimrod-starkware

crates/starknet_committer/src/db/trie_traversal.rs

@@ -434,29 +434,17 @@ where
 {
     let mut storage_tries = HashMap::new();
     for (address, updates) in actual_storage_updates {
-        let sorted_leaf_indices = storage_tries_sorted_indices
-            .get(address)
-            .ok_or(ForestError::MissingSortedLeafIndices(*address))?;
-        let contract_state = original_contracts_trie_leaves
-            .get(&contract_address_into_node_index(address))
-            .ok_or(ForestError::MissingContractCurrentState(*address))?;
-        let config = OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(
-            config.warn_on_trivial_modifications(),
-        );
-
-        let original_skeleton = create_original_skeleton_tree::<Layout::DbLeaf, Layout>(
+        let (address, original_skeleton) = create_storage_trie_for_contract::<Layout>(
             storage,
-            contract_state.storage_root_hash,
-            *sorted_leaf_indices,
-            &config,
-            // TODO(Ariel): Change `LeafModifications` in `actual_storage_updates` to be an
-            // iterator over borrowed data so that the conversion below is costless.
-            &updates.iter().map(|(idx, value)| (*idx, Layout::DbLeaf::from(*value))).collect(),
-            None,
-            address,
+            *address,
+            updates,
+            original_contracts_trie_leaves,
+            storage_tries_sorted_indices,
+            config.warn_on_trivial_modifications(),
         )
         .await?;
-        storage_tries.insert(*address, original_skeleton);
+
+        storage_tries.insert(address, original_skeleton);
     }
     Ok(storage_tries)
 }
@@ -478,45 +466,70 @@ where
     let mut storage_tries = HashMap::new();
 
     for (address, updates) in actual_storage_updates {
-        // Extract data needed for this contract.
-        let sorted_leaf_indices = *storage_tries_sorted_indices
-            .get(address)
-            .ok_or(ForestError::MissingSortedLeafIndices(*address))?;
-        let contract_state = original_contracts_trie_leaves
-            .get(&contract_address_into_node_index(address))
-            .ok_or(ForestError::MissingContractCurrentState(*address))?;
-        let trie_config = OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(
-            config.warn_on_trivial_modifications(),
-        );
-
-        // TODO(Ariel): Change `LeafModifications` in `actual_storage_updates` to be an
-        // iterator over borrowed data so that the conversion below is costless.
-        let leaf_modifications: HashMap<
-            NodeIndex,
-            <Layout as NodeLayoutFor<StarknetStorageValue>>::DbLeaf,
-        > = updates.iter().map(|(idx, value)| (*idx, Layout::DbLeaf::from(*value))).collect();
-
         // Create the future - tokio will poll all futures concurrently.
         futures.push(async move {
-            let original_skeleton = create_original_skeleton_tree::<Layout::DbLeaf, Layout>(
+            create_storage_trie_for_contract::<Layout>(
                 storage,
-                contract_state.storage_root_hash,
-                sorted_leaf_indices,
-                &trie_config,
-                &leaf_modifications,
-                None,
-                address,
+                *address,
+                updates,
+                original_contracts_trie_leaves,
+                storage_tries_sorted_indices,
+                config.warn_on_trivial_modifications(),
             )
-            .await?;
-            Ok::<_, ForestError>((address, original_skeleton))
+            .await
         });
     }
 
     // Collect all results as they complete.
     while let Some(result) = futures.next().await {
         let (address, original_skeleton) = result?;
-        storage_tries.insert(*address, original_skeleton);
+        storage_tries.insert(address, original_skeleton);
     }
 
     Ok(storage_tries)
 }
+
+/// Helper function to create a storage trie for a single contract.
+async fn create_storage_trie_for_contract<'a, Layout: NodeLayoutFor<StarknetStorageValue>>(
+    storage: &impl Storage,
+    address: ContractAddress,
+    updates: &LeafModifications<StarknetStorageValue>,
+    original_contracts_trie_leaves: &HashMap<NodeIndex, ContractState>,
+    storage_tries_sorted_indices: &HashMap<ContractAddress, SortedLeafIndices<'a>>,
+    warn_on_trivial_modifications: bool,
+) -> ForestResult<(ContractAddress, OriginalSkeletonTreeImpl<'a>)>
+where
+    <Layout as NodeLayoutFor<StarknetStorageValue>>::DbLeaf:
+        HasStaticPrefix<KeyContext = ContractAddress>,
+{
+    // Extract data needed for this contract.
+    let sorted_leaf_indices = *storage_tries_sorted_indices
+        .get(&address)
+        .ok_or(ForestError::MissingSortedLeafIndices(address))?;
+    let contract_state = original_contracts_trie_leaves
+        .get(&contract_address_into_node_index(&address))
+        .ok_or(ForestError::MissingContractCurrentState(address))?;
+
+    let trie_config =
+        OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(warn_on_trivial_modifications);
+
+    // TODO(Ariel): Change `LeafModifications` in `actual_storage_updates` to be an
+    // iterator over borrowed data so that the conversion below is costless.
+    let leaf_modifications: HashMap<
+        NodeIndex,
+        <Layout as NodeLayoutFor<StarknetStorageValue>>::DbLeaf,
+    > = updates.iter().map(|(idx, value)| (*idx, Layout::DbLeaf::from(*value))).collect();
+
+    let original_skeleton = create_original_skeleton_tree::<Layout::DbLeaf, Layout>(
+        storage,
+        contract_state.storage_root_hash,
+        sorted_leaf_indices,
+        &trie_config,
+        &leaf_modifications,
+        None,
+        &address,
+    )
+    .await?;
+
+    Ok((address, original_skeleton))
+}