starknet_committer: flush siblings into storage

Author: nimrod-starkware

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

@@ -113,7 +113,7 @@ impl<S: Storage> ForestReader for FactsDb<S> {
         config: ReaderConfig,
     ) -> ForestResult<(OriginalSkeletonForest<'a>, HashMap<NodeIndex, ContractState>)> {
         read_forest::<S, FactsNodeLayout>(
-            &self.storage,
+            &mut self.storage,
             roots,
             storage_updates,
             classes_updates,

crates/starknet_committer/src/db/forest_trait.rs

@@ -87,7 +87,7 @@ pub trait ForestReader {
 
 /// Helper function containing layout-common read logic.
 pub(crate) async fn read_forest<'a, S, Layout>(
-    storage: &S,
+    storage: &mut S,
     roots: StateRoots,
     storage_updates: &'a HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
     classes_updates: &'a LeafModifications<CompiledClassHash>,
@@ -98,22 +98,22 @@ where
     S: Storage,
     Layout: DbLayout,
 {
-    let (contracts_trie, original_contracts_trie_leaves) =
+    let (contracts_trie, original_contracts_trie_leaves, contracts_trie_siblings) =
         create_contracts_trie::<Layout::NodeLayout>(
             storage,
             roots.contracts_trie_root_hash,
             forest_sorted_indices.contracts_trie_sorted_indices,
         )
         .await?;
-    let (storage_tries, _storage_tries_siblings) = create_storage_tries::<Layout::NodeLayout>(
+    let (storage_tries, storage_tries_siblings) = create_storage_tries::<Layout::NodeLayout>(
         storage,
         storage_updates,
         &original_contracts_trie_leaves,
         &config,
         &forest_sorted_indices.storage_tries_sorted_indices,
     )
     .await?;
-    let classes_trie = create_classes_trie::<Layout::NodeLayout>(
+    let (classes_trie, classes_trie_siblings) = create_classes_trie::<Layout::NodeLayout>(
         storage,
         classes_updates,
         roots.classes_trie_root_hash,
@@ -122,6 +122,15 @@ where
     )
     .await?;
 
+    // After all siblings are collected - flush them to storage.
+    let all_siblings = contracts_trie_siblings
+        .into_iter()
+        .chain(storage_tries_siblings.into_iter())
+        .chain(classes_trie_siblings.into_iter())
+        .collect();
+
+    storage.flush_to_cache(all_siblings)?;
+
     Ok((
         OriginalSkeletonForest { classes_trie, contracts_trie, storage_tries },
         original_contracts_trie_leaves,

crates/starknet_committer/src/db/index_db/db.rs

@@ -183,7 +183,7 @@ impl<S: Storage> ForestReader for IndexDb<S> {
         config: ReaderConfig,
     ) -> ForestResult<(OriginalSkeletonForest<'a>, HashMap<NodeIndex, ContractState>)> {
         read_forest::<S, IndexNodeLayout>(
-            &self.storage,
+            &mut self.storage,
             roots,
             storage_updates,
             classes_updates,

crates/starknet_committer/src/db/trie_traversal.rs

@@ -393,29 +393,30 @@ pub async fn create_contracts_trie<'a, Layout: NodeLayoutFor<ContractState>>(
     storage: &impl Storage,
     contracts_trie_root_hash: HashOutput,
     contracts_trie_sorted_indices: SortedLeafIndices<'a>,
-) -> ForestResult<(OriginalSkeletonTreeImpl<'a>, HashMap<NodeIndex, ContractState>)>
+) -> ForestResult<(OriginalSkeletonTreeImpl<'a>, HashMap<NodeIndex, ContractState>, DbHashMap)>
 where
     <Layout as NodeLayoutFor<ContractState>>::DbLeaf: HasStaticPrefix<KeyContext = EmptyKeyContext>,
 {
     let config = OriginalSkeletonTrieConfig::new_for_contracts_trie();
 
     let mut leaves = HashMap::new();
+    let mut siblings_map = DbHashMap::new();
     let skeleton_tree = create_original_skeleton_tree::<Layout::DbLeaf, Layout>(
         storage,
         contracts_trie_root_hash,
         contracts_trie_sorted_indices,
         &config,
         &HashMap::new(),
         Some(&mut leaves),
-        None,
+        Some(&mut siblings_map),
         &EmptyKeyContext,
     )
     .await?;
 
     let leaves: HashMap<NodeIndex, ContractState> =
         leaves.into_iter().map(|(idx, leaf)| (idx, leaf.into())).collect();
 
-    Ok((skeleton_tree, leaves))
+    Ok((skeleton_tree, leaves, siblings_map))
 }
 
 pub async fn create_classes_trie<'a, Layout: NodeLayoutFor<CompiledClassHash>>(
@@ -424,31 +425,35 @@ pub async fn create_classes_trie<'a, Layout: NodeLayoutFor<CompiledClassHash>>(
     classes_trie_root_hash: HashOutput,
     config: &ReaderConfig,
     contracts_trie_sorted_indices: SortedLeafIndices<'a>,
-) -> ForestResult<OriginalSkeletonTreeImpl<'a>>
+) -> ForestResult<(OriginalSkeletonTreeImpl<'a>, DbHashMap)>
 where
     <Layout as NodeLayoutFor<CompiledClassHash>>::DbLeaf:
         HasStaticPrefix<KeyContext = EmptyKeyContext>,
 {
     let config = OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(
         config.warn_on_trivial_modifications(),
     );
+    let mut siblings_map = DbHashMap::new();
 
-    Ok(create_original_skeleton_tree::<Layout::DbLeaf, Layout>(
-        storage,
-        classes_trie_root_hash,
-        contracts_trie_sorted_indices,
-        &config,
-        // TODO(Ariel): Change `actual_classes_updates` to be an iterator over borrowed data so
-        // that the conversion below is costless.
-        &actual_classes_updates
-            .iter()
-            .map(|(idx, value)| (*idx, Layout::DbLeaf::from(*value)))
-            .collect(),
-        None,
-        None,
-        &EmptyKeyContext,
-    )
-    .await?)
+    Ok((
+        create_original_skeleton_tree::<Layout::DbLeaf, Layout>(
+            storage,
+            classes_trie_root_hash,
+            contracts_trie_sorted_indices,
+            &config,
+            // TODO(Ariel): Change `actual_classes_updates` to be an iterator over borrowed data so
+            // that the conversion below is costless.
+            &actual_classes_updates
+                .iter()
+                .map(|(idx, value)| (*idx, Layout::DbLeaf::from(*value)))
+                .collect(),
+            None,
+            Some(&mut siblings_map),
+            &EmptyKeyContext,
+        )
+        .await?,
+        siblings_map,
+    ))
 }
 
 async fn create_storage_tries_sequentially<'a, Layout: NodeLayoutFor<StarknetStorageValue>>(