chore(blockfier): share the utility next_storage_key to starknet api

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<u128, SierraTypeError> {
     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<StorageKey, StarknetApiError> {
-    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<Self> {
         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() })
     }

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

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)

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)?;
 

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,15 +630,15 @@ 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]);
 
     // Verify the ec_point values in storage.
     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]);
 

crates/starknet_api/src/state.rs

@@ -201,6 +201,12 @@ impl From<u128> for StorageKey {
     }
 }
 
+impl StorageKey {
+    pub fn next_storage_key(&self) -> Result<StorageKey, StarknetApiError> {
+        Ok(StorageKey(PatriciaKey::try_from(*self.0.key() + Felt::ONE)?))
+    }
+}
+
 impl_from_through_intermediate!(u128, StorageKey, u8, u16, u32, u64);
 
 /// A contract class.