From 091490ff5102bbf2e8d409fbea9ca91910204b15 Mon Sep 17 00:00:00 2001 From: amosStarkware <88497213+amosStarkware@users.noreply.github.com> Date: Tue, 20 Aug 2024 12:47:16 +0300 Subject: [PATCH] feat(committer): added FilledTree::create function (#494) --- .../filled_tree/errors.rs | 6 +- .../patricia_merkle_tree/filled_tree/tree.rs | 216 ++++++++++++++---- .../filled_tree/tree_test.rs | 194 +++++++++++++--- .../internal_test_utils.rs | 12 +- .../patricia_merkle_tree/node_data/leaf.rs | 3 +- .../original_skeleton_tree/create_tree.rs | 2 +- .../updated_skeleton_tree/tree.rs | 1 + 7 files changed, 343 insertions(+), 91 deletions(-) diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/errors.rs b/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/errors.rs index 2a3189f339..6379ffeff0 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/errors.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/errors.rs @@ -12,8 +12,10 @@ pub enum FilledTreeError { DoubleUpdate { index: NodeIndex, existing_value_as_string: String }, #[error("Got the following error at leaf index {leaf_index:?}: {leaf_error:?}")] Leaf { leaf_error: LeafError, leaf_index: NodeIndex }, - #[error("Missing node at index {0:?}.")] - MissingNode(NodeIndex), + #[error("Missing node placeholder at index {0:?}.")] + MissingNodePlaceholder(NodeIndex), + #[error("Missing leaf input for index {0:?}.")] + MissingLeafInput(NodeIndex), #[error("Missing root.")] MissingRoot, #[error("Poisoned lock: {0}.")] diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/tree.rs b/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/tree.rs index aee60f4eae..db98b66a1b 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/tree.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/tree.rs @@ -25,8 +25,15 @@ pub(crate) type FilledTreeResult = Result; /// Consider a Patricia-Merkle Tree which has been updated with new leaves. /// FilledTree consists of all nodes which were modified in the update, including their updated /// data and hashes. -pub trait FilledTree: Sized { - /// Computes and returns the filled tree. +pub trait FilledTree: Sized + Send { + /// Computes and returns the filled tree and the leaf output map. + fn create<'a, TH: TreeHashFunction + 'static>( + updated_skeleton: impl UpdatedSkeletonTree<'a> + 'static, + leaf_index_to_leaf_input: HashMap, + ) -> impl Future)>> + Send; + + /// Computes and returns the filled tree using the provided leaf modifications. Since the + /// leaves are not computed, no leaf output will be returned. fn create_with_existing_leaves<'a, TH: TreeHashFunction + 'static>( updated_skeleton: impl UpdatedSkeletonTree<'a> + 'static, leaf_modifications: LeafModifications, @@ -47,16 +54,22 @@ pub struct FilledTreeImpl { } impl FilledTreeImpl { - fn initialize_filled_tree_map_with_placeholders<'a>( + fn initialize_filled_tree_output_map_with_placeholders<'a>( updated_skeleton: &impl UpdatedSkeletonTree<'a>, ) -> HashMap>>> { - let mut filled_tree_map = HashMap::new(); + let mut filled_tree_output_map = HashMap::new(); for (index, node) in updated_skeleton.get_nodes() { if !matches!(node, UpdatedSkeletonNode::UnmodifiedSubTree(_)) { - filled_tree_map.insert(index, Mutex::new(None)); + filled_tree_output_map.insert(index, Mutex::new(None)); } } - filled_tree_map + filled_tree_output_map + } + + fn initialize_leaf_output_map_with_placeholders( + leaf_index_to_leaf_input: &HashMap, + ) -> Arc>>> { + Arc::new(leaf_index_to_leaf_input.keys().map(|index| (*index, Mutex::new(None))).collect()) } pub(crate) fn get_all_nodes(&self) -> &HashMap> { @@ -86,34 +99,88 @@ impl FilledTreeImpl { } } } - None => Err(FilledTreeError::MissingNode(index)), + None => Err(FilledTreeError::MissingNodePlaceholder(index)), } } - fn remove_arc_mutex_and_option( - hash_map_in: Arc>>>>, - ) -> FilledTreeResult>> { + // Removes the `Arc` from the map and unwraps the `Mutex` and `Option` from the value. + // If `panic_if_empty_placeholder` is `true`, will panic if an empty placeholder is found. + fn remove_arc_mutex_and_option_from_output_map( + output_map: Arc>>>, + panic_if_empty_placeholder: bool, + ) -> FilledTreeResult> { let mut hash_map_out = HashMap::new(); - for (key, value) in hash_map_in.iter() { + for (key, value) in Arc::into_inner(output_map) + .unwrap_or_else(|| panic!("Cannot retrieve output map from Arc.")) + { let mut value = value .lock() .map_err(|_| FilledTreeError::PoisonedLock("Cannot lock node.".to_owned()))?; match value.take() { - Some(value) => { - hash_map_out.insert(*key, value); + Some(unwrapped_value) => { + hash_map_out.insert(key, unwrapped_value); + } + None => { + if panic_if_empty_placeholder { + panic!("Empty placeholder in an output map for index {0:?}.", key); + } } - None => return Err(FilledTreeError::MissingNode(*key)), } } Ok(hash_map_out) } + fn wrap_leaf_inputs_for_interior_mutability( + leaf_index_to_leaf_input: HashMap, + ) -> Arc>>> { + Arc::new( + leaf_index_to_leaf_input.into_iter().map(|(k, v)| (k, Mutex::new(Some(v)))).collect(), + ) + } + + // If leaf modifications are `None`, will compute the leaf from the corresponding leaf input + // and return the leaf output. Otherwise, will retrieve the leaf from the leaf modifications + // and return `None` in place of the leaf output (ignoring the leaf input). + async fn get_or_compute_leaf( + leaf_modifications: Option>>, + leaf_index_to_leaf_input: Arc>>>, + index: NodeIndex, + ) -> FilledTreeResult<(L, Option)> { + match leaf_modifications { + Some(leaf_modifications) => { + let leaf_data = + L::from_modifications(&index, &leaf_modifications).map_err(|leaf_err| { + FilledTreeError::Leaf { leaf_error: leaf_err, leaf_index: index } + })?; + Ok((leaf_data, None)) + } + None => { + let leaf_input = leaf_index_to_leaf_input + .get(&index) + .ok_or(FilledTreeError::MissingLeafInput(index))? + .lock() + .map_err(|_| FilledTreeError::PoisonedLock("Cannot lock node.".to_owned()))? + .take() + .unwrap_or_else(|| panic!("Leaf input is None for index {0:?}.", index)); + let (leaf_data, leaf_output) = L::create(leaf_input).await.map_err(|leaf_err| { + FilledTreeError::Leaf { leaf_error: leaf_err, leaf_index: index } + })?; + Ok((leaf_data, Some(leaf_output))) + } + } + } + + // Recursively computes the filled tree. If leaf modifications are `None`, will compute the + // leaves from the leaf inputs and fill the leaf output map. Otherwise, will retrieve the + // leaves from the leaf modifications map and ignore the input and output maps. #[async_recursion] async fn compute_filled_tree_rec<'a, TH>( updated_skeleton: Arc + 'async_recursion + 'static>, index: NodeIndex, - leaf_modifications: Arc>, - output_map: Arc>>>>, + leaf_modifications: Option>>, + leaf_index_to_leaf_input: Arc>>>, + filled_tree_output_map: Arc>>>>, + leaf_index_to_leaf_output: Arc>>>, ) -> FilledTreeResult where TH: TreeHashFunction + 'static, @@ -128,14 +195,18 @@ impl FilledTreeImpl { tokio::spawn(Self::compute_filled_tree_rec::( Arc::clone(&updated_skeleton), left_index, - Arc::clone(&leaf_modifications), - Arc::clone(&output_map), + leaf_modifications.as_ref().map(Arc::clone), + Arc::clone(&leaf_index_to_leaf_input), + Arc::clone(&filled_tree_output_map), + Arc::clone(&leaf_index_to_leaf_output), )), tokio::spawn(Self::compute_filled_tree_rec::( Arc::clone(&updated_skeleton), right_index, - Arc::clone(&leaf_modifications), - Arc::clone(&output_map), + leaf_modifications.as_ref().map(Arc::clone), + Arc::clone(&leaf_index_to_leaf_input), + Arc::clone(&filled_tree_output_map), + Arc::clone(&leaf_index_to_leaf_output), )), ); @@ -144,51 +215,54 @@ impl FilledTreeImpl { right_hash: right_hash.await??, }); - let hash_value = TH::compute_node_hash(&data); + let hash = TH::compute_node_hash(&data); Self::write_to_output_map( - output_map, + filled_tree_output_map, index, - FilledNode { hash: hash_value, data }, + FilledNode { hash, data }, )?; - Ok(hash_value) + Ok(hash) } UpdatedSkeletonNode::Edge(path_to_bottom) => { let bottom_node_index = NodeIndex::compute_bottom_index(index, path_to_bottom); let bottom_hash = Self::compute_filled_tree_rec::( Arc::clone(&updated_skeleton), bottom_node_index, - leaf_modifications, - Arc::clone(&output_map), + leaf_modifications.as_ref().map(Arc::clone), + Arc::clone(&leaf_index_to_leaf_input), + Arc::clone(&filled_tree_output_map), + Arc::clone(&leaf_index_to_leaf_output), ) .await?; let data = NodeData::Edge(EdgeData { path_to_bottom: *path_to_bottom, bottom_hash }); - let hash_value = TH::compute_node_hash(&data); + let hash = TH::compute_node_hash(&data); Self::write_to_output_map( - output_map, + filled_tree_output_map, index, - FilledNode { hash: hash_value, data }, + FilledNode { hash, data }, )?; - Ok(hash_value) + Ok(hash) } UpdatedSkeletonNode::UnmodifiedSubTree(hash_result) => Ok(*hash_result), UpdatedSkeletonNode::Leaf => { - // TODO(Amos): use `L::create` when no leaf modifications are provided. - let leaf_data = - L::from_modifications(&index, Arc::clone(&leaf_modifications)).map_err( - |error| FilledTreeError::Leaf { leaf_error: error, leaf_index: index }, - )?; + let (leaf_data, leaf_output) = + Self::get_or_compute_leaf(leaf_modifications, leaf_index_to_leaf_input, index) + .await?; if leaf_data.is_empty() { return Err(FilledTreeError::DeletedLeafInSkeleton(index)); } - let node_data = NodeData::Leaf(leaf_data); - let hash_value = TH::compute_node_hash(&node_data); + let data = NodeData::Leaf(leaf_data); + let hash = TH::compute_node_hash(&data); Self::write_to_output_map( - output_map, + filled_tree_output_map, index, - FilledNode { hash: hash_value, data: node_data }, + FilledNode { hash, data }, )?; - Ok(hash_value) + if let Some(output) = leaf_output { + Self::write_to_output_map(leaf_index_to_leaf_output, index, output)? + }; + Ok(hash) } } } @@ -209,32 +283,82 @@ impl FilledTreeImpl { } impl FilledTree for FilledTreeImpl { + async fn create<'a, TH: TreeHashFunction + 'static>( + updated_skeleton: impl UpdatedSkeletonTree<'a> + 'static, + leaf_index_to_leaf_input: HashMap, + ) -> Result<(Self, HashMap), FilledTreeError> { + // Handle edge cases of no leaf modifications. + if leaf_index_to_leaf_input.is_empty() { + let unmodified = Self::create_unmodified(&updated_skeleton)?; + return Ok((unmodified, HashMap::new())); + } + if updated_skeleton.is_empty() { + return Ok((Self::create_empty(), HashMap::new())); + } + + // Wrap values in `Mutex>` for interior mutability. + let filled_tree_output_map = + Arc::new(Self::initialize_filled_tree_output_map_with_placeholders(&updated_skeleton)); + let leaf_index_to_leaf_output = + Self::initialize_leaf_output_map_with_placeholders(&leaf_index_to_leaf_input); + let wrapped_leaf_index_to_leaf_input = + Self::wrap_leaf_inputs_for_interior_mutability(leaf_index_to_leaf_input); + + // Compute the filled tree. + let root_hash = Self::compute_filled_tree_rec::( + Arc::new(updated_skeleton), + NodeIndex::ROOT, + None, + Arc::clone(&wrapped_leaf_index_to_leaf_input), + Arc::clone(&filled_tree_output_map), + Arc::clone(&leaf_index_to_leaf_output), + ) + .await?; + + Ok(( + FilledTreeImpl { + tree_map: Self::remove_arc_mutex_and_option_from_output_map( + filled_tree_output_map, + true, + )?, + root_hash, + }, + Self::remove_arc_mutex_and_option_from_output_map(leaf_index_to_leaf_output, false)?, + )) + } + async fn create_with_existing_leaves<'a, TH: TreeHashFunction + 'static>( updated_skeleton: impl UpdatedSkeletonTree<'a> + 'static, leaf_modifications: LeafModifications, ) -> FilledTreeResult { + // Handle edge case of no modifications. if leaf_modifications.is_empty() { return Self::create_unmodified(&updated_skeleton); } - if updated_skeleton.is_empty() { return Ok(Self::create_empty()); } // Wrap values in `Mutex>`` for interior mutability. - let filled_tree_map = - Arc::new(Self::initialize_filled_tree_map_with_placeholders(&updated_skeleton)); + let filled_tree_output_map = + Arc::new(Self::initialize_filled_tree_output_map_with_placeholders(&updated_skeleton)); + // Compute the filled tree. let root_hash = Self::compute_filled_tree_rec::( Arc::new(updated_skeleton), NodeIndex::ROOT, - Arc::new(leaf_modifications), - Arc::clone(&filled_tree_map), + Some(leaf_modifications.into()), + Arc::new(HashMap::new()), + Arc::clone(&filled_tree_output_map), + Arc::new(HashMap::new()), ) .await?; Ok(FilledTreeImpl { - tree_map: Self::remove_arc_mutex_and_option(filled_tree_map)?, + tree_map: Self::remove_arc_mutex_and_option_from_output_map( + filled_tree_output_map, + true, + )?, root_hash, }) } diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/tree_test.rs b/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/tree_test.rs index cb8be41992..f564700146 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/tree_test.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/tree_test.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use crate::felt::Felt; use crate::hash::hash_trait::HashOutput; +use crate::patricia_merkle_tree::filled_tree::errors::FilledTreeError; use crate::patricia_merkle_tree::filled_tree::node::FilledNode; use crate::patricia_merkle_tree::filled_tree::tree::{FilledTree, FilledTreeImpl}; use crate::patricia_merkle_tree::internal_test_utils::{ @@ -9,6 +10,7 @@ use crate::patricia_merkle_tree::internal_test_utils::{ OriginalSkeletonMockTrieConfig, TestTreeHashFunction, }; +use crate::patricia_merkle_tree::node_data::errors::LeafError; use crate::patricia_merkle_tree::node_data::inner_node::{ BinaryData, EdgeData, @@ -16,7 +18,7 @@ use crate::patricia_merkle_tree::node_data::inner_node::{ NodeData, PathToBottom, }; -use crate::patricia_merkle_tree::node_data::leaf::SkeletonLeaf; +use crate::patricia_merkle_tree::node_data::leaf::{LeafModifications, SkeletonLeaf}; use crate::patricia_merkle_tree::original_skeleton_tree::tree::OriginalSkeletonTreeImpl; use crate::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices}; use crate::patricia_merkle_tree::updated_skeleton_tree::node::UpdatedSkeletonNode; @@ -67,28 +69,8 @@ async fn test_filled_tree_sanity() { /// i=35: leaf i=36: leaf i=63: leaf /// v=1 v=2 v=3 async fn test_small_filled_tree() { - // Set up the updated skeleton tree. - let new_leaves = [(35, "0x1"), (36, "0x2"), (63, "0x3")]; - let nodes_in_skeleton_tree: Vec<(NodeIndex, UpdatedSkeletonNode)> = [ - create_binary_updated_skeleton_node_for_testing(1), - create_path_to_bottom_edge_updated_skeleton_node_for_testing(2, 0, 1), - create_path_to_bottom_edge_updated_skeleton_node_for_testing(3, 15, 4), - create_binary_updated_skeleton_node_for_testing(4), - create_path_to_bottom_edge_updated_skeleton_node_for_testing(8, 3, 2), - create_path_to_bottom_edge_updated_skeleton_node_for_testing(9, 0, 2), - ] - .into_iter() - .chain( - new_leaves.iter().map(|(index, _)| create_leaf_updated_skeleton_node_for_testing(*index)), - ) - .collect(); - let skeleton_tree: UpdatedSkeletonNodeMap = nodes_in_skeleton_tree.into_iter().collect(); - - let updated_skeleton_tree = UpdatedSkeletonTreeImpl { skeleton_tree }; - let modifications = new_leaves - .iter() - .map(|(index, value)| (NodeIndex::from(*index), MockLeaf(Felt::from_hex(value).unwrap()))) - .collect(); + let (updated_skeleton_tree, modifications) = + get_small_tree_updated_skeleton_and_leaf_modifications(); // Compute the hash values. let filled_tree = FilledTreeImpl::create_with_existing_leaves::( @@ -100,21 +82,117 @@ async fn test_small_filled_tree() { let filled_tree_map = filled_tree.get_all_nodes(); let root_hash = filled_tree.get_root_hash(); - // The expected hash values were computed separately. - let expected_root_hash = HashOutput(Felt::from_hex("0x21").unwrap()); - let expected_filled_tree_map = HashMap::from([ - create_mock_binary_entry_for_testing(1, "0x21", "0xb", "0x16"), - create_mock_edge_entry_for_testing(2, "0xb", 0, 1, "0xa"), - create_mock_edge_entry_for_testing(3, "0x16", 15, 4, "0x3"), - create_mock_binary_entry_for_testing(4, "0xa", "0x6", "0x4"), - create_mock_edge_entry_for_testing(8, "0x6", 3, 2, "0x1"), - create_mock_edge_entry_for_testing(9, "0x4", 0, 2, "0x2"), - create_mock_leaf_entry_for_testing(35, "0x1"), - create_mock_leaf_entry_for_testing(36, "0x2"), - create_mock_leaf_entry_for_testing(63, "0x3"), - ]); + let (expected_filled_tree_map, expected_root_hash) = + get_small_tree_expected_filled_tree_map_and_root_hash(); + assert_eq!(filled_tree_map, &expected_filled_tree_map); + assert_eq!(root_hash, expected_root_hash, "Root hash mismatch"); +} + +#[tokio::test(flavor = "multi_thread")] +/// Similar to `test_small_filled_tree`, except the tree is created via `FilledTree:create()`. +async fn test_small_filled_tree_create() { + let (updated_skeleton_tree, modifications) = + get_small_tree_updated_skeleton_and_leaf_modifications(); + let expected_leaf_index_to_leaf_output: HashMap = + modifications.iter().map(|(index, leaf)| (*index, leaf.0.to_hex())).collect(); + let leaf_index_to_leaf_input: HashMap = + modifications.into_iter().map(|(index, leaf)| (index, leaf.0)).collect(); + + // Compute the hash values. + let (filled_tree, leaf_index_to_leaf_output) = FilledTreeImpl::create::( + updated_skeleton_tree, + leaf_index_to_leaf_input, + ) + .await + .unwrap(); + let filled_tree_map = filled_tree.get_all_nodes(); + let root_hash = filled_tree.get_root_hash(); + + let (expected_filled_tree_map, expected_root_hash) = + get_small_tree_expected_filled_tree_map_and_root_hash(); assert_eq!(filled_tree_map, &expected_filled_tree_map); assert_eq!(root_hash, expected_root_hash, "Root hash mismatch"); + assert_eq!( + leaf_index_to_leaf_output, expected_leaf_index_to_leaf_output, + "Leaf output mismatch" + ); +} + +#[tokio::test(flavor = "multi_thread")] +/// Test the edge case of creating a tree with no leaf modifications. +async fn test_empty_leaf_modifications() { + let root_hash = HashOutput(Felt::ONE); + let unmodified_updated_skeleton_tree_map = + HashMap::from([(NodeIndex::ROOT, UpdatedSkeletonNode::UnmodifiedSubTree(root_hash))]); + + // Test `create_with_existing_leaves`. + let filled_tree = FilledTreeImpl::create_with_existing_leaves::( + UpdatedSkeletonTreeImpl { skeleton_tree: unmodified_updated_skeleton_tree_map.clone() }, + HashMap::new(), + ) + .await + .unwrap(); + assert_eq!(filled_tree.get_root_hash(), root_hash); + assert!(filled_tree.get_all_nodes().is_empty()); + + // Test `create`. + let (filled_tree, leaf_index_to_leaf_output) = FilledTreeImpl::create::( + UpdatedSkeletonTreeImpl { skeleton_tree: unmodified_updated_skeleton_tree_map }, + HashMap::new(), + ) + .await + .unwrap(); + assert_eq!(filled_tree.get_root_hash(), root_hash); + assert!(filled_tree.get_all_nodes().is_empty()); + assert!(leaf_index_to_leaf_output.is_empty()); +} + +#[tokio::test(flavor = "multi_thread")] +/// Test the edge case of creating a tree with an empty updated skeleton and non-empty leaf +/// modifications. This can only happen when the leaf modifications don't actually modify any nodes. +async fn test_empty_updated_skeleton() { + let leaf_modifications = HashMap::from([(NodeIndex::FIRST_LEAF, Felt::ONE)]); + + // Test `create`. + let (filled_tree, leaf_index_to_leaf_output) = FilledTreeImpl::create::( + UpdatedSkeletonTreeImpl { skeleton_tree: HashMap::new() }, + leaf_modifications, + ) + .await + .unwrap(); + assert_eq!(filled_tree.get_root_hash(), HashOutput::ROOT_OF_EMPTY_TREE); + assert!(filled_tree.get_all_nodes().is_empty()); + assert!(leaf_index_to_leaf_output.is_empty()); + + // `create_with_existing_leaves` is tested in `test_delete_leaf_from_empty_tree`. +} + +#[tokio::test(flavor = "multi_thread")] +/// Tests the case of a leaf computation error. +async fn test_leaf_computation_error() { + let (first_leaf_index, second_leaf_index) = (NodeIndex(2_u32.into()), NodeIndex(3_u32.into())); + let leaf_input_map = + HashMap::from([(first_leaf_index, 1_u128.into()), (second_leaf_index, Felt::MAX)]); + let skeleton_tree = HashMap::from([ + (NodeIndex::ROOT, UpdatedSkeletonNode::Binary), + (first_leaf_index, UpdatedSkeletonNode::Leaf), + (second_leaf_index, UpdatedSkeletonNode::Leaf), + ]); + + let result = FilledTreeImpl::create::( + UpdatedSkeletonTreeImpl { skeleton_tree }, + leaf_input_map, + ) + .await; + match result { + Err(FilledTreeError::Leaf { + leaf_error: LeafError::LeafComputationError(_), + leaf_index, + }) => { + assert_eq!(leaf_index, second_leaf_index); + } + _ => panic!("Expected leaf computation error."), + }; } #[tokio::test(flavor = "multi_thread")] @@ -217,6 +295,50 @@ async fn test_delete_leaf_from_empty_tree() { assert!(root_hash == HashOutput::ROOT_OF_EMPTY_TREE); } +fn get_small_tree_updated_skeleton_and_leaf_modifications() +-> (UpdatedSkeletonTreeImpl, LeafModifications) { + // Set up the updated skeleton tree. + let new_leaves = [(35, "0x1"), (36, "0x2"), (63, "0x3")]; + let nodes_in_skeleton_tree: Vec<(NodeIndex, UpdatedSkeletonNode)> = [ + create_binary_updated_skeleton_node_for_testing(1), + create_path_to_bottom_edge_updated_skeleton_node_for_testing(2, 0, 1), + create_path_to_bottom_edge_updated_skeleton_node_for_testing(3, 15, 4), + create_binary_updated_skeleton_node_for_testing(4), + create_path_to_bottom_edge_updated_skeleton_node_for_testing(8, 3, 2), + create_path_to_bottom_edge_updated_skeleton_node_for_testing(9, 0, 2), + ] + .into_iter() + .chain( + new_leaves.iter().map(|(index, _)| create_leaf_updated_skeleton_node_for_testing(*index)), + ) + .collect(); + let skeleton_tree: UpdatedSkeletonNodeMap = nodes_in_skeleton_tree.into_iter().collect(); + + let updated_skeleton_tree = UpdatedSkeletonTreeImpl { skeleton_tree }; + let modifications = new_leaves + .iter() + .map(|(index, value)| (NodeIndex::from(*index), MockLeaf(Felt::from_hex(value).unwrap()))) + .collect(); + (updated_skeleton_tree, modifications) +} + +fn get_small_tree_expected_filled_tree_map_and_root_hash() +-> (HashMap>, HashOutput) { + let expected_root_hash = HashOutput(Felt::from_hex("0x21").unwrap()); + let expected_filled_tree_map = HashMap::from([ + create_mock_binary_entry_for_testing(1, "0x21", "0xb", "0x16"), + create_mock_edge_entry_for_testing(2, "0xb", 0, 1, "0xa"), + create_mock_edge_entry_for_testing(3, "0x16", 15, 4, "0x3"), + create_mock_binary_entry_for_testing(4, "0xa", "0x6", "0x4"), + create_mock_edge_entry_for_testing(8, "0x6", 3, 2, "0x1"), + create_mock_edge_entry_for_testing(9, "0x4", 0, 2, "0x2"), + create_mock_leaf_entry_for_testing(35, "0x1"), + create_mock_leaf_entry_for_testing(36, "0x2"), + create_mock_leaf_entry_for_testing(63, "0x3"), + ]); + (expected_filled_tree_map, expected_root_hash) +} + fn create_binary_updated_skeleton_node_for_testing( index: u128, ) -> (NodeIndex, UpdatedSkeletonNode) { diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/internal_test_utils.rs b/crates/starknet_patricia/src/patricia_merkle_tree/internal_test_utils.rs index aea5d7af10..285cd6f088 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/internal_test_utils.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/internal_test_utils.rs @@ -7,7 +7,7 @@ use crate::generate_trie_config; use crate::hash::hash_trait::HashOutput; use crate::patricia_merkle_tree::external_test_utils::get_random_u256; use crate::patricia_merkle_tree::filled_tree::tree::FilledTreeImpl; -use crate::patricia_merkle_tree::node_data::errors::LeafResult; +use crate::patricia_merkle_tree::node_data::errors::{LeafError, LeafResult}; use crate::patricia_merkle_tree::node_data::inner_node::{EdgePathLength, NodeData, PathToBottom}; use crate::patricia_merkle_tree::node_data::leaf::{Leaf, SkeletonLeaf}; use crate::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonTreeConfig; @@ -50,15 +50,19 @@ impl Deserializable for MockLeaf { } impl Leaf for MockLeaf { - type Input = Self; - type Output = (); + type Input = Felt; + type Output = String; fn is_empty(&self) -> bool { self.0 == Felt::ZERO } + // Create a leaf with value equal to input. If input is `Felt::MAX`, returns an error. async fn create(input: Self::Input) -> LeafResult<(Self, Self::Output)> { - Ok((input, ())) + if input == Felt::MAX { + return Err(LeafError::LeafComputationError("Leaf computation error".to_string())); + } + Ok((Self(input), input.to_hex())) } } diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/node_data/leaf.rs b/crates/starknet_patricia/src/patricia_merkle_tree/node_data/leaf.rs index a277b6ab28..aafd891c52 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/node_data/leaf.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/node_data/leaf.rs @@ -1,7 +1,6 @@ use std::collections::HashMap; use std::fmt::Debug; use std::future::Future; -use std::sync::Arc; use crate::felt::Felt; use crate::patricia_merkle_tree::node_data::errors::{LeafError, LeafResult}; @@ -28,7 +27,7 @@ pub trait Leaf: Clone + Sync + Send + DBObject + Deserializable + Default + Debu /// missing. fn from_modifications( index: &NodeIndex, - leaf_modifications: Arc>, + leaf_modifications: &LeafModifications, ) -> LeafResult { let leaf_data = leaf_modifications .get(index) diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/original_skeleton_tree/create_tree.rs b/crates/starknet_patricia/src/patricia_merkle_tree/original_skeleton_tree/create_tree.rs index 08f859cf13..c17cf47243 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/original_skeleton_tree/create_tree.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/original_skeleton_tree/create_tree.rs @@ -294,7 +294,7 @@ impl<'a> OriginalSkeletonTreeImpl<'a> { } } - fn create_empty(sorted_leaf_indices: SortedLeafIndices<'a>) -> Self { + pub(crate) fn create_empty(sorted_leaf_indices: SortedLeafIndices<'a>) -> Self { Self { nodes: HashMap::new(), sorted_leaf_indices } } diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs b/crates/starknet_patricia/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs index 6afc0cf02f..679da99832 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs @@ -37,6 +37,7 @@ pub trait UpdatedSkeletonTree<'a>: Sized + Send + Sync { fn get_node(&self, index: NodeIndex) -> UpdatedSkeletonTreeResult<&UpdatedSkeletonNode>; } // TODO(Dori, 1/7/2024): Make this a tuple struct. +#[derive(Debug)] pub struct UpdatedSkeletonTreeImpl { pub(crate) skeleton_tree: UpdatedSkeletonNodeMap, }