starknet_os: aliases test (#7680)

nimrod-starkware · web-flow · commit 2cac124d59ba · 2025-07-02T08:38:03.000Z
diff --git a/crates/blockifier/src/state/stateful_compression_test.rs b/crates/blockifier/src/state/stateful_compression_test.rs
@@ -1,9 +1,8 @@
 use std::collections::{HashMap, HashSet};
-use std::sync::LazyLock;
 
 use assert_matches::assert_matches;
 use rstest::rstest;
-use starknet_api::core::{ClassHash, CompiledClassHash, ContractAddress, Nonce, PatriciaKey};
+use starknet_api::core::{ClassHash, CompiledClassHash, ContractAddress, Nonce};
 use starknet_api::felt;
 use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
@@ -23,9 +22,7 @@ use crate::state::cached_state::{CachedState, StateMaps, StorageEntry};
 use crate::state::state_api::{State, StateReader};
 use crate::state::stateful_compression::{AliasCompressor, CompressionError};
 use crate::test_utils::dict_state_reader::DictStateReader;
-
-static ALIAS_CONTRACT_ADDRESS: LazyLock<ContractAddress> =
-    LazyLock::new(|| ContractAddress(PatriciaKey::try_from(Felt::TWO).unwrap()));
+use crate::test_utils::ALIAS_CONTRACT_ADDRESS;
 
 /// Decompresses the state diff by replacing the aliases with addresses and storage keys.
 fn decompress<S: StateReader>(
diff --git a/crates/blockifier/src/test_utils.rs b/crates/blockifier/src/test_utils.rs
@@ -10,14 +10,15 @@ pub mod test_templates;
 pub mod transfers_generator;
 use std::collections::HashMap;
 use std::slice::Iter;
+use std::sync::LazyLock;
 
 use blockifier_test_utils::cairo_versions::{CairoVersion, RunnableCairo1};
 use blockifier_test_utils::contracts::FeatureContract;
 use cairo_vm::types::builtin_name::BuiltinName;
 use cairo_vm::vm::runners::cairo_runner::ExecutionResources;
 use starknet_api::abi::abi_utils::{get_fee_token_var_address, selector_from_name};
 use starknet_api::block::{BlockHash, BlockHashAndNumber, BlockNumber};
-use starknet_api::core::{ClassHash, ContractAddress};
+use starknet_api::core::{ClassHash, ContractAddress, PatriciaKey};
 use starknet_api::executable_transaction::TransactionType;
 use starknet_api::execution_resources::{GasAmount, GasVector};
 use starknet_api::hash::StarkHash;
@@ -65,6 +66,9 @@ pub const ERC20_CONTRACT_PATH: &str = "../blockifier_test_utils/resources/ERC20/
                                        ERC20_without_some_syscalls/ERC20/\
                                        erc20_contract_without_some_syscalls_compiled.json";
 
+pub static ALIAS_CONTRACT_ADDRESS: LazyLock<ContractAddress> =
+    LazyLock::new(|| ContractAddress(PatriciaKey::try_from(Felt::TWO).unwrap()));
+
 #[derive(Clone, Copy, EnumCountMacro, PartialEq, Eq, Debug)]
 pub enum CompilerBasedVersion {
     CairoVersion(CairoVersion),
diff --git a/crates/starknet_os/src/hint_processor/snos_hint_processor.rs b/crates/starknet_os/src/hint_processor/snos_hint_processor.rs
@@ -341,15 +341,17 @@ impl<'a> SnosHintProcessor<'a, DictStateReader> {
         let state_inputs = vec![os_state_input.unwrap_or_default()];
         let os_hints_config = os_hints_config.unwrap_or_default();
 
-        SnosHintProcessor::new(
+        let mut hint_processor = SnosHintProcessor::new(
             os_program,
             os_hints_config,
             block_inputs,
             state_inputs,
             BTreeMap::new(),
             BTreeMap::new(),
             vec![state_reader],
-        )
+        )?;
+        hint_processor.execution_helpers_manager.increment_current_helper_index();
+        Ok(hint_processor)
     }
 }
 
diff --git a/crates/starknet_os/src/test_utils/cairo_runner.rs b/crates/starknet_os/src/test_utils/cairo_runner.rs
@@ -2,6 +2,7 @@ use std::any::Any;
 use std::collections::{HashMap, HashSet};
 
 use blockifier::blockifier_versioned_constants::VersionedConstants;
+use blockifier::test_utils::dict_state_reader::DictStateReader;
 use cairo_vm::serde::deserialize_program::Member;
 use cairo_vm::types::builtin_name::BuiltinName;
 use cairo_vm::types::layout_name::LayoutName;
@@ -553,6 +554,7 @@ fn get_return_values(
 /// Hint locals are added to the outermost exec scope.
 /// If the endpoint used builtins, the respective returned (implicit) arg is the builtin instance
 /// usage, unless the builtin is the output builtin, in which case the arg is the output.
+#[allow(clippy::too_many_arguments)]
 pub fn run_cairo_0_entry_point(
     runner_config: &EntryPointRunnerConfig,
     program_bytes: &[u8],
@@ -561,6 +563,7 @@ pub fn run_cairo_0_entry_point(
     implicit_args: &[ImplicitArg],
     expected_explicit_return_values: &[EndpointArg],
     hint_locals: HashMap<String, Box<dyn Any>>,
+    state_reader: Option<DictStateReader>,
 ) -> Cairo0EntryPointRunnerResult<(Vec<EndpointArg>, Vec<EndpointArg>, CairoRunner)> {
     let mut entrypoint = entrypoint.to_string();
     if runner_config.add_main_prefix_to_entrypoint {
@@ -571,7 +574,7 @@ pub fn run_cairo_0_entry_point(
     let program = inject_builtins(program_bytes, implicit_args)?;
     info!("Successfully injected builtins into program.");
 
-    let (state_reader, os_hints_config, os_state_input) = (None, None, None);
+    let (os_hints_config, os_state_input) = (None, None);
     let os_block_input = OsBlockInput::default();
     let mut hint_processor = SnosHintProcessor::new_for_testing(
         state_reader,
diff --git a/crates/starknet_os/src/test_utils/utils.rs b/crates/starknet_os/src/test_utils/utils.rs
@@ -2,6 +2,8 @@ use std::any::Any;
 use std::collections::HashMap;
 use std::sync::LazyLock;
 
+use cairo_vm::hint_processor::builtin_hint_processor::dict_hint_utils::DICT_ACCESS_SIZE;
+use cairo_vm::types::layout_name::LayoutName;
 use ethnum::U256;
 use num_bigint::{BigInt, Sign};
 use rand::rngs::StdRng;
@@ -30,6 +32,7 @@ pub fn run_cairo_function_and_check_result(
     expected_implicit_retdata: &[EndpointArg],
     hint_locals: HashMap<String, Box<dyn Any>>,
 ) -> Cairo0EntryPointRunnerResult<()> {
+    let state_reader = None;
     let (actual_implicit_retdata, actual_explicit_retdata, _) = run_cairo_0_entry_point(
         runner_config,
         program_bytes,
@@ -38,6 +41,7 @@ pub fn run_cairo_function_and_check_result(
         implicit_args,
         expected_explicit_retdata,
         hint_locals,
+        state_reader,
     )?;
     assert_eq!(expected_explicit_retdata, &actual_explicit_retdata);
     assert_eq!(expected_implicit_retdata, &actual_implicit_retdata);
@@ -62,6 +66,16 @@ pub fn create_squashed_cairo_dict(
     PointerArg::Composed(squashed_dict)
 }
 
+pub fn parse_squashed_cairo_dict(squashed_dict: &[Felt]) -> HashMap<Felt, Felt> {
+    assert!(squashed_dict.len() % DICT_ACCESS_SIZE == 0, "Invalid squashed dict length");
+    let key_offset = 0;
+    let new_val_offset = 2;
+    squashed_dict
+        .chunks(DICT_ACCESS_SIZE)
+        .map(|chunk| (chunk[key_offset], chunk[new_val_offset]))
+        .collect()
+}
+
 // 2**251 + 17 * 2**192 + 1
 pub static DEFAULT_PRIME: LazyLock<BigInt> = LazyLock::new(|| {
     BigInt::from_bytes_be(
@@ -123,3 +137,11 @@ pub fn pack_bigint3(limbs: &[Felt]) -> BigInt {
         acc + as_int(&limb, &DEFAULT_PRIME) * BASE.pow(i.try_into().unwrap())
     })
 }
+
+pub(crate) fn get_entrypoint_runner_config() -> EntryPointRunnerConfig {
+    EntryPointRunnerConfig {
+        layout: LayoutName::small,
+        add_main_prefix_to_entrypoint: false,
+        ..Default::default()
+    }
+}
diff --git a/crates/starknet_os/src/tests/aliases.rs b/crates/starknet_os/src/tests/aliases.rs
@@ -1,9 +1,31 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use apollo_starknet_os_program::test_programs::ALIASES_TEST_BYTES;
+use blockifier::state::stateful_compression::{ALIAS_COUNTER_STORAGE_KEY, INITIAL_AVAILABLE_ALIAS};
+use blockifier::test_utils::dict_state_reader::DictStateReader;
+use blockifier::test_utils::ALIAS_CONTRACT_ADDRESS;
+use cairo_vm::hint_processor::builtin_hint_processor::dict_hint_utils::DICT_ACCESS_SIZE;
+use cairo_vm::hint_processor::hint_processor_utils::felt_to_usize;
+use cairo_vm::types::builtin_name::BuiltinName;
+use rstest::rstest;
+use starknet_api::core::L2_ADDRESS_UPPER_BOUND;
+use starknet_api::state::StorageKey;
+use starknet_types_core::felt::Felt;
+
+use crate::test_utils::cairo_runner::{
+    run_cairo_0_entry_point,
+    EndpointArg,
+    EntryPointRunnerConfig,
+    ImplicitArg,
+    PointerArg,
+    ValueArg,
+};
+use crate::test_utils::utils::{
+    get_entrypoint_runner_config,
+    parse_squashed_cairo_dict,
+    test_cairo_function,
+};
 
-use crate::test_utils::cairo_runner::EntryPointRunnerConfig;
-use crate::test_utils::utils::test_cairo_function;
 // TODO(Nimrod): Move this next to the stateful compression hints implementation.
 // TODO(Amos): This test is incomplete. Add the rest of the test cases and remove this todo.
 
@@ -29,3 +51,153 @@ fn test_constants() {
         HashMap::new(),
     )
 }
+
+#[rstest]
+#[case(
+    Vec::new(),
+    Vec::new(),
+    HashMap::from([(0.into(), 128.into())])
+)]
+#[case(
+    vec![Felt::from(&*L2_ADDRESS_UPPER_BOUND)],
+    vec![128],
+    HashMap::from([
+        (0.into(), 129.into()),
+        (Felt::from(&*L2_ADDRESS_UPPER_BOUND), 128.into())
+    ])
+)]
+#[case(
+    vec![2000.into(), 1999999999.into(), 3000.into(), 2000.into()],
+    vec![128, 129, 130, 128],
+    HashMap::from([
+        (0.into(), 131.into()),
+        (2000.into(), 128.into()),
+        (3000.into(), 130.into()),
+        (1999999999.into(), 129.into())
+    ])
+)]
+#[case(
+    Vec::from_iter((0..128).map(Felt::from)),
+    (0..128).collect::<Vec<_>>(),
+    HashMap::from_iter([(0.into(), 128.into())])
+)]
+#[case(
+    Vec::from_iter((0..129).map(Felt::from)),
+    (0..129).collect::<Vec<_>>(),
+    HashMap::from_iter([
+        (0.into(), 129.into()),
+        (128.into(), 128.into())
+    ])
+)]
+#[case(
+    vec![
+        13.into(),
+        500.into(),
+        11.into(),
+        2000.into(),
+        2001.into(),
+        13.into(),
+        501.into(),
+        98.into(),
+        222.into(),
+        2000.into(),
+        127.into(),
+        128.into()
+    ],
+    vec![13, 128, 11, 129, 130, 13, 131, 98, 132, 129, 127, 133],
+    HashMap::from([
+        (0.into(), 134.into()),
+        (128.into(), 133.into()),
+        (222.into(), 132.into()),
+        (500.into(), 128.into()),
+        (501.into(), 131.into()),
+        (2000.into(), 129.into()),
+        (2001.into(), 130.into())
+    ])
+)]
+#[case(
+    (0..150_u8)
+        .map(|i| Felt::from(128) + Felt::TWO.pow(i))
+        .chain((0..150_u8).map(|i| Felt::from(128) + Felt::TWO.pow(i)))
+        .collect::<Vec<_>>(),
+    (0..150_u128)
+        .map(|i| i + 128)
+        .chain((0..150_u128).map(|i| i + 128))
+        .collect::<Vec<_>>(),
+    HashMap::from_iter(
+        (0..150_u128)
+            .map(|i| (Felt::from(128) + Felt::TWO.pow(i), Felt::from(i + 128)))
+            .chain([(0.into(), (128 + 150).into())])
+    )
+)]
+fn allocate_and_replace_keys_from_empty_storage(
+    #[case] keys: Vec<Felt>,
+    #[case] expected_alias_per_key: Vec<u128>,
+    #[case] expected_alias_storage: HashMap<Felt, Felt>,
+) {
+    let expected_alias_per_key: Vec<_> =
+        expected_alias_per_key.into_iter().map(Felt::from).collect();
+    let (actual_alias_storage, actual_alias_per_key) =
+        allocate_aliases_for_keys_and_replace(keys, HashMap::new());
+    assert_eq!(actual_alias_storage, expected_alias_storage);
+    assert_eq!(actual_alias_per_key, expected_alias_per_key);
+}
+
+fn allocate_aliases_for_keys_and_replace(
+    keys: Vec<Felt>,
+    initial_storage: HashMap<StorageKey, Felt>,
+) -> (HashMap<Felt, Felt>, Vec<Felt>) {
+    let runner_config = get_entrypoint_runner_config();
+    let entrypoint = "__main__.allocate_alias_for_keys_and_replace";
+    let implicit_args = [ImplicitArg::Builtin(BuiltinName::range_check)];
+    let unique_keys: HashSet<Felt> = HashSet::from_iter(
+        keys.iter()
+            .filter(|key| key >= &&INITIAL_AVAILABLE_ALIAS)
+            .copied()
+            .chain([*ALIAS_COUNTER_STORAGE_KEY.key()]),
+    );
+    let expected_explicit_return_values = vec![
+        EndpointArg::Value(ValueArg::Single(Felt::ZERO)), // Aliases.len
+        EndpointArg::Pointer(PointerArg::Array(vec![ // Aliases.ptr
+            Felt::ZERO;
+            (unique_keys.len()) * DICT_ACCESS_SIZE
+            ])),
+        // Aliases per-key ptr.
+        EndpointArg::Pointer(PointerArg::Array(vec![Felt::ZERO; keys.len()])),
+    ];
+    let n_keys_arg = EndpointArg::Value(ValueArg::Single(keys.len().into()));
+    let keys_arg = EndpointArg::Pointer(PointerArg::Array(keys));
+    let explicit_args = vec![n_keys_arg, keys_arg];
+    let storage_view = initial_storage
+        .into_iter()
+        .map(|(key, value)| ((*ALIAS_CONTRACT_ADDRESS, key), value))
+        .collect();
+
+    let state_reader = DictStateReader { storage_view, ..Default::default() };
+    let (_, explicit_return_values, _) = run_cairo_0_entry_point(
+        &runner_config,
+        ALIASES_TEST_BYTES,
+        entrypoint,
+        &explicit_args,
+        &implicit_args,
+        &expected_explicit_return_values,
+        HashMap::new(),
+        Some(state_reader),
+    )
+    .unwrap();
+    if let [
+        EndpointArg::Value(ValueArg::Single(n_aliases)),
+        EndpointArg::Pointer(PointerArg::Array(aliases_storage_updates)),
+        EndpointArg::Pointer(PointerArg::Array(alias_per_key)),
+    ] = explicit_return_values.as_slice()
+    {
+        let n_aliases = felt_to_usize(n_aliases).unwrap();
+        assert_eq!(n_aliases, aliases_storage_updates.len() / DICT_ACCESS_SIZE);
+        let actual_alias_storage = parse_squashed_cairo_dict(aliases_storage_updates);
+        (actual_alias_storage, alias_per_key.clone().to_vec())
+    } else {
+        panic!(
+            "The return value doesn't match the given format.\n Got: {explicit_return_values:?}"
+        );
+    }
+}
diff --git a/crates/starknet_os/src/tests/bls_field.rs b/crates/starknet_os/src/tests/bls_field.rs
@@ -3,7 +3,6 @@ use std::collections::HashMap;
 
 use apollo_starknet_os_program::OS_PROGRAM_BYTES;
 use cairo_vm::types::builtin_name::BuiltinName;
