fix(katana): separate deprecated declared class

crates/katana/executor/src/implementation/blockifier/utils.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::num::NonZeroU128;
 use std::sync::Arc;
 
@@ -432,9 +432,21 @@ pub(super) fn state_update_from_cached_state<S: StateDb>(
         katana_primitives::class::ContractClass,
     > = BTreeMap::new();
 
-    for class_hash in state_diff.compiled_class_hashes.keys() {
+    let mut declared_classes = BTreeMap::new();
+    let mut deprecated_declared_classes = BTreeSet::new();
+
+    // TODO: Legacy class shouldn't have a compiled class hash. This is a hack we added
+    // in our fork of `blockifier. Check if it's possible to remove it now.
+    for (class_hash, compiled_hash) in state_diff.compiled_class_hashes {
         let hash = class_hash.0;
         let class = state.class(hash).unwrap().expect("must exist if declared");
+
+        if class.is_legacy() {
+            deprecated_declared_classes.insert(hash);
+        } else {
+            declared_classes.insert(hash, compiled_hash.0);
+        }
+
         declared_contract_classes.insert(hash, class);
     }
 
@@ -470,24 +482,15 @@ pub(super) fn state_update_from_cached_state<S: StateDb>(
                 katana_primitives::class::ClassHash,
             >>();
 
-    let declared_classes =
-        state_diff
-            .compiled_class_hashes
-            .into_iter()
-            .map(|(key, value)| (key.0, value.0))
-            .collect::<BTreeMap<
-                katana_primitives::class::ClassHash,
-                katana_primitives::class::CompiledClassHash,
-            >>();
-
     StateUpdatesWithClasses {
         classes: declared_contract_classes,
         state_updates: StateUpdates {
             nonce_updates,
             storage_updates,
-            deployed_contracts,
             declared_classes,
-            ..Default::default()
+            deployed_contracts,
+            deprecated_declared_classes,
+            replaced_classes: BTreeMap::default(),
         },
     }
 }

crates/katana/primitives/src/chain_spec.rs

@@ -14,8 +14,7 @@ use crate::genesis::allocation::{DevAllocationsGenerator, GenesisAllocation};
 use crate::genesis::constant::{
     get_fee_token_balance_base_storage_address, DEFAULT_ACCOUNT_CLASS_PUBKEY_STORAGE_SLOT,
     DEFAULT_ETH_FEE_TOKEN_ADDRESS, DEFAULT_LEGACY_ERC20_CLASS, DEFAULT_LEGACY_ERC20_CLASS_HASH,
-    DEFAULT_LEGACY_UDC_CLASS, DEFAULT_LEGACY_UDC_CLASS_HASH,
-    DEFAULT_LEGACY_UDC_COMPILED_CLASS_HASH, DEFAULT_PREFUNDED_ACCOUNT_BALANCE,
+    DEFAULT_LEGACY_UDC_CLASS, DEFAULT_LEGACY_UDC_CLASS_HASH, DEFAULT_PREFUNDED_ACCOUNT_BALANCE,
     DEFAULT_STRK_FEE_TOKEN_ADDRESS, DEFAULT_UDC_ADDRESS, ERC20_DECIMAL_STORAGE_SLOT,
     ERC20_NAME_STORAGE_SLOT, ERC20_SYMBOL_STORAGE_SLOT, ERC20_TOTAL_SUPPLY_STORAGE_SLOT,
 };
@@ -233,11 +232,7 @@ fn add_default_udc(states: &mut StateUpdatesWithClasses) {
         .entry(DEFAULT_LEGACY_UDC_CLASS_HASH)
         .or_insert_with(|| DEFAULT_LEGACY_UDC_CLASS.clone());
 
-    states
-        .state_updates
-        .declared_classes
-        .entry(DEFAULT_LEGACY_UDC_CLASS_HASH)
-        .or_insert(DEFAULT_LEGACY_UDC_COMPILED_CLASS_HASH);
+    states.state_updates.deprecated_declared_classes.insert(DEFAULT_LEGACY_UDC_CLASS_HASH);
 
     // deploy UDC contract
     states

crates/katana/rpc/rpc-types/src/state_update.rs

@@ -1,3 +1,4 @@
+use katana_primitives::class::ClassHash;
 use serde::{Deserialize, Serialize};
 use starknet::core::types::{
     ContractStorageDiffItem, DeclaredClassItem, DeployedContractItem, NonceUpdate, StorageEntry,
@@ -49,6 +50,9 @@
             .map(|(addr, nonce)| NonceUpdate { nonce, contract_address: addr.into() })
             .collect();
 
+        let deprecated_declared_classes: Vec<ClassHash> =
+            value.deprecated_declared_classes.into_iter().collect();
+
         let declared_classes: Vec<DeclaredClassItem> = value
             .declared_classes
             .into_iter()
@@ -81,8 +85,8 @@
             storage_diffs,
             declared_classes,
             deployed_contracts,
+            deprecated_declared_classes,
             replaced_classes: Default::default(),
-            deprecated_declared_classes: Default::default(),
         })
     }
 }

crates/katana/storage/provider/src/providers/db/mod.rs

@@ -1,7 +1,7 @@
 pub mod state;
 pub mod trie;
 
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::fmt::Debug;
 use std::ops::{Range, RangeInclusive};
 
@@ -301,20 +301,24 @@
                 Ok((contract_address, class_hash))
             })?;
 
-            let declared_classes = dup_entries::<
-                Db,
-                tables::ClassDeclarations,
-                BTreeMap<ClassHash, CompiledClassHash>,
-                _,
-            >(&db_tx, block_num, |entry| {
-                let (_, class_hash) = entry?;
-
-                let compiled_hash = db_tx
-                    .get::<tables::CompiledClassHashes>(class_hash)?
-                    .ok_or(ProviderError::MissingCompiledClassHash(class_hash))?;
+            let mut declared_classes = BTreeMap::new();
+            let mut deprecated_declared_classes = BTreeSet::new();
 
-                Ok((class_hash, compiled_hash))
-            })?;
+            if let Some(block_entries) =
+                db_tx.cursor_dup::<tables::ClassDeclarations>()?.walk_dup(Some(block_num), None)?
+            {
+                for entry in block_entries {
+                    let (_, class_hash) = entry?;
+                    match db_tx.get::<tables::CompiledClassHashes>(class_hash)? {
+                        Some(compiled_hash) => {
+                            declared_classes.insert(class_hash, compiled_hash);
+                        }
+                        None => {
+                            deprecated_declared_classes.insert(class_hash);
+                        }
+                    }
+                }
+            }
 
             let storage_updates = {
                 let entries = dup_entries::<
@@ -337,12 +341,14 @@
             };
 
             db_tx.commit()?;
+
             Ok(Some(StateUpdates {
                 nonce_updates,
                 storage_updates,
                 deployed_contracts,
                 declared_classes,
-                ..Default::default()
+                deprecated_declared_classes,
+                replaced_classes: BTreeMap::default(),
             }))
         } else {
             Ok(None)