starknet_committer: split commit_block into read and compute phases

crates/starknet_committer/src/block_committer/commit.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 
 use starknet_api::core::{ClassHash, ContractAddress, Nonce};
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::LeafModifications;
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
 use starknet_types_core::felt::Felt;
 use tracing::{debug, warn};
@@ -21,11 +22,11 @@ use crate::db::forest_trait::ForestReader;
 use crate::forest::deleted_nodes::{find_deleted_nodes, DeletedNodes};
 use crate::forest::filled_forest::FilledForest;
 use crate::forest::forest_errors::ForestError;
-use crate::forest::original_skeleton_forest::ForestSortedIndices;
+use crate::forest::original_skeleton_forest::{ForestSortedIndices, OriginalSkeletonForest};
 use crate::forest::updated_skeleton_forest::UpdatedSkeletonForest;
 use crate::hash_function::hash::TreeHashFunctionImpl;
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
-use crate::patricia_merkle_tree::types::class_hash_into_node_index;
+use crate::patricia_merkle_tree::types::{class_hash_into_node_index, CompiledClassHash};
 
 pub type BlockCommitmentResult<T> = Result<T, BlockCommitmentError>;
 
@@ -54,16 +55,50 @@ pub async fn commit_block<Reader: ForestReader + Send, M: MeasurementsTrait + Se
         n_contracts_trie_modifications,
         actual_classes_updates.len(),
     );
-    // Reads - fetch_nodes.
+
+    // Phase 1 - read the original forest from the DB.
+    let (original_forest, original_contracts_trie_leaves) = read_original_forest(
+        &input,
+        &actual_storage_updates,
+        &actual_classes_updates,
+        &forest_sorted_indices,
+        trie_reader,
+        measurements,
+    )
+    .await?;
+
+    // Phase 2 - compute the new forest.
+    compute_updated_forest(
+        original_forest,
+        original_contracts_trie_leaves,
+        actual_storage_updates,
+        actual_classes_updates,
+        &input.state_diff,
+        measurements,
+    )
+    .await
+}
+
+/// Reads the original forest from the DB.
+async fn read_original_forest<'a, Reader: ForestReader + Send, M: MeasurementsTrait + Send>(
+    input: &Input<Reader::InitialReadContext>,
+    actual_storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
+    actual_classes_updates: &LeafModifications<CompiledClassHash>,
+    forest_sorted_indices: &'a ForestSortedIndices<'a>,
+    trie_reader: &mut Reader,
+    measurements: &mut M,
+) -> BlockCommitmentResult<(OriginalSkeletonForest<'a>, HashMap<NodeIndex, ContractState>)> {
     measurements.start_measurement(Action::Read);
-    let roots =
-        trie_reader.read_roots(input.initial_read_context).await.map_err(ForestError::from)?;
+    let roots = trie_reader
+        .read_roots(input.initial_read_context.clone())
+        .await
+        .map_err(ForestError::from)?;
     let (original_forest, original_contracts_trie_leaves) = trie_reader
         .read(
             roots,
-            &actual_storage_updates,
-            &actual_classes_updates,
-            &forest_sorted_indices,
+            actual_storage_updates,
+            actual_classes_updates,
+            forest_sorted_indices,
             input.config.clone(),
         )
         .await?;
@@ -79,15 +114,28 @@ pub async fn commit_block<Reader: ForestReader + Send, M: MeasurementsTrait + Se
         );
     }
 
-    // Compute the new topology.
+    Ok((original_forest, original_contracts_trie_leaves))
+}
+
+/// Computes the updated forest topology, its new hashes, and the nodes deleted by the state diff.
+async fn compute_updated_forest<M: MeasurementsTrait + Send>(
+    original_forest: OriginalSkeletonForest<'_>,
+    original_contracts_trie_leaves: HashMap<NodeIndex, ContractState>,
+    actual_storage_updates: HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
+    actual_classes_updates: LeafModifications<CompiledClassHash>,
+    state_diff: &StateDiff,
+    measurements: &mut M,
+) -> BlockCommitmentResult<(FilledForest, DeletedNodes)> {
     measurements.start_measurement(Action::Compute);
+
+    // Compute the new topology.
     let updated_forest = UpdatedSkeletonForest::create(
         &original_forest,
-        &input.state_diff.skeleton_classes_updates(),
-        &input.state_diff.skeleton_storage_updates(),
+        &state_diff.skeleton_classes_updates(),
+        &state_diff.skeleton_storage_updates(),
         &original_contracts_trie_leaves,
-        &input.state_diff.address_to_class_hash,
-        &input.state_diff.address_to_nonce,
+        &state_diff.address_to_class_hash,
+        &state_diff.address_to_nonce,
     )?;
     debug!("Updated skeleton forest created successfully.");
 
@@ -106,8 +154,8 @@ pub async fn commit_block<Reader: ForestReader + Send, M: MeasurementsTrait + Se
         actual_storage_updates,
         actual_classes_updates,
         &original_contracts_trie_leaves,
-        &input.state_diff.address_to_class_hash,
-        &input.state_diff.address_to_nonce,
+        &state_diff.address_to_class_hash,
+        &state_diff.address_to_nonce,
     )
     .await?;
     measurements.attempt_to_stop_measurement(Action::Compute, 0).ok();

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

