fix(katana): separate deprecated declared class (#2903)

Currently, we're including both legacy and non-legacy into the same map in the `StateUpdates` struct, while instead the legacy should be put into a different map ie `deprecated_declared_classes`.

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 @@ impl From<katana_primitives::state::StateUpdates> for StateDiff {
             .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 @@ impl From<katana_primitives::state::StateUpdates> for StateDiff {
             storage_diffs,
             declared_classes,
             deployed_contracts,
+            deprecated_declared_classes,
             replaced_classes: Default::default(),
-            deprecated_declared_classes: Default::default(),
         })
     }
 }

crates/katana/rpc/rpc/tests/starknet.rs

@@ -24,8 +24,9 @@ use starknet::core::types::contract::legacy::LegacyContractClass;
 use starknet::core::types::{
     BlockId, BlockTag, Call, DeclareTransactionReceipt, DeployAccountTransactionReceipt,
     EventFilter, EventsPage, ExecutionResult, Felt, MaybePendingBlockWithReceipts,
-    MaybePendingBlockWithTxHashes, MaybePendingBlockWithTxs, StarknetError,
-    TransactionExecutionStatus, TransactionFinalityStatus, TransactionReceipt, TransactionTrace,
+    MaybePendingBlockWithTxHashes, MaybePendingBlockWithTxs, MaybePendingStateUpdate,
+    StarknetError, TransactionExecutionStatus, TransactionFinalityStatus, TransactionReceipt,
+    TransactionTrace,
 };
 use starknet::core::utils::get_contract_address;
 use starknet::macros::{felt, selector};
@@ -56,6 +57,18 @@ async fn declare_and_deploy_contract() -> Result<()> {
     // check that the class is actually declared
     assert!(provider.get_class(BlockId::Tag(BlockTag::Pending), class_hash).await.is_ok());
 
+    // check state update includes class in declared_classes
+    let state_update = provider.get_state_update(BlockId::Tag(BlockTag::Latest)).await?;
+    match state_update {
+        MaybePendingStateUpdate::Update(update) => {
+            assert!(
+                update.state_diff.declared_classes.iter().any(|item| item.class_hash == class_hash
+                    && item.compiled_class_hash == compiled_class_hash)
+            );
+        }
+        _ => panic!("Expected Update, got PendingUpdate"),
+    }
+
     let ctor_args = vec![Felt::ONE, Felt::TWO];
     let calldata = [
         vec![
@@ -110,6 +123,15 @@ async fn declare_and_deploy_legacy_contract() -> Result<()> {
     // check that the class is actually declared
     assert!(provider.get_class(BlockId::Tag(BlockTag::Pending), class_hash).await.is_ok());
 
+    // check state update includes class in deprecated_declared_classes
+    let state_update = provider.get_state_update(BlockId::Tag(BlockTag::Latest)).await?;
+    match state_update {
+        MaybePendingStateUpdate::Update(update) => {
+            assert!(update.state_diff.deprecated_declared_classes.contains(&class_hash));
+        }
+        _ => panic!("Expected Update, got PendingUpdate"),
+    }
+
     let ctor_args = vec![Felt::ONE];
     let calldata = [
         vec![

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 @@ impl<Db: Database> StateUpdateProvider for DbProvider<Db> {
                 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 @@ impl<Db: Database> StateUpdateProvider for DbProvider<Db> {
             };
 
             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)