starknet_os: allocate and replace aliases test

nimrod-starkware · nimrod-starkware · commit 6aec9d93823e · 2025-07-06T15:26:10.000+03:00
diff --git a/crates/starknet_os/src/test_utils/utils.rs b/crates/starknet_os/src/test_utils/utils.rs
@@ -4,6 +4,8 @@ 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 cairo_vm::types::relocatable::{MaybeRelocatable, Relocatable};
+use cairo_vm::vm::vm_core::VirtualMachine;
 use ethnum::U256;
 use num_bigint::{BigInt, Sign};
 use rand::rngs::StdRng;
@@ -17,8 +19,6 @@ use crate::test_utils::cairo_runner::{
     EndpointArg,
     EntryPointRunnerConfig,
     ImplicitArg,
-    PointerArg,
-    ValueArg,
 };
 
 #[allow(clippy::too_many_arguments)]
@@ -49,22 +49,33 @@ pub fn run_cairo_function_and_check_result(
     Ok(())
 }
 
-pub fn create_squashed_cairo_dict(
-    prev_values: &HashMap<Felt, EndpointArg>,
-    new_values: &HashMap<Felt, EndpointArg>,
-) -> PointerArg {
-    let mut squashed_dict: Vec<EndpointArg> = vec![];
+/// Creates a squashed dict from previous and new values, and stores it in a new memory segment.
+pub fn allocate_squashed_cairo_dict(
+    prev_values: &HashMap<Felt, MaybeRelocatable>,
+    new_values: &HashMap<Felt, MaybeRelocatable>,
+    vm: &mut VirtualMachine,
+) -> (Relocatable, Relocatable) {
+    let squashed_dict = flatten_cairo_dict(prev_values, new_values);
+    let dict_segment_start = vm.add_memory_segment();
+    let dict_segment_end = vm.load_data(dict_segment_start, &squashed_dict).unwrap();
+    (dict_segment_start, dict_segment_end)
+}
+
+pub fn flatten_cairo_dict(
+    prev_values: &HashMap<Felt, MaybeRelocatable>,
+    new_values: &HashMap<Felt, MaybeRelocatable>,
+) -> Vec<MaybeRelocatable> {
+    let mut squashed_dict = vec![];
     let mut sorted_new_values: Vec<_> = new_values.iter().collect();
     sorted_new_values.sort_by_key(|(key, _)| *key);
 
     for (key, value) in sorted_new_values {
-        let prev_value: &EndpointArg =
-            prev_values.get(key).unwrap_or(&EndpointArg::Value(ValueArg::Single(Felt::ZERO)));
+        let prev_value = prev_values.get(key).unwrap_or(&MaybeRelocatable::Int(Felt::ZERO));
         squashed_dict.push((*key).into());
         squashed_dict.push(prev_value.clone());
         squashed_dict.push(value.clone());
     }
-    PointerArg::Composed(squashed_dict)
+    squashed_dict
 }
 
 pub fn parse_squashed_cairo_dict(squashed_dict: &[Felt]) -> HashMap<Felt, Felt> {
diff --git a/crates/starknet_os/src/tests/aliases.rs b/crates/starknet_os/src/tests/aliases.rs
@@ -1,4 +1,5 @@
-use std::collections::{HashMap, HashSet};
+use core::panic;
+use std::collections::{BTreeSet, HashMap, HashSet};
 
 use apollo_starknet_os_program::test_programs::ALIASES_TEST_BYTES;
 use blockifier::state::stateful_compression::{ALIAS_COUNTER_STORAGE_KEY, INITIAL_AVAILABLE_ALIAS};
@@ -15,20 +16,25 @@ use starknet_types_core::felt::Felt;
 
 use crate::test_utils::cairo_runner::{
     initialize_and_run_cairo_0_entry_point,
+    initialize_cairo_runner,
+    run_cairo_0_entrypoint,
     EndpointArg,
     EntryPointRunnerConfig,
     ImplicitArg,
     PointerArg,
     ValueArg,
 };
 use crate::test_utils::utils::{
+    allocate_squashed_cairo_dict,
+    flatten_cairo_dict,
     get_entrypoint_runner_config,
     parse_squashed_cairo_dict,
     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.
+
+const DEFAULT_CLASS_HASH: u128 = 7777;
 
 #[test]
 fn test_constants() {
@@ -290,22 +296,258 @@ fn allocate_aliases_for_keys_and_replace(
         Some(state_reader),
     )
     .unwrap();
-    if let [
+    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 aliases_storage_updates_as_felts: Vec<Felt> =
-            aliases_storage_updates.iter().map(|f| f.get_int().unwrap()).collect();
-        let actual_alias_storage = parse_squashed_cairo_dict(&aliases_storage_updates_as_felts);
-        let alias_per_key: Vec<Felt> = alias_per_key.iter().map(|f| f.get_int().unwrap()).collect();
-        (actual_alias_storage, alias_per_key)
-    } else {
+    else {
         panic!(
             "The return value doesn't match the given format.\n Got: {explicit_return_values:?}"
         );
+    };
+    let n_aliases = felt_to_usize(n_aliases).unwrap();
+    assert_eq!(n_aliases, aliases_storage_updates.len() / DICT_ACCESS_SIZE);
+    let aliases_storage_updates_as_felts: Vec<Felt> =
+        aliases_storage_updates.iter().map(|f| f.get_int().unwrap()).collect();
+    let actual_alias_storage = parse_squashed_cairo_dict(&aliases_storage_updates_as_felts);
+    let alias_per_key: Vec<Felt> = alias_per_key.iter().map(|f| f.get_int().unwrap()).collect();
+    (actual_alias_storage, alias_per_key)
+}
+
+#[rstest]
+#[case::non_allocation_of_address_lt_16_from_empty_storage(
+    HashMap::from([
+        (
+            15,
+            HashMap::from([(5534, 1), (98435, 1), (99999, 1)])
+        ),
+        (
+            16,
+            HashMap::from([(11, 1), (127, 1), (128, 1), (129, 1), (225, 1), (7659, 1)])
+        ),
+        (
+            7659,
+            HashMap::from([(12, 0), (200, 1), (300, 1), (1111, 1)])
+        ),
+        (
+            99999,
+            HashMap::from([(225, 1)])
+        )
+    ]),
+    HashMap::new(),
+    HashMap::new(),
+    HashMap::new(),
+    HashMap::from([(0, 136), (128, 128), (129, 129), (200, 132), (225, 130), (300, 133), (1111, 134), (7659, 131), (99999, 135)])
+)]
+#[case::non_allocation_of_address_lt_16_from_non_empty_storage(
+    HashMap::from([
+        (
+            9,
+            HashMap::from([(5534, 1), (98435, 1), (99999, 1)])
+        ),
+        (
+            44,
+            HashMap::from([(11, 1), (129, 1), (225, 1), (7659, 1)])
+        ),
+        (
+            400,
+            HashMap::from([(225, 1), (400, 1), (700, 1), (701, 1), (1111, 1)])
+        ),
+    ]),
+    HashMap::new(),
+    HashMap::new(),
+    HashMap::from([(0, 135), (129, 128), (225, 129), (7659, 130), (200, 131), (300, 132), (1111, 133), (99999, 134)]),
+    HashMap::from([(0, 138), (129, 128), (225, 129), (400, 135), (700, 136), (701, 137), (1111, 133), (7659, 130)]
+    )
+)]
+#[case::non_allocation_with_only_trivial_updates(
+    HashMap::from([
+        (
+            11,
+            HashMap::from([(5534, 1), (98435, 1), (99999, 1)])
+        ),
+        (
+            44,
+            HashMap::from([(11, 0), (129, 1), (225, 0), (400, 1), (7659, 1)])
+        ),
+        (
+            400,
+            HashMap::from([(225, 0), (406, 0), (700, 1), (701, 1), (1111, 1)])
+        ),
+        (
+            598,
+            HashMap::from([(2255, 0), (7008, 0)]) // Trivial update.
+        )
+    ]),
+    HashMap::new(),
+    HashMap::new(),
+    HashMap::new(),
+    HashMap::from([(0, 134), (129, 128), (400, 129), (700, 131), (701, 132), (1111, 133), (7659, 130)]
+    )
+)]
+#[case::allocation_with_only_nonce_change(
+    HashMap::new(),
+    HashMap::from([
+        (13, 1),
+        (58, 1),
+        (11111, DEFAULT_CLASS_HASH), // Gets a new nonce.
+        (222222, 1),
+        (3333333, 1),
+        (3333336, DEFAULT_CLASS_HASH), // Nothing changed.
+    ]),
+    HashMap::from([(11111, 1)]),
+    HashMap::new(),
+    HashMap::from([(0, 131), (11111, 128), (222222, 129), (3333333, 130)]
+    )
+)]
+#[case::non_allocation_with_trivial_class_hash_update(
+    HashMap::new(),
+    HashMap::from([(24, 1), (5000, 1), (6666, 1), (9999, 1), (11111, DEFAULT_CLASS_HASH),
+    ]),
+    HashMap::new(),
+    HashMap::from([(0, 133), (5000, 128), (11111, 129), (222222, 130), (3333333, 131), (87777, 132)]),
+    HashMap::from([(0, 135), (5000, 128), (6666, 133), (9999, 134)]
+    )
+)]
+#[case::allocation_with_partially_trivial_updates(
+    HashMap::from([
+        (
+            1,
+            HashMap::from([(777, 1), (8888, 1), (9999, 1)]) // No aliases.
+        ),
+        (
+            100,
+            HashMap::from([(200, 1), (777, 1), (888, 0)]) // Aliases for non-trivial diffs.
+        ),
+        (
+            600,
+            HashMap::from([(2000, 1), (3000, 1)])
+        ),
+        (
+            800,
+            HashMap::from([(700, 1), (701, 1)])
+        ),
+        (
+            3000,
+            HashMap::from([(600, 1), (2000, 1)])
+        ),
+        (
+            10000,
+            HashMap::from([(34567, 0), (435, 0)])  // No aliases (all diffs trivial).
+        )
+    ]),
+    HashMap::from([(200, 1), (500, 1), (700, 1), (800, DEFAULT_CLASS_HASH)]),
+    HashMap::from([(700, 1), (10000, 0)]),
+    HashMap::new(),
+    HashMap::from([(0, 137), (200, 128), (500, 130), (600, 133), (700, 134), (701, 135), (777, 129), (800, 136), (2000, 131), (3000, 132)])
+)]
+fn test_allocate_addresses_for_state_diff_and_replace(
+    #[case] storage_updates: HashMap<u128, HashMap<u128, u128>>,
+    #[case] address_to_class_hash: HashMap<u128, u128>,
+    #[case] address_to_nonce: HashMap<u128, u128>,
+    #[case] initial_alias_storage: HashMap<u128, u128>,
+    #[case] expected_alias_storage: HashMap<u128, u128>,
+) {
+    let runner_config = get_entrypoint_runner_config();
+    let entrypoint = "__main__.allocate_aliases_and_replace";
+    let implicit_args = [ImplicitArg::Builtin(BuiltinName::range_check)];
+    let modified_contracts: BTreeSet<_> = storage_updates
+        .keys()
+        .chain(address_to_class_hash.keys().chain(address_to_nonce.keys()))
+        .collect();
+
+    // Initialize the runner to be able to allocate segments.
+    let (mut cairo_runner, program, entrypoint) = initialize_cairo_runner(
+        &runner_config,
+        ALIASES_TEST_BYTES,
+        entrypoint,
+        &implicit_args,
+        HashMap::new(),
+    )
+    .unwrap();
+
+    // Construct the contract state changes.
+    let mut prev_state_entries = HashMap::new();
+    let mut new_state_entries = HashMap::new();
+    let n_contracts = modified_contracts.len();
+    for address in modified_contracts {
+        let inner_updates = storage_updates
+            .get(address)
+            .unwrap_or(&HashMap::new())
+            .iter()
+            .map(|(k, v)| ((*k).into(), Felt::from(*v).into()))
+            .collect();
+        let (new_nonce, prev_nonce) = (address_to_nonce.get(address).copied().unwrap_or(0), 0);
+        let (new_class_hash, prev_class_hash) = (
+            address_to_class_hash.get(address).copied().unwrap_or(DEFAULT_CLASS_HASH),
+            DEFAULT_CLASS_HASH,
+        );
+        let (prev_storage_ptr, new_storage_ptr) =
+            allocate_squashed_cairo_dict(&HashMap::new(), &inner_updates, &mut cairo_runner.vm);
+        let new_state_entry: Vec<MaybeRelocatable> = vec![
+            Felt::from(new_class_hash).into(),
+            new_storage_ptr.into(),
+            Felt::from(new_nonce).into(),
+        ];
+        let prev_state_entry: Vec<MaybeRelocatable> = vec![
+            Felt::from(prev_class_hash).into(),
+            prev_storage_ptr.into(),
+            Felt::from(prev_nonce).into(),
+        ];
+        new_state_entries
+            .insert((*address).into(), cairo_runner.vm.gen_arg(&new_state_entry).unwrap());
+        prev_state_entries
+            .insert((*address).into(), cairo_runner.vm.gen_arg(&prev_state_entry).unwrap());
     }
+    let flat_contract_state_changes = flatten_cairo_dict(&prev_state_entries, &new_state_entries);
+    let explicit_args = vec![
+        EndpointArg::Value(ValueArg::Single(n_contracts.into())),
+        EndpointArg::Pointer(PointerArg::Array(flat_contract_state_changes)),
+    ];
+    let storage_view = initial_alias_storage
+        .into_iter()
+        .map(|(key, value)| ((*ALIAS_CONTRACT_ADDRESS, key.into()), value.into()))
+        .collect();
+    let state_reader = DictStateReader { storage_view, ..Default::default() };
+    let expected_aliases_storage_flat_length = expected_alias_storage.len() * DICT_ACCESS_SIZE;
+    let expected_explicit_return_values = vec![
+        EndpointArg::Pointer(PointerArg::Array(vec![
+            MaybeRelocatable::Int(Felt::ZERO);
+            expected_aliases_storage_flat_length
+        ])),
+        EndpointArg::Pointer(PointerArg::Array(vec![])),
+        EndpointArg::Pointer(PointerArg::Array(vec![])),
+    ];
+
+    let (_, explicit_return_values) = run_cairo_0_entrypoint(
+        entrypoint,
+        &explicit_args,
+        &implicit_args,
+        Some(state_reader),
+        &mut cairo_runner,
+        &program,
+        &runner_config,
+        &expected_explicit_return_values,
+    )
+    .unwrap();
+
+    // TODO(Nimrod): Complete this test to also compare the other return values.
+    let [
+        EndpointArg::Pointer(PointerArg::Array(aliases_storage_updates)),
+        EndpointArg::Pointer(PointerArg::Array(_)),
+        EndpointArg::Pointer(PointerArg::Array(_)),
+    ] = explicit_return_values.as_slice()
+    else {
+        panic!(
+            "The return value doesn't match the given format.\n Got: {explicit_return_values:?}"
+        );
+    };
+
+    let aliases_storage_updates_as_felts: Vec<Felt> =
+        aliases_storage_updates.iter().map(|f| f.get_int().unwrap()).collect();
+    let actual_alias_storage = parse_squashed_cairo_dict(&aliases_storage_updates_as_felts);
+    let expected_alias_storage: HashMap<Felt, Felt> =
+        expected_alias_storage.into_iter().map(|(key, value)| (key.into(), value.into())).collect();
+    assert_eq!(actual_alias_storage, expected_alias_storage);
 }
