From 7bcd113e8e51a1de02dde32e6a276f02de689b64 Mon Sep 17 00:00:00 2001 From: amos rothberg Date: Sun, 18 Aug 2024 17:00:15 +0300 Subject: [PATCH] feat(committer): use FilledTree::create() in filled forest --- .../src/block_committer/input.rs | 10 ++ .../src/forest/filled_forest.rs | 122 +++++++++--------- .../src/forest/forest_errors.rs | 5 +- .../src/patricia_merkle_tree/types.rs | 2 +- 4 files changed, 79 insertions(+), 60 deletions(-) diff --git a/crates/starknet_committer/src/block_committer/input.rs b/crates/starknet_committer/src/block_committer/input.rs index 74830e7806e..efc53e0c597 100644 --- a/crates/starknet_committer/src/block_committer/input.rs +++ b/crates/starknet_committer/src/block_committer/input.rs @@ -14,6 +14,16 @@ use crate::patricia_merkle_tree::types::{ClassHash, CompiledClassHash, Nonce}; // TODO(Nimrod, 1/6/2025): Use the ContractAddress defined in starknet-types-core when available. pub struct ContractAddress(pub Felt); +impl From<&NodeIndex> for ContractAddress { + fn from(node_index: &NodeIndex) -> ContractAddress { + assert!(node_index.is_leaf(), "NodeIndex {:?} is not a leaf.", node_index); + ContractAddress( + Felt::try_from(*node_index - NodeIndex::FIRST_LEAF) + .expect("Unable to convert node index to felt."), + ) + } +} + impl From<&ContractAddress> for NodeIndex { fn from(address: &ContractAddress) -> NodeIndex { NodeIndex::from_leaf_felt(&address.0) diff --git a/crates/starknet_committer/src/forest/filled_forest.rs b/crates/starknet_committer/src/forest/filled_forest.rs index a1e28e01325..e43374be192 100644 --- a/crates/starknet_committer/src/forest/filled_forest.rs +++ b/crates/starknet_committer/src/forest/filled_forest.rs @@ -2,11 +2,10 @@ 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 crate::block_committer::input::{ContractAddress, StarknetStorageValue}; use crate::forest::forest_errors::{ForestError, ForestResult}; @@ -19,7 +18,6 @@ use crate::patricia_merkle_tree::types::{ CompiledClassHash, ContractsTrie, Nonce, - StorageTrie, StorageTrieMap, }; @@ -52,8 +50,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 +64,73 @@ 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 (contracts_trie, storage_tries) = + contracts_trie_task.await?.map_err(ForestError::ContractsTrie)?; let classes_trie = classes_trie_task.await?.map_err(ForestError::ClassesTrie)?; - 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::from(&node_index), 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() + ); + // `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 82da32069fd..51c866d24eb 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 c70b49af1c6..370731e0bd1 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 }