feature: compute Storage Tries and Contracts Trie together

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<TH: ForestHashFunction + 'static>(
-        contract_address: ContractAddress,
-        new_nonce: Nonce,
-        new_class_hash: ClassHash,
-        updated_storage_trie: UpdatedSkeletonTreeImpl,
-        inner_updates: LeafModifications<StarknetStorageValue>,
-    ) -> ForestResult<(ContractAddress, ContractState, StorageTrie)> {
-        let filled_storage_trie = StorageTrie::create_no_additional_output::<TH>(
-            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,
-        ))
-    }
 }

crates/committer/src/patricia_merkle_tree/filled_tree/tree.rs

@@ -153,17 +153,24 @@ impl<L: Leaf + 'static> FilledTreeImpl<L> {
 
     fn remove_arc_mutex_and_option<V>(
         hash_map_in: Arc<HashMap<NodeIndex, Mutex<Option<V>>>>,
+        fail_on_none_value: bool,
     ) -> FilledTreeResult<HashMap<NodeIndex, V>, 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::<L>::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::<L>::MissingNode(key));
+                    }
                 }
-                None => return Err(FilledTreeError::<L>::MissingNode(*key)),
             }
         }
         Ok(hash_map_out)
@@ -181,15 +188,6 @@ impl<L: Leaf + 'static> FilledTreeImpl<L> {
         Arc::new(res)
     }
 
-    fn clone_arc_if_needed(
-        wrapped_input: &Option<Arc<HashMap<NodeIndex, Mutex<L::I>>>>,
-    ) -> Option<Arc<HashMap<NodeIndex, Mutex<L::I>>>> {
-        match wrapped_input {
-            Some(input) => Some(Arc::clone(&input)),
-            None => None,
-        }
-    }
-
     #[async_recursion]
     async fn compute_filled_tree_rec<'a, TH>(
         updated_skeleton: Arc<impl UpdatedSkeletonTree<'a> + 'async_recursion + 'static>,
@@ -271,13 +269,8 @@ impl<L: Leaf + 'static> FilledTreeImpl<L> {
                 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<L: Leaf + 'static> FilledTree<L> for FilledTreeImpl<L> {
         // 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<L: Leaf + 'static> FilledTree<L> for FilledTreeImpl<L> {
         )
         .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),
         ))
     }
 

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<Self::O>)> {
         let (leaf_index, nonce, class_hash, updated_skeleton, leaf_modifications) = input;
 
-        match FilledTreeImpl::<StarknetStorageValue>::create_no_additional_output::<
-            TreeHashFunctionImpl,
-        >(updated_skeleton.into(), leaf_modifications.into())
+        match FilledTreeImpl::<StarknetStorageValue>::create::<TreeHashFunctionImpl>(
+            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(),

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."),
         )
     }
 

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,
 }