feat(starknet_batcher): implement sync_block

crates/starknet_batcher/src/batcher.rs

@@ -1,13 +1,15 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use blockifier::state::global_cache::GlobalContractCache;
 #[cfg(test)]
 use mockall::automock;
 use papyrus_storage::state::{StateStorageReader, StateStorageWriter};
 use starknet_api::block::BlockNumber;
+use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::executable_transaction::Transaction;
 use starknet_api::state::ThinStateDiff;
+use starknet_api::transaction::TransactionHash;
 use starknet_batcher_types::batcher_types::{
     BatcherResult,
     DecisionReachedInput,
@@ -351,9 +353,22 @@ impl Batcher {
         Ok(GetProposalContentResponse { content: GetProposalContent::Finished(commitment) })
     }
 
-    // TODO(Arni): Impl add sync block
-    pub async fn add_sync_block(&mut self, _sync_block: SyncBlock) -> BatcherResult<()> {
-        todo!("Implement add sync block");
+    pub async fn add_sync_block(&mut self, sync_block: SyncBlock) -> BatcherResult<()> {
+        if let Some(height) = self.active_height {
+            info!("Aborting all work on height {} due to state sync.", height);
+            self.abort_active_height().await;
+            self.active_height = None;
+        }
+
+        let SyncBlock { state_diff, transaction_hashes } = sync_block;
+        let address_to_nonce = state_diff.nonces.iter().map(|(k, v)| (*k, *v)).collect();
+        let tx_hashes = transaction_hashes.into_iter().collect();
+
+        let info_callback = |height: BlockNumber| {
+            info!("Syncing block at height {} and notifying mempool of the block.", height);
+        };
+        // TODO(Arni): Assert the input `sync_block` corresponds to this `height`.
+        self.commit_proposal_and_block(info_callback, state_diff, address_to_nonce, tx_hashes).await
     }
 
     #[instrument(skip(self), err)]
@@ -367,12 +382,26 @@ impl Batcher {
             .map_err(|_| BatcherError::InternalError)?;
         let ProposalOutput { state_diff, nonces: address_to_nonce, tx_hashes, .. } =
             proposal_output;
+        let info_callback = |height: BlockNumber| {
+            info!(
+                "Committing proposal {} at height {} and notifying mempool of the block.",
+                proposal_id, height
+            );
+        };
+        self.commit_proposal_and_block(info_callback, state_diff, address_to_nonce, tx_hashes).await
+    }
+
+    async fn commit_proposal_and_block(
+        &mut self,
+        info_callback: impl Fn(BlockNumber),
+        state_diff: ThinStateDiff,
+        address_to_nonce: HashMap<ContractAddress, Nonce>,
+        tx_hashes: HashSet<TransactionHash>,
+    ) -> BatcherResult<()> {
         // TODO: Keep the height from start_height or get it from the input.
         let height = self.get_height_from_storage()?;
-        info!(
-            "Committing proposal {} at height {} and notifying mempool of the block.",
-            proposal_id, height
-        );
+        info_callback(height);
+
         trace!("Transactions: {:#?}, State diff: {:#?}.", tx_hashes, state_diff);
         self.storage_writer.commit_proposal(height, state_diff).map_err(|err| {
             error!("Failed to commit proposal to storage: {}", err);

crates/starknet_batcher/src/batcher_test.rs

@@ -8,6 +8,7 @@ use blockifier::test_utils::struct_impls::BlockInfoExt;
 use chrono::Utc;
 use futures::future::BoxFuture;
 use futures::FutureExt;
+use indexmap::indexmap;
 use mockall::automock;
 use mockall::predicate::{always, eq};
 use rstest::rstest;
@@ -37,6 +38,7 @@ use starknet_batcher_types::batcher_types::{
 use starknet_batcher_types::errors::BatcherError;
 use starknet_mempool_types::communication::MockMempoolClient;
 use starknet_mempool_types::mempool_types::CommitBlockArgs;
+use starknet_state_sync_types::state_sync_types::SyncBlock;
 use tokio::sync::Mutex;
 
 use crate::batcher::{Batcher, MockBatcherStorageReaderTrait, MockBatcherStorageWriterTrait};
@@ -476,6 +478,38 @@ async fn get_content_from_unknown_proposal() {
     assert_eq!(result, Err(BatcherError::ProposalNotFound { proposal_id: PROPOSAL_ID }));
 }
 
+#[rstest]
+#[tokio::test]
+async fn add_sync_block() {
+    let mut mock_dependencies = MockDependencies::default();
+
+    mock_dependencies
+        .storage_writer
+        .expect_commit_proposal()
+        .with(eq(INITIAL_HEIGHT), eq(test_state_diff()))
+        .returning(|_, _| Ok(()));
+
+    mock_dependencies
+        .mempool_client
+        .expect_commit_block()
+        .with(eq(CommitBlockArgs {
+            address_to_nonce: test_contract_nonces(),
+            tx_hashes: test_tx_hashes(),
+        }))
+        .returning(|_| Ok(()));
+
+    let mut batcher = create_batcher(mock_dependencies);
+
+    let sync_block = SyncBlock {
+        state_diff: test_state_diff(),
+        transaction_hashes: test_tx_hashes().into_iter().collect(),
+    };
+    batcher.add_sync_block(sync_block).await.unwrap();
+
+    let get_height_response = batcher.get_height().await.unwrap();
+    assert_eq!(get_height_response.height, INITIAL_HEIGHT);
+}
+
 #[rstest]
 #[tokio::test]
 async fn decision_reached() {
@@ -613,3 +647,19 @@ fn test_tx_hashes() -> HashSet<TransactionHash> {
 fn test_contract_nonces() -> HashMap<ContractAddress, Nonce> {
     HashMap::from_iter((0..3u8).map(|i| (contract_address!(i + 33), nonce!(i + 9))))
 }
+
+pub fn test_state_diff() -> ThinStateDiff {
+    ThinStateDiff {
+        storage_diffs: indexmap! {
+            4u64.into() => indexmap! {
+                5u64.into() => 6u64.into(),
+                7u64.into() => 8u64.into(),
+            },
+            9u64.into() => indexmap! {
+                10u64.into() => 11u64.into(),
+            },
+        },
+        nonces: test_contract_nonces().into_iter().collect(),
+        ..Default::default()
+    }
+}