From 7f4e7b3a05ddd3a9ddad265c62993f63489448a4 Mon Sep 17 00:00:00 2001 From: amosStarkware <88497213+amosStarkware@users.noreply.github.com> Date: Tue, 20 Aug 2024 15:45:50 +0300 Subject: [PATCH] feat(committer): use FilledTree::create() in filled forest (#495) --- .../src/block_committer/input.rs | 22 +++ .../src/block_committer/input_test.rs | 25 +++ .../src/forest/filled_forest.rs | 144 +++++++++++------- .../src/forest/forest_errors.rs | 5 +- .../src/patricia_merkle_tree/types.rs | 2 +- 5 files changed, 138 insertions(+), 60 deletions(-) create mode 100644 crates/starknet_committer/src/block_committer/input_test.rs diff --git a/crates/starknet_committer/src/block_committer/input.rs b/crates/starknet_committer/src/block_committer/input.rs index 74830e7806..62590b2465 100644 --- a/crates/starknet_committer/src/block_committer/input.rs +++ b/crates/starknet_committer/src/block_committer/input.rs @@ -10,10 +10,32 @@ use tracing::level_filters::LevelFilter; use crate::patricia_merkle_tree::types::{ClassHash, CompiledClassHash, Nonce}; +#[cfg(test)] +#[path = "input_test.rs"] +pub mod input_test; + #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] // TODO(Nimrod, 1/6/2025): Use the ContractAddress defined in starknet-types-core when available. pub struct ContractAddress(pub Felt); +impl TryFrom<&NodeIndex> for ContractAddress { + type Error = String; + + fn try_from(node_index: &NodeIndex) -> Result { + if !node_index.is_leaf() { + return Err("NodeIndex is not a leaf.".to_string()); + } + let result = Felt::try_from(*node_index - NodeIndex::FIRST_LEAF); + match result { + Ok(felt) => Ok(ContractAddress(felt)), + Err(error) => Err(format!( + "Tried to convert node index to felt and got the following error: {:?}", + error.to_string() + )), + } + } +} + impl From<&ContractAddress> for NodeIndex { fn from(address: &ContractAddress) -> NodeIndex { NodeIndex::from_leaf_felt(&address.0) diff --git a/crates/starknet_committer/src/block_committer/input_test.rs b/crates/starknet_committer/src/block_committer/input_test.rs new file mode 100644 index 0000000000..5838d8b018 --- /dev/null +++ b/crates/starknet_committer/src/block_committer/input_test.rs @@ -0,0 +1,25 @@ +use rstest::rstest; +use starknet_patricia::felt::Felt; +use starknet_patricia::patricia_merkle_tree::types::NodeIndex; + +use crate::block_committer::input::ContractAddress; + +#[rstest] +fn test_node_index_to_contract_address_conversion() { + // Positive flow. + assert_eq!(ContractAddress::try_from(&NodeIndex::FIRST_LEAF), Ok(ContractAddress(Felt::ZERO))); + assert_eq!( + ContractAddress::try_from(&(NodeIndex::FIRST_LEAF + NodeIndex(1_u32.into()))), + Ok(ContractAddress(Felt::ONE)) + ); + assert_eq!( + ContractAddress::try_from(&NodeIndex::MAX), + Ok(ContractAddress(Felt::try_from(NodeIndex::MAX - NodeIndex::FIRST_LEAF).unwrap())) + ); + + // Negative flow. + assert_eq!( + ContractAddress::try_from(&(NodeIndex::FIRST_LEAF - NodeIndex(1_u32.into()))), + Err("NodeIndex is not a leaf.".to_string()) + ); +} diff --git a/crates/starknet_committer/src/forest/filled_forest.rs b/crates/starknet_committer/src/forest/filled_forest.rs index a1e28e0132..6ee31f0123 100644 --- a/crates/starknet_committer/src/forest/filled_forest.rs +++ b/crates/starknet_committer/src/forest/filled_forest.rs @@ -2,11 +2,11 @@ use std::collections::HashMap; use starknet_patricia::hash::hash_trait::HashOutput; use starknet_patricia::patricia_merkle_tree::filled_tree::tree::FilledTree; -use starknet_patricia::patricia_merkle_tree::node_data::leaf::LeafModifications; +use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications}; use starknet_patricia::patricia_merkle_tree::types::NodeIndex; use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::tree::UpdatedSkeletonTreeImpl; use starknet_patricia::storage::storage_trait::Storage; -use tokio::task::JoinSet; +use tracing::info; use crate::block_committer::input::{ContractAddress, StarknetStorageValue}; use crate::forest::forest_errors::{ForestError, ForestResult}; @@ -19,7 +19,6 @@ use crate::patricia_merkle_tree::types::{ CompiledClassHash, ContractsTrie, Nonce, - StorageTrie, StorageTrieMap, }; @@ -52,8 +51,10 @@ impl FilledForest { self.classes_trie.get_root_hash() } + /// Creates a filled forest. Assumes the storage updates and the updated skeletons of the + /// storage tries include all modified contracts, including those with unmodified storage. pub(crate) async fn create( - mut updated_forest: UpdatedSkeletonForest, + updated_forest: UpdatedSkeletonForest, storage_updates: HashMap>, classes_updates: LeafModifications, original_contracts_trie_leaves: &HashMap, @@ -64,67 +65,94 @@ impl FilledForest { updated_forest.classes_trie, classes_updates, )); - let mut contracts_trie_modifications = HashMap::new(); - let mut filled_storage_tries = HashMap::new(); - let mut contracts_state_tasks = JoinSet::new(); - for (address, inner_updates) in storage_updates { - let updated_storage_trie = updated_forest - .storage_tries - .remove(&address) - .ok_or(ForestError::MissingUpdatedSkeleton(address))?; - - let original_contract_state = original_contracts_trie_leaves - .get(&((&address).into())) - .ok_or(ForestError::MissingContractCurrentState(address))?; - contracts_state_tasks.spawn(Self::new_contract_state::( - address, - *(address_to_nonce.get(&address).unwrap_or(&original_contract_state.nonce)), - *(address_to_class_hash - .get(&address) - .unwrap_or(&original_contract_state.class_hash)), - updated_storage_trie, - inner_updates, - )); - } - - while let Some(result) = contracts_state_tasks.join_next().await { - let (address, new_contract_state, filled_storage_trie) = result??; - contracts_trie_modifications.insert((&address).into(), new_contract_state); - filled_storage_tries.insert(address, filled_storage_trie); - } - - let contracts_trie_task = tokio::spawn(ContractsTrie::create_with_existing_leaves::( + let contracts_trie_task = tokio::task::spawn(ContractsTrie::create::( updated_forest.contracts_trie, - contracts_trie_modifications, + FilledForest::get_contracts_trie_leaf_input( + original_contracts_trie_leaves, + storage_updates, + updated_forest.storage_tries, + address_to_class_hash, + address_to_nonce, + )?, )); - let contracts_trie = contracts_trie_task.await?.map_err(ForestError::ContractsTrie)?; let classes_trie = classes_trie_task.await?.map_err(ForestError::ClassesTrie)?; + info!( + "Classes trie update complete; {:?} new facts computed.", + classes_trie.tree_map.len() + ); + let (contracts_trie, storage_tries) = + contracts_trie_task.await?.map_err(ForestError::ContractsTrie)?; + info!( + "Contracts trie update complete; {:?} new facts computed.", + contracts_trie.tree_map.len() + ); - Ok(Self { storage_tries: filled_storage_tries, contracts_trie, classes_trie }) + Ok(Self { + storage_tries: storage_tries + .into_iter() + .map(|(node_index, storage_trie)| { + ( + ContractAddress::try_from(&node_index).unwrap_or_else(|error| { + panic!( + "Got the following error when trying to convert node index \ + {node_index:?} to a contract address: {error:?}", + ) + }), + storage_trie, + ) + }) + .collect(), + contracts_trie, + classes_trie, + }) } - async fn new_contract_state( - contract_address: ContractAddress, - new_nonce: Nonce, - new_class_hash: ClassHash, - updated_storage_trie: UpdatedSkeletonTreeImpl, - inner_updates: LeafModifications, - ) -> ForestResult<(ContractAddress, ContractState, StorageTrie)> { - let filled_storage_trie = - StorageTrie::create_with_existing_leaves::(updated_storage_trie, inner_updates) - .await - .map_err(ForestError::StorageTrie)?; - let new_root_hash = filled_storage_trie.get_root_hash(); - Ok(( - contract_address, - ContractState { - nonce: new_nonce, - storage_root_hash: new_root_hash, - class_hash: new_class_hash, - }, - filled_storage_trie, - )) + fn get_contracts_trie_leaf_input( + original_contracts_trie_leaves: &HashMap, + contract_address_to_storage_updates: HashMap< + ContractAddress, + LeafModifications, + >, + mut contract_address_to_storage_skeleton: HashMap, + address_to_class_hash: &HashMap, + address_to_nonce: &HashMap, + ) -> ForestResult::Input>> { + let mut leaf_index_to_leaf_input = HashMap::new(); + assert_eq!( + contract_address_to_storage_updates.len(), + contract_address_to_storage_skeleton.len(), + "Mismatch between number of updated storage trie skeletons and number of storage \ + leaf-modification maps. Number of updated storage trie skeletons: {0:?}, number of \ + storage leaf-modification maps: {1:?}", + contract_address_to_storage_skeleton.len(), + contract_address_to_storage_updates.len() + ); + // `contract_address_to_storage_updates` includes all modified contracts, even those with + // unmodified storage, see StateDiff::actual_storage_updates(). + for (contract_address, storage_updates) in contract_address_to_storage_updates { + let node_index = NodeIndex::from(&contract_address); + let original_contract_state = original_contracts_trie_leaves + .get(&node_index) + .ok_or(ForestError::MissingContractCurrentState(contract_address))?; + leaf_index_to_leaf_input.insert( + node_index, + ( + node_index, + *(address_to_nonce + .get(&contract_address) + .unwrap_or(&original_contract_state.nonce)), + *(address_to_class_hash + .get(&contract_address) + .unwrap_or(&original_contract_state.class_hash)), + contract_address_to_storage_skeleton + .remove(&contract_address) + .ok_or(ForestError::MissingUpdatedSkeleton(contract_address))?, + storage_updates, + ), + ); + } + Ok(leaf_index_to_leaf_input) } } diff --git a/crates/starknet_committer/src/forest/forest_errors.rs b/crates/starknet_committer/src/forest/forest_errors.rs index 82da32069f..51c866d24e 100644 --- a/crates/starknet_committer/src/forest/forest_errors.rs +++ b/crates/starknet_committer/src/forest/forest_errors.rs @@ -25,7 +25,10 @@ pub enum ForestError { address {0:?}" )] MissingOriginalSkeleton(ContractAddress), - #[error("Can't fill storage trie, because there is no updated skeleton at address {0:?}")] + #[error( + "Can't create Contracts trie, because there is no updated skeleton for storage trie at \ + address {0:?}" + )] MissingUpdatedSkeleton(ContractAddress), #[error( "Can't build storage trie, because there are no sorted leaf indices of the contract at \ diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/types.rs b/crates/starknet_patricia/src/patricia_merkle_tree/types.rs index c70b49af1c..370731e0bd 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/types.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/types.rs @@ -55,7 +55,7 @@ impl NodeIndex { Self(index) } - pub(crate) fn is_leaf(&self) -> bool { + pub fn is_leaf(&self) -> bool { Self::FIRST_LEAF <= *self && *self <= Self::MAX }