@@ -125,8 +125,8 @@ impl<S: Storage> ForestReader for FactsDb<S> {
     async fn read<'a>(
         &mut self,
         roots: StateRoots,
-        storage_updates: &'a HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
-        classes_updates: &'a LeafModifications<CompiledClassHash>,
+        storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
+        classes_updates: &LeafModifications<CompiledClassHash>,
         forest_sorted_indices: &'a ForestSortedIndices<'a>,
         config: ReaderConfig,
     ) -> ForestResult<(OriginalSkeletonForest<'a>, HashMap<NodeIndex, ContractState>)> {

crates/starknet_committer/src/db/forest_trait.rs

@@ -70,13 +70,13 @@ pub trait ForestMetadata {
 #[async_trait]
 pub trait ForestReader {
     /// Input required to start reading the storage trie.
-    type InitialReadContext: InputContext + Send;
+    type InitialReadContext: InputContext + Send + Sync;
 
     async fn read<'a>(
         &mut self,
         roots: StateRoots,
-        storage_updates: &'a HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
-        classes_updates: &'a LeafModifications<CompiledClassHash>,
+        storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
+        classes_updates: &LeafModifications<CompiledClassHash>,
         forest_sorted_indices: &'a ForestSortedIndices<'a>,
         config: ReaderConfig,
     ) -> ForestResult<(OriginalSkeletonForest<'a>, HashMap<NodeIndex, ContractState>)>;
@@ -91,8 +91,8 @@ pub trait ForestReader {
 pub(crate) async fn read_forest<'a, S, Layout>(
     storage: &mut S,
     roots: StateRoots,
-    storage_updates: &'a HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
-    classes_updates: &'a LeafModifications<CompiledClassHash>,
+    storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
+    classes_updates: &LeafModifications<CompiledClassHash>,
     forest_sorted_indices: &'a ForestSortedIndices<'a>,
     config: ReaderConfig,
 ) -> ForestResult<(OriginalSkeletonForest<'a>, HashMap<NodeIndex, ContractState>)>

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

@@ -212,8 +212,8 @@ where
     async fn read<'a>(
         &mut self,
         roots: StateRoots,
-        storage_updates: &'a HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
-        classes_updates: &'a LeafModifications<CompiledClassHash>,
+        storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
+        classes_updates: &LeafModifications<CompiledClassHash>,
         forest_sorted_indices: &'a ForestSortedIndices<'a>,
         config: ReaderConfig,
     ) -> ForestResult<(OriginalSkeletonForest<'a>, HashMap<NodeIndex, ContractState>)> {

crates/starknet_committer/src/db/trie_traversal.rs

@@ -662,7 +662,7 @@ pub async fn create_original_skeleton_tree<'a, L: Leaf, Layout: NodeLayout<'a, L
 /// sequentially.
 pub async fn create_storage_tries<'a, Layout>(
     storage: &mut impl Storage,
-    actual_storage_updates: &'a HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
+    actual_storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
     original_contracts_trie_leaves: &HashMap<NodeIndex, ContractState>,
     config: &ReaderConfig,
     storage_tries_sorted_indices: &'a HashMap<ContractAddress, SortedLeafIndices<'a>>,
@@ -796,17 +796,25 @@ where
 }
 
 /// Holds all data needed to build one storage trie concurrently.
-/// Implements [StorageTask] so that [GatherableStorage::gather] can drive it.
-struct TrieReadTask<'a, Layout: NodeLayoutFor<StarknetStorageValue>> {
+///
+/// Two distinct lifetimes are needed:
+/// - `'a` is the lifetime of the caller-owned index vectors that produced `sorted_leaf_indices`.
+/// - `'u` is the borrow of `updates`, used only during construction of the original skeleton tree,
+///   but does not retain a reference into the output `OriginalSkeletonTreeImpl<'a>`.
+///
+/// Unifying the lifetimes would make the returned `OriginalSkeletonForest<'a>` to
+/// borrow the storage update maps too, which would prevent `commit_block` from later moving those
+/// maps by value into the compute phase.
+struct TrieReadTask<'a, 'u, Layout: NodeLayoutFor<StarknetStorageValue>> {
     address: ContractAddress,
-    updates: &'a LeafModifications<StarknetStorageValue>,
+    updates: &'u LeafModifications<StarknetStorageValue>,
     storage_root_hash: HashOutput,
     sorted_leaf_indices: SortedLeafIndices<'a>,
     warn_on_trivial_modifications: bool,
     _layout: PhantomData<Layout>,
 }
 
-impl<'a, S, Layout> StorageTaskOutput<S> for TrieReadTask<'a, Layout>
+impl<'a, 'u, S, Layout> StorageTaskOutput<S> for TrieReadTask<'a, 'u, Layout>
 where
     S: ImmutableReadOnlyStorage,
     Layout: NodeLayoutFor<StarknetStorageValue> + Send + 'static,
@@ -815,7 +823,7 @@ where
 }
 
 #[async_trait]
-impl<'a, 's, S, Layout> StorageTask<'s, S> for TrieReadTask<'a, Layout>
+impl<'a, 'u, 's, S, Layout> StorageTask<'s, S> for TrieReadTask<'a, 'u, Layout>
 where
     S: ImmutableReadOnlyStorage + 's,
     Layout: NodeLayoutFor<StarknetStorageValue> + Send + 'static,
@@ -838,7 +846,7 @@ where
 
 async fn create_storage_tries_concurrently<'a, S, Layout>(
     storage: &mut S,
-    actual_storage_updates: &'a HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
+    actual_storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
     original_contracts_trie_leaves: &HashMap<NodeIndex, ContractState>,
     warn_on_trivial_modifications: bool,
     storage_tries_sorted_indices: &'a HashMap<ContractAddress, SortedLeafIndices<'a>>,