fix: kakarot sequencer (kkrt-labs#468)

* add fee token and storate + nonce helpers

* update based on comments

crates/ef-testing/src/evm_sequencer/constants.rs

@@ -9,6 +9,12 @@ use starknet_api::{
     hash::StarkFelt,
 };
 
+fn load_legacy_contract_class(path: &str) -> Result<LegacyContractClass, eyre::Error> {
+    let file = std::fs::File::open(path)?;
+    let class = serde_json::from_reader::<_, LegacyContractClass>(file)?;
+    Ok(class)
+}
+
 lazy_static! {
     // Chain params
     pub static ref CHAIN_ID: u64 = 0x4b4b5254;
@@ -68,16 +74,18 @@ lazy_static! {
         ContractAddress(TryInto::<PatriciaKey>::try_into(StarkFelt::from(2_u8)).unwrap());
 
     // Main contract classes
-    pub static ref KAKAROT_CLASS: LegacyContractClass = serde_json::from_reader::<_, LegacyContractClass>(std::fs::File::open("../../lib/kakarot/build/kakarot.json").unwrap()).unwrap();
-    pub static ref CONTRACT_ACCOUNT_CLASS: LegacyContractClass = serde_json::from_reader::<_, LegacyContractClass>(std::fs::File::open("../../lib/kakarot/build/contract_account.json").unwrap()).unwrap();
-    pub static ref EOA_CLASS: LegacyContractClass = serde_json::from_reader::<_, LegacyContractClass>(std::fs::File::open("../../lib/kakarot/build/externally_owned_account.json").unwrap()).unwrap();
-    pub static ref PROXY_CLASS: LegacyContractClass = serde_json::from_reader::<_, LegacyContractClass>(std::fs::File::open("../../lib/kakarot/build/proxy.json").unwrap()).unwrap();
+    pub static ref KAKAROT_CLASS: LegacyContractClass = load_legacy_contract_class("../../lib/kakarot/build/kakarot.json").expect("Failed to load Kakarot contract class");
+    pub static ref CONTRACT_ACCOUNT_CLASS: LegacyContractClass = load_legacy_contract_class("../../lib/kakarot/build/contract_account.json").expect("Failed to load ContractAccount contract class");
+    pub static ref EOA_CLASS: LegacyContractClass = load_legacy_contract_class("../../lib/kakarot/build/externally_owned_account.json").expect("Failed to load EOA contract class");
+    pub static ref PROXY_CLASS: LegacyContractClass = load_legacy_contract_class("../../lib/kakarot/build/proxy.json").expect("Failed to load Proxy contract class");
+    pub static ref FEE_TOKEN_CLASS: LegacyContractClass = load_legacy_contract_class("../../lib/kakarot/build/fixtures/ERC20.json").expect("Failed to load FeeToken contract class");
 
     // Main class hashes
     pub static ref KAKAROT_CLASS_HASH: ClassHash = ClassHash(KAKAROT_CLASS.class_hash().unwrap().into());
     pub static ref CONTRACT_ACCOUNT_CLASS_HASH: ClassHash = ClassHash(CONTRACT_ACCOUNT_CLASS.class_hash().unwrap().into());
     pub static ref EOA_CLASS_HASH: ClassHash = ClassHash(EOA_CLASS.class_hash().unwrap().into());
     pub static ref PROXY_CLASS_HASH: ClassHash = ClassHash(PROXY_CLASS.class_hash().unwrap().into());
+    pub static ref FEE_TOKEN_CLASS_HASH: ClassHash = ClassHash(FEE_TOKEN_CLASS.class_hash().unwrap().into());
 
 }
 

crates/ef-testing/src/evm_sequencer/setup.rs → crates/ef-testing/src/evm_sequencer/evm_state.rs

@@ -1,9 +1,10 @@
 use blockifier::abi::abi_utils::{
     get_erc20_balance_var_addresses, get_storage_var_address, get_uint256_storage_var_addresses,
 };
-use blockifier::state::state_api::{State, StateResult};
+use blockifier::state::state_api::{State, StateReader, StateResult};
 use reth_primitives::{Address, Bytes};
 use revm_primitives::U256;
+use starknet::core::types::FieldElement;
 use starknet_api::core::Nonce;
 use starknet_api::hash::StarkFelt;
 use starknet_api::StarknetApiError;
@@ -26,6 +27,10 @@ pub trait EvmState {
     ) -> StateResult<()>;
 
     fn fund(&mut self, evm_address: &Address, balance: U256) -> StateResult<()>;
+
+    fn get_storage_at(&mut self, evm_address: &Address, key: U256) -> StateResult<U256>;
+
+    fn get_nonce(&mut self, evm_address: &Address) -> StateResult<U256>;
 }
 
 impl EvmState for KakarotSequencer {
@@ -102,14 +107,10 @@ impl EvmState for KakarotSequencer {
         (&mut self.0.state).set_class_hash_at(starknet_address, *PROXY_CLASS_HASH)?;
 
         // Add the address to the Kakarot evm to starknet mapping
-        let evm_starknet_address_mapping_storage = (
-            get_storage_var_address("evm_to_starknet_address", &[evm_address]).unwrap(), // safe unwrap: var is ASCII
-            *starknet_address.0.key(),
-        );
         (&mut self.0.state).set_storage_at(
             *KAKAROT_ADDRESS,
-            evm_starknet_address_mapping_storage.0,
-            evm_starknet_address_mapping_storage.1,
+            get_storage_var_address("evm_to_starknet_address", &[evm_address]).unwrap(),
+            *starknet_address.0.key(),
         );
         Ok(())
     }
@@ -120,7 +121,7 @@ impl EvmState for KakarotSequencer {
         let mut storage = vec![];
 
         // Initialize the balance storage var.
-        let balance_keys = get_erc20_balance_var_addresses(&FEE_TOKEN_ADDRESS)?;
+        let balance_keys = get_erc20_balance_var_addresses(&starknet_address.try_into()?)?;
         let balance_keys = [balance_keys.0, balance_keys.1];
         let balance_storage = &mut balance_keys
             .into_iter()
@@ -132,7 +133,7 @@ impl EvmState for KakarotSequencer {
         // Initialize the allowance storage var.
         let allowance_keys = get_uint256_storage_var_addresses(
             "ERC20_allowances",
-            &[*FEE_TOKEN_ADDRESS.0.key(), starknet_address.into()],
+            &[starknet_address.into(), *KAKAROT_ADDRESS.0.key()],
         )?;
         let allowance_keys = [allowance_keys.0, allowance_keys.1];
         let allowance_storage = &mut allowance_keys
@@ -142,12 +143,53 @@ impl EvmState for KakarotSequencer {
         storage.append(allowance_storage);
 
         // Write all the storage vars to the sequencer state.
-        let starknet_address = starknet_address.try_into()?;
         for (k, v) in storage {
-            (&mut self.0.state).set_storage_at(starknet_address, k, v);
+            (&mut self.0.state).set_storage_at(*FEE_TOKEN_ADDRESS, k, v);
         }
         Ok(())
     }
+
+    fn get_storage_at(&mut self, evm_address: &Address, key: U256) -> StateResult<U256> {
+        let keys = split_u256(key).map(Into::into);
+        let keys = get_uint256_storage_var_addresses("storage_", &keys).unwrap(); // safe unwrap: all vars are ASCII
+
+        let starknet_address = compute_starknet_address(evm_address);
+
+        let low = (&mut self.0.state).get_storage_at(starknet_address.try_into()?, keys.0)?;
+        let high = (&mut self.0.state).get_storage_at(starknet_address.try_into()?, keys.1)?;
+
+        let low = U256::from_be_bytes(Into::<FieldElement>::into(low).to_bytes_be());
+        let high = U256::from_be_bytes(Into::<FieldElement>::into(high).to_bytes_be());
+
+        Ok(high << 128 | low)
+    }
+
+    fn get_nonce(&mut self, evm_address: &Address) -> StateResult<U256> {
+        let starknet_address = compute_starknet_address(evm_address);
+
+        let implementation = (&mut self.0.state)
+            .get_storage_at(
+                starknet_address.try_into()?,
+                get_storage_var_address("_implementation", &[])?,
+            )
+            .unwrap();
+
+        let nonce = if implementation == EOA_CLASS_HASH.0 {
+            (&mut self.0.state)
+                .get_nonce_at(starknet_address.try_into()?)?
+                .0
+        } else if implementation == CONTRACT_ACCOUNT_CLASS_HASH.0 {
+            let key = get_storage_var_address("nonce", &[])?;
+            (&mut self.0.state).get_storage_at(starknet_address.try_into()?, key)?
+        } else {
+            // We can't throw an error here, because it could just be an uninitialized account.
+            StarkFelt::from(0_u8)
+        };
+
+        Ok(U256::from_be_bytes(
+            Into::<FieldElement>::into(nonce).to_bytes_be(),
+        ))
+    }
 }
 
 #[cfg(test)]

crates/ef-testing/src/evm_sequencer/mod.rs

@@ -1,20 +1,23 @@
 pub mod constants;
-pub mod setup;
+pub mod evm_state;
 pub mod types;
 pub mod utils;
 
 use blockifier::abi::abi_utils::get_storage_var_address;
 use blockifier::execution::contract_class::{ContractClass, ContractClassV0};
 use blockifier::state::errors::StateError;
 use blockifier::state::state_api::{State as BlockifierState, StateResult};
+use blockifier::transaction::errors::TransactionExecutionError;
+use blockifier::transaction::transaction_execution::Transaction;
 use cairo_vm::types::errors::program_errors::ProgramError;
+use sequencer::execution::Execution;
 use sequencer::sequencer::Sequencer;
 use sequencer::state::State;
 
 use self::constants::{
     BLOCK_CONTEXT, CONTRACT_ACCOUNT_CLASS, CONTRACT_ACCOUNT_CLASS_HASH, EOA_CLASS, EOA_CLASS_HASH,
-    FEE_TOKEN_ADDRESS, KAKAROT_ADDRESS, KAKAROT_CLASS, KAKAROT_CLASS_HASH, KAKAROT_OWNER_ADDRESS,
-    PROXY_CLASS, PROXY_CLASS_HASH,
+    FEE_TOKEN_ADDRESS, FEE_TOKEN_CLASS, FEE_TOKEN_CLASS_HASH, KAKAROT_ADDRESS, KAKAROT_CLASS,
+    KAKAROT_CLASS_HASH, KAKAROT_OWNER_ADDRESS, PROXY_CLASS, PROXY_CLASS_HASH,
 };
 
 pub(crate) struct KakarotSequencer(Sequencer<State>);
@@ -54,7 +57,7 @@ impl KakarotSequencer {
             )?),
         )?;
 
-        // Write proxy, eoa and contract account classes and class hashes.
+        // Write proxy, eoa, contract account and erc20 classes and class hashes.
         (&mut self.0.state).set_contract_class(
             &PROXY_CLASS_HASH,
             ContractClass::V0(ContractClassV0::try_from_json_string(
@@ -76,11 +79,25 @@ impl KakarotSequencer {
                     .map_err(|err| StateError::ProgramError(ProgramError::Parse(err)))?,
             )?),
         )?;
+        (&mut self.0.state).set_contract_class(
+            &FEE_TOKEN_CLASS_HASH,
+            ContractClass::V0(ContractClassV0::try_from_json_string(
+                &serde_json::to_string(&*FEE_TOKEN_CLASS)
+                    .map_err(|err| StateError::ProgramError(ProgramError::Parse(err)))?,
+            )?),
+        )?;
+        (&mut self.0.state).set_class_hash_at(*FEE_TOKEN_ADDRESS, *FEE_TOKEN_CLASS_HASH)?;
 
         Ok(self)
     }
 }
 
+impl Execution for KakarotSequencer {
+    fn execute(&mut self, transaction: Transaction) -> Result<(), TransactionExecutionError> {
+        self.0.execute(transaction)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;