From 8cd1275d6732fb2a23ab41e5ecd39b87acf6033a Mon Sep 17 00:00:00 2001 From: Nimrod Weiss Date: Tue, 9 Jul 2024 12:25:31 +0300 Subject: [PATCH] refactor: sort modifications once --- .../external_test_utils.rs | 5 +- .../filled_tree/tree_test.rs | 9 ++-- .../create_tree_test.rs | 9 ++++ .../original_skeleton_tree/skeleton_forest.rs | 54 ++++++++++++------- .../skeleton_forest_test.rs | 23 ++++++-- .../src/patricia_merkle_tree/types.rs | 5 ++ .../create_tree_helper.rs | 5 +- .../create_tree_helper_test.rs | 2 +- .../updated_skeleton_tree/skeleton_forest.rs | 11 ++-- .../updated_skeleton_tree/tree.rs | 4 +- .../updated_skeleton_tree/tree_test.rs | 13 +++-- 11 files changed, 99 insertions(+), 41 deletions(-) diff --git a/crates/committer/src/patricia_merkle_tree/external_test_utils.rs b/crates/committer/src/patricia_merkle_tree/external_test_utils.rs index fbc35ae6..29c7ee68 100644 --- a/crates/committer/src/patricia_merkle_tree/external_test_utils.rs +++ b/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 = leaf_modifications.keys().copied().collect(); - let sorted_leaf_indices = SortedLeafIndices::new(&mut sorted_leaf_indices); + let mut leaf_indices: Vec = 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"); diff --git a/crates/committer/src/patricia_merkle_tree/filled_tree/tree_test.rs b/crates/committer/src/patricia_merkle_tree/filled_tree/tree_test.rs index 01326569..d2795fa1 100644 --- a/crates/committer/src/patricia_merkle_tree/filled_tree/tree_test.rs +++ b/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))]); diff --git a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/create_tree_test.rs b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/create_tree_test.rs index 82f5ae0f..606f9c1f 100644 --- a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/create_tree_test.rs +++ b/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 { + let mut leaf_indices: Vec = indices + .iter() + .map(|idx| NodeIndex::FIRST_LEAF + NodeIndex::from(*idx)) + .collect(); + leaf_indices.sort(); + leaf_indices +} diff --git a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest.rs b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest.rs index 17e60235..3b0c584f 100644 --- a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest.rs +++ b/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, pub(crate) contracts_trie: OriginalSkeletonTreeImpl, - pub(crate) storage_tries: HashMap, + pub(crate) contracts_trie_sorted_indices: Vec, + pub(crate) storage_tries: HashMap)>, } impl OriginalSkeletonForest { @@ -45,7 +47,7 @@ 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(), @@ -53,7 +55,7 @@ impl OriginalSkeletonForest { &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,17 +80,24 @@ impl OriginalSkeletonForest { accessed_addresses: &HashSet<&ContractAddress>, contracts_trie_root_hash: HashOutput, storage: &impl Storage, - ) -> ForestResult<(OriginalSkeletonTreeImpl, HashMap)> { - let mut sorted_leaf_indices: Vec = accessed_addresses + ) -> ForestResult<( + OriginalSkeletonTreeImpl, + HashMap, + Vec, + )> { + let mut leaf_indices: Vec = 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( @@ -94,7 +105,7 @@ impl OriginalSkeletonForest { original_contracts_trie_leaves: &HashMap, storage: &impl Storage, config: &impl Config, - ) -> ForestResult> { + ) -> ForestResult)>> { let mut storage_tries = HashMap::new(); for (address, updates) in actual_storage_updates { let mut sorted_leaf_indices: Vec = 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,7 +133,7 @@ impl OriginalSkeletonForest { classes_trie_root_hash: HashOutput, storage: &impl Storage, config: &impl Config, - ) -> ForestResult { + ) -> ForestResult<(OriginalSkeletonTreeImpl, Vec)> { let config = OriginalSkeletonClassesTrieConfig::new( actual_classes_updates, config.warn_on_trivial_modifications(), @@ -130,11 +141,14 @@ impl OriginalSkeletonForest { let mut sorted_leaf_indices: Vec = 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, + )) } } diff --git a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest_test.rs b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest_test.rs index b3f73a0c..a984f9c1 100644 --- a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/skeleton_forest_test.rs +++ b/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), diff --git a/crates/committer/src/patricia_merkle_tree/types.rs b/crates/committer/src/patricia_merkle_tree/types.rs index b9b0ea1e..62c27650 100644 --- a/crates/committer/src/patricia_merkle_tree/types.rs +++ b/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]) diff --git a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/create_tree_helper.rs b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/create_tree_helper.rs index 1584e3b2..92e08794 100644 --- a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/create_tree_helper.rs +++ b/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, + sorted_leaf_indices: &[NodeIndex], ) -> TempSkeletonNode { - let mut leaf_indices: Vec = 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 { diff --git a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/create_tree_helper_test.rs b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/create_tree_helper_test.rs index 5eb63c15..d59dec7b 100644 --- a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/create_tree_helper_test.rs +++ b/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::(updated, Arc::new(empty_map)) .await .unwrap(); diff --git a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/skeleton_forest.rs b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/skeleton_forest.rs index 59917678..fa162c9d 100644 --- a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/skeleton_forest.rs +++ b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/skeleton_forest.rs @@ -34,6 +34,7 @@ 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. @@ -41,13 +42,16 @@ impl UpdatedSkeletonForest { 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 { diff --git a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs index a0c5c44f..fedb0e0d 100644 --- a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs +++ b/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, + sorted_leaf_indices: &[NodeIndex], ) -> UpdatedSkeletonTreeResult; /// 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, + sorted_leaf_indices: &[NodeIndex], ) -> UpdatedSkeletonTreeResult { 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()), diff --git a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree_test.rs b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree_test.rs index 183fbefd..31dc6038 100644 --- a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree_test.rs +++ b/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 = 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