From db277ff50b32be4251e32a1bf6fec19b6307cca1 Mon Sep 17 00:00:00 2001
From: Noam Spiegelstein <noam.s@starkware.co>
Date: Mon, 16 Dec 2024 17:48:04 +0200
Subject: [PATCH] feat(starknet_state_sync): implement sync state reader get
 nonce at

---
 .../starknet_gateway/src/sync_state_reader.rs |  8 +++-
 crates/starknet_state_sync/src/lib.rs         | 23 ++++++++++-
 .../src/communication.rs                      | 41 ++++++++++++++++++-
 3 files changed, 67 insertions(+), 5 deletions(-)

diff --git a/crates/starknet_gateway/src/sync_state_reader.rs b/crates/starknet_gateway/src/sync_state_reader.rs
index 8f706bdbfb..9db7c3e78a 100644
--- a/crates/starknet_gateway/src/sync_state_reader.rs
+++ b/crates/starknet_gateway/src/sync_state_reader.rs
@@ -45,8 +45,12 @@ impl BlockifierStateReader for SyncStateReader {
         Ok(res.unwrap_or_default())
     }
 
-    fn get_nonce_at(&self, _contract_address: ContractAddress) -> StateResult<Nonce> {
-        todo!()
+    fn get_nonce_at(&self, contract_address: ContractAddress) -> StateResult<Nonce> {
+        let res = block_on(
+            self.shared_state_sync_client.get_nonce_at(self.block_number, contract_address),
+        )?;
+
+        Ok(res.unwrap_or_default())
     }
 
     fn get_compiled_class(&self, _class_hash: ClassHash) -> StateResult<RunnableCompiledClass> {
diff --git a/crates/starknet_state_sync/src/lib.rs b/crates/starknet_state_sync/src/lib.rs
index daa6207682..2bb9a89057 100644
--- a/crates/starknet_state_sync/src/lib.rs
+++ b/crates/starknet_state_sync/src/lib.rs
@@ -6,7 +6,7 @@ use papyrus_storage::body::BodyStorageReader;
 use papyrus_storage::state::StateStorageReader;
 use papyrus_storage::StorageReader;
 use starknet_api::block::BlockNumber;
-use starknet_api::core::{ContractAddress, BLOCK_HASH_TABLE_ADDRESS};
+use starknet_api::core::{ContractAddress, Nonce, BLOCK_HASH_TABLE_ADDRESS};
 use starknet_api::state::{StateNumber, StorageKey};
 use starknet_sequencer_infra::component_definitions::{ComponentRequestHandler, ComponentStarter};
 use starknet_sequencer_infra::component_server::{LocalComponentServer, RemoteComponentServer};
@@ -50,6 +50,9 @@ impl ComponentRequestHandler<StateSyncRequest, StateSyncResponse> for StateSync
                     storage_key,
                 ))
             }
+            StateSyncRequest::GetNonceAt(block_number, contract_address) => {
+                StateSyncResponse::GetNonceAt(self.get_nonce_at(block_number, contract_address))
+            }
         }
     }
 }
@@ -97,6 +100,24 @@ impl StateSync {
 
         Ok(None)
     }
+
+    fn get_nonce_at(
+        &self,
+        block_number: BlockNumber,
+        contract_address: ContractAddress,
+    ) -> StateSyncResult<Option<Nonce>> {
+        let txn = self.storage_reader.begin_ro_txn()?;
+
+        if let Some(state_number) = StateNumber::right_after_block(block_number) {
+            let res = txn
+                .get_state_reader()?
+                .get_nonce_at(state_number, &contract_address)?
+                .ok_or_else(|| StateSyncError::StorageError("Contract not found".to_string()))?;
+            return Ok(Some(res));
+        }
+
+        Ok(None)
+    }
 }
 
 pub type LocalStateSyncServer =
diff --git a/crates/starknet_state_sync_types/src/communication.rs b/crates/starknet_state_sync_types/src/communication.rs
index 62d017e790..579a527ef4 100644
--- a/crates/starknet_state_sync_types/src/communication.rs
+++ b/crates/starknet_state_sync_types/src/communication.rs
@@ -4,7 +4,7 @@ use async_trait::async_trait;
 use papyrus_proc_macros::handle_response_variants;
 use serde::{Deserialize, Serialize};
 use starknet_api::block::BlockNumber;
-use starknet_api::core::ContractAddress;
+use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::state::StorageKey;
 use starknet_sequencer_infra::component_client::{
     ClientError,
@@ -44,7 +44,12 @@ pub trait StateSyncClient: Send + Sync {
         storage_key: StorageKey,
     ) -> StateSyncClientResult<Option<Felt>>;
 
-    // TODO: add get_nonce_at for BlockifierStateReader trait
+    async fn get_nonce_at(
+        &self,
+        block_number: BlockNumber,
+        contract_address: ContractAddress,
+    ) -> StateSyncClientResult<Option<Nonce>>;
+
     // TODO: add get_compiled_class for BlockifierStateReader trait
     // TODO: add get_class_hash_at for BlockifierStateReader trait
     // TODO: add get_compiled_class_hash for BlockifierStateReader trait
@@ -73,6 +78,7 @@ pub enum StateSyncRequest {
     GetBlock(BlockNumber),
     AddNewBlock(BlockNumber, SyncBlock),
     GetStorageAt(BlockNumber, ContractAddress, StorageKey),
+    GetNonceAt(BlockNumber, ContractAddress),
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -80,6 +86,7 @@ pub enum StateSyncResponse {
     GetBlock(StateSyncResult<Option<SyncBlock>>),
     AddNewBlock(StateSyncResult<()>),
     GetStorageAt(StateSyncResult<Option<Felt>>),
+    GetNonceAt(StateSyncResult<Option<Nonce>>),
 }
 
 #[async_trait]
@@ -123,6 +130,21 @@ impl StateSyncClient for LocalStateSyncClient {
             StateSyncError
         )
     }
+
+    async fn get_nonce_at(
+        &self,
+        block_number: BlockNumber,
+        contract_address: ContractAddress,
+    ) -> StateSyncClientResult<Option<Nonce>> {
+        let request = StateSyncRequest::GetNonceAt(block_number, contract_address);
+        let response = self.send(request).await;
+        handle_response_variants!(
+            StateSyncResponse,
+            GetNonceAt,
+            StateSyncClientError,
+            StateSyncError
+        )
+    }
 }
 
 #[async_trait]
@@ -166,4 +188,19 @@ impl StateSyncClient for RemoteStateSyncClient {
             StateSyncError
         )
     }
+
+    async fn get_nonce_at(
+        &self,
+        block_number: BlockNumber,
+        contract_address: ContractAddress,
+    ) -> StateSyncClientResult<Option<Nonce>> {
+        let request = StateSyncRequest::GetNonceAt(block_number, contract_address);
+        let response = self.send(request).await;
+        handle_response_variants!(
+            StateSyncResponse,
+            GetNonceAt,
+            StateSyncClientError,
+            StateSyncError
+        )
+    }
 }