test: add single tree flow test

crates/committer/src/patricia_merkle_tree/external_test_utils.rs

@@ -1,9 +1,21 @@
+use std::collections::HashMap;
+
 use ethnum::U256;
+use serde_json::json;
 
 use crate::felt::Felt;
+use crate::hash::hash_trait::HashOutput;
 use crate::patricia_merkle_tree::errors::TypesError;
+use crate::storage::map_storage::MapStorage;
 use rand::Rng;
 
+use super::filled_tree::tree::{FilledTree, FilledTreeImpl};
+use super::node_data::leaf::{LeafData, LeafDataImpl, LeafModifications, SkeletonLeaf};
+use super::original_skeleton_tree::tree::{OriginalSkeletonTree, OriginalSkeletonTreeImpl};
+use super::types::NodeIndex;
+use super::updated_skeleton_tree::hash_function::TreeHashFunctionImpl;
+use super::updated_skeleton_tree::tree::{UpdatedSkeletonTree, UpdatedSkeletonTreeImpl};
+
 impl TryFrom<&U256> for Felt {
     type Error = TypesError<U256>;
     fn try_from(value: &U256) -> Result<Self, Self::Error> {
@@ -51,3 +63,56 @@ pub fn get_random_u256<R: Rng>(rng: &mut R, low: U256, high: U256) -> U256 {
     }
     result
 }
+
+pub async fn single_tree_flow_test(
+    leaf_modifications: LeafModifications<LeafDataImpl>,
+    storage: MapStorage,
+    root_hash: HashOutput,
+) -> String {
+    // Move from leaf number to actual index.
+    let leaf_modifications = leaf_modifications
+        .into_iter()
+        .map(|(k, v)| (NodeIndex::FIRST_LEAF + k, v))
+        .collect::<HashMap<NodeIndex, LeafDataImpl>>();
+    let mut sorted_leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().copied().collect();
+    sorted_leaf_indices.sort();
+
+    // Build the original tree.
+    let mut original_skeleton: OriginalSkeletonTreeImpl =
+        OriginalSkeletonTree::create(&storage, &sorted_leaf_indices, root_hash)
+            .expect("Failed to create the original skeleton tree");
+
+    // Update the tree with the new data.
+    let updated_skeleton: UpdatedSkeletonTreeImpl = UpdatedSkeletonTree::create(
+        &mut original_skeleton,
+        &leaf_modifications
+            .iter()
+            .map(|(index, data)| {
+                (
+                    *index,
+                    match data.is_empty() {
+                        true => SkeletonLeaf::Zero,
+                        false => SkeletonLeaf::NonZero,
+                    },
+                )
+            })
+            .collect(),
+    )
+    .expect("Failed to create the updated skeleton tree");
+
+    // Compute the hash.
+    let filled_tree: FilledTreeImpl =
+        FilledTreeImpl::create::<TreeHashFunctionImpl>(updated_skeleton, leaf_modifications)
+            .await
+            .expect("Failed to create the filled tree");
+    let hash_result = filled_tree.get_root_hash();
+
+    let mut result_map = HashMap::new();
+    // Serialize the hash result.
+    let json_hash = &json!(hash_result.0.to_hex());
+    result_map.insert("root_hash", json_hash);
+    // Serlialize the storage modifications.
+    let json_storage = &json!(filled_tree.serialize());
+    result_map.insert("storage_changes", json_storage);
+    serde_json::to_string(&result_map).expect("serialization failed")
+}

crates/committer/src/patricia_merkle_tree/internal_test_utils.rs

@@ -1,7 +1,11 @@
 use crate::felt::Felt;
 use crate::patricia_merkle_tree::external_test_utils::get_random_u256;
+
 use crate::patricia_merkle_tree::node_data::inner_node::{EdgePathLength, PathToBottom};
 use crate::patricia_merkle_tree::node_data::leaf::SkeletonLeaf;
+
+use crate::patricia_merkle_tree::types::{NodeIndex, SubTreeHeight};
+
 use ethnum::U256;
 use rand::rngs::ThreadRng;
 use rstest::{fixture, rstest};
@@ -34,8 +38,6 @@ pub(crate) fn random() -> ThreadRng {
     rand::thread_rng()
 }
 
-use crate::patricia_merkle_tree::types::{NodeIndex, SubTreeHeight};
-
 impl NodeIndex {
     /// Assumes self represents an index in a smaller tree height. Returns a node index represents
     /// the same index in the starknet state tree as if the smaller tree was 'planted' at the lowest

crates/committer_cli/src/main.rs

@@ -79,6 +79,7 @@ async fn main() {
             // Run relevant test.
             let output = test
                 .run(inputs.as_deref())
+                .await
                 .unwrap_or_else(|error| panic!("Failed to run test: {}", error));
 
             // Print test's output.

crates/committer_cli/src/tests/python_tests.rs

@@ -17,6 +17,8 @@ use committer::patricia_merkle_tree::node_data::inner_node::{
 };
 use committer::patricia_merkle_tree::node_data::leaf::{ContractState, LeafDataImpl};
 use committer::patricia_merkle_tree::types::SubTreeHeight;
+
+use committer::patricia_merkle_tree::external_test_utils::single_tree_flow_test;
 use committer::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunctionImpl;
 use committer::storage::db_object::DBObject;
 use committer::storage::errors::{DeserializationError, SerializationError};
@@ -50,6 +52,7 @@ pub(crate) enum PythonTest {
     ParseTxOutput,
     ParseStateDiff,
     SerializeForRustCommitterFlowTest,
+    ComputeHashSingleTree,
 }
 
 /// Error type for PythonTest enum.
@@ -98,6 +101,7 @@ impl TryFrom<String> for PythonTest {
             "parse_tx_output_test" => Ok(Self::ParseTxOutput),
             "parse_state_diff_test" => Ok(Self::ParseStateDiff),
             "serialize_to_rust_committer_flow_test" => Ok(Self::SerializeForRustCommitterFlowTest),
+            "tree_test" => Ok(Self::ComputeHashSingleTree),
             _ => Err(PythonTestError::UnknownTestName(value)),
         }
     }
@@ -110,7 +114,7 @@ impl PythonTest {
     }
 
     /// Runs the test with the given arguments.
-    pub(crate) fn run(&self, input: Option<&str>) -> Result<String, PythonTestError> {
+    pub(crate) async fn run(&self, input: Option<&str>) -> Result<String, PythonTestError> {
         match self {
             Self::ExampleTest => {
                 let example_input: HashMap<String, String> =
@@ -161,14 +165,25 @@ impl PythonTest {
                     serde_json::from_str(Self::non_optional_input(input)?)?;
                 Ok(serialize_for_rust_committer_flow_test(input))
             }
+            Self::ComputeHashSingleTree => {
+                // 1. Get and deserialize input.
+                let input: HashMap<String, String> =
+                    serde_json::from_str(Self::non_optional_input(input)?)?;
+                let (leaf_modifications, storage, root_hash) =
+                    parse_input_single_storage_tree_flow_test(input);
+                // 2. Run the test.
+                let output = single_tree_flow_test(leaf_modifications, storage, root_hash).await;
+                // 3. Serialize and return output.
+                Ok(output)
+            }
         }
     }
 }
 
 // Test that the fetching of the input to flow test is working.
 fn serialize_for_rust_committer_flow_test(input: HashMap<String, String>) -> String {
     let (leaf_modifications, storage, root_hash) = parse_input_single_storage_tree_flow_test(input);
-    // Serialize the leaf modifications to an object that can JSON-serialized.
+    // Serialize the leaf modifications to an object that can be JSON-serialized.
     let leaf_modifications_to_print: HashMap<String, Vec<u8>> = leaf_modifications
         .into_iter()
         .map(|(k, v)| (k.0.to_string(), v.serialize().0))