feat(katana): simplify genesis class (#2948)

dojoengine · Jan 24, 2025 · 894478e · 894478e
1 parent 4ffbac0
commit 894478e
Show file tree

Hide file tree

Showing 8 changed files with 4,903 additions and 4,891 deletions.
diff --git a/crates/katana/chain-spec/src/lib.rs b/crates/katana/chain-spec/src/lib.rs
@@ -124,13 +124,14 @@ impl ChainSpec {
         for (class_hash, class) in &self.genesis.classes {
             let class_hash = *class_hash;
 
-            if class.class.is_legacy() {
+            if 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);
+                let casm_hash = class.as_ref().clone().compile().unwrap().class_hash().unwrap();
+                states.state_updates.declared_classes.insert(class_hash, casm_hash);
             }
 
-            states.classes.insert(class_hash, class.class.as_ref().clone());
+            states.classes.insert(class_hash, class.as_ref().clone());
         }
 
         for (address, alloc) in &self.genesis.allocations {
@@ -313,7 +314,6 @@ mod tests {
         DEFAULT_LEGACY_ERC20_CLASS, DEFAULT_LEGACY_ERC20_COMPILED_CLASS_HASH,
         DEFAULT_LEGACY_UDC_CLASS, DEFAULT_LEGACY_UDC_COMPILED_CLASS_HASH,
     };
-    use katana_primitives::genesis::GenesisClass;
     use katana_primitives::version::CURRENT_STARKNET_VERSION;
     use starknet::macros::felt;
 
@@ -324,35 +324,11 @@ mod tests {
         // setup initial states to test
 
         let classes = BTreeMap::from([
-            (
-                DEFAULT_LEGACY_UDC_CLASS_HASH,
-                GenesisClass {
-                    class: DEFAULT_LEGACY_UDC_CLASS.clone().into(),
-                    compiled_class_hash: DEFAULT_LEGACY_UDC_COMPILED_CLASS_HASH,
-                },
-            ),
-            (
-                DEFAULT_LEGACY_ERC20_CLASS_HASH,
-                GenesisClass {
-                    class: DEFAULT_LEGACY_ERC20_CLASS.clone().into(),
-                    compiled_class_hash: DEFAULT_LEGACY_ERC20_COMPILED_CLASS_HASH,
-                },
-            ),
-            (
-                DEFAULT_ACCOUNT_CLASS_HASH,
-                GenesisClass {
-                    compiled_class_hash: DEFAULT_ACCOUNT_COMPILED_CLASS_HASH,
-                    class: DEFAULT_ACCOUNT_CLASS.clone().into(),
-                },
-            ),
+            (DEFAULT_LEGACY_UDC_CLASS_HASH, DEFAULT_LEGACY_UDC_CLASS.clone().into()),
+            (DEFAULT_LEGACY_ERC20_CLASS_HASH, DEFAULT_LEGACY_ERC20_CLASS.clone().into()),
+            (DEFAULT_ACCOUNT_CLASS_HASH, DEFAULT_ACCOUNT_CLASS.clone().into()),
             #[cfg(feature = "controller")]
-            (
-                CONTROLLER_CLASS_HASH,
-                GenesisClass {
-                    compiled_class_hash: CONTROLLER_CLASS_HASH,
-                    class: CONTROLLER_ACCOUNT_CLASS.clone().into(),
-                },
-            ),
+            (CONTROLLER_CLASS_HASH, CONTROLLER_ACCOUNT_CLASS.clone().into()),
         ]);
 
         let allocations = [
@@ -522,7 +498,7 @@ mod tests {
         {
             assert_eq!(
                 actual_state_updates.state_updates.declared_classes.get(&CONTROLLER_CLASS_HASH),
-                Some(&CONTROLLER_CLASS_HASH),
+                Some(&CONTROLLER_ACCOUNT_CLASS.clone().compile().unwrap().class_hash().unwrap()),
                 "The controller account class should be declared"
             );
 

diff --git a/crates/katana/controller/src/lib.rs b/crates/katana/controller/src/lib.rs
@@ -115,11 +115,8 @@ pub mod json {
         // parse the controller class json file
         let json = serde_json::from_str::<Value>(CONTROLLER_SIERRA_ARTIFACT)?;
 
-        let class = GenesisClassJson {
-            class_hash: None,
-            class: json.into(),
-            name: Some(CONTROLLER_CLASS_NAME.to_string()),
-        };
+        let class =
+            GenesisClassJson { class: json.into(), name: Some(CONTROLLER_CLASS_NAME.to_string()) };
 
         genesis.classes.push(class);
 

diff --git a/crates/katana/primitives/src/class.rs b/crates/katana/primitives/src/class.rs
@@ -67,7 +67,21 @@ impl ContractClass {
     ///
     /// Returns `true` if the contract class is a legacy class, `false` otherwise.
     pub fn is_legacy(&self) -> bool {
-        matches!(self, Self::Legacy(_))
+        self.as_legacy().is_some()
+    }
+
+    pub fn as_legacy(&self) -> Option<&LegacyContractClass> {
+        match self {
+            Self::Legacy(class) => Some(class),
+            _ => None,
+        }
+    }
+
+    pub fn as_sierra(&self) -> Option<&SierraContractClass> {
+        match self {
+            Self::Class(class) => Some(class),
+            _ => None,
+        }
     }
 }
 
@@ -126,6 +140,13 @@ impl CompiledClass {
             Self::Legacy(class) => Ok(compute_legacy_class_hash(class)?),
         }
     }
+
+    /// Checks if the compiled contract class is a legacy (Cairo 0) class.
+    ///
+    /// Returns `true` if the compiled contract class is a legacy class, `false` otherwise.
+    pub fn is_legacy(&self) -> bool {
+        matches!(self, Self::Legacy(_))
+    }
 }
 
 #[derive(Debug, thiserror::Error)]