From 16a4f8893eaee76eefb3bb4fd8352a96eed52dee Mon Sep 17 00:00:00 2001 From: Arni Hod Date: Sun, 17 Nov 2024 09:14:47 +0200 Subject: [PATCH] chore(blockfier): share the utility next_storage_key to starknet api --- crates/blockifier/src/abi/sierra_types.rs | 9 ++------- crates/blockifier/src/concurrency/worker_logic_test.rs | 3 +-- crates/blockifier/src/fee/fee_utils.rs | 3 +-- crates/blockifier/src/state/state_api.rs | 3 +-- crates/blockifier/src/transaction/transactions_test.rs | 10 +++++----- crates/starknet_api/src/state.rs | 6 ++++++ 6 files changed, 16 insertions(+), 18 deletions(-) diff --git a/crates/blockifier/src/abi/sierra_types.rs b/crates/blockifier/src/abi/sierra_types.rs index 1def7724f1c..58677f36fc4 100644 --- a/crates/blockifier/src/abi/sierra_types.rs +++ b/crates/blockifier/src/abi/sierra_types.rs @@ -4,7 +4,7 @@ use cairo_vm::vm::errors::memory_errors::MemoryError; use cairo_vm::vm::vm_core::VirtualMachine; use num_bigint::{BigUint, ToBigUint}; use num_traits::ToPrimitive; -use starknet_api::core::{ContractAddress, PatriciaKey}; +use starknet_api::core::ContractAddress; use starknet_api::state::StorageKey; use starknet_api::StarknetApiError; use starknet_types_core::felt::Felt; @@ -45,11 +45,6 @@ pub fn felt_to_u128(felt: &Felt) -> Result { felt.to_u128().ok_or_else(|| SierraTypeError::ValueTooLargeForType { val: *felt, ty: "u128" }) } -// TODO(barak, 01/10/2023): Move to starknet_api under StorageKey implementation. -pub fn next_storage_key(key: &StorageKey) -> Result { - Ok(StorageKey(PatriciaKey::try_from(*key.0.key() + Felt::ONE)?)) -} - // Implementations. // We implement the trait SierraType for SierraU128 and not for u128 since it's not guaranteed that @@ -111,7 +106,7 @@ impl SierraType for SierraU256 { key: &StorageKey, ) -> SierraTypeResult { let low_val = SierraU128::from_storage(state, contract_address, key)?; - let high_key = next_storage_key(key)?; + let high_key = key.next_storage_key()?; let high_val = SierraU128::from_storage(state, contract_address, &high_key)?; Ok(Self { low_val: low_val.as_value(), high_val: high_val.as_value() }) } diff --git a/crates/blockifier/src/concurrency/worker_logic_test.rs b/crates/blockifier/src/concurrency/worker_logic_test.rs index 3bb1dad82fa..9798dd2f485 100644 --- a/crates/blockifier/src/concurrency/worker_logic_test.rs +++ b/crates/blockifier/src/concurrency/worker_logic_test.rs @@ -11,7 +11,6 @@ use starknet_types_core::felt::Felt; use super::WorkerExecutor; use crate::abi::abi_utils::get_fee_token_var_address; -use crate::abi::sierra_types::next_storage_key; use crate::bouncer::Bouncer; use crate::concurrency::fee_utils::STORAGE_READ_SEQUENCER_BALANCE_INDICES; use crate::concurrency::scheduler::{Task, TransactionStatus}; @@ -345,7 +344,7 @@ fn test_worker_execute(default_all_resource_bounds: ValidResourceBounds) { let erc20 = FeatureContract::ERC20(CairoVersion::Cairo0); let erc_contract_address = contract_address!(TEST_ERC20_CONTRACT_ADDRESS2); let account_balance_key_low = get_fee_token_var_address(account_address); - let account_balance_key_high = next_storage_key(&account_balance_key_low).unwrap(); + let account_balance_key_high = account_balance_key_low.next_storage_key().unwrap(); // Both in write and read sets, only the account balance appear, and not the sequencer balance. // This is because when executing transaction in concurrency mode on, we manually remove the // writes and reads to and from the sequencer balance (to avoid the inevitable dependency diff --git a/crates/blockifier/src/fee/fee_utils.rs b/crates/blockifier/src/fee/fee_utils.rs index 7c73ace01da..6dba8b39daa 100644 --- a/crates/blockifier/src/fee/fee_utils.rs +++ b/crates/blockifier/src/fee/fee_utils.rs @@ -11,7 +11,6 @@ use starknet_api::transaction::fields::{Fee, GasVectorComputationMode, Resource} use starknet_types_core::felt::Felt; use crate::abi::abi_utils::get_fee_token_var_address; -use crate::abi::sierra_types::next_storage_key; use crate::blockifier::block::BlockInfo; use crate::context::{BlockContext, TransactionContext}; use crate::fee::resources::TransactionFeeResult; @@ -145,7 +144,7 @@ pub fn get_sequencer_balance_keys(block_context: &BlockContext) -> (StorageKey, pub fn get_address_balance_keys(address: ContractAddress) -> (StorageKey, StorageKey) { let balance_key_low = get_fee_token_var_address(address); - let balance_key_high = next_storage_key(&balance_key_low).unwrap_or_else(|_| { + let balance_key_high = balance_key_low.next_storage_key().unwrap_or_else(|_| { panic!("Failed to get balance_key_high for address: {:?}", address.0); }); (balance_key_low, balance_key_high) diff --git a/crates/blockifier/src/state/state_api.rs b/crates/blockifier/src/state/state_api.rs index 8d430a18929..4844511f2b5 100644 --- a/crates/blockifier/src/state/state_api.rs +++ b/crates/blockifier/src/state/state_api.rs @@ -6,7 +6,6 @@ use starknet_types_core::felt::Felt; use super::cached_state::{ContractClassMapping, StateMaps}; use crate::abi::abi_utils::get_fee_token_var_address; -use crate::abi::sierra_types::next_storage_key; use crate::execution::contract_class::RunnableContractClass; use crate::state::errors::StateError; @@ -60,7 +59,7 @@ pub trait StateReader { fee_token_address: ContractAddress, ) -> Result<(Felt, Felt), StateError> { let low_key = get_fee_token_var_address(contract_address); - let high_key = next_storage_key(&low_key)?; + let high_key = low_key.next_storage_key()?; let low = self.get_storage_at(fee_token_address, low_key)?; let high = self.get_storage_at(fee_token_address, high_key)?; diff --git a/crates/blockifier/src/transaction/transactions_test.rs b/crates/blockifier/src/transaction/transactions_test.rs index 3044e969ee1..6df3cf60e5f 100644 --- a/crates/blockifier/src/transaction/transactions_test.rs +++ b/crates/blockifier/src/transaction/transactions_test.rs @@ -54,7 +54,6 @@ use crate::abi::abi_utils::{ selector_from_name, }; use crate::abi::constants as abi_constants; -use crate::abi::sierra_types::next_storage_key; use crate::context::{BlockContext, ChainInfo, FeeTokenAddresses, TransactionContext}; use crate::execution::call_info::{ CallExecution, @@ -304,9 +303,10 @@ fn expected_fee_transfer_call_info( let sender_balance_key_low = get_fee_token_var_address(account_address); let sender_balance_key_high = - next_storage_key(&sender_balance_key_low).expect("Cannot get sender balance high key."); + sender_balance_key_low.next_storage_key().expect("Cannot get sender balance high key."); let sequencer_balance_key_low = get_fee_token_var_address(expected_sequencer_address); - let sequencer_balance_key_high = next_storage_key(&sequencer_balance_key_low) + let sequencer_balance_key_high = sequencer_balance_key_low + .next_storage_key() .expect("Cannot get sequencer balance high key."); Some(CallInfo { call: expected_fee_transfer_call, @@ -630,7 +630,7 @@ fn verify_storage_after_invoke_advanced_operations( let key = get_storage_var_address("two_counters", &[index]); let value = state.get_storage_at(contract_address, key).unwrap(); assert_eq!(value, expected_counters[0]); - let key = next_storage_key(&key).unwrap(); + let key = key.next_storage_key().unwrap(); let value = state.get_storage_at(contract_address, key).unwrap(); assert_eq!(value, expected_counters[1]); @@ -638,7 +638,7 @@ fn verify_storage_after_invoke_advanced_operations( let key = get_storage_var_address("ec_point", &[]); let value = state.get_storage_at(contract_address, key).unwrap(); assert_eq!(value, expected_ec_point[0]); - let key = next_storage_key(&key).unwrap(); + let key = key.next_storage_key().unwrap(); let value = state.get_storage_at(contract_address, key).unwrap(); assert_eq!(value, expected_ec_point[1]); diff --git a/crates/starknet_api/src/state.rs b/crates/starknet_api/src/state.rs index d7351617fcd..289ce13663b 100644 --- a/crates/starknet_api/src/state.rs +++ b/crates/starknet_api/src/state.rs @@ -201,6 +201,12 @@ impl From for StorageKey { } } +impl StorageKey { + pub fn next_storage_key(&self) -> Result { + Ok(StorageKey(PatriciaKey::try_from(*self.0.key() + Felt::ONE)?)) + } +} + impl_from_through_intermediate!(u128, StorageKey, u8, u16, u32, u64); /// A contract class.