starknet_committer: add utility that returns sorted leaf indices

Author: ArielElp

crates/starknet_committer/src/patricia_merkle_tree/tree.rs

@@ -4,7 +4,8 @@ use starknet_api::core::{ClassHash, ContractAddress};
 use starknet_api::hash::HashOutput;
 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, SortedLeafIndices};
+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;
 
@@ -50,10 +51,64 @@ impl OriginalSkeletonTreeConfig for OriginalSkeletonTrieConfig {
     }
 }
 
+/// Requested trie leaves for Patricia witness collection (classes trie, contracts trie, and
+/// per-contract storage leaves). Built via [`LeavesRequest::from`].
+#[derive(Clone)]
+pub struct LeavesRequest {
+    pub class_leaf_indices: Vec<NodeIndex>,
+    pub contract_leaf_indices: Vec<NodeIndex>,
+    pub contract_storage_leaf_indices: HashMap<NodeIndex, Vec<NodeIndex>>,
+}
+
+pub struct SortedLeavesRequest<'a> {
+    pub class_sorted: SortedLeafIndices<'a>,
+    pub contract_sorted: SortedLeafIndices<'a>,
+    pub storage_sorted: HashMap<NodeIndex, SortedLeafIndices<'a>>,
+}
+
+impl LeavesRequest {
+    /// Builds index buffers expected by [`fetch_all_patricia_paths`].
+    pub fn from(
+        class_hashes: &[ClassHash],
+        contract_addresses: &[ContractAddress],
+        contract_storage_keys: &HashMap<ContractAddress, Vec<StarknetStorageKey>>,
+    ) -> Self {
+        let contract_leaf_indices: Vec<NodeIndex> =
+            contract_addresses.iter().map(contract_address_into_node_index).collect();
+        let contract_storage_leaf_indices: HashMap<NodeIndex, Vec<NodeIndex>> =
+            contract_storage_keys
+                .iter()
+                .map(|(address, keys)| {
+                    let node_index = contract_address_into_node_index(address);
+                    let leaf_indices: Vec<_> = keys.iter().map(NodeIndex::from).collect();
+                    (node_index, leaf_indices)
+                })
+                .collect();
+        Self {
+            class_leaf_indices: class_hashes.iter().map(class_hash_into_node_index).collect(),
+            contract_leaf_indices,
+            contract_storage_leaf_indices,
+        }
+    }
+
+    /// Sorts class, contract, and per-contract storage leaf index buffers in place, then returns
+    /// a single [`SortedLeavesRequest`].
+    pub fn get_sorted(&mut self) -> SortedLeavesRequest<'_> {
+        let class_sorted = SortedLeafIndices::new(&mut self.class_leaf_indices);
+        let contract_sorted = SortedLeafIndices::new(&mut self.contract_leaf_indices);
+        let storage_sorted: HashMap<_, _> = self
+            .contract_storage_leaf_indices
+            .iter_mut()
+            .map(|(address, leaf_indices)| (*address, SortedLeafIndices::new(leaf_indices)))
+            .collect();
+        SortedLeavesRequest { class_sorted, contract_sorted, storage_sorted }
+    }
+}
+
 /// Fetch all tries patricia paths given the modified leaves.
 /// Fetch the leaves in the contracts trie only, to be able to get the storage root hashes.
-/// Assumption: `contract_sorted_leaf_indices` contains all `contract_storage_sorted_leaf_indices`
-/// keys.
+/// Assumption: `contract_sorted_leaf_indices` lists every contract that appears in
+/// `contract_storage_sorted_leaf_indices`.
 pub async fn fetch_all_patricia_paths<Layout>(
     storage: &mut impl ReadOnlyStorage,
     classes_trie_root_hash: HashOutput,
@@ -78,7 +133,8 @@ where
         address_counter,
         contract_storage_sorted_leaf_indices.len(),
         "contract_sorted_leaf_indices is missing an address with requested storage witnesses. \
-         contract_sorted_leaf_indices: {contract_sorted_leaf_indices:?}, storage addresses: {:?}",
+         contract_sorted_leaf_indices: {:?}, storage addresses: {:?}",
+        contract_sorted_leaf_indices,
         contract_storage_sorted_leaf_indices.keys()
     );
 
@@ -179,43 +235,27 @@ pub async fn fetch_previous_and_new_patricia_paths(
     contract_addresses: &[ContractAddress],
     contract_storage_keys: &HashMap<ContractAddress, Vec<StarknetStorageKey>>,
 ) -> TraversalResult<StarknetForestProofs> {
-    let mut class_leaf_indices: Vec<NodeIndex> =
-        class_hashes.iter().map(class_hash_into_node_index).collect();
-    let class_sorted_leaf_indices = SortedLeafIndices::new(&mut class_leaf_indices);
-
-    let mut contract_leaf_indices: Vec<NodeIndex> =
-        contract_addresses.iter().map(contract_address_into_node_index).collect();
-    let contract_sorted_leaf_indices = SortedLeafIndices::new(&mut contract_leaf_indices);
-
-    let mut contract_storage_leaf_indices: HashMap<NodeIndex, Vec<NodeIndex>> =
-        contract_storage_keys
-            .iter()
-            .map(|(address, keys)| {
-                let node_index = contract_address_into_node_index(address);
-                let leaf_indices: Vec<_> = keys.iter().map(NodeIndex::from).collect();
-                (node_index, leaf_indices)
-            })
-            .collect();
-    let contract_storage_sorted_leaf_indices = &contract_storage_leaf_indices
-        .iter_mut()
-        .map(|(address, leaf_indices)| (*address, SortedLeafIndices::new(leaf_indices)))
-        .collect();
+    let mut leaves_request =
+        LeavesRequest::from(class_hashes, contract_addresses, contract_storage_keys);
+
+    let sorted_leaves = leaves_request.get_sorted();
     let prev_proofs = fetch_all_patricia_paths::<FactsNodeLayout>(
         storage,
         classes_trie_root_hashes.previous_root_hash,
         contracts_trie_root_hashes.previous_root_hash,
-        class_sorted_leaf_indices,
-        contract_sorted_leaf_indices,
-        contract_storage_sorted_leaf_indices,
+        sorted_leaves.class_sorted,
+        sorted_leaves.contract_sorted,
+        &sorted_leaves.storage_sorted,
     )
     .await?;
+
     let new_proofs = fetch_all_patricia_paths::<FactsNodeLayout>(
         storage,
         classes_trie_root_hashes.new_root_hash,
         contracts_trie_root_hashes.new_root_hash,
-        class_sorted_leaf_indices,
-        contract_sorted_leaf_indices,
-        contract_storage_sorted_leaf_indices,
+        sorted_leaves.class_sorted,
+        sorted_leaves.contract_sorted,
+        &sorted_leaves.storage_sorted,
     )
     .await?;