From df672bc5a755e00cf30232464de6268082e78904 Mon Sep 17 00:00:00 2001
From: amosStarkware <88497213+amosStarkware@users.noreply.github.com>
Date: Tue, 20 Aug 2024 16:35:06 +0300
Subject: [PATCH] chore(committer): add struct for contract state leaf input
 (#530)

---
 .../src/forest/filled_forest.rs               | 18 ++++++++--------
 .../patricia_merkle_tree/leaf/leaf_impl.rs    | 21 +++++++++++--------
 2 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/crates/starknet_committer/src/forest/filled_forest.rs b/crates/starknet_committer/src/forest/filled_forest.rs
index 6ee31f0123..e1d0037112 100644
--- a/crates/starknet_committer/src/forest/filled_forest.rs
+++ b/crates/starknet_committer/src/forest/filled_forest.rs
@@ -2,7 +2,7 @@ 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::{Leaf, LeafModifications};
+use starknet_patricia::patricia_merkle_tree::node_data::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;
@@ -12,7 +12,7 @@ use crate::block_committer::input::{ContractAddress, StarknetStorageValue};
 use crate::forest::forest_errors::{ForestError, ForestResult};
 use crate::forest::updated_skeleton_forest::UpdatedSkeletonForest;
 use crate::hash_function::hash::ForestHashFunction;
-use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
+use crate::patricia_merkle_tree::leaf::leaf_impl::{ContractState, ContractStateInput};
 use crate::patricia_merkle_tree::types::{
     ClassHash,
     ClassesTrie,
@@ -118,7 +118,7 @@ impl FilledForest {
         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>> {
+    ) -> ForestResult<HashMap<NodeIndex, ContractStateInput>> {
         let mut leaf_index_to_leaf_input = HashMap::new();
         assert_eq!(
             contract_address_to_storage_updates.len(),
@@ -138,19 +138,19 @@ impl FilledForest {
                 .ok_or(ForestError::MissingContractCurrentState(contract_address))?;
             leaf_index_to_leaf_input.insert(
                 node_index,
-                (
-                    node_index,
-                    *(address_to_nonce
+                ContractStateInput {
+                    leaf_index: node_index,
+                    nonce: *(address_to_nonce
                         .get(&contract_address)
                         .unwrap_or(&original_contract_state.nonce)),
-                    *(address_to_class_hash
+                    class_hash: *(address_to_class_hash
                         .get(&contract_address)
                         .unwrap_or(&original_contract_state.class_hash)),
-                    contract_address_to_storage_skeleton
+                    updated_skeleton: 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/patricia_merkle_tree/leaf/leaf_impl.rs b/crates/starknet_committer/src/patricia_merkle_tree/leaf/leaf_impl.rs
index 02af6f26e2..0e7d3df7ce 100644
--- a/crates/starknet_committer/src/patricia_merkle_tree/leaf/leaf_impl.rs
+++ b/crates/starknet_committer/src/patricia_merkle_tree/leaf/leaf_impl.rs
@@ -44,13 +44,7 @@ impl Leaf for CompiledClassHash {
 }
 
 impl Leaf for ContractState {
-    type Input = (
-        NodeIndex,
-        Nonce,
-        ClassHash,
-        UpdatedSkeletonTreeImpl,
-        LeafModifications<StarknetStorageValue>,
-    );
+    type Input = ContractStateInput;
     type Output = FilledTreeImpl<StarknetStorageValue>;
 
     fn is_empty(&self) -> bool {
@@ -60,11 +54,12 @@ impl Leaf for ContractState {
     }
 
     async fn create(input: Self::Input) -> LeafResult<(Self, Self::Output)> {
-        let (leaf_index, nonce, class_hash, updated_skeleton, storage_modifications) = input;
+        let ContractStateInput { leaf_index, nonce, class_hash, updated_skeleton, storage_updates } =
+            input;
 
         let storage_trie = FilledTreeImpl::<StarknetStorageValue>::create_with_existing_leaves::<
             TreeHashFunctionImpl,
-        >(updated_skeleton, storage_modifications)
+        >(updated_skeleton, storage_updates)
         .await
         .map_err(|storage_error| {
             LeafError::LeafComputationError(format!(
@@ -79,3 +74,11 @@ impl Leaf for ContractState {
         ))
     }
 }
+
+pub struct ContractStateInput {
+    pub leaf_index: NodeIndex,
+    pub nonce: Nonce,
+    pub class_hash: ClassHash,
+    pub updated_skeleton: UpdatedSkeletonTreeImpl,
+    pub storage_updates: LeafModifications<StarknetStorageValue>,
+}