wip

kariy · kariy · commit afef89113682 · 2025-10-17T22:02:25.000-04:00
diff --git a/crates/storage/db/src/models/contract.rs b/crates/storage/db/src/models/contract.rs
@@ -12,30 +12,108 @@ pub struct ContractInfoChangeList {
     pub nonce_change_list: BlockList,
 }
 
+/// The type of event that triggered the class change.
+#[derive(Debug, Default, PartialEq, Eq)]
+#[cfg_attr(test, derive(::arbitrary::Arbitrary))]
+pub enum ContractClassChangeType {
+    /// The contract was deployed with the given class hash.
+    #[default]
+    Deployed,
+    /// The class change is made using the `replace_class` syscall.
+    Replaced,
+}
+
 #[derive(Debug, Default, PartialEq, Eq)]
 #[cfg_attr(test, derive(::arbitrary::Arbitrary))]
 pub struct ContractClassChange {
+    /// The type of class change.
+    pub r#type: ContractClassChangeType,
+    /// The address of the contract whose class has changed.
     pub contract_address: ContractAddress,
     /// The updated class hash of `contract_address`.
     pub class_hash: ClassHash,
 }
 
+impl ContractClassChange {
+    /// Creates a new `ContractClassChange` instance representing a deployed contract.
+    pub fn deployed(contract_address: ContractAddress, class_hash: ClassHash) -> Self {
+        Self { r#type: ContractClassChangeType::Deployed, contract_address, class_hash }
+    }
+
+    /// Creates a new `ContractClassChange` instance representing a replaced class
+    pub fn replaced(contract_address: ContractAddress, new_class_hash: ClassHash) -> Self {
+        Self {
+            r#type: ContractClassChangeType::Replaced,
+            contract_address,
+            class_hash: new_class_hash,
+        }
+    }
+
+    fn decompress_current(bytes: &[u8]) -> Result<Self, CodecError> {
+        let r#type = ContractClassChangeType::decompress(&bytes[0..1])?;
+        let contract_address = ContractAddress::decode(&bytes[1..33])?;
+        let class_hash = ClassHash::decompress(&bytes[33..])?;
+        Ok(Self { r#type, contract_address, class_hash })
+    }
+
+    /// Backward compatibility purposes.
+    ///
+    /// For old format, ContractClassChangeType::Deployed type is assumed.
+    #[cold]
+    fn decompress_legacy(bytes: &[u8]) -> Result<Self, CodecError> {
+        let contract_address = ContractAddress::decode(&bytes[0..32])?;
+        let class_hash = ClassHash::decompress(&bytes[32..])?;
+        Ok(Self { r#type: ContractClassChangeType::Deployed, contract_address, class_hash })
+    }
+}
+
+impl Compress for ContractClassChangeType {
+    type Compressed = Vec<u8>;
+    fn compress(self) -> Result<Self::Compressed, CodecError> {
+        let byte: u8 = match self {
+            ContractClassChangeType::Deployed => 0,
+            ContractClassChangeType::Replaced => 1,
+        };
+        Ok(vec![byte])
+    }
+}
+
+impl Decompress for ContractClassChangeType {
+    fn decompress<B: AsRef<[u8]>>(bytes: B) -> Result<Self, CodecError> {
+        let bytes = bytes.as_ref();
+        if bytes.is_empty() {
+            return Err(CodecError::Decompress("can't decompress empty bytes".to_string()));
+        }
+
+        match bytes[0] {
+            0 => Ok(ContractClassChangeType::Deployed),
+            1 => Ok(ContractClassChangeType::Replaced),
+            _ => Err(CodecError::Decompress("unknown ContractClassChangeType variant".to_string())),
+        }
+    }
+}
+
 impl Compress for ContractClassChange {
     type Compressed = Vec<u8>;
+
     fn compress(self) -> Result<Self::Compressed, CodecError> {
         let mut buf = Vec::new();
-        buf.extend_from_slice(self.contract_address.encode().as_ref());
-        buf.extend_from_slice(self.class_hash.compress()?.as_ref());
+        buf.extend(self.r#type.compress()?);
+        buf.extend(self.contract_address.encode());
+        buf.extend(self.class_hash.compress()?);
         Ok(buf)
     }
 }
 
 impl Decompress for ContractClassChange {
     fn decompress<B: AsRef<[u8]>>(bytes: B) -> Result<Self, crate::error::CodecError> {
         let bytes = bytes.as_ref();
-        let contract_address = ContractAddress::decode(&bytes[0..32])?;
-        let class_hash = ClassHash::decompress(&bytes[32..])?;
-        Ok(Self { contract_address, class_hash })
+
+        if let Ok(result) = Self::decompress_current(bytes) {
+            return Ok(result);
+        }
+
+        ContractClassChange::decompress_legacy(bytes)
     }
 }
 
diff --git a/crates/storage/provider/provider/src/providers/db/mod.rs b/crates/storage/provider/provider/src/providers/db/mod.rs
@@ -9,7 +9,7 @@ use katana_db::abstraction::{Database, DbCursor, DbCursorMut, DbDupSortCursor, D
 use katana_db::error::DatabaseError;
 use katana_db::models::block::StoredBlockBodyIndices;
 use katana_db::models::contract::{
-    ContractClassChange, ContractInfoChangeList, ContractNonceChange,
+    ContractClassChange, ContractClassChangeType, ContractInfoChangeList, ContractNonceChange,
 };
 use katana_db::models::list::BlockList;
 use katana_db::models::stage::StageCheckpoint;
@@ -301,8 +301,28 @@ impl<Db: Database> StateUpdateProvider for DbProvider<Db> {
                 BTreeMap<ContractAddress, ClassHash>,
                 _,
             >(&db_tx, block_num, |entry| {
-                let (_, ContractClassChange { contract_address, class_hash }) = entry?;
-                Ok(Some((contract_address, class_hash)))
+                let (_, ContractClassChange { r#type, contract_address, class_hash }) = entry?;
+
+                if r#type == ContractClassChangeType::Deployed {
+                    Ok(Some((contract_address, class_hash)))
+                } else {
+                    Ok(None)
+                }
+            })?;
+
+            let replaced_classes = dup_entries::<
+                Db,
+                tables::ClassChangeHistory,
+                BTreeMap<ContractAddress, ClassHash>,
+                _,
+            >(&db_tx, block_num, |entry| {
+                let (_, ContractClassChange { r#type, contract_address, class_hash }) = entry?;
+
+                if r#type == ContractClassChangeType::Replaced {
+                    Ok(Some((contract_address, class_hash)))
+                } else {
+                    Ok(None)
+                }
             })?;
 
             let mut declared_classes = BTreeMap::new();
@@ -351,8 +371,8 @@ impl<Db: Database> StateUpdateProvider for DbProvider<Db> {
                 storage_updates,
                 deployed_contracts,
                 declared_classes,
+                replaced_classes,
                 deprecated_declared_classes,
-                replaced_classes: BTreeMap::default(),
             }))
         } else {
             Ok(None)
@@ -431,8 +451,12 @@ impl<Db: Database> StateUpdateProvider for DbProvider<Db> {
                 BTreeMap<ContractAddress, ClassHash>,
                 _,
             >(&db_tx, block_num, |entry| {
-                let (_, ContractClassChange { contract_address, class_hash }) = entry?;
-                Ok(Some((contract_address, class_hash)))
+                let (_, ContractClassChange { r#type, contract_address, class_hash }) = entry?;
+                if r#type == ContractClassChangeType::Deployed {
+                    Ok(Some((contract_address, class_hash)))
+                } else {
+                    Ok(None)
+                }
             })?;
 
             db_tx.commit()?;
@@ -835,11 +859,24 @@ impl<Db: Database> BlockWriter for DbProvider<Db> {
 
                 db_tx.put::<tables::ContractInfo>(addr, value)?;
 
-                let class_change_key = ContractClassChange { contract_address: addr, class_hash };
+                let class_change_key = ContractClassChange::deployed(addr, class_hash);
                 db_tx.put::<tables::ClassChangeHistory>(block_number, class_change_key)?;
                 db_tx.put::<tables::ContractInfoChangeSet>(addr, new_change_set)?;
             }
 
+            for (addr, new_class_hash) in states.state_updates.replaced_classes {
+                let mut info = db_tx.get::<tables::ContractInfo>(addr)?.unwrap();
+                info.class_hash = new_class_hash;
+                db_tx.put::<tables::ContractInfo>(addr, info)?;
+
+                let mut change_set = db_tx.get::<tables::ContractInfoChangeSet>(addr)?.unwrap();
+                change_set.class_change_list.insert(block_number);
+                db_tx.put::<tables::ContractInfoChangeSet>(addr, change_set)?;
+
+                let class_change_key = ContractClassChange::replaced(addr, new_class_hash);
+                db_tx.put::<tables::ClassChangeHistory>(block_number, class_change_key)?;
+            }
+
             for (addr, nonce) in states.state_updates.nonce_updates {
                 let value = if let Some(info) = db_tx.get::<tables::ContractInfo>(addr)? {
                     GenericContractInfo { nonce, ..info }
diff --git a/crates/storage/provider/provider/src/providers/fork/state.rs b/crates/storage/provider/provider/src/providers/fork/state.rs
@@ -2,7 +2,9 @@ use std::cmp::Ordering;
 use std::sync::Arc;
 
 use katana_db::abstraction::{Database, DbTx, DbTxMut};
-use katana_db::models::contract::{ContractClassChange, ContractNonceChange};
+use katana_db::models::contract::{
+    ContractClassChange, ContractClassChangeType, ContractNonceChange,
+};
 use katana_db::models::storage::{ContractStorageEntry, ContractStorageKey, StorageEntry};
 use katana_db::tables;
 use katana_fork::BackendClient;
@@ -270,7 +272,11 @@ impl<Db: Database> StateProvider for HistoricalStateProvider<Db> {
             Ok(res)
         } else if let res @ Some(class_hash) = self.backend.get_class_hash_at(address)? {
             let block = self.provider.block();
-            let entry = ContractClassChange { contract_address: address, class_hash };
+            let entry = ContractClassChange {
+                r#type: ContractClassChangeType::Deployed,
+                contract_address: address,
+                class_hash,
+            };
 
             self.db.db().tx_mut()?.put::<tables::ClassChangeHistory>(block, entry)?;
             Ok(res)