-use cairo_vm::types::layout_name::LayoutName;
 use cairo_vm::types::program::Program;
 use ethnum::U256;
 use num_bigint::{BigInt, BigUint, RandBigInt, RandomBits, Sign, ToBigInt};
@@ -16,12 +15,12 @@ use crate::hints::hint_implementation::kzg::utils::{split_bigint3, BASE, BLS_PRI
 use crate::test_utils::cairo_runner::{
     run_cairo_0_entry_point,
     EndpointArg,
-    EntryPointRunnerConfig,
     ImplicitArg,
     PointerArg,
     ValueArg,
 };
 use crate::test_utils::utils::{
+    get_entrypoint_runner_config,
     pack_bigint3,
     seeded_random_prng,
     test_cairo_function,
@@ -32,14 +31,6 @@ const REDUCED_MUL_LIMB_BOUND: i128 = 2_i128.pow(104);
 
 // TODO(Nimrod): Move this next to the BLS hints implementation.
 
-fn get_entrypoint_runner_config() -> EntryPointRunnerConfig {
-    EntryPointRunnerConfig {
-        layout: LayoutName::small,
-        add_main_prefix_to_entrypoint: false,
-        ..Default::default()
-    }
-}
-
 fn run_reduced_mul_test(a_split: &[Felt], b_split: &[Felt]) {
     let explicit_args = [
         EndpointArg::Value(ValueArg::Array(a_split.to_vec())),
@@ -151,6 +142,7 @@ fn test_horner_eval() {
         )));
         let implicit_args = [ImplicitArg::Builtin(BuiltinName::range_check)];
 
+        let state_reader = None;
         let (_, explicit_retdata, _) = run_cairo_0_entry_point(
             &entrypoint_runner_config,
             OS_PROGRAM_BYTES,
@@ -159,6 +151,7 @@ fn test_horner_eval() {
             &implicit_args,
             &[EndpointArg::Value(ValueArg::Array(vec![Felt::ZERO, Felt::ZERO, Felt::ZERO]))],
             HashMap::new(),
+            state_reader,
         )
         .unwrap();
 
