From ef3461a0f032f42c7c6cb1d2dc65b9d290ebce82 Mon Sep 17 00:00:00 2001
From: Arni Hod <arnon@starkware.co>
Date: Mon, 9 Dec 2024 20:11:08 +0200
Subject: [PATCH] feat(starknet_batcher): implement sync_block

---
 crates/starknet_batcher/src/batcher.rs      | 39 +++++++++++++++++---
 crates/starknet_batcher/src/batcher_test.rs | 40 +++++++++++++++++++++
 2 files changed, 75 insertions(+), 4 deletions(-)

diff --git a/crates/starknet_batcher/src/batcher.rs b/crates/starknet_batcher/src/batcher.rs
index 9f47cbc3d2d..49b2126506d 100644
--- a/crates/starknet_batcher/src/batcher.rs
+++ b/crates/starknet_batcher/src/batcher.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use blockifier::abi::constants;
@@ -8,8 +8,10 @@ use chrono::Utc;
 use mockall::automock;
 use papyrus_storage::state::{StateStorageReader, StateStorageWriter};
 use starknet_api::block::{BlockHashAndNumber, 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,
@@ -353,9 +355,28 @@ 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.clear_previous_proposals().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();
+
+        // TODO(Arni): Assert the input sync_block corresponds to this `height`.
+        let height = match self.get_height_from_storage()?.next() {
+            Some(height) => Ok(height),
+            None => {
+                // We use an explicit match pattern instead of 'ok_or' for improved logging.
+                info!("Height in storage is at the upper limit of block number.");
+                Err(BatcherError::InternalError)
+            }
+        }?;
+        info!("Syncing block at height {} and notifying mempool of the block.", height);
+        self.commit_proposal_and_block(height, state_diff, address_to_nonce, tx_hashes).await
     }
 
     #[instrument(skip(self), err)]
@@ -374,6 +395,16 @@ impl Batcher {
             "Committing proposal {} at height {} and notifying mempool of the block.",
             proposal_id, height
         );
+        self.commit_proposal_and_block(height, state_diff, address_to_nonce, tx_hashes).await
+    }
+
+    async fn commit_proposal_and_block(
+        &mut self,
+        height: BlockNumber,
+        state_diff: ThinStateDiff,
+        address_to_nonce: HashMap<ContractAddress, Nonce>,
+        tx_hashes: HashSet<TransactionHash>,
+    ) -> BatcherResult<()> {
         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);
diff --git a/crates/starknet_batcher/src/batcher_test.rs b/crates/starknet_batcher/src/batcher_test.rs
index 0b968dfaddf..99cd4815bec 100644
--- a/crates/starknet_batcher/src/batcher_test.rs
+++ b/crates/starknet_batcher/src/batcher_test.rs
@@ -37,6 +37,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};
@@ -483,6 +484,41 @@ 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(ThinStateDiff::default()))
+        .returning(|_, _| Ok(()));
+
+    mock_dependencies
+        .storage_writer
+        .expect_commit_proposal()
+        .with(eq(INITIAL_HEIGHT.unchecked_next()), eq(get_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: get_state_diff(),
+        transaction_hashes: test_tx_hashes().into_iter().collect(),
+    };
+    batcher.add_sync_block(sync_block).await.unwrap();
+}
+
 #[rstest]
 #[tokio::test]
 async fn decision_reached() {
@@ -620,3 +656,7 @@ 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 get_state_diff() -> ThinStateDiff {
+    ThinStateDiff { nonces: test_contract_nonces().into_iter().collect(), ..Default::default() }
+}