fix(katana): genesis deprecated declared classes in state updates (#2893

)
dojoengine · Jan 11, 2025 · b5e45c3 · b5e45c3
1 parent dd1ccce
commit b5e45c3
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 27 deletions.
diff --git a/crates/katana/core/Cargo.toml b/crates/katana/core/Cargo.toml
@@ -28,7 +28,6 @@ reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 starknet.workspace = true
-starknet-crypto.workspace = true
 starknet-types-core.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
@@ -37,6 +36,7 @@ url.workspace = true
 
 alloy-primitives = { workspace = true, features = [ "serde" ] }
 alloy-sol-types = { workspace = true, default-features = false, features = [ "json" ] }
+starknet-crypto = { workspace = true, optional = true }
 
 alloy-contract = { workspace = true, default-features = false }
 alloy-network = { workspace = true, default-features = false }
@@ -50,4 +50,4 @@ hex.workspace = true
 tempfile.workspace = true
 
 [features]
-starknet-messaging = [  ]
+starknet-messaging = [ "dep:starknet-crypto" ]
diff --git a/crates/katana/primitives/src/chain_spec.rs b/crates/katana/primitives/src/chain_spec.rs
@@ -80,7 +80,12 @@ impl ChainSpec {
         for (class_hash, class) in &self.genesis.classes {
             let class_hash = *class_hash;
 
-            states.state_updates.declared_classes.insert(class_hash, class.compiled_class_hash);
+            if class.class.is_legacy() {
+                states.state_updates.deprecated_declared_classes.insert(class_hash);
+            } else {
+                states.state_updates.declared_classes.insert(class_hash, class.compiled_class_hash);
+            }
+
             states.classes.insert(class_hash, class.class.as_ref().clone());
         }
 
@@ -412,7 +417,7 @@ mod tests {
         assert_eq!(
             actual_state_updates
                 .state_updates
-                .declared_classes
+                .deprecated_declared_classes
                 .get(&DEFAULT_LEGACY_ERC20_CLASS_HASH),
             Some(&DEFAULT_LEGACY_ERC20_COMPILED_CLASS_HASH),
         );
@@ -444,7 +449,10 @@ mod tests {
         );
 
         assert_eq!(
-            actual_state_updates.state_updates.declared_classes.get(&DEFAULT_LEGACY_UDC_CLASS_HASH),
+            actual_state_updates
+                .state_updates
+                .deprecated_declared_classes
+                .get(&DEFAULT_LEGACY_UDC_CLASS_HASH),
             Some(&DEFAULT_LEGACY_UDC_COMPILED_CLASS_HASH),
             "The default universal deployer class should be declared"
         );

diff --git a/crates/katana/primitives/src/class.rs b/crates/katana/primitives/src/class.rs
@@ -57,20 +57,11 @@ impl ContractClass {
         }
     }
 
-    /// Returns the class as a Sierra class, if any.
-    pub fn as_class(&self) -> Option<&SierraContractClass> {
-        match self {
-            Self::Class(class) => Some(class),
-            _ => None,
-        }
-    }
-
-    /// Returns the class as a legacy class, if any.
-    pub fn as_legacy(&self) -> Option<&LegacyContractClass> {
-        match self {
-            Self::Legacy(class) => Some(class),
-            _ => None,
-        }
+    /// Checks if this contract class is a Cairo 0 legacy class.
+    ///
+    /// Returns `true` if the contract class is a legacy class, `false` otherwise.
+    pub fn is_legacy(&self) -> bool {
+        matches!(self, Self::Legacy(_))
     }
 }
 

diff --git a/crates/katana/storage/provider/src/providers/db/mod.rs b/crates/katana/storage/provider/src/providers/db/mod.rs
@@ -702,6 +702,11 @@ impl<Db: Database> BlockWriter for DbProvider<Db> {
                 db_tx.put::<tables::ClassDeclarations>(block_number, class_hash)?
             }
 
+            for class_hash in states.state_updates.deprecated_declared_classes {
+                db_tx.put::<tables::ClassDeclarationBlock>(class_hash, block_number)?;
+                db_tx.put::<tables::ClassDeclarations>(block_number, class_hash)?
+            }
+
             for (class_hash, class) in states.classes {
                 // generate the compiled class
                 let compiled = class.clone().compile()?;

diff --git a/crates/katana/storage/provider/src/providers/db/state.rs b/crates/katana/storage/provider/src/providers/db/state.rs
@@ -225,22 +225,29 @@ impl<Tx: DbTx + fmt::Debug> HistoricalStateProvider<Tx> {
     pub fn new(tx: Tx, block_number: u64) -> Self {
         Self { tx, block_number }
     }
+
+    /// Check if the class was declared before the pinned block number.
+    fn is_class_declared_before_block(&self, hash: ClassHash) -> ProviderResult<bool> {
+        let decl_block_num = self.tx.get::<tables::ClassDeclarationBlock>(hash)?;
+        let is_declared = decl_block_num.is_some_and(|num| num <= self.block_number);
+        Ok(is_declared)
+    }
 }
 
 impl<Tx> ContractClassProvider for HistoricalStateProvider<Tx>
 where
     Tx: DbTx + fmt::Debug + Send + Sync,
 {
     fn class(&self, hash: ClassHash) -> ProviderResult<Option<ContractClass>> {
-        if self.compiled_class_hash_of_class_hash(hash)?.is_some() {
+        if self.is_class_declared_before_block(hash)? {
             Ok(self.tx.get::<tables::Classes>(hash)?)
         } else {
             Ok(None)
         }
     }
 
     fn compiled_class(&self, hash: ClassHash) -> ProviderResult<Option<CompiledClass>> {
-        if self.compiled_class_hash_of_class_hash(hash)?.is_some() {
+        if self.is_class_declared_before_block(hash)? {
             Ok(self.tx.get::<tables::CompiledClasses>(hash)?)
         } else {
             Ok(None)
@@ -251,12 +258,7 @@ where
         &self,
         hash: ClassHash,
     ) -> ProviderResult<Option<CompiledClassHash>> {
-        // check that the requested class hash was declared before the pinned block number
-        if self
-            .tx
-            .get::<tables::ClassDeclarationBlock>(hash)?
-            .is_some_and(|num| num <= self.block_number)
-        {
+        if self.is_class_declared_before_block(hash)? {
             Ok(self.tx.get::<tables::CompiledClassHashes>(hash)?)
         } else {
             Ok(None)