Pre-forking state refactor

crates/starknet-devnet-core/src/account.rs

@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use blockifier::state::state_api::{State, StateReader};
 use starknet_api::core::{calculate_contract_address, PatriciaKey};
 use starknet_api::hash::{StarkFelt, StarkHash};
 use starknet_api::transaction::{Calldata, ContractAddressSalt};
@@ -17,7 +18,8 @@ use crate::constants::{
     CHARGEABLE_ACCOUNT_PUBLIC_KEY,
 };
 use crate::error::DevnetResult;
-use crate::traits::{Accounted, Deployed, StateChanger, StateExtractor};
+use crate::state::{CustomState, StarknetState};
+use crate::traits::{Accounted, Deployed};
 use crate::utils::get_storage_var_address;
 
 /// data taken from https://github.com/0xSpaceShard/starknet-devnet/blob/fb96e0cc3c1c31fb29892ecefd2a670cf8a32b51/starknet_devnet/account.py
@@ -114,19 +116,21 @@ impl Account {
 }
 
 impl Deployed for Account {
-    fn deploy(&self, state: &mut (impl StateChanger + StateExtractor)) -> DevnetResult<()> {
-        // declare if not declared
-        if !state.is_contract_declared(&self.class_hash) {
-            state.declare_contract_class(self.class_hash, self.contract_class.clone())?;
-        }
+    fn deploy(&self, state: &mut StarknetState) -> DevnetResult<()> {
+        self.declare_if_undeclared(state, self.class_hash, &self.contract_class)?;
 
-        // deploy
-        state.deploy_contract(self.account_address, self.class_hash)?;
+        state.predeploy_contract(self.account_address, self.class_hash)?;
 
-        // set public key
+        // set public key directly in the most underlying state
         let public_key_storage_var = get_storage_var_address("Account_public_key", &[])?;
-        let storage_key = ContractStorageKey::new(self.account_address, public_key_storage_var);
-        state.change_storage(storage_key, self.public_key)?;
+        state.state.state.set_storage_at(
+            self.account_address.try_into()?,
+            public_key_storage_var.try_into()?,
+            self.public_key.into(),
+        )?;
+
+        // set balance directly in the most underlying state
+        self.set_initial_balance(&mut state.state.state)?;
 
         Ok(())
     }
@@ -137,29 +141,31 @@ impl Deployed for Account {
 }
 
 impl Accounted for Account {
-    fn set_initial_balance(&self, state: &mut impl StateChanger) -> DevnetResult<()> {
+    fn set_initial_balance(&self, state: &mut impl State) -> DevnetResult<()> {
         let storage_var_address =
             get_storage_var_address("ERC20_balances", &[Felt::from(self.account_address)])?;
 
         for fee_token_address in [self.eth_fee_token_address, self.strk_fee_token_address] {
-            let storage_key = ContractStorageKey::new(fee_token_address, storage_var_address);
-
-            state.change_storage(storage_key, self.initial_balance)?;
+            state.set_storage_at(
+                fee_token_address.try_into()?,
+                storage_var_address.try_into()?,
+                self.initial_balance.into(),
+            )?;
         }
 
         Ok(())
     }
 
-    fn get_balance(
-        &self,
-        state: &mut impl StateExtractor,
-        token: FeeToken,
-    ) -> DevnetResult<Balance> {
+    fn get_balance(&self, state: &mut impl StateReader, token: FeeToken) -> DevnetResult<Balance> {
         let balance_storage_key = match token {
             FeeToken::ETH => self.eth_balance_storage_key()?,
             FeeToken::STRK => self.strk_balance_storage_key()?,
         };
-        state.get_storage(balance_storage_key)
+        let balance = state.get_storage_at(
+            (*balance_storage_key.get_contract_address()).try_into()?,
+            (*balance_storage_key.get_storage_key()).try_into()?,
+        )?;
+        Ok(balance.into())
     }
 }
 
@@ -172,8 +178,8 @@ mod tests {
     use super::Account;
     use crate::account::FeeToken;
     use crate::constants::CAIRO_1_ERC20_CONTRACT_CLASS_HASH;
-    use crate::state::StarknetState;
-    use crate::traits::{Accounted, Deployed, StateChanger};
+    use crate::state::{CustomState, StarknetState};
+    use crate::traits::{Accounted, Deployed};
     use crate::utils::exported_test_utils::dummy_cairo_0_contract_class;
     use crate::utils::get_storage_var_address;
     use crate::utils::test_utils::{dummy_contract_address, dummy_felt};
@@ -253,25 +259,12 @@ mod tests {
         assert!(account.deploy(&mut state).is_ok());
     }
 
-    #[test]
-    fn account_get_balance_should_return_zero_because_balance_was_not_set() {
-        let (account, mut state) = setup();
-
-        account.deploy(&mut state).unwrap();
-        let balance = account.get_balance(&mut state, FeeToken::ETH).unwrap();
-        assert_eq!(balance, Felt::from(0));
-
-        let balance = account.get_balance(&mut state, FeeToken::STRK).unwrap();
-        assert_eq!(balance, Felt::from(0));
-    }
-
     #[test]
     fn account_get_balance_should_return_correct_value() {
         let (mut account, mut state) = setup();
         let expected_balance = Felt::from(100);
         account.initial_balance = expected_balance;
         account.deploy(&mut state).unwrap();
-        account.set_initial_balance(&mut state).unwrap();
         let generated_balance = account.get_balance(&mut state, FeeToken::ETH).unwrap();
 
         assert_eq!(expected_balance, generated_balance);
@@ -301,7 +294,7 @@ mod tests {
 
         // deploy the erc20 contract
         state
-            .deploy_contract(
+            .predeploy_contract(
                 fee_token_address,
                 Felt::from_prefixed_hex_str(CAIRO_1_ERC20_CONTRACT_CLASS_HASH).unwrap(),
             )

crates/starknet-devnet-core/src/starknet/add_declare_transaction.rs

@@ -126,10 +126,15 @@ pub fn add_declare_transaction_v1(
 
 #[cfg(test)]
 mod tests {
+    use blockifier::state::state_api::StateReader;
     use nonzero_ext::nonzero;
     use starknet_api::block::BlockNumber;
+    use starknet_api::core::CompiledClassHash;
+    use starknet_api::hash::StarkHash;
     use starknet_api::transaction::Fee;
-    use starknet_rs_core::types::{TransactionExecutionStatus, TransactionFinalityStatus};
+    use starknet_rs_core::types::{
+        BlockId, BlockTag, TransactionExecutionStatus, TransactionFinalityStatus,
+    };
     use starknet_types::contract_address::ContractAddress;
     use starknet_types::contract_class::{Cairo0Json, ContractClass};
     use starknet_types::felt::Felt;
@@ -143,7 +148,8 @@ mod tests {
     };
     use crate::starknet::predeployed::create_erc20_at_address;
     use crate::starknet::{predeployed, Starknet};
-    use crate::traits::{Accounted, Deployed, HashIdentifiedMut, StateExtractor};
+    use crate::state::CustomStateReader;
+    use crate::traits::{Deployed, HashIdentifiedMut};
     use crate::utils::exported_test_utils::dummy_cairo_0_contract_class;
     use crate::utils::test_utils::{
         convert_broadcasted_declare_v2_to_v3, dummy_broadcasted_declare_transaction_v2,
@@ -248,7 +254,7 @@ mod tests {
         // check if txn is with status accepted
         assert_eq!(tx.finality_status, TransactionFinalityStatus::AcceptedOnL2);
         assert_eq!(tx.execution_result.status(), TransactionExecutionStatus::Succeeded);
-        assert!(starknet.state.contract_classes.get(&class_hash).is_some());
+        starknet.state.get_rpc_contract_class(&class_hash).unwrap();
     }
 
     #[test]
@@ -262,14 +268,17 @@ mod tests {
         let tx = starknet.transactions.get_by_hash_mut(&tx_hash).unwrap();
 
         // check if generated class hash is expected one
-        assert_eq!(
-            class_hash,
-            ContractClass::Cairo1(declare_txn.contract_class).generate_hash().unwrap()
-        );
+        let generated_hash =
+            ContractClass::Cairo1(declare_txn.contract_class.clone()).generate_hash().unwrap();
+        assert_eq!(class_hash, generated_hash);
+
         // check if txn is with status accepted
         assert_eq!(tx.finality_status, TransactionFinalityStatus::AcceptedOnL2);
         assert_eq!(tx.execution_result.status(), TransactionExecutionStatus::Succeeded);
-        assert!(starknet.state.contract_classes.get(&class_hash).is_some());
+        assert_eq!(
+            starknet.get_class(&BlockId::Tag(BlockTag::Latest), class_hash).unwrap(),
+            declare_txn.contract_class.into()
+        );
     }
 
     #[test]
@@ -281,15 +290,12 @@ mod tests {
         let expected_compiled_class_hash = declare_txn.compiled_class_hash;
 
         // check if contract is not declared
-        assert!(!starknet.state.is_contract_declared(&expected_class_hash));
-        assert!(
-            !starknet
-                .state
-                .state
-                .state
-                .class_hash_to_compiled_class
-                .contains_key(&expected_compiled_class_hash)
+        assert!(!starknet.state.is_contract_declared(expected_class_hash));
+        assert_eq!(
+            starknet.state.get_compiled_class_hash(expected_class_hash.into()).unwrap(),
+            CompiledClassHash(StarkHash::ZERO)
         );
+        assert!(starknet.get_class(&BlockId::Tag(BlockTag::Latest), expected_class_hash).is_err());
 
         let (tx_hash, retrieved_class_hash) =
             starknet.add_declare_transaction_v2(declare_txn).unwrap();
@@ -301,7 +307,11 @@ mod tests {
         // check if txn is with status accepted
         assert_eq!(retrieved_txn.finality_status, TransactionFinalityStatus::AcceptedOnL2);
         assert_eq!(retrieved_txn.execution_result.status(), TransactionExecutionStatus::Succeeded);
-        assert!(starknet.state.is_contract_declared(&expected_class_hash));
+        assert!(starknet.state.is_contract_declared(expected_class_hash));
+        assert_eq!(
+            starknet.state.get_compiled_class_hash(expected_class_hash.into()).unwrap(),
+            expected_compiled_class_hash.into()
+        );
     }
 
     #[test]
@@ -374,7 +384,7 @@ mod tests {
         assert_eq!(tx.finality_status, TransactionFinalityStatus::AcceptedOnL2);
         assert_eq!(tx.execution_result.status(), TransactionExecutionStatus::Succeeded);
         // check if contract is successfully declared
-        assert!(starknet.state.is_contract_declared(&class_hash));
+        assert!(starknet.state.is_contract_declared(class_hash));
         // check if pending block is resetted
         assert!(starknet.pending_block().get_transactions().is_empty());
         // check if there is generated block
@@ -400,7 +410,7 @@ mod tests {
 
         let expected_class_hash = declare_txn.contract_class.generate_hash().unwrap();
         // check if contract is not declared
-        assert!(!starknet.state.is_contract_declared(&expected_class_hash));
+        assert!(!starknet.state.is_contract_declared(expected_class_hash));
 
         let (tx_hash, class_hash) = starknet.add_declare_transaction_v1(declare_txn).unwrap();
 
@@ -411,7 +421,7 @@ mod tests {
         assert_eq!(tx.execution_result.status(), TransactionExecutionStatus::Succeeded);
 
         // check if contract is declared
-        assert!(starknet.state.is_contract_declared(&class_hash));
+        assert!(starknet.state.is_contract_declared(class_hash));
     }
 
     /// Initializes starknet with 1 account - account without validations
@@ -442,9 +452,7 @@ mod tests {
         .unwrap();
 
         acc.deploy(&mut starknet.state).unwrap();
-        acc.set_initial_balance(&mut starknet.state).unwrap();
 
-        starknet.state.clear_dirty_state();
         starknet.block_context = Starknet::init_block_context(
             nonzero!(1u128),
             constants::ETH_ERC20_CONTRACT_ADDRESS,