From 7ce138b58af8b354e57868a535d850044438a24a Mon Sep 17 00:00:00 2001
From: Arni Hod <arnon@starkware.co>
Date: Sun, 7 Dec 2025 16:48:09 +0200
Subject: [PATCH] apollo_state_sync: align functionality with apollo state
 reader

---
 .../apollo_state_reader/src/apollo_state.rs   |  7 ++-
 crates/apollo_state_sync/src/lib.rs           | 44 ++++++++++---------
 2 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/crates/apollo_state_reader/src/apollo_state.rs b/crates/apollo_state_reader/src/apollo_state.rs
index c468f64f05e..64d8f1da176 100644
--- a/crates/apollo_state_reader/src/apollo_state.rs
+++ b/crates/apollo_state_reader/src/apollo_state.rs
@@ -294,14 +294,13 @@ impl FetchCompiledClasses for ApolloReader {
     }
 
     fn is_declared(&self, class_hash: ClassHash) -> StateResult<bool> {
-        let state_number = StateNumber(self.latest_block);
+        let state_number = self.latest_block;
         let class_declaration_block_number = self
             .reader()?
             .get_state_reader()
             .and_then(|sr| sr.get_class_definition_block_number(&class_hash))
             .map_err(|err| StateError::StateReadError(err.to_string()))?;
-        Ok(
-            matches!(class_declaration_block_number, Some(block_number) if block_number <= state_number.0),
-        )
+        Ok(matches!(class_declaration_block_number, Some(block_number)
+            if block_number <= state_number))
     }
 }
diff --git a/crates/apollo_state_sync/src/lib.rs b/crates/apollo_state_sync/src/lib.rs
index ea43643fbf2..36beb00ac8d 100644
--- a/crates/apollo_state_sync/src/lib.rs
+++ b/crates/apollo_state_sync/src/lib.rs
@@ -248,19 +248,32 @@ impl StateSync {
 
     async fn is_cairo_1_class_declared_at(
         &self,
-        block_number: BlockNumber,
+        state_number: BlockNumber,
         class_hash: ClassHash,
     ) -> StateSyncResult<bool> {
         let storage_reader = self.storage_reader.clone();
-        let class_definition_block_number_opt = storage_reader
+        let class_definition_block_number = storage_reader
             .begin_ro_txn()?
-            .get_state_reader()?
-            .get_class_definition_block_number(&class_hash)?;
-        if let Some(class_definition_block_number) = class_definition_block_number_opt {
-            Ok(class_definition_block_number <= block_number)
-        } else {
-            Ok(false)
-        }
+            .get_state_reader()
+            .and_then(|sr| sr.get_class_definition_block_number(&class_hash))?;
+        Ok(matches!(class_definition_block_number, Some(block_number)
+            if block_number <= state_number))
+    }
+
+    async fn is_cairo_0_class_declared_at(
+        &self,
+        state_number: BlockNumber,
+        class_hash: ClassHash,
+    ) -> StateSyncResult<bool> {
+        let storage_reader = self.storage_reader.clone();
+        // TODO(noamsp): Add unit testing for cairo0
+        let deprecated_class_definition_block_number = storage_reader
+            .begin_ro_txn()?
+            .get_state_reader()
+            .and_then(|sr| sr.get_deprecated_class_definition_block_number(&class_hash))?;
+
+        Ok(matches!(deprecated_class_definition_block_number, Some(block_number)
+            if block_number <= state_number))
     }
 
     async fn is_class_declared_at(
@@ -272,18 +285,7 @@ impl StateSync {
             return Ok(true);
         }
 
-        let storage_reader = self.storage_reader.clone();
-        // TODO(noamsp): Add unit testing for cairo0
-        let deprecated_class_definition_block_number_opt = storage_reader
-            .begin_ro_txn()?
-            .get_state_reader()?
-            .get_deprecated_class_definition_block_number(&class_hash)?;
-
-        Ok(deprecated_class_definition_block_number_opt.is_some_and(
-            |deprecated_class_definition_block_number| {
-                deprecated_class_definition_block_number <= block_number
-            },
-        ))
+        self.is_cairo_0_class_declared_at(block_number, class_hash).await
     }
 }