starknet_committer: collect siblings when traversing the tree

Author: nimrod-starkware

crates/starknet_committer/src/db/create_original_skeleton_tests.rs

@@ -188,6 +188,7 @@ pub(crate) async fn test_create_original_skeleton<L, Layout>(
         &config,
         &leaf_modifications,
         None,
+        None,
         &EmptyKeyContext,
     )
     .await

crates/starknet_committer/src/db/external_test_utils.rs

@@ -42,6 +42,7 @@ where
         &config,
         &leaf_modifications.iter().map(|(k, v)| (*k, v.clone().into())).collect(),
         None,
+        None,
         key_context,
     )
     .await

crates/starknet_committer/src/db/facts_db/create_facts_tree.rs

@@ -34,6 +34,7 @@ pub async fn get_leaves<'a, L: Leaf>(
         &config,
         &leaf_modifications,
         Some(&mut previous_leaves),
+        None,
         key_context,
     )
     .await?;

crates/starknet_committer/src/db/facts_db/traversal.rs

@@ -69,9 +69,13 @@ pub(crate) async fn fetch_patricia_paths_inner<'a, L: Leaf>(
     let mut current_subtrees = subtrees;
     let mut next_subtrees = Vec::new();
     while !current_subtrees.is_empty() {
-        let filled_roots =
-            get_roots_from_storage::<L, FactsNodeLayout>(&current_subtrees, storage, key_context)
-                .await?;
+        let filled_roots = get_roots_from_storage::<L, FactsNodeLayout>(
+            &current_subtrees,
+            storage,
+            key_context,
+            None,
+        )
+        .await?;
         for (filled_root, subtree) in filled_roots.into_iter().zip(current_subtrees.iter()) {
             match filled_root.data {
                 // Binary node.

crates/starknet_committer/src/db/trie_traversal.rs

@@ -27,7 +27,7 @@ use starknet_patricia::patricia_merkle_tree::traversal::{
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
 use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::errors::StorageError;
-use starknet_patricia_storage::storage_trait::{AsyncStorage, DbKey, Storage};
+use starknet_patricia_storage::storage_trait::{AsyncStorage, DbHashMap, DbKey, Storage};
 use tracing::warn;
 
 use crate::block_committer::input::{
@@ -47,6 +47,7 @@ macro_rules! log_trivial_modification {
     };
 }
 
+#[allow(clippy::too_many_arguments)]
 /// Fetches the Patricia witnesses, required to build the original skeleton tree from storage.
 ///
 /// Given a list of subtrees, traverses towards their leaves and fetches all non-empty,
@@ -62,6 +63,7 @@ pub(crate) async fn fetch_nodes<'a, L, Layout>(
     leaf_modifications: &LeafModifications<L>,
     config: &impl OriginalSkeletonTreeConfig,
     mut previous_leaves: Option<&mut HashMap<NodeIndex, L>>,
+    mut siblings_map: Option<&mut DbHashMap>,
     key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> OriginalSkeletonTreeResult<()>
 where
@@ -73,8 +75,13 @@ where
     let should_fetch_modified_leaves =
         config.compare_modified_leaves() || previous_leaves.is_some();
     while !current_subtrees.is_empty() {
-        let filled_roots =
-            get_roots_from_storage::<L, Layout>(&current_subtrees, storage, key_context).await?;
+        let filled_roots = get_roots_from_storage::<L, Layout>(
+            &current_subtrees,
+            storage,
+            key_context,
+            siblings_map.as_deref_mut(),
+        )
+        .await?;
         for (filled_root, subtree) in filled_roots.into_iter().zip(current_subtrees.into_iter()) {
             if subtree.is_unmodified() {
                 handle_unmodified_subtree(skeleton_tree, &mut next_subtrees, filled_root, subtree);
@@ -244,17 +251,21 @@ pub async fn get_roots_from_storage<'a, L: Leaf, Layout: NodeLayout<'a, L>>(
     subtrees: &[Layout::SubTree],
     storage: &impl Storage,
     key_context: &<L as HasStaticPrefix>::KeyContext,
+    mut siblings_map: Option<&mut DbHashMap>,
 ) -> TraversalResult<Vec<FilledNode<L, Layout::NodeData>>> {
     let mut subtrees_roots = vec![];
     let db_keys: Vec<DbKey> =
         subtrees.iter().map(|subtree| subtree.get_root_db_key::<L>(key_context)).collect();
 
     let db_vals = storage.mget(&db_keys.iter().collect::<Vec<&DbKey>>()).await?;
-    for ((subtree, optional_val), db_key) in subtrees.iter().zip(db_vals.iter()).zip(db_keys) {
+    for ((subtree, optional_val), db_key) in subtrees.iter().zip(db_vals.into_iter()).zip(db_keys) {
         let Some(val) = optional_val else { Err(StorageError::MissingKey(db_key))? };
         let filled_node =
-            Layout::NodeDbObject::deserialize(val, &subtree.get_root_context())?.into();
+            Layout::NodeDbObject::deserialize(&val, &subtree.get_root_context())?.into();
         subtrees_roots.push(filled_node);
+        if subtree.is_unmodified() {
+            siblings_map.as_mut().map(|m| m.insert(db_key, val));
+        }
     }
     Ok(subtrees_roots)
 }
@@ -279,6 +290,7 @@ pub(crate) fn log_warning_for_empty_leaves<L: Leaf, T: Borrow<NodeIndex> + Debug
     Ok(())
 }
 
+#[allow(clippy::too_many_arguments)]
 /// Creates an original skeleton tree by fetching Patricia nodes from storage.
 ///
 /// Traverses the trie from the root towards the modified leaves, collecting all nodes needed
@@ -300,6 +312,7 @@ pub async fn create_original_skeleton_tree<'a, L: Leaf, Layout: NodeLayout<'a, L
     config: &impl OriginalSkeletonTreeConfig,
     leaf_modifications: &LeafModifications<L>,
     previous_leaves: Option<&mut HashMap<NodeIndex, L>>,
+    siblings_map: Option<&mut DbHashMap>,
     key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> OriginalSkeletonTreeResult<OriginalSkeletonTreeImpl<'a>> {
     if sorted_leaf_indices.is_empty() {
@@ -332,6 +345,7 @@ pub async fn create_original_skeleton_tree<'a, L: Leaf, Layout: NodeLayout<'a, L
         leaf_modifications,
         config,
         previous_leaves,
+        siblings_map,
         key_context,
     )
     .await?;
@@ -379,6 +393,7 @@ where
         &config,
         &HashMap::new(),
         Some(&mut leaves),
+        None,
         &EmptyKeyContext,
     )
     .await?;
@@ -416,6 +431,7 @@ where
             .map(|(idx, value)| (*idx, Layout::DbLeaf::from(*value)))
             .collect(),
         None,
+        None,
         &EmptyKeyContext,
     )
     .await?)
@@ -527,6 +543,7 @@ where
         &trie_config,
         &leaf_modifications,
         None,
+        None,
         &address,
     )
     .await?;