diff --git a/crates/committer/src/patricia_merkle_tree/external_test_utils.rs b/crates/committer/src/patricia_merkle_tree/external_test_utils.rs index 14233a51..31d7f1ee 100644 --- a/crates/committer/src/patricia_merkle_tree/external_test_utils.rs +++ b/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; fn try_from(value: &U256) -> Result { @@ -51,3 +63,56 @@ pub fn get_random_u256(rng: &mut R, low: U256, high: U256) -> U256 { } result } + +pub async fn single_tree_flow_test( + leaf_modifications: LeafModifications, + 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::>(); + let mut sorted_leaf_indices: Vec = 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::(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") +} diff --git a/crates/committer/src/patricia_merkle_tree/internal_test_utils.rs b/crates/committer/src/patricia_merkle_tree/internal_test_utils.rs index bf2d2b6c..e72423d7 100644 --- a/crates/committer/src/patricia_merkle_tree/internal_test_utils.rs +++ b/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 diff --git a/crates/committer_cli/src/main.rs b/crates/committer_cli/src/main.rs index 21dc4924..c5fd7c4c 100644 --- a/crates/committer_cli/src/main.rs +++ b/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. diff --git a/crates/committer_cli/src/tests/python_tests.rs b/crates/committer_cli/src/tests/python_tests.rs index ad600e7b..5991b76f 100644 --- a/crates/committer_cli/src/tests/python_tests.rs +++ b/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 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 { + pub(crate) async fn run(&self, input: Option<&str>) -> Result { match self { Self::ExampleTest => { let example_input: HashMap = @@ -161,6 +165,17 @@ 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 = + 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) + } } } } @@ -168,7 +183,7 @@ impl PythonTest { // Test that the fetching of the input to flow test is working. fn serialize_for_rust_committer_flow_test(input: HashMap) -> 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> = leaf_modifications .into_iter() .map(|(k, v)| (k.0.to_string(), v.serialize().0))