feat(committer): use FilledTree::create() in filled forest (#495)

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<ContractAddress, Self::Error> {
+        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)

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())
+    );
+}

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<TH: ForestHashFunction + 'static>(
-        mut updated_forest: UpdatedSkeletonForest,
+        updated_forest: UpdatedSkeletonForest,
         storage_updates: HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
         classes_updates: LeafModifications<CompiledClassHash>,
         original_contracts_trie_leaves: &HashMap<NodeIndex, ContractState>,
@@ -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::<TH>(
-                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::<TH>(
+        let contracts_trie_task = tokio::task::spawn(ContractsTrie::create::<TH>(
             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<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_with_existing_leaves::<TH>(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<NodeIndex, ContractState>,
+        contract_address_to_storage_updates: HashMap<
+            ContractAddress,
+            LeafModifications<StarknetStorageValue>,
+        >,
+        mut contract_address_to_storage_skeleton: HashMap<ContractAddress, UpdatedSkeletonTreeImpl>,
+        address_to_class_hash: &HashMap<ContractAddress, ClassHash>,
+        address_to_nonce: &HashMap<ContractAddress, Nonce>,
+    ) -> ForestResult<HashMap<NodeIndex, <ContractState as Leaf>::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)
     }
 }

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 \

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
     }