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 53d42849..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,9 +4,7 @@ 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::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; @@ -129,28 +127,4 @@ impl FilledForest { } Ok(leaf_index_to_leaf_input) } - - 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_no_additional_output::( - 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, - )) - } } 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 71728d48..2f6b5d5d 100644 --- a/crates/committer/src/patricia_merkle_tree/filled_tree/tree.rs +++ b/crates/committer/src/patricia_merkle_tree/filled_tree/tree.rs @@ -153,17 +153,24 @@ impl FilledTreeImpl { 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) @@ -181,15 +188,6 @@ impl FilledTreeImpl { Arc::new(res) } - fn clone_arc_if_needed( - wrapped_input: &Option>>>, - ) -> Option>>> { - match wrapped_input { - Some(input) => Some(Arc::clone(&input)), - None => None, - } - } - #[async_recursion] async fn compute_filled_tree_rec<'a, TH>( updated_skeleton: Arc + 'async_recursion + 'static>, @@ -271,13 +269,8 @@ impl FilledTreeImpl { 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)?; - match additional_output { - Some(output) => Self::write_to_additional_output_map( - leaf_index_to_leaf_output, - index, - output, - )?, - None => (), + if let Some(output) = additional_output { + Self::write_to_additional_output_map(leaf_index_to_leaf_output, index, output)? }; Ok(hash_value) } @@ -313,15 +306,13 @@ impl FilledTree for FilledTreeImpl { // Compute the filled tree in two steps: // 1. Create a map containing the tree structure without hash values. // 2. Fill in the hash values. - // TODO(Amos, 1/8/2024): How should this work for Contracts Trie? Probably need to return - // All the unmodified filled storage tries. if leaf_index_to_leaf_input.is_empty() { let unmodified = Self::create_unmodified(&updated_skeleton)?; - return Ok((unmodified, None)); + return Ok((unmodified, Some(HashMap::new()))); } if updated_skeleton.is_empty() { - return Ok((Self::create_empty(), None)); + return Ok((Self::create_empty(), Some(HashMap::new()))); } let filled_tree_map = Arc::new(Self::initialize_filled_tree_map_with_placeholders( @@ -340,18 +331,14 @@ impl FilledTree for FilledTreeImpl { ) .await?; let unwrapped_leaf_index_to_leaf_output = - Self::remove_arc_mutex_and_option(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)?, + tree_map: Self::remove_arc_mutex_and_option(filled_tree_map, true)?, root_hash, }, - match unwrapped_leaf_index_to_leaf_output.is_empty() { - true => None, - false => Some(unwrapped_leaf_index_to_leaf_output), - }, + Some(unwrapped_leaf_index_to_leaf_output), )) } 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 8a2aaa26..70a28708 100644 --- a/crates/committer/src/patricia_merkle_tree/node_data/leaf.rs +++ b/crates/committer/src/patricia_merkle_tree/node_data/leaf.rs @@ -78,12 +78,13 @@ impl Leaf for ContractState { async fn create(input: Self::I) -> LeafResult<(Self, Option)> { let (leaf_index, nonce, class_hash, updated_skeleton, leaf_modifications) = input; - match FilledTreeImpl::::create_no_additional_output::< - TreeHashFunctionImpl, - >(updated_skeleton.into(), leaf_modifications.into()) + match FilledTreeImpl::::create::( + updated_skeleton.into(), + leaf_modifications.into(), + ) .await { - Ok(storage_trie) => Ok(( + Ok((storage_trie, _)) => Ok(( Self { nonce, storage_root_hash: storage_trie.get_root_hash(), diff --git a/crates/committer/src/patricia_merkle_tree/types.rs b/crates/committer/src/patricia_merkle_tree/types.rs index c5cc7cc9..5eeae994 100644 --- a/crates/committer/src/patricia_merkle_tree/types.rs +++ b/crates/committer/src/patricia_merkle_tree/types.rs @@ -141,11 +141,10 @@ 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); + 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."), + Felt::try_from(self - Self::FIRST_LEAF).expect("Unable to convert node index to felt."), ) } 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, }