feat(blockifier): count the number of charged storage keys

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 charged new allocated aliases.
+/// Counts declared classes, deployed contracts, and storage keys that were previously empty and
+/// are now filled.
 pub fn n_charged_invoke_aliases<S: StateReader>(
-    _state: &CachedState<S>,
-    _state_changes: &StateChanges,
+    state: &CachedState<S>,
+    state_changes: &StateChanges,
 ) -> StateResult<usize> {
-    // 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.iter() {
+        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)
 }

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<DictStateReader> {
+    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<StorageEntry, Felt>,
+    #[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<ClassHash, bool>,
+    #[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<ContractAddress, ClassHash>,
+    #[case] charged_change: bool,
+) {
+    assert_number_of_charged_changes(
+        StateChanges(StateMaps { class_hashes, ..Default::default() }),
+        charged_change,
+    );
+}

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 =