starknet_os: aliases test

Author: nimrod-starkware

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)
     }
 }
 

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,

crates/starknet_os/src/test_utils/utils.rs

@@ -2,6 +2,7 @@ use std::any::Any;
 use std::collections::HashMap;
 use std::sync::LazyLock;
 
+use cairo_vm::types::layout_name::LayoutName;
 use ethnum::U256;
 use num_bigint::{BigInt, Sign};
 use rand::rngs::StdRng;
@@ -38,6 +39,7 @@ pub fn run_cairo_function_and_check_result(
         implicit_args,
         expected_explicit_retdata,
         hint_locals,
+        None,
     )?;
     assert_eq!(expected_explicit_retdata, &actual_explicit_retdata);
     assert_eq!(expected_implicit_retdata, &actual_implicit_retdata);
@@ -122,3 +124,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()
+    }
+}

crates/starknet_os/src/tests/aliases.rs

@@ -1,9 +1,25 @@
-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 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::{ContractAddress, L2_ADDRESS_UPPER_BOUND};
+use starknet_api::state::StorageKey;
+use starknet_types_core::felt::Felt;
 
-use crate::test_utils::cairo_runner::EntryPointRunnerConfig;
-use crate::test_utils::utils::test_cairo_function;
+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, test_cairo_function};
 
 // TODO(Amos): This test is incomplete. Add the rest of the test cases and remove this todo.
 
@@ -29,3 +45,91 @@ 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())]))]
+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
+            ])),
+        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 alias_contract_address: ContractAddress = Felt::TWO.try_into().unwrap();
+    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();
+    let mut actual_alias_storage = HashMap::new();
+    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!(aliases_storage_updates.len() % DICT_ACCESS_SIZE == 0);
+        assert!(aliases_storage_updates.len() / DICT_ACCESS_SIZE == n_aliases);
+        let key_offset = 0;
+        let new_value_offset = 2;
+        for i in 0..n_aliases {
+            let key = aliases_storage_updates[i * DICT_ACCESS_SIZE + key_offset];
+            let new_value = aliases_storage_updates[i * DICT_ACCESS_SIZE + new_value_offset];
+            actual_alias_storage.insert(key, new_value);
+        }
+        (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:?}"
+        );
+    }
+}

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};
@@ -15,12 +14,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,
@@ -29,14 +28,6 @@ use crate::test_utils::utils::{
 
 const REDUCED_MUL_LIMB_BOUND: i128 = 2_i128.pow(104);
 
-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())),
@@ -155,6 +146,7 @@ fn test_horner_eval() {
             &implicit_args,
             &[EndpointArg::Value(ValueArg::Array(vec![Felt::ZERO, Felt::ZERO, Felt::ZERO]))],
             HashMap::new(),
+            None,
         )
         .unwrap();