diff --git a/crates/starknet_committer/src/db/facts_db.rs b/crates/starknet_committer/src/db/facts_db.rs index a9fdfb378d1..07069b3466c 100644 --- a/crates/starknet_committer/src/db/facts_db.rs +++ b/crates/starknet_committer/src/db/facts_db.rs @@ -1,7 +1,6 @@ pub mod create_facts_tree; pub mod db; pub mod node_serde; -pub mod traversal; pub mod types; pub use db::{FactDbFilledNode, FactsDb, FactsNodeLayout}; diff --git a/crates/starknet_committer/src/db/facts_db/traversal.rs b/crates/starknet_committer/src/db/facts_db/traversal.rs deleted file mode 100644 index b8c6956f245..00000000000 --- a/crates/starknet_committer/src/db/facts_db/traversal.rs +++ /dev/null @@ -1,133 +0,0 @@ -use std::collections::HashMap; - -use starknet_api::hash::HashOutput; -use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{ - NodeData, - Preimage, - PreimageMap, -}; -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::storage_trait::ReadOnlyStorage; - -use crate::db::facts_db::types::FactsSubTree; -use crate::db::facts_db::FactsNodeLayout; -use crate::db::trie_traversal::get_roots_from_storage; - -#[cfg(test)] -#[path = "traversal_test.rs"] -pub mod traversal_test; - -/// Returns the Patricia inner nodes ([PreimageMap]) in the paths to the given `leaf_indices` in the -/// given tree according to the `root_hash`. -/// If `leaves` is not `None`, it also fetches the modified leaves and inserts them into the -/// provided map. -pub async fn fetch_patricia_paths( - storage: &mut impl ReadOnlyStorage, - root_hash: HashOutput, - sorted_leaf_indices: SortedLeafIndices<'_>, - leaves: Option<&mut HashMap>, - key_context: &::KeyContext, -) -> TraversalResult { - let mut witnesses = PreimageMap::new(); - - if sorted_leaf_indices.is_empty() || root_hash == HashOutput::ROOT_OF_EMPTY_TREE { - return Ok(witnesses); - } - - let main_subtree = FactsSubTree::create(sorted_leaf_indices, NodeIndex::ROOT, root_hash); - - fetch_patricia_paths_inner::( - storage, - vec![main_subtree], - &mut witnesses, - leaves, - key_context, - ) - .await?; - Ok(witnesses) -} - -// TODO(Rotem): Match the python logic of fetching the nodes. -/// Fetches the inner nodes [HashOutput] and [Preimage] in the paths to modified leaves. -/// Required for `patricia_update` function in Cairo. -/// Extra preimages (more than the data required to verify merkle paths) are required to verify -/// correctness of final tree topology; for more details see 'traverse_edge' in 'patricia.cairo'. -/// Given a list of subtrees, traverses towards their leaves and fetches all non-empty, -/// inner nodes in their paths and their siblings. -/// If `leaves` is not `None`, it also fetches the modified leaves and inserts them into the -/// provided map. -pub(crate) async fn fetch_patricia_paths_inner<'a, L: Leaf>( - storage: &mut impl ReadOnlyStorage, - 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 = - get_roots_from_storage::(¤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. - // If it's the root: It's a modified subtree. - // Otherwise: - // If it was inserted as a child of a binary node - it's a modified subtree. - // If it was inserted as a child of an edge node - it should be fetched anyway - // (modified or unmodified). - NodeData::Binary(binary_data) => { - witnesses.insert(subtree.root_hash, Preimage::Binary(binary_data.clone())); - let (left_subtree, right_subtree) = subtree - .get_children_subtrees(binary_data.left_data, binary_data.right_data); - - if !left_subtree.is_unmodified() { - next_subtrees.push(left_subtree); - } - if !right_subtree.is_unmodified() { - next_subtrees.push(right_subtree); - } - } - // Edge node. - // If it's the root: it's not necessarily a modified tree, because the modification - // might be a deletion. In this case, we want to fetch the bottom node just if it's - // a binary node (and not a leaf). Otherwise: It was inserted as a child of a binary - // node, so it's a modified subtree. - NodeData::Edge(edge_data) => { - witnesses.insert(subtree.root_hash, Preimage::Edge(edge_data)); - // Parse bottom. - let (bottom_subtree, empty_leaves_indices) = subtree - .get_bottom_subtree(&edge_data.path_to_bottom, edge_data.bottom_data); - if let Some(ref mut leaves_map) = leaves { - // Insert empty leaves descendent of the current subtree, that are not - // descendents of the bottom node. - for index in empty_leaves_indices { - leaves_map.insert(*index, L::default()); - } - } - // Insert the bottom subtree if it's modified or a binary node. - if !bottom_subtree.is_unmodified() || !bottom_subtree.is_leaf() { - next_subtrees.push(bottom_subtree); - } - } - // Leaf node. - // If it was inserted as a child of a binary node - it's a modified leaf. - // If it was inserted as a child of an edge node - it means the edge node is - // modified, meaning the leaf is also modified. - NodeData::Leaf(leaf_data) => { - // Fetch the leaf if it's modified and should be fetched. - if let Some(ref mut leaves_map) = leaves { - leaves_map.insert(subtree.root_index, leaf_data); - } - } - } - } - current_subtrees = next_subtrees; - next_subtrees = Vec::new(); - } - Ok(()) -} 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 4cc52973d5b..774fa2589d2 100644 --- a/crates/starknet_committer/src/db/facts_db/traversal_test.rs +++ b/crates/starknet_committer/src/db/facts_db/traversal_test.rs @@ -33,8 +33,9 @@ use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue}; use starknet_types_core::felt::Felt; use starknet_types_core::hash::Pedersen; -use crate::db::facts_db::traversal::fetch_patricia_paths_inner; +use super::fetch_patricia_paths_inner; use crate::db::facts_db::types::FactsSubTree; +use crate::db::facts_db::FactsNodeLayout; fn to_preimage_map(raw_preimages: HashMap>) -> PreimageMap { raw_preimages @@ -83,7 +84,7 @@ async fn test_fetch_patricia_paths_inner_impl( let mut nodes = HashMap::new(); let mut fetched_leaves = HashMap::new(); - fetch_patricia_paths_inner::( + fetch_patricia_paths_inner::( &mut storage, vec![main_subtree], &mut nodes, diff --git a/crates/starknet_committer/src/db/trie_traversal.rs b/crates/starknet_committer/src/db/trie_traversal.rs index 8bbe2397363..9fe5bd83907 100644 --- a/crates/starknet_committer/src/db/trie_traversal.rs +++ b/crates/starknet_committer/src/db/trie_traversal.rs @@ -12,6 +12,9 @@ use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{ BinaryData, EdgeData, NodeData, + PathToBottom, + Preimage, + PreimageMap, }; use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications}; use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonTreeConfig; @@ -50,6 +53,311 @@ use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState; use crate::patricia_merkle_tree::tree::OriginalSkeletonTrieConfig; use crate::patricia_merkle_tree::types::CompiledClassHash; +#[cfg(test)] +#[path = "facts_db/traversal_test.rs"] +mod traversal_test; + +/// Returns the Patricia inner nodes ([PreimageMap]) in the paths to the given `leaf_indices` in the +/// given tree according to the `root_hash` (including siblings). +/// If `leaves` is not `None`, it also fetches the modified leaves and inserts them into the +/// provided map. +pub async fn fetch_patricia_paths<'a, L: Leaf, Layout: NodeLayout<'a, L>>( + storage: &mut impl ReadOnlyStorage, + root_hash: HashOutput, + sorted_leaf_indices: SortedLeafIndices<'a>, + leaves: Option<&mut HashMap>, + key_context: &::KeyContext, +) -> TraversalResult { + let mut witnesses = PreimageMap::new(); + + if sorted_leaf_indices.is_empty() || root_hash == HashOutput::ROOT_OF_EMPTY_TREE { + return Ok(witnesses); + } + + let main_subtree = + Layout::SubTree::create(sorted_leaf_indices, NodeIndex::ROOT, root_hash.into()); + + fetch_patricia_paths_inner::( + storage, + vec![main_subtree], + &mut witnesses, + leaves, + key_context, + ) + .await?; + Ok(witnesses) +} + +/// Fetches the inner nodes [HashOutput] and [Preimage] in the paths to modified leaves. +/// Required for `patricia_update` function in Cairo. +/// Extra preimages (more than the data required to verify merkle paths) are required to verify +/// correctness of final tree topology; for more details see 'traverse_edge' in 'patricia.cairo'. +/// Given a list of subtrees, traverses towards their leaves and fetches all non-empty, +/// inner nodes in their paths and their siblings. +/// If `leaves` is not `None`, it also fetches the modified leaves and inserts them into the +/// provided map. +pub(crate) async fn fetch_patricia_paths_inner<'a, L: Leaf, Layout: NodeLayout<'a, L>>( + storage: &mut impl ReadOnlyStorage, + subtrees: Vec, + witnesses: &mut PreimageMap, + mut leaves: Option<&mut HashMap>, + key_context: &::KeyContext, +) -> TraversalResult<()> { + // Hashes collected so far, keyed by node index. + let mut hash_by_index: HashMap = HashMap::new(); + // Pending binary preimage entries: (node_hash, left_index, right_index). + let mut pending_binary: Vec<(HashOutput, NodeIndex, NodeIndex)> = Vec::new(); + // Pending edge preimage entries: (node_hash, path, bottom_index). + let mut pending_edge: Vec<(HashOutput, PathToBottom, NodeIndex)> = Vec::new(); + + // `current_traversal`: subtrees requiring full processing (preimage building + traversal). + // `next_hash_only`: unmodified subtrees queued to load their hash (no traversal) before the + // next round's pendings flush. + let mut current_traversal = subtrees; + let mut next_traversal: Vec = Vec::new(); + let mut next_hash_only: Vec = Vec::new(); + + while !current_traversal.is_empty() { + let filled_roots = + get_roots_from_storage::(¤t_traversal, storage, key_context).await?; + for (filled_root, subtree) in filled_roots.into_iter().zip(current_traversal.iter()) { + hash_by_index.insert(subtree.get_root_index(), filled_root.hash); + match filled_root.data { + // Binary node. + // If it's the root: It's a modified subtree. + // Otherwise: + // If it was inserted as a child of a binary node - it's a modified subtree. + // If it was inserted as a child of an edge node - it should be fetched anyway + // (modified or unmodified). + NodeData::Binary(BinaryData { left_data, right_data }) => { + let (left_subtree, right_subtree) = + subtree.get_children_subtrees(left_data.clone(), right_data.clone()); + let left_index = left_subtree.get_root_index(); + let right_index = right_subtree.get_root_index(); + + let left_hash = schedule_binary_child::( + left_data, + left_subtree, + &mut hash_by_index, + &mut next_traversal, + &mut next_hash_only, + ); + let right_hash = schedule_binary_child::( + right_data, + right_subtree, + &mut hash_by_index, + &mut next_traversal, + &mut next_hash_only, + ); + + if let (Some(left_hash), Some(right_hash)) = (left_hash, right_hash) { + witnesses.insert( + filled_root.hash, + Preimage::Binary(BinaryData { + left_data: left_hash, + right_data: right_hash, + }), + ); + } else { + pending_binary.push((filled_root.hash, left_index, right_index)); + } + } + // Edge node. + // If it's the root: it's not necessarily a modified tree, because the modification + // might be a deletion. In this case, we want to fetch the bottom node only if it's + // a binary node (and not a leaf). Otherwise: It was inserted as a child of a binary + // node, so it's a modified subtree. + NodeData::Edge(EdgeData { bottom_data, path_to_bottom }) => { + let (bottom_subtree, empty_leaves_indices) = + subtree.get_bottom_subtree(&path_to_bottom, bottom_data.clone()); + let bottom_index = bottom_subtree.get_root_index(); + + if let Some(ref mut leaves_map) = leaves { + // Insert empty leaves descendent of the current subtree, that are not + // descendents of the bottom node. + for index in empty_leaves_indices { + leaves_map.insert(*index, L::default()); + } + } + + let bottom_hash = schedule_edge_child::( + bottom_data, + bottom_subtree, + &mut hash_by_index, + &mut next_traversal, + ); + + match bottom_hash { + Some(hash) => { + witnesses.insert( + filled_root.hash, + Preimage::Edge(EdgeData { bottom_data: hash, path_to_bottom }), + ); + } + None => { + pending_edge.push((filled_root.hash, path_to_bottom, bottom_index)); + } + } + } + // Leaf node. + // If it was inserted as a child of a binary node - it's a modified leaf. + // If it was inserted as a child of an edge node - it means the edge node is + // modified, meaning the leaf is also modified. + NodeData::Leaf(leaf_data) => { + if let Some(ref mut leaves_map) = leaves { + if !subtree.is_unmodified() { + leaves_map.insert(subtree.get_root_index(), leaf_data); + } + } + } + } + } + + clear_pending_nodes::( + &mut pending_binary, + &mut pending_edge, + &next_hash_only, + storage, + key_context, + &mut hash_by_index, + witnesses, + ) + .await?; + + current_traversal = next_traversal; + next_traversal = Vec::new(); + next_hash_only = Vec::new(); + } + + Ok(()) +} + +/// Loads hashes for `hash_only_subtrees` into `hash_by_index`. +async fn read_hashes<'a, L: Leaf, Layout: NodeLayout<'a, L>>( + hash_only_subtrees: &[Layout::SubTree], + storage: &mut impl ReadOnlyStorage, + key_context: &::KeyContext, + hash_by_index: &mut HashMap, +) -> TraversalResult<()> { + let filled_roots = + get_roots_from_storage::(hash_only_subtrees, storage, key_context).await?; + for (filled_root, subtree) in filled_roots.into_iter().zip(hash_only_subtrees.iter()) { + hash_by_index.insert(subtree.get_root_index(), filled_root.hash); + } + Ok(()) +} + +/// Flushes [`PreimageMap`] entries that were waiting on child hashes from storage. +async fn clear_pending_nodes<'a, L: Leaf, Layout: NodeLayout<'a, L>>( + pending_binary: &mut Vec<(HashOutput, NodeIndex, NodeIndex)>, + pending_edge: &mut Vec<(HashOutput, PathToBottom, NodeIndex)>, + hash_only_subtrees: &[Layout::SubTree], + storage: &mut impl ReadOnlyStorage, + key_context: &::KeyContext, + hash_by_index: &mut HashMap, + witnesses: &mut PreimageMap, +) -> TraversalResult<()> { + read_hashes::(hash_only_subtrees, storage, key_context, hash_by_index).await?; + + pending_binary.retain(|&(node_hash, left_index, right_index)| { + match (hash_by_index.get(&left_index), hash_by_index.get(&right_index)) { + (Some(&left_hash), Some(&right_hash)) => { + witnesses.insert( + node_hash, + Preimage::Binary(BinaryData { left_data: left_hash, right_data: right_hash }), + ); + false + } + // Children hashes are not yet available, so we keep the entry in the pending queue. + // Not possible in facts layout. In index layout, this occurs when at least one of the + // children has modified leaves, hence its hash will only be available in + // the next iteration. + _ => true, + } + }); + pending_edge.retain(|&(node_hash, path_to_bottom, bottom_index)| { + if let Some(&bottom_hash) = hash_by_index.get(&bottom_index) { + witnesses.insert( + node_hash, + Preimage::Edge(EdgeData { bottom_data: bottom_hash, path_to_bottom }), + ); + false + } + // Bottom hash is not yet available, so we keep the entry in the pending queue. + // Not possible in facts layout. In index layout, edge nodes remain pending until the next + // iteration, when the bottom node is traversed. + else { + true + } + }); + + Ok(()) +} + +/// Handles a left/right child under a binary node in [`fetch_patricia_paths_inner`]. +/// +/// On [`UnmodifiedChildTraversal::Skip`], pushes the child onto `next_traversal` when the subtree +/// is modified. +/// +/// On [`UnmodifiedChildTraversal::Traverse`], unmodified subtrees added to the `next_hash_only` +/// queue; modified subtrees use enqueued for further traversal. +/// +/// Returns the child hash if immediately available, or `None` if the child must be read first. +fn schedule_binary_child<'a, L: Leaf, Layout: NodeLayout<'a, L>>( + node_data: Layout::NodeData, + subtree: Layout::SubTree, + hash_by_index: &mut HashMap, + next_traversal: &mut Vec, + next_hash_only: &mut Vec, +) -> Option { + match Layout::SubTree::should_traverse_unmodified_child(node_data) { + UnmodifiedChildTraversal::Skip(hash) => { + hash_by_index.insert(subtree.get_root_index(), hash); + if !subtree.is_unmodified() { + next_traversal.push(subtree); + } + Some(hash) + } + UnmodifiedChildTraversal::Traverse => { + if subtree.is_unmodified() { + next_hash_only.push(subtree); + } else { + next_traversal.push(subtree); + } + None + } + } +} + +/// Handles the bottom child under an edge node in [`fetch_patricia_paths_inner`]. +/// +/// On [`UnmodifiedChildTraversal::Skip`], pushes onto `next_traversal` (unless it's an unmodified +/// leaf), as we always need the bottom preimage. +/// +/// On [`UnmodifiedChildTraversal::Traverse`], enqueues for further traversal. +/// +/// Returns the child hash if immediately available, or `None` if the child must be read first. +fn schedule_edge_child<'a, L: Leaf, Layout: NodeLayout<'a, L>>( + node_data: Layout::NodeData, + subtree: Layout::SubTree, + hash_by_index: &mut HashMap, + next_traversal: &mut Vec, +) -> Option { + match Layout::SubTree::should_traverse_unmodified_child(node_data) { + UnmodifiedChildTraversal::Skip(hash) => { + hash_by_index.insert(subtree.get_root_index(), hash); + if !subtree.is_unmodified() || !subtree.is_leaf() { + next_traversal.push(subtree); + } + Some(hash) + } + UnmodifiedChildTraversal::Traverse => { + next_traversal.push(subtree); + None + } + } +} + /// Logs out a warning of a trivial modification. macro_rules! log_trivial_modification { ($index:expr, $value:expr) => { diff --git a/crates/starknet_committer/src/patricia_merkle_tree/tree.rs b/crates/starknet_committer/src/patricia_merkle_tree/tree.rs index 744a6d3cc9b..09ae29099dc 100644 --- a/crates/starknet_committer/src/patricia_merkle_tree/tree.rs +++ b/crates/starknet_committer/src/patricia_merkle_tree/tree.rs @@ -14,7 +14,8 @@ use crate::block_committer::input::{ StarknetStorageKey, StarknetStorageValue, }; -use crate::db::facts_db::traversal::fetch_patricia_paths; +use crate::db::facts_db::FactsNodeLayout; +use crate::db::trie_traversal::fetch_patricia_paths; use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState; use crate::patricia_merkle_tree::types::{ class_hash_into_node_index, @@ -80,7 +81,7 @@ async fn fetch_all_patricia_paths( // Classes trie - no need to fetch the leaves. let leaves = None; - let classes_trie_proof = fetch_patricia_paths::( + let classes_trie_proof = fetch_patricia_paths::( storage, classes_trie_root_hash, class_sorted_leaf_indices, @@ -91,7 +92,7 @@ async fn fetch_all_patricia_paths( // Contracts trie - the leaves are required. let mut leaves = HashMap::new(); - let contracts_proof_nodes = fetch_patricia_paths::( + let contracts_proof_nodes = fetch_patricia_paths::( storage, contracts_trie_root_hash, contract_sorted_leaf_indices, @@ -122,7 +123,7 @@ async fn fetch_all_patricia_paths( }; // No need to fetch the leaves. let leaves = None; - let proof = fetch_patricia_paths::( + let proof = fetch_patricia_paths::( storage, storage_root_hash, *sorted_leaf_indices,