refactor: account balances (#569)

* refactor: account balances

fix: test_kakarot_core_compute_starknet_address

chore: update compute sn address script

* fix test

.tool-versions

@@ -1 +1 @@
-scarb nightly-2023-11-22
+scarb nightly-2023-11-22

crates/contracts/src/kakarot_core/kakarot.cairo

@@ -203,7 +203,7 @@ mod KakarotCore {
             let address = Address {
                 evm: evm_address, starknet: self.compute_starknet_address(evm_address)
             };
-            address.balance().unwrap()
+            address.fetch_balance()
         }
 
         fn contract_account_storage_at(

crates/contracts/src/tests/test_eoa.cairo

@@ -14,8 +14,9 @@ mod test_external_owned_account {
         IMockContractUpgradeableDispatcher, IMockContractUpgradeableDispatcherTrait,
         MockContractUpgradeableV1
     };
-    use contracts::tests::test_utils::deploy_contract_account;
-    use contracts::tests::test_utils::{setup_contracts_for_testing};
+    use contracts::tests::test_utils::{
+        setup_contracts_for_testing, deploy_eoa, deploy_contract_account
+    };
     use contracts::uninitialized_account::{
         IUninitializedAccountDispatcher, IUninitializedAccountDispatcherTrait, UninitializedAccount,
         IUninitializedAccount
@@ -40,25 +41,6 @@ mod test_external_owned_account {
     };
     use utils::helpers::{U8SpanExTrait, u256_to_bytes_array};
 
-    fn deploy_eoa(eoa_address: EthAddress) -> IExternallyOwnedAccountDispatcher {
-        let kakarot_address = get_contract_address();
-        let calldata: Span<felt252> = array![kakarot_address.into(), eoa_address.into()].span();
-
-        let (starknet_address, _) = deploy_syscall(
-            UninitializedAccount::TEST_CLASS_HASH.try_into().unwrap(),
-            evm_address().into(),
-            calldata,
-            false
-        )
-            .expect('failed to deploy EOA');
-
-        let account = IUninitializedAccountDispatcher { contract_address: starknet_address };
-
-        account.initialize(ExternallyOwnedAccount::TEST_CLASS_HASH.try_into().unwrap());
-        let eoa = IExternallyOwnedAccountDispatcher { contract_address: starknet_address };
-        eoa.set_chain_id(chain_id());
-        eoa
-    }
 
     #[test]
     #[available_gas(2000000000)]

crates/contracts/src/tests/test_utils.cairo

@@ -3,6 +3,7 @@ use contracts::contract_account::{
 };
 
 use contracts::eoa::{ExternallyOwnedAccount};
+use contracts::eoa::{IExternallyOwnedAccountDispatcher, IExternallyOwnedAccountDispatcherTrait};
 use contracts::kakarot_core::{interface::IExtendedKakarotCoreDispatcher, KakarotCore};
 use contracts::uninitialized_account::{
     IUninitializedAccountDispatcher, IUninitializedAccountDispatcherTrait, UninitializedAccount
@@ -111,6 +112,27 @@ fn deploy_contract_account(evm_address: EthAddress, bytecode: Span<u8>) -> Addre
 }
 
 
+fn deploy_eoa(eoa_address: EthAddress) -> IExternallyOwnedAccountDispatcher {
+    let kakarot_address = get_contract_address();
+    let calldata: Span<felt252> = array![kakarot_address.into(), eoa_address.into()].span();
+
+    let (starknet_address, _) = deploy_syscall(
+        UninitializedAccount::TEST_CLASS_HASH.try_into().unwrap(),
+        eoa_address.into(),
+        calldata,
+        false
+    )
+        .expect('failed to deploy EOA');
+
+    let account = IUninitializedAccountDispatcher { contract_address: starknet_address };
+
+    account.initialize(ExternallyOwnedAccount::TEST_CLASS_HASH.try_into().unwrap());
+    let eoa = IExternallyOwnedAccountDispatcher { contract_address: starknet_address };
+    eoa.set_chain_id(chain_id());
+    eoa
+}
+
+
 fn fund_account_with_native_token(
     contract_address: ContractAddress, native_token: IERC20CamelDispatcher, amount: u256,
 ) {

crates/evm/src/create_helpers.cairo

@@ -73,7 +73,7 @@ impl MachineCreateHelpersImpl of MachineCreateHelpers {
         let caller = self.address();
         let mut caller_account = self.state.get_account(caller.evm);
         let caller_current_nonce = caller_account.nonce();
-        let caller_balance = self.state.read_balance(caller.evm)?;
+        let caller_balance = caller_account.balance();
         if caller_balance < create_args.value
             || target_account.nonce() == integer::BoundedInt::<u64>::max() {
             return self.stack.push(0);

crates/evm/src/instructions/block_information.cairo

@@ -106,7 +106,7 @@ impl BlockInformation of BlockInformationTrait {
     fn exec_selfbalance(ref self: Machine) -> Result<(), EVMError> {
         let evm_address = self.address().evm;
 
-        let balance = self.state.read_balance(evm_address)?;
+        let balance = self.state.get_account(evm_address).balance;
 
         self.stack.push(balance)
     }

crates/evm/src/instructions/environmental_information.cairo

@@ -36,8 +36,7 @@ impl EnvironmentInformationImpl of EnvironmentInformationTrait {
     fn exec_balance(ref self: Machine) -> Result<(), EVMError> {
         let evm_address = self.stack.pop_eth_address()?;
 
-        let balance = self.state.read_balance(evm_address)?;
-
+        let balance = self.state.get_account(evm_address).balance();
         self.stack.push(balance)
     }
 

crates/evm/src/instructions/system_operations.cairo

@@ -92,7 +92,7 @@ impl SystemOperations of SystemOperationsTrait {
         // If sender_balance < value, return early, pushing
         // 0 on the stack to indicate call failure.
         let caller_address = self.address();
-        let sender_balance = self.state.read_balance(caller_address.evm)?;
+        let sender_balance = self.state.get_account(caller_address.evm).balance();
         if sender_balance < value {
             return self.stack.push(0);
         }
@@ -161,7 +161,7 @@ impl SystemOperations of SystemOperationsTrait {
                 Transfer {
                     sender: account.address(),
                     recipient,
-                    amount: self.state.read_balance(account.address().evm)?
+                    amount: self.state.get_account(account.address().evm).balance
                 }
             );
 

crates/evm/src/model.cairo

@@ -9,9 +9,7 @@ use evm::model::account::{Account, AccountTrait};
 use evm::model::contract_account::{ContractAccountTrait};
 use evm::model::eoa::EOATrait;
 use evm::state::State;
-use openzeppelin::token::erc20::interface::{
-    IERC20CamelSafeDispatcher, IERC20CamelSafeDispatcherTrait
-};
+use openzeppelin::token::erc20::interface::{IERC20CamelDispatcher, IERC20CamelDispatcherTrait};
 use starknet::{EthAddress, get_contract_address, ContractAddress};
 use utils::helpers::{ResultExTrait};
 use utils::traits::{EthAddressDefault, ContractAddressDefault};
@@ -40,13 +38,11 @@ impl AddressImpl of AddressTrait {
         }
     }
 
-    fn balance(self: @Address) -> Result<u256, EVMError> {
+    fn fetch_balance(self: @Address) -> u256 {
         let kakarot_state = KakarotCore::unsafe_new_contract_state();
         let native_token_address = kakarot_state.native_token();
-        let native_token = IERC20CamelSafeDispatcher { contract_address: native_token_address };
-        native_token
-            .balanceOf(*self.starknet)
-            .map_err(EVMError::SyscallFailed(CONTRACT_SYSCALL_FAILED))
+        let native_token = IERC20CamelDispatcher { contract_address: native_token_address };
+        native_token.balanceOf(*self.starknet)
     }
 }
 

crates/evm/src/model/account.cairo

@@ -7,9 +7,7 @@ use contracts::kakarot_core::{KakarotCore, IKakarotCore};
 use evm::errors::{EVMError, CONTRACT_SYSCALL_FAILED};
 use evm::model::contract_account::{ContractAccountTrait};
 use evm::model::{Address, AddressTrait, AccountType};
-use openzeppelin::token::erc20::interface::{
-    IERC20CamelSafeDispatcher, IERC20CamelSafeDispatcherTrait
-};
+use openzeppelin::token::erc20::interface::{IERC20CamelDispatcher, IERC20CamelDispatcherTrait};
 use starknet::{ContractAddress, EthAddress, get_contract_address};
 use utils::helpers::{ResultExTrait, ByteArrayExTrait, compute_starknet_address};
 
@@ -19,42 +17,70 @@ struct Account {
     address: Address,
     code: Span<u8>,
     nonce: u64,
+    balance: u256,
     selfdestruct: bool,
 }
 
-struct ContractAccountBuilder {
+#[derive(Drop)]
+struct AccountBuilder {
     account: Account
 }
 
 #[generate_trait]
-impl ContractAccountBuilderImpl of ContractAccountBuilderTrait {
-    fn new(address: Address) -> ContractAccountBuilder {
-        ContractAccountBuilder {
+impl AccountBuilderImpl of AccountBuilderTrait {
+    fn new(address: Address) -> AccountBuilder {
+        AccountBuilder {
             account: Account {
-                account_type: AccountType::ContractAccount,
+                account_type: AccountType::Unknown,
                 address: address,
                 code: Default::default().span(),
                 nonce: 0,
+                balance: 0,
                 selfdestruct: false,
             }
         }
     }
 
     #[inline(always)]
-    fn fetch_nonce(mut self: ContractAccountBuilder) -> ContractAccountBuilder {
+    fn set_type(mut self: AccountBuilder, account_type: AccountType) -> AccountBuilder {
+        self.account.account_type = account_type;
+        self
+    }
+
+    #[inline(always)]
+    fn fetch_balance(mut self: AccountBuilder) -> AccountBuilder {
+        let kakarot_state = KakarotCore::unsafe_new_contract_state();
+        let native_token_address = kakarot_state.native_token();
+        let native_token = IERC20CamelDispatcher { contract_address: native_token_address };
+        self.account.balance = self.account.address.fetch_balance();
+        self
+    }
+
+    #[inline(always)]
+    fn fetch_nonce(mut self: AccountBuilder) -> AccountBuilder {
+        assert!(
+            self.account.account_type == AccountType::ContractAccount,
+            "Cannot fetch nonce of an EOA"
+        );
         let contract_account = IContractAccountDispatcher {
             contract_address: self.account.address.starknet
         };
         self.account.nonce = contract_account.nonce();
         self
     }
 
+    #[inline(always)]
+    fn set_nonce(mut self: AccountBuilder, nonce: u64) -> AccountBuilder {
+        self.account.nonce = nonce;
+        self
+    }
+
     /// Loads the bytecode of a ContractAccount from Kakarot Core's contract storage into a Span<u8>.
     /// # Arguments
     /// * `self` - The address of the Contract Account to load the bytecode from
     /// # Returns
     /// * The bytecode of the Contract Account as a ByteArray
-    fn fetch_bytecode(mut self: ContractAccountBuilder) -> ContractAccountBuilder {
+    fn fetch_bytecode(mut self: AccountBuilder) -> AccountBuilder {
         let contract_account = IContractAccountDispatcher {
             contract_address: self.account.address.starknet
         };
@@ -64,7 +90,7 @@ impl ContractAccountBuilderImpl of ContractAccountBuilderTrait {
     }
 
     #[inline(always)]
-    fn build(self: ContractAccountBuilder) -> Account {
+    fn build(self: AccountBuilder) -> Account {
         self.account
     }
 }
@@ -89,13 +115,9 @@ impl AccountImpl of AccountTrait {
                 // If no account exists at `address`, then we are trying to
                 // access an undeployed account (CA or EOA). We create an
                 // empty account with the correct address and return it.
-                Account {
-                    account_type: AccountType::Unknown,
-                    address: Address { starknet: starknet_address, evm: evm_address, },
-                    code: Default::default().span(),
-                    nonce: 0,
-                    selfdestruct: false,
-                }
+                AccountBuilderTrait::new(Address { starknet: starknet_address, evm: evm_address })
+                    .fetch_balance()
+                    .build()
             }
         }
     }
@@ -117,21 +139,21 @@ impl AccountImpl of AccountTrait {
             Option::Some((
                 account_type, starknet_address
             )) => {
+                let address = Address { evm: evm_address, starknet: starknet_address };
                 match account_type {
                     AccountType::EOA => Option::Some(
-                        Account {
-                            account_type: AccountType::EOA,
-                            address: Address { evm: evm_address, starknet: starknet_address },
-                            code: Default::default().span(),
-                            nonce: 1,
-                            selfdestruct: false,
-                        }
+                        AccountBuilderTrait::new(address)
+                            .set_type(AccountType::EOA)
+                            .set_nonce(1)
+                            .fetch_balance()
+                            .build()
                     ),
                     AccountType::ContractAccount => {
-                        let address = Address { evm: evm_address, starknet: starknet_address };
-                        let account = ContractAccountBuilderTrait::new(address)
+                        let account = AccountBuilderTrait::new(address)
+                            .set_type(AccountType::ContractAccount)
                             .fetch_nonce()
                             .fetch_bytecode()
+                            .fetch_balance()
                             .build();
                         Option::Some(account)
                     },
@@ -227,6 +249,16 @@ impl AccountImpl of AccountTrait {
         }
     }
 
+    #[inline(always)]
+    fn set_balance(ref self: Account, value: u256) {
+        self.balance = value;
+    }
+
+    #[inline(always)]
+    fn balance(self: @Account) -> u256 {
+        *self.balance
+    }
+
     #[inline(always)]
     fn address(self: @Account) -> Address {
         *self.address
@@ -280,20 +312,6 @@ impl AccountImpl of AccountTrait {
         self.address().evm
     }
 
-    /// Returns the balance in native token for a given EVM account (EOA or CA)
-    /// This is equivalent to checking the balance in native coin, i.e. ETHER of an account in Ethereum
-    #[inline(always)]
-    fn balance(self: @Account) -> Result<u256, EVMError> {
-        let kakarot_state = KakarotCore::unsafe_new_contract_state();
-        let native_token_address = kakarot_state.native_token();
-        let native_token = IERC20CamelSafeDispatcher { contract_address: native_token_address };
-        //Note: Starknet OS doesn't allow error management of failed syscalls yet.
-        // If this call fails, the entire transaction will revert.
-        native_token
-            .balanceOf(self.address().starknet)
-            .map_err(EVMError::SyscallFailed(CONTRACT_SYSCALL_FAILED))
-    }
-
     /// Returns the bytecode of the EVM account (EOA or CA)
     #[inline(always)]
     fn bytecode(self: @Account) -> Span<u8> {

crates/evm/src/model/contract_account.cairo

@@ -23,9 +23,7 @@ use evm::errors::{
 use evm::execution::execute;
 use evm::model::{Address, Account, AccountType, AccountTrait};
 use hash::{HashStateTrait, HashStateExTrait};
-use openzeppelin::token::erc20::interface::{
-    IERC20CamelSafeDispatcher, IERC20CamelSafeDispatcherTrait
-};
+use openzeppelin::token::erc20::interface::{IERC20CamelDispatcher, IERC20CamelDispatcherTrait};
 use poseidon::PoseidonTrait;
 use starknet::{
     deploy_syscall, StorageBaseAddress, storage_base_address_from_felt252, Store, EthAddress,