From 688648811632f38047d4311b0195d611dfaee188 Mon Sep 17 00:00:00 2001 From: amos rothberg Date: Mon, 29 Jul 2024 13:12:36 +0300 Subject: [PATCH 1/2] feat: compute Storage Tries and Contracts Trie together --- .../external_test_utils.rs | 9 +- .../filled_tree/errors.rs | 5 +- .../filled_tree/forest.rs | 136 +++++++-------- .../patricia_merkle_tree/filled_tree/tree.rs | 157 ++++++++++++++---- .../filled_tree/tree_test.rs | 32 ++-- .../internal_test_utils.rs | 12 +- .../patricia_merkle_tree/node_data/errors.rs | 3 + .../patricia_merkle_tree/node_data/leaf.rs | 92 +++++----- .../src/patricia_merkle_tree/types.rs | 7 + .../create_tree_helper_test.rs | 9 +- .../updated_skeleton_tree/tree.rs | 3 +- 11 files changed, 297 insertions(+), 168 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 ca7472c3..1f4a8dcc 100644 --- a/crates/committer/src/patricia_merkle_tree/external_test_utils.rs +++ b/crates/committer/src/patricia_merkle_tree/external_test_utils.rs @@ -96,9 +96,12 @@ pub async fn tree_computation_flow( ) .expect("Failed to create the updated skeleton tree"); - StorageTrie::create::(updated_skeleton.into(), leaf_modifications) - .await - .expect("Failed to create the filled tree") + StorageTrie::create_no_additional_output::( + updated_skeleton.into(), + leaf_modifications, + ) + .await + .expect("Failed to create the filled tree") } pub async fn single_tree_flow_test( diff --git a/crates/committer/src/patricia_merkle_tree/filled_tree/errors.rs b/crates/committer/src/patricia_merkle_tree/filled_tree/errors.rs index 073a2d37..ae8d54c1 100644 --- a/crates/committer/src/patricia_merkle_tree/filled_tree/errors.rs +++ b/crates/committer/src/patricia_merkle_tree/filled_tree/errors.rs @@ -13,10 +13,13 @@ pub enum FilledTreeError { #[error("Deleted leaf at index {0:?} appears in the updated skeleton tree.")] DeletedLeafInSkeleton(NodeIndex), #[error("Double update at node {index:?}. Existing value: {existing_value:?}.")] - DoubleUpdate { + DoubleOutputUpdate { index: NodeIndex, existing_value: Box>, }, + #[error("Double update at node {index:?}.")] + //TODO(Amos): Add the existing value to the error message. + DoubleAdditionalOutputUpdate { index: NodeIndex }, #[error(transparent)] Leaf(#[from] LeafError), #[error("Missing node at index {0:?}.")] diff --git a/crates/committer/src/patricia_merkle_tree/filled_tree/forest.rs b/crates/committer/src/patricia_merkle_tree/filled_tree/forest.rs index cadbe2a7..78744710 100644 --- a/crates/committer/src/patricia_merkle_tree/filled_tree/forest.rs +++ b/crates/committer/src/patricia_merkle_tree/filled_tree/forest.rs @@ -4,10 +4,8 @@ use crate::hash::hash_trait::HashOutput; use crate::patricia_merkle_tree::filled_tree::node::CompiledClassHash; use crate::patricia_merkle_tree::filled_tree::node::{ClassHash, Nonce}; use crate::patricia_merkle_tree::filled_tree::tree::FilledTree; -use crate::patricia_merkle_tree::filled_tree::tree::{ - ClassesTrie, ContractsTrie, StorageTrie, StorageTrieMap, -}; -use crate::patricia_merkle_tree::node_data::leaf::{ContractState, LeafModifications}; +use crate::patricia_merkle_tree::filled_tree::tree::{ClassesTrie, ContractsTrie, StorageTrieMap}; +use crate::patricia_merkle_tree::node_data::leaf::{ContractState, Leaf, LeafModifications}; use crate::patricia_merkle_tree::types::NodeIndex; use crate::patricia_merkle_tree::updated_skeleton_tree::hash_function::ForestHashFunction; use crate::patricia_merkle_tree::updated_skeleton_tree::skeleton_forest::UpdatedSkeletonForest; @@ -16,7 +14,6 @@ use crate::storage::storage_trait::Storage; use std::collections::HashMap; use std::sync::Arc; -use tokio::task::JoinSet; pub struct FilledForest { pub storage_tries: StorageTrieMap, @@ -48,83 +45,86 @@ impl FilledForest { } pub(crate) async fn create( - mut updated_forest: UpdatedSkeletonForest, + updated_forest: UpdatedSkeletonForest, storage_updates: HashMap>, classes_updates: LeafModifications, original_contracts_trie_leaves: &HashMap, address_to_class_hash: &HashMap, address_to_nonce: &HashMap, ) -> ForestResult { - let classes_trie_task = tokio::spawn(ClassesTrie::create::( + let classes_trie = ClassesTrie::create_no_additional_output::( Arc::new(updated_forest.classes_trie), Arc::new(classes_updates), - )); - let mut contracts_trie_modifications = HashMap::new(); - let mut filled_storage_tries = HashMap::new(); - let mut contracts_state_tasks = JoinSet::new(); + ) + .await?; - 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(&NodeIndex::from_contract_address(&address)) - .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( - NodeIndex::from_contract_address(&address), - new_contract_state, - ); - filled_storage_tries.insert(address, filled_storage_trie); - } - - let contracts_trie_task = tokio::spawn(ContractsTrie::create::( + let (contracts_trie, storage_tries) = ContractsTrie::create::( Arc::new(updated_forest.contracts_trie), - Arc::new(contracts_trie_modifications), - )); + Arc::new(FilledForest::get_contracts_trie_leaf_input( + original_contracts_trie_leaves, + storage_updates, + updated_forest.storage_tries, + address_to_class_hash, + address_to_nonce, + )?), + ) + .await?; Ok(Self { - storage_tries: filled_storage_tries, - contracts_trie: contracts_trie_task.await??, - classes_trie: classes_trie_task.await??, + storage_tries: storage_tries + .unwrap_or_else(|| panic!("Missing storage tries.")) + .into_iter() + .map(|(node_index, storage_trie)| (node_index.to_contract_address(), 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::(Arc::new(updated_storage_trie), Arc::new(inner_updates)) - .await?; - 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, - )) + // TODO(Amos, 1/8/2024): Can this be done more efficiently? + // should error be returned if keys are missing? + fn get_contracts_trie_leaf_input( + original_contracts_trie_leaves: &HashMap, + mut storage_updates: HashMap>, + mut storage_tries: HashMap, + address_to_class_hash: &HashMap, + address_to_nonce: &HashMap, + ) -> ForestResult::I>> { + let mut leaf_index_to_leaf_input = HashMap::new(); + assert_eq!(storage_updates.len(), storage_tries.len()); + // `storage_updates` includes all modified contracts, see + // StateDiff::actual_storage_updates(). + for contract_address in storage_updates.keys().cloned().collect::>() { + let original_contract_state = original_contracts_trie_leaves + .get(&NodeIndex::from_contract_address(&contract_address)) + .ok_or(ForestError::MissingContractCurrentState(contract_address))?; + leaf_index_to_leaf_input.insert( + NodeIndex::from_contract_address(&contract_address), + ( + NodeIndex::from_contract_address(&contract_address), + *(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)), + storage_tries.remove(&contract_address).unwrap_or_else(|| { + panic!( + "Missing update skeleton tree for contract {:?}", + contract_address + ) + }), + storage_updates + .remove(&contract_address) + .unwrap_or_else(|| { + panic!( + "Missing storage updates for contract {:?}", + contract_address + ) + }), + ), + ); + } + Ok(leaf_index_to_leaf_input) } } diff --git a/crates/committer/src/patricia_merkle_tree/filled_tree/tree.rs b/crates/committer/src/patricia_merkle_tree/filled_tree/tree.rs index 6bc032f4..1bdde902 100644 --- a/crates/committer/src/patricia_merkle_tree/filled_tree/tree.rs +++ b/crates/committer/src/patricia_merkle_tree/filled_tree/tree.rs @@ -10,12 +10,12 @@ use crate::hash::hash_trait::HashOutput; use crate::patricia_merkle_tree::filled_tree::errors::FilledTreeError; use crate::patricia_merkle_tree::filled_tree::node::CompiledClassHash; use crate::patricia_merkle_tree::filled_tree::node::FilledNode; +use crate::patricia_merkle_tree::node_data::errors::LeafError; use crate::patricia_merkle_tree::node_data::inner_node::BinaryData; use crate::patricia_merkle_tree::node_data::inner_node::EdgeData; use crate::patricia_merkle_tree::node_data::inner_node::NodeData; use crate::patricia_merkle_tree::node_data::leaf::ContractState; use crate::patricia_merkle_tree::node_data::leaf::Leaf; -use crate::patricia_merkle_tree::node_data::leaf::LeafModifications; use crate::patricia_merkle_tree::types::NodeIndex; use crate::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunction; use crate::patricia_merkle_tree::updated_skeleton_tree::node::UpdatedSkeletonNode; @@ -32,12 +32,22 @@ pub(crate) type FilledTreeResult = Result>; /// Consider a Patricia-Merkle Tree which has been updated with new leaves. /// FilledTree consists of all nodes which were modified in the update, including their updated /// data and hashes. -pub(crate) trait FilledTree: Sized { +pub(crate) trait FilledTree: Sized + Send { /// Computes and returns the filled tree. + // TODO(Amos, 1/8/2024): Implement and use a LeafIndex type for input & output. It should probably + // also be used in other places across the code. async fn create<'a, TH: TreeHashFunction + 'static>( updated_skeleton: Arc + 'static>, - leaf_modifications: Arc>, - ) -> FilledTreeResult; + leaf_index_to_leaf_input: Arc>, + ) -> FilledTreeResult<(Self, Option>), L>; + + async fn create_no_additional_output<'a, TH: TreeHashFunction + 'static>( + updated_skeleton: Arc + 'static>, + leaf_index_to_leaf_input: Arc>, + ) -> FilledTreeResult { + let (result, _) = Self::create::(updated_skeleton, leaf_index_to_leaf_input).await?; + Ok(result) + } /// Serializes the current state of the tree into a hashmap, /// where each key-value pair corresponds @@ -59,7 +69,7 @@ pub type ContractsTrie = FilledTreeImpl; pub type StorageTrieMap = HashMap; impl FilledTreeImpl { - fn initialize_with_placeholders<'a>( + fn initialize_filled_tree_map_with_placeholders<'a>( updated_skeleton: &Arc>, ) -> HashMap>>> { let mut filled_tree_map = HashMap::new(); @@ -71,6 +81,17 @@ impl FilledTreeImpl { filled_tree_map } + fn initialize_leaf_output_map_with_placeholders( + leaf_index_to_leaf_input: &Arc>, + ) -> Arc>>> { + Arc::new( + leaf_index_to_leaf_input + .keys() + .map(|index| (*index, Mutex::new(None))) + .collect(), + ) + } + pub(crate) fn get_all_nodes(&self) -> &HashMap> { &self.tree_map } @@ -89,7 +110,7 @@ impl FilledTreeImpl { .lock() .map_err(|_| FilledTreeError::PoisonedLock("Cannot lock node.".to_owned()))?; match node.take() { - Some(existing_node) => Err(FilledTreeError::DoubleUpdate { + Some(existing_node) => Err(FilledTreeError::DoubleOutputUpdate { index, existing_value: Box::new(existing_node), }), @@ -103,30 +124,77 @@ impl FilledTreeImpl { } } - fn remove_arc_mutex_and_option( - hash_map_in: Arc>>>>, - ) -> FilledTreeResult>, L> { + /// Similar to `write_to_output_map`, but for the additional output map. + fn write_to_leaf_output_map( + output_map: Arc>>>, + index: NodeIndex, + data: L::O, + ) -> FilledTreeResult<(), L> { + match output_map.get(&index) { + Some(node) => { + let mut node = node + .lock() + // TODO(Amos): add error for this, instead of using incorrect error. + .map_err(|_| FilledTreeError::PoisonedLock("Cannot lock node.".to_owned()))?; + match node.take() { + Some(_existing_output_data) => { + Err(FilledTreeError::DoubleAdditionalOutputUpdate { index }) + } + None => { + *node = Some(data); + Ok(()) + } + } + } + // TODO(Amos): add error for this, instead of using incorrect error. + None => Err(FilledTreeError::::MissingNode(index)), + } + } + + fn remove_arc_mutex_and_option( + hash_map_in: Arc>>>, + fail_on_none_value: bool, + ) -> FilledTreeResult, L> { let mut hash_map_out = HashMap::new(); - for (key, value) in hash_map_in.iter() { + for (key, value) in Arc::into_inner(hash_map_in) + .unwrap_or_else(|| panic!("Cannot retrieve hashmap from Arc.")) + { let mut value = value .lock() .map_err(|_| FilledTreeError::::PoisonedLock("Cannot lock node.".to_owned()))?; match value.take() { Some(value) => { - hash_map_out.insert(*key, value); + hash_map_out.insert(key, value); + } + None => { + if fail_on_none_value { + return Err(FilledTreeError::::MissingNode(key)); + } } - None => return Err(FilledTreeError::::MissingNode(*key)), } } Ok(hash_map_out) } + fn wrap_leaf_input_keys_in_mutex( + leaf_index_to_leaf_input: Arc>, + ) -> Arc>>> { + // TODO(Amos, 1/8/2024): Can this be done without exiting the Arc? + let res = Arc::into_inner(leaf_index_to_leaf_input) + .expect("Cannot retrieve hashmap from Arc.") + .into_iter() + .map(|(k, v)| (k, Mutex::new(Some(v)))) + .collect(); + Arc::new(res) + } + #[async_recursion] async fn compute_filled_tree_rec<'a, TH>( updated_skeleton: Arc + 'async_recursion + 'static>, index: NodeIndex, - leaf_modifications: Arc>, + leaf_index_to_leaf_input: Arc>>>, output_map: Arc>>>>, + leaf_index_to_leaf_output: Arc>>>, ) -> FilledTreeResult where TH: TreeHashFunction + 'static, @@ -141,14 +209,16 @@ impl FilledTreeImpl { tokio::spawn(Self::compute_filled_tree_rec::( Arc::clone(&updated_skeleton), left_index, - Arc::clone(&leaf_modifications), + Arc::clone(&leaf_index_to_leaf_input), Arc::clone(&output_map), + Arc::clone(&leaf_index_to_leaf_output), )), tokio::spawn(Self::compute_filled_tree_rec::( Arc::clone(&updated_skeleton), right_index, - Arc::clone(&leaf_modifications), + Arc::clone(&leaf_index_to_leaf_input), Arc::clone(&output_map), + Arc::clone(&leaf_index_to_leaf_output), )), ); @@ -166,8 +236,9 @@ impl FilledTreeImpl { let bottom_hash = Self::compute_filled_tree_rec::( Arc::clone(&updated_skeleton), bottom_node_index, - leaf_modifications, + Arc::clone(&leaf_index_to_leaf_input), Arc::clone(&output_map), + Arc::clone(&leaf_index_to_leaf_output), ) .await?; let data = NodeData::Edge(EdgeData { @@ -180,13 +251,27 @@ impl FilledTreeImpl { } UpdatedSkeletonNode::UnmodifiedSubTree(hash_result) => Ok(*hash_result), UpdatedSkeletonNode::Leaf => { - let leaf_data = L::create(&index, leaf_modifications).await?; + // TODO(Amos): Use correct error types, and consider returning error instead of + // panic. + let leaf_input = leaf_index_to_leaf_input + .get(&index) + .ok_or(LeafError::MissingLeafModificationData(index))? + .lock() + .map_err(|_| { + FilledTreeError::::PoisonedLock("Cannot lock node.".to_owned()) + })? + .take() + .unwrap_or_else(|| panic!("Missing input for leaf {:?}.", index)); + let (leaf_data, leaf_output) = L::create(leaf_input).await?; if leaf_data.is_empty() { return Err(FilledTreeError::::DeletedLeafInSkeleton(index)); } let node_data = NodeData::Leaf(leaf_data); let hash_value = TH::compute_node_hash(&node_data); Self::write_to_output_map(output_map, index, hash_value, node_data)?; + if let Some(output) = leaf_output { + Self::write_to_leaf_output_map(leaf_index_to_leaf_output, index, output)? + }; Ok(hash_value) } } @@ -216,33 +301,45 @@ impl FilledTreeImpl { impl FilledTree for FilledTreeImpl { async fn create<'a, TH: TreeHashFunction + 'static>( updated_skeleton: Arc + 'static>, - leaf_modifications: Arc>, - ) -> Result> { + leaf_index_to_leaf_input: Arc>, + ) -> Result<(Self, Option>), FilledTreeError> { // Compute the filled tree in two steps: // 1. Create a map containing the tree structure without hash values. // 2. Fill in the hash values. - if leaf_modifications.is_empty() { - return Self::create_unmodified(&updated_skeleton); + if leaf_index_to_leaf_input.is_empty() { + let unmodified = Self::create_unmodified(&updated_skeleton)?; + return Ok((unmodified, Some(HashMap::new()))); } if updated_skeleton.is_empty() { - return Ok(Self::create_empty()); + return Ok((Self::create_empty(), Some(HashMap::new()))); } - let filled_tree_map = Arc::new(Self::initialize_with_placeholders(&updated_skeleton)); + let filled_tree_map = Arc::new(Self::initialize_filled_tree_map_with_placeholders( + &updated_skeleton, + )); + let leaf_index_to_leaf_output = + Self::initialize_leaf_output_map_with_placeholders(&leaf_index_to_leaf_input); + let wrapped_leaf_index_to_leaf_input = + Self::wrap_leaf_input_keys_in_mutex(leaf_index_to_leaf_input); let root_hash = Self::compute_filled_tree_rec::( - updated_skeleton, + Arc::clone(&updated_skeleton), NodeIndex::ROOT, - leaf_modifications, + Arc::clone(&wrapped_leaf_index_to_leaf_input), Arc::clone(&filled_tree_map), + Arc::clone(&leaf_index_to_leaf_output), ) .await?; + let unwrapped_leaf_index_to_leaf_output = + Self::remove_arc_mutex_and_option(leaf_index_to_leaf_output, false)?; - // Create and return a new FilledTreeImpl from the hashmap. - Ok(FilledTreeImpl { - tree_map: Self::remove_arc_mutex_and_option(filled_tree_map)?, - root_hash, - }) + Ok(( + FilledTreeImpl { + tree_map: Self::remove_arc_mutex_and_option(filled_tree_map, true)?, + root_hash, + }, + Some(unwrapped_leaf_index_to_leaf_output), + )) } fn serialize(&self) -> HashMap { 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 26f61139..f7ec66ee 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 @@ -28,13 +28,14 @@ async fn test_filled_tree_sanity() { skeleton_tree.insert(new_leaf_index, UpdatedSkeletonNode::Leaf); let modifications = HashMap::from([(new_leaf_index, new_filled_leaf)]); let updated_skeleton_tree = UpdatedSkeletonTreeImpl { skeleton_tree }; - let root_hash = FilledTreeImpl::create::( - Arc::new(updated_skeleton_tree), - Arc::new(modifications), - ) - .await - .unwrap() - .get_root_hash(); + let root_hash = + FilledTreeImpl::::create_no_additional_output::( + Arc::new(updated_skeleton_tree), + Arc::new(modifications), + ) + .await + .unwrap() + .get_root_hash(); assert_eq!(root_hash, HashOutput(Felt::ONE), "Root hash mismatch"); } @@ -89,7 +90,7 @@ async fn test_small_filled_tree() { .collect(); // Compute the hash values. - let filled_tree = FilledTreeImpl::create::( + let filled_tree = FilledTreeImpl::create_no_additional_output::( Arc::new(updated_skeleton_tree), Arc::new(modifications), ) @@ -152,7 +153,7 @@ async fn test_small_tree_with_unmodified_nodes() { )]); // Compute the hash values. - let filled_tree = FilledTreeImpl::create::( + let filled_tree = FilledTreeImpl::create_no_additional_output::( Arc::new(updated_skeleton_tree), Arc::new(modifications), ) @@ -203,12 +204,13 @@ async fn test_delete_leaf_from_empty_tree() { let leaf_modifications = HashMap::from([(NodeIndex::FIRST_LEAF, MockLeaf(Felt::ZERO))]); // Compute the filled tree. - let filled_tree = FilledTreeImpl::create::( - updated_skeleton_tree.into(), - leaf_modifications.into(), - ) - .await - .unwrap(); + let filled_tree = + FilledTreeImpl::::create_no_additional_output::( + updated_skeleton_tree.into(), + leaf_modifications.into(), + ) + .await + .unwrap(); // The filled tree should be empty. let filled_tree_map = filled_tree.get_all_nodes(); diff --git a/crates/committer/src/patricia_merkle_tree/internal_test_utils.rs b/crates/committer/src/patricia_merkle_tree/internal_test_utils.rs index 2cf06aaa..f440fb09 100644 --- a/crates/committer/src/patricia_merkle_tree/internal_test_utils.rs +++ b/crates/committer/src/patricia_merkle_tree/internal_test_utils.rs @@ -1,5 +1,3 @@ -use std::sync::Arc; - use crate::felt::Felt; use crate::generate_trie_config; use crate::hash::hash_trait::HashOutput; @@ -52,15 +50,15 @@ impl Deserializable for MockLeaf { } impl Leaf for MockLeaf { + type I = Self; + type O = (); + fn is_empty(&self) -> bool { self.0 == Felt::ZERO } - async fn create( - index: &NodeIndex, - leaf_modifications: Arc>, - ) -> LeafResult { - Self::from_modifications(index, leaf_modifications) + async fn create(input: Self::I) -> LeafResult<(Self, Option)> { + Ok((input, None)) } } diff --git a/crates/committer/src/patricia_merkle_tree/node_data/errors.rs b/crates/committer/src/patricia_merkle_tree/node_data/errors.rs index 4d1a729b..cd63434e 100644 --- a/crates/committer/src/patricia_merkle_tree/node_data/errors.rs +++ b/crates/committer/src/patricia_merkle_tree/node_data/errors.rs @@ -1,6 +1,7 @@ use std::fmt::Debug; use thiserror::Error; +use crate::patricia_merkle_tree::filled_tree::errors::StorageTrieError; use crate::patricia_merkle_tree::node_data::inner_node::{EdgePath, EdgePathLength}; use crate::patricia_merkle_tree::types::NodeIndex; @@ -28,6 +29,8 @@ pub enum EdgePathError { pub enum LeafError { #[error("Missing modification data at index {0:?}.")] MissingLeafModificationData(NodeIndex), + #[error("While computing the storage trie at index {0:?} got the following error: {1:?}")] + StorageTrieComputationFailed(Box, NodeIndex), } pub type LeafResult = Result; diff --git a/crates/committer/src/patricia_merkle_tree/node_data/leaf.rs b/crates/committer/src/patricia_merkle_tree/node_data/leaf.rs index 7cdde637..a97ac8e8 100644 --- a/crates/committer/src/patricia_merkle_tree/node_data/leaf.rs +++ b/crates/committer/src/patricia_merkle_tree/node_data/leaf.rs @@ -1,43 +1,30 @@ -use std::collections::HashMap; -use std::fmt::Debug; -use std::future::Future; -use std::sync::Arc; - use crate::block_committer::input::StarknetStorageValue; use crate::felt::Felt; use crate::hash::hash_trait::HashOutput; use crate::patricia_merkle_tree::filled_tree::node::{ClassHash, CompiledClassHash, Nonce}; +use crate::patricia_merkle_tree::filled_tree::tree::{FilledTree, FilledTreeImpl}; use crate::patricia_merkle_tree::node_data::errors::{LeafError, LeafResult}; use crate::patricia_merkle_tree::types::NodeIndex; +use crate::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunctionImpl; +use crate::patricia_merkle_tree::updated_skeleton_tree::tree::UpdatedSkeletonTreeImpl; use crate::storage::db_object::{DBObject, Deserializable}; +use std::collections::HashMap; +use std::fmt::Debug; +use std::future::Future; pub trait Leaf: Clone + Sync + Send + DBObject + Deserializable + Default + Debug + Eq { + // TODO(Amos, 1/1/2025): Add default values when it is stable. + type I: Send + Sync + 'static; + type O: Send + 'static; + /// Returns true if leaf is empty. fn is_empty(&self) -> bool; /// Creates a leaf. - /// This function is async to allow computation of a leaf on the fly; in simple cases, it can - /// be enough to return the leaf data directly using [Self::from_modifications]. // Use explicit desugaring of `async fn` to allow adding trait bounds to the return type, see // https://blog.rust-lang.org/2023/12/21/async-fn-rpit-in-traits.html#async-fn-in-public-traits // for details. - fn create( - index: &NodeIndex, - leaf_modifications: Arc>, - ) -> impl Future> + Send; - - /// Extracts the leaf data from the leaf modifications. Returns an error if the leaf data is - /// missing. - fn from_modifications( - index: &NodeIndex, - leaf_modifications: Arc>, - ) -> LeafResult { - let leaf_data = leaf_modifications - .get(index) - .ok_or(LeafError::MissingLeafModificationData(*index))? - .clone(); - Ok(leaf_data) - } + fn create(input: Self::I) -> impl Future)>> + Send; } #[derive(Clone, Debug, Default, Eq, PartialEq)] @@ -48,43 +35,68 @@ pub struct ContractState { } impl Leaf for StarknetStorageValue { + type I = Self; + type O = (); + fn is_empty(&self) -> bool { self.0 == Felt::ZERO } - - async fn create( - index: &NodeIndex, - leaf_modifications: Arc>, - ) -> LeafResult { - Self::from_modifications(index, leaf_modifications) + async fn create(input: Self::I) -> LeafResult<(Self, Option)> { + Ok((input, None)) } } impl Leaf for CompiledClassHash { + type I = Self; + type O = (); + fn is_empty(&self) -> bool { self.0 == Felt::ZERO } - async fn create( - index: &NodeIndex, - leaf_modifications: Arc>, - ) -> LeafResult { - Self::from_modifications(index, leaf_modifications) + async fn create(input: Self::I) -> LeafResult<(Self, Option)> { + Ok((input, None)) } } impl Leaf for ContractState { + type I = ( + NodeIndex, + Nonce, + ClassHash, + UpdatedSkeletonTreeImpl, + LeafModifications, + ); + type O = FilledTreeImpl; + fn is_empty(&self) -> bool { self.nonce.0 == Felt::ZERO && self.class_hash.0 == Felt::ZERO && self.storage_root_hash.0 == Felt::ZERO } - async fn create( - index: &NodeIndex, - leaf_modifications: Arc>, - ) -> LeafResult { - Self::from_modifications(index, leaf_modifications) + async fn create(input: Self::I) -> LeafResult<(Self, Option)> { + let (leaf_index, nonce, class_hash, updated_skeleton, storage_modifications) = input; + + match FilledTreeImpl::::create::( + updated_skeleton.into(), + storage_modifications.into(), + ) + .await + { + Ok((storage_trie, _)) => Ok(( + Self { + nonce, + storage_root_hash: storage_trie.get_root_hash(), + class_hash, + }, + Some(storage_trie), + )), + Err(storage_error) => Err(LeafError::StorageTrieComputationFailed( + storage_error.into(), + leaf_index, + )), + } } } diff --git a/crates/committer/src/patricia_merkle_tree/types.rs b/crates/committer/src/patricia_merkle_tree/types.rs index c014d3ae..5eeae994 100644 --- a/crates/committer/src/patricia_merkle_tree/types.rs +++ b/crates/committer/src/patricia_merkle_tree/types.rs @@ -141,6 +141,13 @@ impl NodeIndex { Self::from_leaf_felt(&address.0) } + pub(crate) fn to_contract_address(self) -> ContractAddress { + assert!(self.is_leaf(), "NodeIndex {:?} is not a leaf.", self); + ContractAddress( + Felt::try_from(self - Self::FIRST_LEAF).expect("Unable to convert node index to felt."), + ) + } + pub(crate) fn from_class_hash(class_hash: &ClassHash) -> Self { Self::from_leaf_felt(&class_hash.0) } 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 a039882f..46a06ea1 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 @@ -512,8 +512,11 @@ async fn test_update_non_modified_storage_tree(#[case] root_hash: HashOutput) { .unwrap(); let updated = UpdatedSkeletonTreeImpl::create(&mut original_skeleton_tree, &HashMap::new()).unwrap(); - let filled = MockTrie::create::(Arc::new(updated), Arc::new(empty_map)) - .await - .unwrap(); + let filled = MockTrie::create_no_additional_output::( + Arc::new(updated), + Arc::new(empty_map), + ) + .await + .unwrap(); assert_eq!(root_hash, filled.get_root_hash()); } 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 5e41f085..0fde1934 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 @@ -37,7 +37,8 @@ pub(crate) trait UpdatedSkeletonTree<'a>: Sized + Send + Sync { fn get_node(&self, index: NodeIndex) -> UpdatedSkeletonTreeResult<&UpdatedSkeletonNode>; } // TODO(Dori, 1/7/2024): Make this a tuple struct. -pub(crate) struct UpdatedSkeletonTreeImpl { +#[derive(Debug)] +pub struct UpdatedSkeletonTreeImpl { pub(crate) skeleton_tree: UpdatedSkeletonNodeMap, } From 8d101ffee93229ab026e9f2c3dd71db4e1a51792 Mon Sep 17 00:00:00 2001 From: amos rothberg Date: Tue, 30 Jul 2024 11:49:22 +0300 Subject: [PATCH 2/2] Pass benchmarks. --- .../external_test_utils.rs | 7 +++--- .../committer_cli/benches/committer_bench.rs | 25 +++++++++++-------- 2 files changed, 17 insertions(+), 15 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 1f4a8dcc..83590a0f 100644 --- a/crates/committer/src/patricia_merkle_tree/external_test_utils.rs +++ b/crates/committer/src/patricia_merkle_tree/external_test_utils.rs @@ -68,7 +68,7 @@ pub fn get_random_u256(rng: &mut R, low: U256, high: U256) -> U256 { } pub async fn tree_computation_flow( - leaf_modifications: Arc>, + leaf_modifications: LeafModifications, storage: &MapStorage, root_hash: HashOutput, ) -> StorageTrie { @@ -98,7 +98,7 @@ pub async fn tree_computation_flow( StorageTrie::create_no_additional_output::( updated_skeleton.into(), - leaf_modifications, + Arc::new(leaf_modifications), ) .await .expect("Failed to create the filled tree") @@ -115,8 +115,7 @@ pub async fn single_tree_flow_test( .map(|(k, v)| (NodeIndex::FIRST_LEAF + k, v)) .collect::>(); - let filled_tree = - tree_computation_flow(Arc::new(leaf_modifications), &storage, root_hash).await; + let filled_tree = tree_computation_flow(leaf_modifications, &storage, root_hash).await; let hash_result = filled_tree.get_root_hash(); diff --git a/crates/committer_cli/benches/committer_bench.rs b/crates/committer_cli/benches/committer_bench.rs index a83442c4..daa31750 100644 --- a/crates/committer_cli/benches/committer_bench.rs +++ b/crates/committer_cli/benches/committer_bench.rs @@ -7,7 +7,7 @@ // Then upload the new files to GCS with this new prefix (run e.g., // gcloud storage cp LOCAL_FILE gs://committer-testing-artifacts/NEW_PREFIX/tree_flow_inputs.json). -use std::{collections::HashMap, sync::Arc}; +use std::collections::HashMap; use committer::{ block_committer::input::StarknetStorageValue, @@ -17,7 +17,7 @@ use committer::{ }, }; use committer_cli::{commands::parse_and_commit, tests::utils::parse_from_python::TreeFlowInput}; -use criterion::{criterion_group, criterion_main, Criterion}; +use criterion::{criterion_group, criterion_main, BatchSize, Criterion}; const CONCURRENCY_MODE: bool = true; const SINGLE_TREE_FLOW_INPUT: &str = include_str!("tree_flow_inputs.json"); @@ -42,16 +42,19 @@ pub fn single_tree_flow_benchmark(criterion: &mut Criterion) { .into_iter() .map(|(k, v)| (NodeIndex::FIRST_LEAF + k, v)) .collect::>(); - let arc_leaf_modifications = Arc::new(leaf_modifications); - criterion.bench_function("tree_computation_flow", |benchmark| { - benchmark.iter(|| { - runtime.block_on(tree_computation_flow( - Arc::clone(&arc_leaf_modifications), - &storage, - root_hash, - )); - }) + criterion.bench_function("tree_computation_flow", move |b| { + b.iter_batched( + || leaf_modifications.clone(), + |leaf_modifications_input| { + runtime.block_on(tree_computation_flow( + leaf_modifications_input, + &storage, + root_hash, + )); + }, + BatchSize::LargeInput, + ) }); }