refactor: sort modifications once

crates/committer/src/patricia_merkle_tree/external_test_utils.rs

@@ -73,8 +73,8 @@ pub async fn tree_computation_flow(
     root_hash: HashOutput,
 ) -> StorageTrie {
     let config = OriginalSkeletonStorageTrieConfig::new(&leaf_modifications, false);
-    let mut sorted_leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().copied().collect();
-    let sorted_leaf_indices = SortedLeafIndices::new(&mut sorted_leaf_indices);
+    let mut leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().copied().collect();
+    let sorted_leaf_indices = SortedLeafIndices::new(&mut leaf_indices);
     let mut original_skeleton: OriginalSkeletonTreeImpl =
         OriginalSkeletonTree::create(storage, root_hash, sorted_leaf_indices, &config)
             .expect("Failed to create the original skeleton tree");
@@ -93,6 +93,7 @@ pub async fn tree_computation_flow(
                 )
             })
             .collect(),
+        &leaf_indices,
     )
     .expect("Failed to create the updated skeleton tree");
 

crates/committer/src/patricia_merkle_tree/filled_tree/tree_test.rs

@@ -263,9 +263,12 @@ async fn test_delete_leaf_from_empty_tree() {
     // Create an updated skeleton tree with a single leaf that is deleted.
     let skeleton_modifications = HashMap::from([(NodeIndex::FIRST_LEAF, SkeletonLeaf::Zero)]);
 
-    let updated_skeleton_tree =
-        UpdatedSkeletonTreeImpl::create(&mut original_skeleton_tree, &skeleton_modifications)
-            .unwrap();
+    let updated_skeleton_tree = UpdatedSkeletonTreeImpl::create(
+        &mut original_skeleton_tree,
+        &skeleton_modifications,
+        &[NodeIndex::FIRST_LEAF],
+    )
+    .unwrap();
 
     let leaf_modifications =
         HashMap::from([(NodeIndex::FIRST_LEAF, StarknetStorageValue(Felt::ZERO))]);

crates/committer/src/patricia_merkle_tree/original_skeleton_tree/create_tree_test.rs

@@ -567,3 +567,12 @@ fn create_previously_empty_leaf_indices<'a>(
         .filter(|idx| !subtree_leaf_indices.contains(idx))
         .collect()
 }
+
+pub(crate) fn create_expected_sorted_leaf_indices(indices: &[u128]) -> Vec<NodeIndex> {
+    let mut leaf_indices: Vec<NodeIndex> = indices
+        .iter()
+        .map(|idx| NodeIndex::FIRST_LEAF + NodeIndex::from(*idx))
+        .collect();
+    leaf_indices.sort();
+    leaf_indices
+}

crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest.rs

