diff --git a/crates/starknet_committer/src/db/external_test_utils.rs b/crates/starknet_committer/src/db/external_test_utils.rs index 6375f4b6fa9..d4b91aa4d2b 100644 --- a/crates/starknet_committer/src/db/external_test_utils.rs +++ b/crates/starknet_committer/src/db/external_test_utils.rs @@ -15,10 +15,12 @@ use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::tree::{ UpdatedSkeletonTree, UpdatedSkeletonTreeImpl, }; +use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix}; use starknet_patricia_storage::map_storage::MapStorage; use crate::db::facts_db::create_facts_tree::create_original_skeleton_tree; +// TODO(Ariel, 14/12/2025): make this generic over the layout. pub async fn tree_computation_flow( leaf_modifications: LeafModifications, storage: &mut MapStorage, @@ -27,7 +29,7 @@ pub async fn tree_computation_flow( ) -> FilledTreeImpl where TH: TreeHashFunction + 'static, - L: Leaf + 'static, + L: Leaf + HasStaticPrefix + 'static, { let mut sorted_leaf_indices: Vec = leaf_modifications.keys().copied().collect(); let sorted_leaf_indices = SortedLeafIndices::new(&mut sorted_leaf_indices); @@ -37,6 +39,7 @@ where sorted_leaf_indices, &config, &leaf_modifications, + &EmptyKeyContext, ) .await .expect("Failed to create the original skeleton tree"); @@ -63,7 +66,10 @@ where .expect("Failed to create the filled tree") } -pub async fn single_tree_flow_test + 'static>( +pub async fn single_tree_flow_test< + L: Leaf + HasStaticPrefix + 'static, + TH: TreeHashFunction + 'static, +>( leaf_modifications: LeafModifications, storage: &mut MapStorage, root_hash: HashOutput, @@ -85,7 +91,7 @@ pub async fn single_tree_flow_test + let json_hash = &json!(hash_result.0.to_hex_string()); result_map.insert("root_hash", json_hash); // Serlialize the storage modifications. - let json_storage = &json!(filled_tree.serialize()); + let json_storage = &json!(filled_tree.serialize(&EmptyKeyContext)); result_map.insert("storage_changes", json_storage); serde_json::to_string(&result_map).expect("serialization failed") } diff --git a/crates/starknet_committer/src/db/facts_db/create_facts_tree.rs b/crates/starknet_committer/src/db/facts_db/create_facts_tree.rs index 50aafe5a021..340ba80387f 100644 --- a/crates/starknet_committer/src/db/facts_db/create_facts_tree.rs +++ b/crates/starknet_committer/src/db/facts_db/create_facts_tree.rs @@ -20,6 +20,7 @@ use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::tree::{ }; use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait; use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices}; +use starknet_patricia_storage::db_object::HasStaticPrefix; use starknet_patricia_storage::storage_trait::Storage; use tracing::warn; @@ -48,13 +49,15 @@ async fn fetch_nodes<'a, L: Leaf>( leaf_modifications: &LeafModifications, config: &impl OriginalSkeletonTreeConfig, mut previous_leaves: Option<&mut HashMap>, + key_context: &::KeyContext, ) -> OriginalSkeletonTreeResult<()> { let mut current_subtrees = subtrees; let mut next_subtrees = Vec::new(); while !current_subtrees.is_empty() { let should_fetch_modified_leaves = config.compare_modified_leaves() || previous_leaves.is_some(); - let filled_roots = calculate_subtrees_roots::(¤t_subtrees, storage).await?; + let filled_roots = + calculate_subtrees_roots::(¤t_subtrees, storage, key_context).await?; for (filled_root, subtree) in filled_roots.into_iter().zip(current_subtrees.iter()) { match filled_root.data { // Binary node. @@ -150,6 +153,7 @@ pub async fn create_original_skeleton_tree<'a, L: Leaf>( sorted_leaf_indices: SortedLeafIndices<'a>, config: &impl OriginalSkeletonTreeConfig, leaf_modifications: &LeafModifications, + key_context: &::KeyContext, ) -> OriginalSkeletonTreeResult> { if sorted_leaf_indices.is_empty() { return Ok(OriginalSkeletonTreeImpl::create_unmodified(root_hash)); @@ -171,6 +175,7 @@ pub async fn create_original_skeleton_tree<'a, L: Leaf>( leaf_modifications, config, None, + key_context, ) .await?; Ok(skeleton_tree) @@ -182,6 +187,7 @@ pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>( sorted_leaf_indices: SortedLeafIndices<'a>, leaf_modifications: &LeafModifications, config: &impl OriginalSkeletonTreeConfig, + key_context: &::KeyContext, ) -> OriginalSkeletonTreeResult<(OriginalSkeletonTreeImpl<'a>, HashMap)> { if sorted_leaf_indices.is_empty() { let unmodified = OriginalSkeletonTreeImpl::create_unmodified(root_hash); @@ -203,15 +209,20 @@ pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>( leaf_modifications, config, Some(&mut leaves), + key_context, ) .await?; Ok((skeleton_tree, leaves)) } +/// Prepares the OS inputs by fetching paths to the given leaves (i.e. their induced Skeleton tree). +/// Note that ATM, the Rust committer does not manage history and is not used for storage proofs; +/// Thus, this function assumes facts layout. pub async fn get_leaves<'a, L: Leaf>( storage: &mut impl Storage, root_hash: HashOutput, sorted_leaf_indices: SortedLeafIndices<'a>, + key_context: &::KeyContext, ) -> OriginalSkeletonTreeResult> { let config = NoCompareOriginalSkeletonTrieConfig::default(); let leaf_modifications = LeafModifications::new(); @@ -221,6 +232,7 @@ pub async fn get_leaves<'a, L: Leaf>( sorted_leaf_indices, &leaf_modifications, &config, + key_context, ) .await?; Ok(previous_leaves) diff --git a/crates/starknet_committer/src/db/facts_db/create_facts_tree_test.rs b/crates/starknet_committer/src/db/facts_db/create_facts_tree_test.rs index 923dc0affbe..5dbc7eed214 100644 --- a/crates/starknet_committer/src/db/facts_db/create_facts_tree_test.rs +++ b/crates/starknet_committer/src/db/facts_db/create_facts_tree_test.rs @@ -18,7 +18,7 @@ use starknet_patricia::patricia_merkle_tree::external_test_utils::{ use starknet_patricia::patricia_merkle_tree::node_data::leaf::LeafModifications; use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::node::OriginalSkeletonNode; use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices, SubTreeHeight}; -use starknet_patricia_storage::db_object::DBObject; +use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext}; use starknet_patricia_storage::map_storage::MapStorage; use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue}; use starknet_types_core::felt::Felt; @@ -217,6 +217,7 @@ async fn test_create_tree( sorted_leaf_indices, &config, &leaf_modifications, + &EmptyKeyContext, ) .await .unwrap(); @@ -225,7 +226,7 @@ async fn test_create_tree( pub(crate) fn create_mock_leaf_entry(val: u128) -> (DbKey, DbValue) { let leaf = MockLeaf(Felt::from(val)); - (leaf.get_db_key(&leaf.0.to_bytes_be()), leaf.serialize()) + (leaf.get_db_key(&EmptyKeyContext, &leaf.0.to_bytes_be()), leaf.serialize()) } fn create_mock_leaf_modifications( diff --git a/crates/starknet_committer/src/db/facts_db/db.rs b/crates/starknet_committer/src/db/facts_db/db.rs index 107e1ac79da..9e75d912323 100644 --- a/crates/starknet_committer/src/db/facts_db/db.rs +++ b/crates/starknet_committer/src/db/facts_db/db.rs @@ -7,6 +7,7 @@ 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::original_skeleton_tree::tree::OriginalSkeletonTreeImpl; use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices}; +use starknet_patricia_storage::db_object::EmptyKeyContext; use starknet_patricia_storage::map_storage::MapStorage; use starknet_patricia_storage::storage_trait::{ create_db_key, @@ -65,6 +66,7 @@ impl FactsDb { contracts_trie_sorted_indices, &HashMap::new(), &OriginalSkeletonContractsTrieConfig::new(), + &EmptyKeyContext, ) .await?) } @@ -93,6 +95,7 @@ impl FactsDb { *sorted_leaf_indices, &config, updates, + &EmptyKeyContext, ) .await?; storage_tries.insert(*address, original_skeleton); @@ -115,6 +118,7 @@ impl FactsDb { contracts_trie_sorted_indices, &config, actual_classes_updates, + &EmptyKeyContext, ) .await?) } @@ -175,9 +179,9 @@ impl ForestWriter for FactsDb { filled_forest .storage_tries .values() - .flat_map(|tree| tree.serialize().into_iter()) - .chain(filled_forest.contracts_trie.serialize()) - .chain(filled_forest.classes_trie.serialize()) + .flat_map(|tree| tree.serialize(&EmptyKeyContext).into_iter()) + .chain(filled_forest.contracts_trie.serialize(&EmptyKeyContext)) + .chain(filled_forest.classes_trie.serialize(&EmptyKeyContext)) .collect() } diff --git a/crates/starknet_committer/src/db/facts_db/traversal.rs b/crates/starknet_committer/src/db/facts_db/traversal.rs index 93e6f36e46c..c260ff22b1b 100644 --- a/crates/starknet_committer/src/db/facts_db/traversal.rs +++ b/crates/starknet_committer/src/db/facts_db/traversal.rs @@ -10,6 +10,7 @@ use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{ use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf; use starknet_patricia::patricia_merkle_tree::traversal::{SubTreeTrait, TraversalResult}; use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices}; +use starknet_patricia_storage::db_object::HasStaticPrefix; use starknet_patricia_storage::errors::StorageError; use starknet_patricia_storage::storage_trait::{create_db_key, DbKey, Storage}; @@ -23,12 +24,16 @@ pub mod traversal_test; pub async fn calculate_subtrees_roots<'a, L: Leaf>( subtrees: &[FactsSubTree<'a>], storage: &mut impl Storage, + key_context: &::KeyContext, ) -> TraversalResult>> { let mut subtrees_roots = vec![]; let db_keys: Vec = subtrees .iter() .map(|subtree| { - create_db_key(subtree.get_root_prefix::(), &subtree.root_hash.0.to_bytes_be()) + create_db_key( + subtree.get_root_prefix::(key_context), + &subtree.root_hash.0.to_bytes_be(), + ) }) .collect(); @@ -49,6 +54,7 @@ pub async fn fetch_patricia_paths( root_hash: HashOutput, sorted_leaf_indices: SortedLeafIndices<'_>, leaves: Option<&mut HashMap>, + key_context: &::KeyContext, ) -> TraversalResult { let mut witnesses = PreimageMap::new(); @@ -58,7 +64,14 @@ pub async fn fetch_patricia_paths( let main_subtree = FactsSubTree::create(sorted_leaf_indices, NodeIndex::ROOT, root_hash); - fetch_patricia_paths_inner::(storage, vec![main_subtree], &mut witnesses, leaves).await?; + fetch_patricia_paths_inner::( + storage, + vec![main_subtree], + &mut witnesses, + leaves, + key_context, + ) + .await?; Ok(witnesses) } @@ -74,11 +87,13 @@ pub(crate) async fn fetch_patricia_paths_inner<'a, L: Leaf>( subtrees: Vec>, witnesses: &mut PreimageMap, mut leaves: Option<&mut HashMap>, + key_context: &::KeyContext, ) -> TraversalResult<()> { let mut current_subtrees = subtrees; let mut next_subtrees = Vec::new(); while !current_subtrees.is_empty() { - let filled_roots = calculate_subtrees_roots::(¤t_subtrees, storage).await?; + let filled_roots = + calculate_subtrees_roots::(¤t_subtrees, storage, key_context).await?; for (filled_root, subtree) in filled_roots.into_iter().zip(current_subtrees.iter()) { // Always insert root. // No need to insert an unmodified node (which is not the root), because its parent is diff --git a/crates/starknet_committer/src/db/facts_db/traversal_test.rs b/crates/starknet_committer/src/db/facts_db/traversal_test.rs index 721b06e7d68..373234f960c 100644 --- a/crates/starknet_committer/src/db/facts_db/traversal_test.rs +++ b/crates/starknet_committer/src/db/facts_db/traversal_test.rs @@ -27,6 +27,7 @@ use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{ }; use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait; use starknet_patricia::patricia_merkle_tree::types::{SortedLeafIndices, SubTreeHeight}; +use starknet_patricia_storage::db_object::EmptyKeyContext; use starknet_patricia_storage::map_storage::MapStorage; use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue}; use starknet_types_core::felt::Felt; @@ -87,6 +88,7 @@ async fn test_fetch_patricia_paths_inner_impl( vec![main_subtree], &mut nodes, Some(&mut fetched_leaves), + &EmptyKeyContext, ) .await .unwrap(); diff --git a/crates/starknet_committer/src/db/facts_db/types.rs b/crates/starknet_committer/src/db/facts_db/types.rs index 3d0a718c683..b8531ab4701 100644 --- a/crates/starknet_committer/src/db/facts_db/types.rs +++ b/crates/starknet_committer/src/db/facts_db/types.rs @@ -3,6 +3,7 @@ use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::PatriciaPr use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf; use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait; use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices}; +use starknet_patricia_storage::db_object::HasStaticPrefix; use starknet_patricia_storage::storage_trait::DbKeyPrefix; #[derive(Debug, PartialEq)] @@ -35,9 +36,12 @@ impl<'a> SubTreeTrait<'a> for FactsSubTree<'a> { false } - fn get_root_prefix(&self) -> DbKeyPrefix { + fn get_root_prefix( + &self, + key_context: &::KeyContext, + ) -> DbKeyPrefix { if self.is_leaf() { - PatriciaPrefix::Leaf(L::get_static_prefix()).into() + PatriciaPrefix::Leaf(L::get_static_prefix(key_context)).into() } else { PatriciaPrefix::InnerNode.into() } diff --git a/crates/starknet_committer/src/forest/filled_forest.rs b/crates/starknet_committer/src/forest/filled_forest.rs index cf2812b0bd9..24dd6618496 100644 --- a/crates/starknet_committer/src/forest/filled_forest.rs +++ b/crates/starknet_committer/src/forest/filled_forest.rs @@ -117,8 +117,8 @@ impl FilledForest { 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(). + // `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 = contract_address_into_node_index(&contract_address); let original_contract_state = original_contracts_trie_leaves diff --git a/crates/starknet_committer/src/forest/skeleton_forest_test.rs b/crates/starknet_committer/src/forest/skeleton_forest_test.rs index 32dd56f34d7..54dbb7c4ff1 100644 --- a/crates/starknet_committer/src/forest/skeleton_forest_test.rs +++ b/crates/starknet_committer/src/forest/skeleton_forest_test.rs @@ -17,7 +17,7 @@ use starknet_patricia::patricia_merkle_tree::external_test_utils::{ }; use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::tree::OriginalSkeletonTreeImpl; use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices, SubTreeHeight}; -use starknet_patricia_storage::db_object::DBObject; +use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext}; use starknet_patricia_storage::map_storage::MapStorage; use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue}; use starknet_types_core::felt::Felt; @@ -50,12 +50,12 @@ macro_rules! compare_skeleton_tree { pub(crate) fn create_storage_leaf_entry(val: u128) -> (DbKey, DbValue) { let leaf = StarknetStorageValue(Felt::from(val)); - (leaf.get_db_key(&leaf.0.to_bytes_be()), leaf.serialize()) + (leaf.get_db_key(&EmptyKeyContext, &leaf.0.to_bytes_be()), leaf.serialize()) } pub(crate) fn create_compiled_class_leaf_entry(val: u128) -> (DbKey, DbValue) { let leaf = CompiledClassHash(Felt::from(val)); - (leaf.get_db_key(&leaf.0.to_bytes_be()), leaf.serialize()) + (leaf.get_db_key(&EmptyKeyContext, &leaf.0.to_bytes_be()), leaf.serialize()) } pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) { @@ -65,7 +65,7 @@ pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) { storage_root_hash: HashOutput(felt), class_hash: ClassHash(felt), }; - (leaf.get_db_key(&felt.to_bytes_be()), leaf.serialize()) + (leaf.get_db_key(&EmptyKeyContext, &felt.to_bytes_be()), leaf.serialize()) } // This test uses addition hash for simplicity (i.e hash(a,b) = a + b). 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 bcaf7a7fd08..0979fe3f6ea 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 @@ -5,7 +5,7 @@ use starknet_patricia::patricia_merkle_tree::node_data::errors::{LeafError, Leaf 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::db_object::HasStaticPrefix; +use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix}; use starknet_patricia_storage::storage_trait::DbKeyPrefix; use starknet_types_core::felt::Felt; @@ -22,7 +22,8 @@ pub struct ContractState { } impl HasStaticPrefix for StarknetStorageValue { - fn get_static_prefix() -> DbKeyPrefix { + type KeyContext = EmptyKeyContext; + fn get_static_prefix(_key_context: &Self::KeyContext) -> DbKeyPrefix { CommitterLeafPrefix::StorageLeaf.into() } } @@ -41,7 +42,8 @@ impl Leaf for StarknetStorageValue { } impl HasStaticPrefix for CompiledClassHash { - fn get_static_prefix() -> DbKeyPrefix { + type KeyContext = EmptyKeyContext; + fn get_static_prefix(_key_context: &Self::KeyContext) -> DbKeyPrefix { CommitterLeafPrefix::CompiledClassLeaf.into() } } @@ -60,7 +62,8 @@ impl Leaf for CompiledClassHash { } impl HasStaticPrefix for ContractState { - fn get_static_prefix() -> DbKeyPrefix { + type KeyContext = EmptyKeyContext; + fn get_static_prefix(_key_context: &Self::KeyContext) -> DbKeyPrefix { CommitterLeafPrefix::StateTreeLeaf.into() } } diff --git a/crates/starknet_committer/src/patricia_merkle_tree/tree.rs b/crates/starknet_committer/src/patricia_merkle_tree/tree.rs index 912babc2b24..2326aa47a7e 100644 --- a/crates/starknet_committer/src/patricia_merkle_tree/tree.rs +++ b/crates/starknet_committer/src/patricia_merkle_tree/tree.rs @@ -6,6 +6,7 @@ use starknet_patricia::generate_trie_config; use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonTreeConfig; use starknet_patricia::patricia_merkle_tree::traversal::TraversalResult; use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices}; +use starknet_patricia_storage::db_object::EmptyKeyContext; use starknet_patricia_storage::storage_trait::Storage; use crate::block_committer::input::{ @@ -76,6 +77,7 @@ async fn fetch_all_patricia_paths( classes_trie_root_hash, class_sorted_leaf_indices, leaves, + &EmptyKeyContext, ) .await?; @@ -86,6 +88,7 @@ async fn fetch_all_patricia_paths( contracts_trie_root_hash, contract_sorted_leaf_indices, Some(&mut leaves), + &EmptyKeyContext, ) .await?; @@ -105,6 +108,7 @@ async fn fetch_all_patricia_paths( storage_root_hash, *sorted_leaf_indices, leaves, + &EmptyKeyContext, ) .await?; contracts_trie_storage_proofs.insert( diff --git a/crates/starknet_committer_and_os_cli/src/committer_cli/tests/python_tests.rs b/crates/starknet_committer_and_os_cli/src/committer_cli/tests/python_tests.rs index 85ea1d3b89b..8f64dfc68e2 100644 --- a/crates/starknet_committer_and_os_cli/src/committer_cli/tests/python_tests.rs +++ b/crates/starknet_committer_and_os_cli/src/committer_cli/tests/python_tests.rs @@ -26,7 +26,7 @@ use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{ PathToBottom, }; use starknet_patricia::patricia_merkle_tree::types::SubTreeHeight; -use starknet_patricia_storage::db_object::DBObject; +use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext}; use starknet_patricia_storage::errors::DeserializationError; use starknet_patricia_storage::map_storage::MapStorage; use starknet_patricia_storage::storage_trait::{DbKey, DbValue, Storage}; @@ -421,17 +421,17 @@ pub(crate) fn test_node_db_key() -> String { data: NodeData::Binary(BinaryData { left_hash: hash, right_hash: hash }), hash, }; - let binary_node_key = binary_node.db_key().0; + let binary_node_key = binary_node.db_key(&EmptyKeyContext).0; let edge_node: FilledNode = FilledNode { data: NodeData::Edge(EdgeData { bottom_hash: hash, path_to_bottom: Default::default() }), hash, }; - let edge_node_key = edge_node.db_key().0; + let edge_node_key = edge_node.db_key(&EmptyKeyContext).0; let storage_leaf = FilledNode { data: NodeData::Leaf(StarknetStorageValue(zero)), hash }; - let storage_leaf_key = storage_leaf.db_key().0; + let storage_leaf_key = storage_leaf.db_key(&EmptyKeyContext).0; let state_tree_leaf = FilledNode { data: NodeData::Leaf(ContractState { @@ -441,10 +441,10 @@ pub(crate) fn test_node_db_key() -> String { }), hash, }; - let state_tree_leaf_key = state_tree_leaf.db_key().0; + let state_tree_leaf_key = state_tree_leaf.db_key(&EmptyKeyContext).0; let compiled_class_leaf = FilledNode { data: NodeData::Leaf(CompiledClassHash(zero)), hash }; - let compiled_class_leaf_key = compiled_class_leaf.db_key().0; + let compiled_class_leaf_key = compiled_class_leaf.db_key(&EmptyKeyContext).0; // Store keys in a HashMap. let mut map: HashMap> = HashMap::new(); @@ -518,7 +518,7 @@ async fn test_storage_node(data: HashMap) -> CommitterPythonTest }; // Store the binary node in the storage. - rust_fact_storage.set(binary_rust.db_key(), binary_rust.serialize()).await?; + rust_fact_storage.set(binary_rust.db_key(&EmptyKeyContext), binary_rust.serialize()).await?; // Parse the edge node data from the input. let edge_json = get_or_key_not_found(&data, "edge")?; @@ -545,7 +545,7 @@ async fn test_storage_node(data: HashMap) -> CommitterPythonTest }; // Store the edge node in the storage. - rust_fact_storage.set(edge_rust.db_key(), edge_rust.serialize()).await?; + rust_fact_storage.set(edge_rust.db_key(&EmptyKeyContext), edge_rust.serialize()).await?; // Parse the storage leaf data from the input. let storage_leaf_json = get_or_key_not_found(&data, "storage")?; @@ -561,7 +561,9 @@ async fn test_storage_node(data: HashMap) -> CommitterPythonTest }; // Store the storage leaf node in the storage. - rust_fact_storage.set(storage_leaf_rust.db_key(), storage_leaf_rust.serialize()).await?; + rust_fact_storage + .set(storage_leaf_rust.db_key(&EmptyKeyContext), storage_leaf_rust.serialize()) + .await?; // Parse the contract state leaf data from the input. let contract_state_leaf = get_or_key_not_found(&data, "contract_state_leaf")?; @@ -587,7 +589,10 @@ async fn test_storage_node(data: HashMap) -> CommitterPythonTest // Store the contract state leaf node in the storage. rust_fact_storage - .set(contract_state_leaf_rust.db_key(), contract_state_leaf_rust.serialize()) + .set( + contract_state_leaf_rust.db_key(&EmptyKeyContext), + contract_state_leaf_rust.serialize(), + ) .await?; // Parse the compiled class leaf data from the input. @@ -606,7 +611,10 @@ async fn test_storage_node(data: HashMap) -> CommitterPythonTest // Store the compiled class leaf node in the storage. rust_fact_storage - .set(compiled_class_leaf_rust.db_key(), compiled_class_leaf_rust.serialize()) + .set( + compiled_class_leaf_rust.db_key(&EmptyKeyContext), + compiled_class_leaf_rust.serialize(), + ) .await?; // Serialize the storage to a JSON string and handle serialization errors. diff --git a/crates/starknet_os_flow_tests/src/utils.rs b/crates/starknet_os_flow_tests/src/utils.rs index 8695d19bf78..6bb066aa2a7 100644 --- a/crates/starknet_os_flow_tests/src/utils.rs +++ b/crates/starknet_os_flow_tests/src/utils.rs @@ -60,6 +60,7 @@ use starknet_os::hints::vars::Const; use starknet_os::io::os_input::{CachedStateInput, CommitmentInfo}; use starknet_patricia::patricia_merkle_tree::node_data::inner_node::flatten_preimages; use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices, SubTreeHeight}; +use starknet_patricia_storage::db_object::EmptyKeyContext; use starknet_patricia_storage::map_storage::MapStorage; use starknet_types_core::felt::Felt; @@ -257,6 +258,7 @@ pub(crate) async fn create_cached_state_input_and_commitment_infos( commitments, previous_state_roots.classes_trie_root_hash, sorted_class_leaf_indices, + &EmptyKeyContext, ) .await .unwrap(); @@ -284,6 +286,7 @@ pub(crate) async fn create_cached_state_input_and_commitment_infos( commitments, address_to_previous_storage_root_hash[&address], sorted_leaf_indices, + &EmptyKeyContext, ) .await .unwrap(); @@ -372,14 +375,22 @@ pub(crate) async fn get_previous_states_and_new_storage_roots< // Get previous contract state leaves. let sorted_contract_leaf_indices = SortedLeafIndices::new(&mut contract_leaf_indices); // Get the previous and the new contract states. - let previous_contract_states = - get_leaves(commitments, previous_contract_trie_root, sorted_contract_leaf_indices) - .await - .unwrap(); - let new_contract_states: HashMap = - get_leaves(commitments, new_contract_trie_root, sorted_contract_leaf_indices) - .await - .unwrap(); + let previous_contract_states = get_leaves( + commitments, + previous_contract_trie_root, + sorted_contract_leaf_indices, + &EmptyKeyContext, + ) + .await + .unwrap(); + let new_contract_states: HashMap = get_leaves( + commitments, + new_contract_trie_root, + sorted_contract_leaf_indices, + &EmptyKeyContext, + ) + .await + .unwrap(); let new_contract_roots: HashMap = new_contract_states .into_iter() .map(|(idx, contract_state)| { diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/external_test_utils.rs b/crates/starknet_patricia/src/patricia_merkle_tree/external_test_utils.rs index c403c1cc825..6834efef190 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/external_test_utils.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/external_test_utils.rs @@ -4,7 +4,12 @@ use ethnum::U256; use num_bigint::{BigUint, RandBigInt}; use rand::Rng; use starknet_api::hash::HashOutput; -use starknet_patricia_storage::db_object::{DBObject, Deserializable, HasStaticPrefix}; +use starknet_patricia_storage::db_object::{ + DBObject, + Deserializable, + EmptyKeyContext, + HasStaticPrefix, +}; use starknet_patricia_storage::errors::DeserializationError; use starknet_patricia_storage::storage_trait::{create_db_key, DbKey, DbKeyPrefix, DbValue}; use starknet_types_core::felt::Felt; @@ -25,7 +30,8 @@ use crate::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeleto pub struct MockLeaf(pub Felt); impl HasStaticPrefix for MockLeaf { - fn get_static_prefix() -> DbKeyPrefix { + type KeyContext = EmptyKeyContext; + fn get_static_prefix(_key_context: &Self::KeyContext) -> DbKeyPrefix { DbKeyPrefix::new(&[0]) } } @@ -112,8 +118,10 @@ fn create_inner_node_patricia_key(val: Felt) -> DbKey { create_db_key(PatriciaPrefix::InnerNode.into(), &val.to_bytes_be()) } -pub fn create_leaf_patricia_key(val: u128) -> DbKey { - create_db_key(L::get_static_prefix(), &U256::from(val).to_be_bytes()) +pub fn create_leaf_patricia_key>( + val: u128, +) -> DbKey { + create_db_key(L::get_static_prefix(&EmptyKeyContext), &U256::from(val).to_be_bytes()) } fn create_binary_val(left: Felt, right: Felt) -> DbValue { @@ -153,7 +161,9 @@ pub fn create_edge_entry_from_u128( create_edge_entry::(Felt::from(hash), path, length) } -pub fn create_leaf_entry(hash: u128) -> (DbKey, DbValue) { +pub fn create_leaf_entry>( + hash: u128, +) -> (DbKey, DbValue) { (create_leaf_patricia_key::(hash), DbValue(create_32_bytes_entry(hash).to_vec())) } diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/node_serde.rs b/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/node_serde.rs index 21d2ffbda5c..0c7c82b0fbe 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/node_serde.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/node_serde.rs @@ -1,6 +1,6 @@ use ethnum::U256; use starknet_api::hash::HashOutput; -use starknet_patricia_storage::db_object::{DBObject, HasDynamicPrefix}; +use starknet_patricia_storage::db_object::{DBObject, HasDynamicPrefix, HasStaticPrefix}; use starknet_patricia_storage::errors::DeserializationError; use starknet_patricia_storage::storage_trait::{DbKey, DbKeyPrefix, DbValue}; use starknet_types_core::felt::Felt; @@ -44,16 +44,19 @@ impl FilledNode { self.hash.0.to_bytes_be() } - pub fn db_key(&self) -> DbKey { - self.get_db_key(&self.suffix()) + pub fn db_key(&self, key_context: &::KeyContext) -> DbKey { + self.get_db_key(key_context, &self.suffix()) } } impl HasDynamicPrefix for FilledNode { - fn get_prefix(&self) -> DbKeyPrefix { + // Inherit the KeyContext from the HasStaticPrefix implementation of the leaf. + type KeyContext = ::KeyContext; + + fn get_prefix(&self, key_context: &Self::KeyContext) -> DbKeyPrefix { match &self.data { NodeData::Binary(_) | NodeData::Edge(_) => PatriciaPrefix::InnerNode, - NodeData::Leaf(_) => PatriciaPrefix::Leaf(L::get_static_prefix()), + NodeData::Leaf(_) => PatriciaPrefix::Leaf(L::get_static_prefix(key_context)), } .into() } 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 4895ab2dbe6..77fb45b3815 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 @@ -5,7 +5,7 @@ use std::sync::{Arc, Mutex}; use async_recursion::async_recursion; use starknet_api::hash::HashOutput; -use starknet_patricia_storage::db_object::DBObject; +use starknet_patricia_storage::db_object::{DBObject, HasStaticPrefix}; use starknet_patricia_storage::storage_trait::DbHashMap; use crate::patricia_merkle_tree::filled_tree::errors::FilledTreeError; @@ -42,7 +42,7 @@ pub trait FilledTree: Sized + Send { /// Serializes the current state of the tree into a hashmap, /// where each key-value pair corresponds /// to a storage key and its serialized storage value. - fn serialize(&self) -> DbHashMap; + fn serialize(&self, key_context: &::KeyContext) -> DbHashMap; fn get_root_hash(&self) -> HashOutput; } @@ -363,10 +363,13 @@ impl FilledTree for FilledTreeImpl { }) } - fn serialize(&self) -> DbHashMap { + fn serialize(&self, key_context: &::KeyContext) -> DbHashMap { // This function iterates over each node in the tree, using the node's `db_key` as the // hashmap key and the result of the node's `serialize` method as the value. - self.get_all_nodes().values().map(|node| (node.db_key(), node.serialize())).collect() + self.get_all_nodes() + .values() + .map(|node| (node.db_key(key_context), node.serialize())) + .collect() } fn get_root_hash(&self) -> HashOutput { diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/traversal.rs b/crates/starknet_patricia/src/patricia_merkle_tree/traversal.rs index 0e5b3993a19..862a07bd345 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/traversal.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/traversal.rs @@ -1,3 +1,4 @@ +use starknet_patricia_storage::db_object::HasStaticPrefix; use starknet_patricia_storage::errors::{DeserializationError, StorageError}; use starknet_patricia_storage::storage_trait::{DbKeyPrefix, PatriciaStorageError}; use thiserror::Error; @@ -27,7 +28,7 @@ pub type TraversalResult = Result; /// storage layout. pub trait SubTreeTrait<'a>: Sized { /// Extra data a node can carry (e.g. its hash). - type NodeData: Copy; + type NodeData; /// Creates a concrete child node given its index and data. fn create( @@ -101,5 +102,8 @@ pub trait SubTreeTrait<'a>: Sized { fn should_traverse_unmodified_children() -> bool; /// Returns the [DbKeyPrefix] of the root node. - fn get_root_prefix(&self) -> DbKeyPrefix; + fn get_root_prefix( + &self, + key_context: &::KeyContext, + ) -> DbKeyPrefix; } diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/traversal_test.rs b/crates/starknet_patricia/src/patricia_merkle_tree/traversal_test.rs index 6ffa3f2e805..528bc826150 100644 --- a/crates/starknet_patricia/src/patricia_merkle_tree/traversal_test.rs +++ b/crates/starknet_patricia/src/patricia_merkle_tree/traversal_test.rs @@ -1,6 +1,7 @@ use ethnum::U256; use pretty_assertions::assert_eq; use rstest::rstest; +use starknet_patricia_storage::db_object::HasStaticPrefix; use starknet_patricia_storage::storage_trait::DbKeyPrefix; use crate::patricia_merkle_tree::external_test_utils::small_tree_index_to_full; @@ -40,7 +41,10 @@ impl<'a> SubTreeTrait<'a> for TestSubTree<'a> { false } - fn get_root_prefix(&self) -> DbKeyPrefix { + fn get_root_prefix( + &self, + _key_context: &::KeyContext, + ) -> DbKeyPrefix { // Dummy prefix for testing purposes (we only need a prefix when interacting with storage). DbKeyPrefix::new(&[0]) } diff --git a/crates/starknet_patricia_storage/src/db_object.rs b/crates/starknet_patricia_storage/src/db_object.rs index 91af298424d..08af754d51b 100644 --- a/crates/starknet_patricia_storage/src/db_object.rs +++ b/crates/starknet_patricia_storage/src/db_object.rs @@ -1,19 +1,32 @@ use crate::errors::DeserializationError; use crate::storage_trait::{create_db_key, DbKey, DbKeyPrefix, DbValue}; +pub struct EmptyKeyContext; + pub trait HasDynamicPrefix { + /// Extra data needed to construct a leaf for node db key prefix. For example, in index layout, + /// we need to know the trie type of inner nodes. + type KeyContext; + /// Returns the storage key prefix of the DB object. - fn get_prefix(&self) -> DbKeyPrefix; + fn get_prefix(&self, key_context: &Self::KeyContext) -> DbKeyPrefix; } pub trait HasStaticPrefix { + /// Extra data needed to construct a leaf for node db key prefix. For example, in index layout, + /// we need to know the trie type of inner nodes. + type KeyContext; + /// Returns the storage key prefix of the DB object. - fn get_static_prefix() -> DbKeyPrefix; + fn get_static_prefix(key_context: &Self::KeyContext) -> DbKeyPrefix; } impl HasDynamicPrefix for T { - fn get_prefix(&self) -> DbKeyPrefix { - T::get_static_prefix() + /// Inherit the KeyContext from the HasStaticPrefix trait. + type KeyContext = T::KeyContext; + + fn get_prefix(&self, key_context: &Self::KeyContext) -> DbKeyPrefix { + T::get_static_prefix(key_context) } } @@ -22,8 +35,8 @@ pub trait DBObject: HasDynamicPrefix { fn serialize(&self) -> DbValue; /// Returns a `DbKey` from a prefix and a suffix. - fn get_db_key(&self, suffix: &[u8]) -> DbKey { - create_db_key(self.get_prefix(), suffix) + fn get_db_key(&self, key_context: &Self::KeyContext, suffix: &[u8]) -> DbKey { + create_db_key(self.get_prefix(key_context), suffix) } }