starknet_os: implement squash_dict and parse_contract_changes test utils (#9881)

Author: dorimedini-starkware

crates/starknet_os/Cargo.toml

@@ -14,7 +14,7 @@ deserialize = [
   "starknet_patricia/deserialize",
 ]
 include_program_output = []
-testing = ["blockifier/testing", "starknet_patricia/testing"]
+testing = ["blockifier/testing", "itertools", "starknet_patricia/testing"]
 
 [dependencies]
 apollo_starknet_os_program.workspace = true
@@ -38,6 +38,7 @@ derive_more.workspace = true
 digest.workspace = true
 indexmap.workspace = true
 indoc.workspace = true
+itertools = { workspace = true, optional = true }
 log.workspace = true
 num-bigint.workspace = true
 num-integer.workspace = true

crates/starknet_os/src/test_utils/cairo_dict.rs

@@ -1,10 +1,18 @@
 use std::collections::HashMap;
 
+use apollo_starknet_os_program::OS_PROGRAM;
 use cairo_vm::hint_processor::builtin_hint_processor::dict_hint_utils::DICT_ACCESS_SIZE;
 use cairo_vm::types::relocatable::{MaybeRelocatable, Relocatable};
 use cairo_vm::vm::vm_core::VirtualMachine;
+use itertools::Itertools;
+use starknet_api::core::{ClassHash, ContractAddress, Nonce, PatriciaKey};
+use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
 
+use crate::hints::vars::CairoStruct;
+use crate::io::os_output_types::{FullContractChanges, FullContractStorageUpdate};
+use crate::vm_utils::get_address_of_nested_fields_from_base_address;
+
 /// 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>,
@@ -43,3 +51,129 @@ pub fn parse_squashed_cairo_dict(squashed_dict: &[Felt]) -> HashMap<Felt, Felt>
         .map(|chunk| (chunk[key_offset], chunk[new_val_offset]))
         .collect()
 }
+
+/// Parses a cairo dictionary from VM memory into a squashed dictionary.
+/// Each entry is a tuple of the form (key, (prev_value, new_value)).
+pub fn squash_dict(
+    vm: &VirtualMachine,
+    dict_start: Relocatable,
+    dict_end: Relocatable,
+) -> Vec<(Felt, (Felt, Felt))> {
+    let mut prev_vals = HashMap::new();
+    let mut new_vals = HashMap::new();
+    let flat_dict: Vec<MaybeRelocatable> = vm
+        .segments
+        .memory
+        .get_range(dict_start, (dict_end - dict_start).unwrap())
+        .into_iter()
+        .map(|item| item.unwrap().into_owned())
+        .collect();
+    for chunk in flat_dict.chunks_exact(DICT_ACCESS_SIZE) {
+        let (key, prev, new) = (
+            chunk.first().unwrap().get_int().unwrap(),
+            chunk.get(1).unwrap().get_int().unwrap(),
+            chunk.get(2).unwrap().get_int().unwrap(),
+        );
+        if !prev_vals.contains_key(&prev) {
+            prev_vals.insert(key, prev);
+        } else {
+            assert_eq!(new_vals.get(&key).unwrap(), &prev);
+        }
+        new_vals.insert(key, new);
+    }
+    prev_vals.into_iter().map(|(key, prev)| (key, (prev, *new_vals.get(&key).unwrap()))).collect()
+}
+
+/// Parses (from VM memory) a squashed cairo dictionary of contract changes.
+/// Squashes the contract changes per contract address.
+pub fn parse_contract_changes(
+    vm: &VirtualMachine,
+    dict_start: Relocatable,
+    dict_end: Relocatable,
+) -> HashMap<ContractAddress, FullContractChanges> {
+    let flat_outer_dict: Vec<MaybeRelocatable> = vm
+        .segments
+        .memory
+        .get_range(dict_start, (dict_end - dict_start).unwrap())
+        .into_iter()
+        .map(|item| item.unwrap().into_owned())
+        .collect();
+    assert!(flat_outer_dict.len() % DICT_ACCESS_SIZE == 0, "Invalid outer dict length");
+    flat_outer_dict
+        .chunks_exact(DICT_ACCESS_SIZE)
+        .map(|chunk| {
+            let (address, prev_state_entry_ptr, new_state_entry_ptr) = (
+                ContractAddress(
+                    PatriciaKey::try_from(chunk.first().unwrap().get_int().unwrap()).unwrap(),
+                ),
+                chunk.get(1).unwrap().get_relocatable().unwrap(),
+                chunk.get(2).unwrap().get_relocatable().unwrap(),
+            );
+
+            // Fetch fields of previous and new state entries.
+            // Note that nonces and class hash addresses point to integers, while storage pointer
+            // points to a relocatable.
+            let (
+                prev_nonce_ptr,
+                new_nonce_ptr,
+                prev_class_hash_ptr,
+                new_class_hash_ptr,
+                prev_storage_ptr,
+                new_storage_ptr,
+            ) = [
+                (prev_state_entry_ptr, "nonce"),
+                (new_state_entry_ptr, "nonce"),
+                (prev_state_entry_ptr, "class_hash"),
+                (new_state_entry_ptr, "class_hash"),
+                (prev_state_entry_ptr, "storage_ptr"),
+                (new_state_entry_ptr, "storage_ptr"),
+            ]
+            .into_iter()
+            .map(|(ptr, field)| {
+                get_address_of_nested_fields_from_base_address(
+                    ptr,
+                    CairoStruct::StateEntry,
+                    vm,
+                    &[field],
+                    &*OS_PROGRAM,
+                )
+                .unwrap()
+            })
+            .collect_tuple()
+            .unwrap();
+
+            let (prev_nonce, new_nonce, prev_class_hash, new_class_hash) =
+                [prev_nonce_ptr, new_nonce_ptr, prev_class_hash_ptr, new_class_hash_ptr]
+                    .into_iter()
+                    .map(|ptr| vm.get_integer(ptr).unwrap())
+                    .collect_tuple()
+                    .unwrap();
+            let (prev_storage_ptr, new_storage_ptr) = [prev_storage_ptr, new_storage_ptr]
+                .into_iter()
+                .map(|ptr| vm.get_relocatable(ptr).unwrap())
+                .collect_tuple()
+                .unwrap();
+
+            let storage_changes = squash_dict(vm, prev_storage_ptr, new_storage_ptr)
+                .into_iter()
+                .map(|(key, (prev_value, new_value))| FullContractStorageUpdate {
+                    key: StorageKey(PatriciaKey::try_from(key).unwrap()),
+                    prev_value,
+                    new_value,
+                })
+                .collect();
+
+            (
+                address,
+                FullContractChanges {
+                    addr: address,
+                    prev_nonce: Nonce(*prev_nonce),
+                    new_nonce: Nonce(*new_nonce),
+                    prev_class_hash: ClassHash(*prev_class_hash),
+                    new_class_hash: ClassHash(*new_class_hash),
+                    storage_changes,
+                },
+            )
+        })
+        .collect()
+}