@@ -26,8 +26,10 @@ pub mod skeleton_forest_test;
 #[derive(Debug, Eq, PartialEq)]
 pub(crate) struct OriginalSkeletonForest {
     pub(crate) classes_trie: OriginalSkeletonTreeImpl,
+    pub(crate) classes_trie_sorted_indices: Vec<NodeIndex>,
     pub(crate) contracts_trie: OriginalSkeletonTreeImpl,
-    pub(crate) storage_tries: HashMap<ContractAddress, OriginalSkeletonTreeImpl>,
+    pub(crate) contracts_trie_sorted_indices: Vec<NodeIndex>,
+    pub(crate) storage_tries: HashMap<ContractAddress, (OriginalSkeletonTreeImpl, Vec<NodeIndex>)>,
 }
 
 impl OriginalSkeletonForest {
@@ -45,15 +47,15 @@ impl OriginalSkeletonForest {
         Self: std::marker::Sized,
     {
         let accessed_addresses = state_diff.accessed_addresses();
-        let (contracts_trie, original_contracts_trie_leaves) =
+        let (contracts_trie, original_contracts_trie_leaves, classes_trie_sorted_indices) =
             Self::create_contracts_trie(&accessed_addresses, contracts_trie_root_hash, &storage)?;
         let storage_tries = Self::create_storage_tries(
             &state_diff.actual_storage_updates(),
             &original_contracts_trie_leaves,
             &storage,
             config,
         )?;
-        let classes_trie = Self::create_classes_trie(
+        let (classes_trie, contracts_trie_sorted_indices) = Self::create_classes_trie(
             &state_diff.actual_classes_updates(),
             classes_trie_root_hash,
             &storage,
@@ -63,7 +65,9 @@ impl OriginalSkeletonForest {
         Ok((
             Self {
                 classes_trie,
+                classes_trie_sorted_indices,
                 contracts_trie,
+                contracts_trie_sorted_indices,
                 storage_tries,
             },
             original_contracts_trie_leaves,
@@ -76,25 +80,32 @@ impl OriginalSkeletonForest {
         accessed_addresses: &HashSet<&ContractAddress>,
         contracts_trie_root_hash: HashOutput,
         storage: &impl Storage,
-    ) -> ForestResult<(OriginalSkeletonTreeImpl, HashMap<NodeIndex, ContractState>)> {
-        let mut sorted_leaf_indices: Vec<NodeIndex> = accessed_addresses
+    ) -> ForestResult<(
+        OriginalSkeletonTreeImpl,
+        HashMap<NodeIndex, ContractState>,
+        Vec<NodeIndex>,
+    )> {
+        let mut leaf_indices: Vec<NodeIndex> = accessed_addresses
             .iter()
             .map(|address| NodeIndex::from_contract_address(address))
             .collect();
-        Ok(OriginalSkeletonTreeImpl::create_and_get_previous_leaves(
-            storage,
-            contracts_trie_root_hash,
-            SortedLeafIndices::new(&mut sorted_leaf_indices),
-            &OriginalSkeletonContractsTrieConfig::new(),
-        )?)
+        let sorted_leaf_indices = SortedLeafIndices::new(&mut leaf_indices);
+        let (contracts_trie, previous_leaves) =
+            OriginalSkeletonTreeImpl::create_and_get_previous_leaves(
+                storage,
+                contracts_trie_root_hash,
+                sorted_leaf_indices,
+                &OriginalSkeletonContractsTrieConfig::new(),
+            )?;
+        Ok((contracts_trie, previous_leaves, leaf_indices))
     }
 
     fn create_storage_tries(
         actual_storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
         original_contracts_trie_leaves: &HashMap<NodeIndex, ContractState>,
         storage: &impl Storage,
         config: &impl Config,
-    ) -> ForestResult<HashMap<ContractAddress, OriginalSkeletonTreeImpl>> {
+    ) -> ForestResult<HashMap<ContractAddress, (OriginalSkeletonTreeImpl, Vec<NodeIndex>)>> {
         let mut storage_tries = HashMap::new();
         for (address, updates) in actual_storage_updates {
             let mut sorted_leaf_indices: Vec<NodeIndex> = updates.keys().copied().collect();
@@ -112,7 +123,7 @@ impl OriginalSkeletonForest {
                 SortedLeafIndices::new(&mut sorted_leaf_indices),
                 &config,
             )?;
-            storage_tries.insert(*address, original_skeleton);
+            storage_tries.insert(*address, (original_skeleton, sorted_leaf_indices));
         }
         Ok(storage_tries)
     }
@@ -122,19 +133,22 @@ impl OriginalSkeletonForest {
         classes_trie_root_hash: HashOutput,
         storage: &impl Storage,
         config: &impl Config,
-    ) -> ForestResult<OriginalSkeletonTreeImpl> {
+    ) -> ForestResult<(OriginalSkeletonTreeImpl, Vec<NodeIndex>)> {
         let config = OriginalSkeletonClassesTrieConfig::new(
             actual_classes_updates,
             config.warn_on_trivial_modifications(),
         );
         let mut sorted_leaf_indices: Vec<NodeIndex> =
             actual_classes_updates.keys().copied().collect();
 
-        Ok(OriginalSkeletonTreeImpl::create(
-            storage,
-            classes_trie_root_hash,
-            SortedLeafIndices::new(&mut sorted_leaf_indices),
-            &config,
-        )?)
+        Ok((
+            OriginalSkeletonTreeImpl::create(
+                storage,
+                classes_trie_root_hash,
+                SortedLeafIndices::new(&mut sorted_leaf_indices),
+                &config,
+            )?,
+            sorted_leaf_indices,
+        ))
     }
 }

crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest_test.rs

@@ -12,7 +12,8 @@ use crate::patricia_merkle_tree::filled_tree::node::{ClassHash, CompiledClassHas
 use crate::patricia_merkle_tree::node_data::leaf::ContractState;
 use crate::patricia_merkle_tree::original_skeleton_tree::create_tree::create_tree_test::{
     create_32_bytes_entry, create_binary_entry, create_binary_skeleton_node, create_edge_entry,
-    create_edge_skeleton_node, create_expected_skeleton, create_unmodified_subtree_skeleton_node,
+    create_edge_skeleton_node, create_expected_skeleton, create_expected_sorted_leaf_indices,
+    create_unmodified_subtree_skeleton_node,
 };
 use crate::patricia_merkle_tree::original_skeleton_tree::create_tree::create_tree_test::{
     create_compiled_class_leaf_entry, create_contract_state_leaf_entry, create_root_edge_entry,
@@ -162,6 +163,7 @@ use crate::storage::map_storage::MapStorage;
             ],
             3
         ),
+        classes_trie_sorted_indices: create_expected_sorted_leaf_indices(&[0, 6, 7]),
         contracts_trie: create_expected_skeleton(
             vec![
                 create_binary_skeleton_node(1),
@@ -173,10 +175,12 @@ use crate::storage::map_storage::MapStorage;
             ],
             3
         ),
+        contracts_trie_sorted_indices: create_expected_sorted_leaf_indices(&[7, 6, 0]),
         storage_tries: HashMap::from([
             (
                 ContractAddress(Felt::from(0_u128)),
-                create_expected_skeleton(
+                (
+                    create_expected_skeleton(
                     vec![
                         create_binary_skeleton_node(1),
                         create_binary_skeleton_node(2),
@@ -188,11 +192,14 @@ use crate::storage::map_storage::MapStorage;
                         create_unmodified_subtree_skeleton_node(11, 16),
                     ],
                     3
+                    ),
+                    create_expected_sorted_leaf_indices(&[2, 5, 0])
                 )
             ),
             (
                 ContractAddress(Felt::from(6_u128)),
-                create_expected_skeleton(
+                (
+                    create_expected_skeleton(
                     vec![
                         create_binary_skeleton_node(1),
                         create_edge_skeleton_node(2, 0, 1),
@@ -203,11 +210,14 @@ use crate::storage::map_storage::MapStorage;
                         create_unmodified_subtree_skeleton_node(9, 2),
                     ],
                     3
+                    ),
+                    create_expected_sorted_leaf_indices(&[0, 5, 3])
                 )
             ),
             (
                 ContractAddress(Felt::from(7_u128)),
-                create_expected_skeleton(
+                (
+                    create_expected_skeleton(
                     vec![
                         create_binary_skeleton_node(1),
                         create_edge_skeleton_node(2, 0, 1),
@@ -218,9 +228,12 @@ use crate::storage::map_storage::MapStorage;
                         create_unmodified_subtree_skeleton_node(9, 2),
                     ],
                     3
+                    ),
+                    create_expected_sorted_leaf_indices(&[0, 5, 3])
                 )
-            )
+            ),
             ]),
+
         },
         create_contract_leaves(&[
             (7, 29 + 248),

crates/committer/src/patricia_merkle_tree/types.rs

@@ -235,6 +235,11 @@ impl<'a> SortedLeafIndices<'a> {
         Self(indices)
     }
 
+    /// Creates a new instance from the given indices. Assumes indices are sorted.
+    pub(crate) fn from_sorted(sorted_indices: &'a [NodeIndex]) -> Self {
+        Self(sorted_indices)
+    }
+
     /// Returns a subslice of the indices stored at self, at the range [leftmost_idx, rightmost_idx).
     pub(crate) fn subslice(&self, leftmost_idx: usize, rightmost_idx: usize) -> Self {
         Self(&self.0[leftmost_idx..rightmost_idx])

crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/create_tree_helper.rs

@@ -98,10 +98,9 @@ impl UpdatedSkeletonTreeImpl {
     pub(crate) fn finalize_middle_layers(
         &mut self,
         original_skeleton: &mut impl OriginalSkeletonTree,
-        leaf_modifications: &LeafModifications<SkeletonLeaf>,
+        sorted_leaf_indices: &[NodeIndex],
     ) -> TempSkeletonNode {
-        let mut leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().cloned().collect();
-        let sorted_leaf_indices = SortedLeafIndices::new(&mut leaf_indices);
+        let sorted_leaf_indices = SortedLeafIndices::from_sorted(sorted_leaf_indices);
         if original_skeleton.get_nodes().is_empty() {
             self.update_node_in_empty_tree(&NodeIndex::ROOT, &sorted_leaf_indices)
         } else {

crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/create_tree_helper_test.rs

@@ -511,7 +511,7 @@ async fn test_update_non_modified_storage_tree(#[case] root_hash: HashOutput) {
     )
     .unwrap();
     let updated =
-        UpdatedSkeletonTreeImpl::create(&mut original_skeleton_tree, &HashMap::new()).unwrap();
+        UpdatedSkeletonTreeImpl::create(&mut original_skeleton_tree, &HashMap::new(), &[]).unwrap();
     let filled = StorageTrie::create::<TreeHashFunctionImpl>(updated, Arc::new(empty_map))
         .await
         .unwrap();

crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/skeleton_forest.rs

@@ -34,20 +34,24 @@ impl UpdatedSkeletonForest {
         let classes_trie = UpdatedSkeletonTreeImpl::create(
             &mut original_skeleton_forest.classes_trie,
             class_hash_leaf_modifications,
+            &original_skeleton_forest.classes_trie_sorted_indices,
         )?;
 
         // Storage tries.
         let mut contracts_trie_leaves = HashMap::new();
         let mut storage_tries = HashMap::new();
 
         for (address, updates) in storage_updates {
-            let original_storage_trie = original_skeleton_forest
+            let (original_storage_trie, sorted_leaf_indices) = original_skeleton_forest
                 .storage_tries
                 .get_mut(address)
                 .ok_or(ForestError::MissingOriginalSkeleton(*address))?;
 
-            let updated_storage_trie =
-                UpdatedSkeletonTreeImpl::create(original_storage_trie, updates)?;
+            let updated_storage_trie = UpdatedSkeletonTreeImpl::create(
+                original_storage_trie,
+                updates,
+                sorted_leaf_indices,
+            )?;
             let storage_trie_becomes_empty = updated_storage_trie.is_empty();
 
             storage_tries.insert(*address, updated_storage_trie);
@@ -69,6 +73,7 @@ impl UpdatedSkeletonForest {
         let contracts_trie = UpdatedSkeletonTreeImpl::create(
             &mut original_skeleton_forest.contracts_trie,
             &contracts_trie_leaves,
+            &original_skeleton_forest.contracts_trie_sorted_indices,
         )?;
 
         Ok(Self {

crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs

@@ -25,6 +25,7 @@ pub(crate) trait UpdatedSkeletonTree: Sized + Send + Sync {
     fn create(
         original_skeleton: &mut impl OriginalSkeletonTree,
         leaf_modifications: &LeafModifications<SkeletonLeaf>,
+        sorted_leaf_indices: &[NodeIndex],
     ) -> UpdatedSkeletonTreeResult<Self>;
 
     /// Does the skeleton represents an empty-tree (i.e. all leaves are empty).
@@ -45,6 +46,7 @@ impl UpdatedSkeletonTree for UpdatedSkeletonTreeImpl {
     fn create(
         original_skeleton: &mut impl OriginalSkeletonTree,
         leaf_modifications: &LeafModifications<SkeletonLeaf>,
+        sorted_leaf_indices: &[NodeIndex],
     ) -> UpdatedSkeletonTreeResult<Self> {
         if leaf_modifications.is_empty() {
             return Self::create_unmodified(original_skeleton);
@@ -54,7 +56,7 @@ impl UpdatedSkeletonTree for UpdatedSkeletonTreeImpl {
         let mut updated_skeleton_tree = UpdatedSkeletonTreeImpl { skeleton_tree };
 
         let temp_root_node =
-            updated_skeleton_tree.finalize_middle_layers(original_skeleton, leaf_modifications);
+            updated_skeleton_tree.finalize_middle_layers(original_skeleton, sorted_leaf_indices);
         // Finalize root.
         match temp_root_node {
             TempSkeletonNode::Empty => assert!(updated_skeleton_tree.skeleton_tree.is_empty()),

crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree_test.rs

@@ -131,8 +131,14 @@ fn test_updated_skeleton_tree_impl_create(
         .iter()
         .map(|(index, val)| (*index, (*val).into()))
         .collect();
-    let updated_skeleton_tree =
-        UpdatedSkeletonTreeImpl::create(&mut original_skeleton, &leaf_modifications).unwrap();
+    let mut sorted_leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().copied().collect();
+    sorted_leaf_indices.sort();
+    let updated_skeleton_tree = UpdatedSkeletonTreeImpl::create(
+        &mut original_skeleton,
+        &leaf_modifications,
+        &sorted_leaf_indices,
+    )
+    .unwrap();
 
     let mut expected_skeleton_tree = initial_updated_skeleton.skeleton_tree.clone();
     expected_skeleton_tree.extend(expected_skeleton_additions.iter().cloned());
@@ -159,6 +165,7 @@ fn test_updated_empty_tree(#[case] modifications: LeafModifications<StarknetStor
         .map(|(idx, leaf)| (idx, leaf.0.into()))
         .collect();
     let updated_skeleton_tree =
-        UpdatedSkeletonTreeImpl::create(&mut original_skeleton, &skeleton_modifications).unwrap();
+        UpdatedSkeletonTreeImpl::create(&mut original_skeleton, &skeleton_modifications, &indices)
+            .unwrap();
     assert!(updated_skeleton_tree.is_empty());
 }