From 96605ab3ba002eb769d39687758a1c562d2615c7 Mon Sep 17 00:00:00 2001 From: Yoav Gross Date: Sun, 10 Nov 2024 15:38:08 +0200 Subject: [PATCH] feat(blockifier): count the number of charged storage keys --- crates/blockifier/src/execution/alias_keys.rs | 54 ++++++++++++- .../src/execution/alias_keys_test.rs | 77 +++++++++++++++++++ crates/starknet_api/src/core.rs | 2 +- 3 files changed, 128 insertions(+), 5 deletions(-) create mode 100644 crates/blockifier/src/execution/alias_keys_test.rs diff --git a/crates/blockifier/src/execution/alias_keys.rs b/crates/blockifier/src/execution/alias_keys.rs index ebf8ce2564..cd596c6067 100644 --- a/crates/blockifier/src/execution/alias_keys.rs +++ b/crates/blockifier/src/execution/alias_keys.rs @@ -1,11 +1,57 @@ +use starknet_api::core::ClassHash; +use starknet_types_core::felt::Felt; + use crate::state::cached_state::{CachedState, StateChanges}; use crate::state::state_api::{StateReader, StateResult}; +#[cfg(test)] +#[path = "alias_keys_test.rs"] +mod test; + +/// Keys in contract addresses up to this address don't get aliases. +pub const N_SAVED_CONTRACT_ADDRESSES: u8 = 16; +/// The alias of a felt X up to this number is X. +/// This trivial mappings don't write to the alias contract. +pub const N_TRIVIAL_SELF_ALIASES: u8 = 128; + /// Returns the number of aliases we charge the transaction for. +/// Counts declared classes, deployed contracts, and storage keys that were previously empty and +/// are now filled. pub fn n_charged_invoke_aliases( - _state: &CachedState, - _state_changes: &StateChanges, + state: &CachedState, + state_changes: &StateChanges, ) -> StateResult { - // TODO: Implement this function - Ok(0) + let n_declared_classes = state_changes + .0 + .declared_contracts + .iter() + .filter(|(class_hash, is_declared)| { + **is_declared && (class_hash.0 >= N_TRIVIAL_SELF_ALIASES.into()) + }) + .count(); + + let mut n_deployed_contracts = 0; + for contract_address in state_changes.0.class_hashes.keys() { + if contract_address.0 >= N_TRIVIAL_SELF_ALIASES.into() + // The contract is deployed, not replaced class. + && state.get_class_hash_at(*contract_address)? == ClassHash(Felt::ZERO) + { + n_deployed_contracts += 1; + } + } + + let storage_changes = &state_changes.0.storage; + let mut n_storage_keys = 0; + for ((contract_address, storage_key), new_value) in storage_changes { + if contract_address.0.0 >= N_SAVED_CONTRACT_ADDRESSES.into() + && storage_key.0.0 >= N_TRIVIAL_SELF_ALIASES.into() + // Zero to non-zero. + && state.get_storage_at(*contract_address, *storage_key)? == Felt::ZERO + && new_value != &Felt::ZERO + { + n_storage_keys += 1; + } + } + + Ok(n_declared_classes + n_deployed_contracts + n_storage_keys) } diff --git a/crates/blockifier/src/execution/alias_keys_test.rs b/crates/blockifier/src/execution/alias_keys_test.rs new file mode 100644 index 0000000000..44cb94bd76 --- /dev/null +++ b/crates/blockifier/src/execution/alias_keys_test.rs @@ -0,0 +1,77 @@ +use std::collections::HashMap; + +use rstest::rstest; +use starknet_api::core::{ClassHash, ContractAddress}; +use starknet_api::state::StorageKey; +use starknet_types_core::felt::Felt; + +use super::{n_charged_invoke_aliases, N_SAVED_CONTRACT_ADDRESSES, N_TRIVIAL_SELF_ALIASES}; +use crate::state::cached_state::{CachedState, StateChanges, StateMaps, StorageEntry}; +use crate::test_utils::dict_state_reader::DictStateReader; + +fn initial_state() -> CachedState { + CachedState::from(DictStateReader { + storage_view: HashMap::from([( + (ContractAddress::from(1000_u32), StorageKey::from(1001_u32)), + Felt::from(1002), + )]), + address_to_class_hash: HashMap::from([( + ContractAddress::from(2000_u32), + ClassHash(Felt::from(2001)), + )]), + ..Default::default() + }) +} + +fn assert_number_of_charged_changes(state_changes: StateChanges, charged_change: bool) { + let state = initial_state(); + let n_charged_aliases = n_charged_invoke_aliases(&state, &state_changes).unwrap(); + assert_eq!(n_charged_aliases, if charged_change { 1 } else { 0 }); +} + +#[rstest] +#[case::empty(HashMap::new(), false)] +#[case::non_zero_to_non_zero(HashMap::from([((ContractAddress::from(1000_u32), StorageKey::from(1001_u32)), Felt::from(1003))]), false)] +#[case::non_zero_to_zero(HashMap::from([((ContractAddress::from(1000_u32), StorageKey::from(1001_u32)), Felt::ZERO)]), false)] +#[case::zero_to_non_zero(HashMap::from([((ContractAddress::from(1000_u32), StorageKey::from(1004_u32)), Felt::from(1005))]), true)] +#[case::low_key(HashMap::from([((ContractAddress::from(1000_u32), StorageKey::from(N_TRIVIAL_SELF_ALIASES - 1)), Felt::from(1005))]), false)] +#[case::low_address(HashMap::from([((ContractAddress::from(N_SAVED_CONTRACT_ADDRESSES - 1), StorageKey::from(1004_u32)), Felt::from(1005))]), false)] +fn test_charged_storage_changes( + #[case] storage_changes: HashMap, + #[case] charged_change: bool, +) { + assert_number_of_charged_changes( + StateChanges(StateMaps { storage: storage_changes, ..Default::default() }), + charged_change, + ); +} + +#[rstest] +#[case::empty(HashMap::new(), false)] +#[case::add_class_hash(HashMap::from([(ClassHash(3000.into()), true)]), true)] +#[case::remove_class_hash(HashMap::from([(ClassHash(3000.into()), false)]), false)] +#[case::low_key(HashMap::from([(ClassHash((N_TRIVIAL_SELF_ALIASES - 1).into()), false)]), false)] +fn test_charged_declared_classes( + #[case] declared_contracts: HashMap, + #[case] charged_change: bool, +) { + assert_number_of_charged_changes( + StateChanges(StateMaps { declared_contracts, ..Default::default() }), + charged_change, + ); +} + +#[rstest] +#[case::empty(HashMap::new(), false)] +#[case::new_contract(HashMap::from([(ContractAddress::from(3000_u32), ClassHash(3001.into()))]), true)] +#[case::replace_class(HashMap::from([(ContractAddress::from(2000_u32), ClassHash(2002.into()))]), false)] +#[case::low_key(HashMap::from([(ContractAddress::from(N_TRIVIAL_SELF_ALIASES - 1), ClassHash(3001.into()))]), false)] +fn test_charged_deployed_contracts( + #[case] class_hashes: HashMap, + #[case] charged_change: bool, +) { + assert_number_of_charged_changes( + StateChanges(StateMaps { class_hashes, ..Default::default() }), + charged_change, + ); +} diff --git a/crates/starknet_api/src/core.rs b/crates/starknet_api/src/core.rs index ec2ac3ed8e..ac7f2e2d4c 100644 --- a/crates/starknet_api/src/core.rs +++ b/crates/starknet_api/src/core.rs @@ -347,7 +347,7 @@ pub struct StateDiffCommitment(pub PoseidonHash); derive_more:: Deref, )] #[display(fmt = "{}", "_0.to_fixed_hex_string()")] -pub struct PatriciaKey(StarkHash); +pub struct PatriciaKey(pub StarkHash); // 2**251 pub const PATRICIA_KEY_UPPER_BOUND: &str =