load cheatnet function (#1350)

<!-- Reference any GitHub issues resolved by this PR -->

Related: #346 

## Introduced changes

<!-- A brief description of the changes -->

- Adds a `load` function which will be used in snforge_std and invoked
in cheatcodes in order to retrieve data from the blockfier state
directly

## Checklist

<!-- Make sure all of these are complete -->

- [x] Linked relevant issue
- [x] Updated relevant documentation
- [x] Added relevant tests
- [x] Performed self-review of the code
- [x] Added changes to `CHANGELOG.md`

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/cheatcodes/load.rs

@@ -0,0 +1,49 @@
+use crate::state::BlockifierState;
+use cairo_felt::Felt252;
+use conversions::{FromConv, IntoConv};
+use num_traits::Pow;
+use starknet::core::types::FieldElement;
+use starknet_api::core::{ContractAddress, PatriciaKey};
+use starknet_api::hash::StarkHash;
+use starknet_api::state::StorageKey;
+
+///
+/// # Arguments
+///
+/// * `blockifier_state`: Blockifier state reader
+/// * `target`: The address of the contract we want to target
+/// * `storage_address`: Beginning of the storage of the variable
+/// * `size`: How many felts we want to read from the calculated offset (`target` + `storage_address`)
+///
+/// returns: Result<Vec<Felt252>, Error> - a result containing the read data  
+///
+pub fn load(
+    blockifier_state: &mut BlockifierState,
+    target: ContractAddress,
+    storage_address: &Felt252,
+    size: &Felt252,
+) -> Result<Vec<Felt252>, anyhow::Error> {
+    let mut values: Vec<Felt252> = vec![];
+    let mut current_slot = storage_address.clone();
+
+    while current_slot < storage_address + size {
+        let storage_value = blockifier_state.blockifier_state.get_storage_at(
+            target,
+            StorageKey(PatriciaKey::try_from(StarkHash::from_(
+                current_slot.clone(),
+            ))?),
+        );
+        values.push(Felt252::from_(storage_value?));
+        current_slot += Felt252::from(1);
+    }
+    Ok(values)
+}
+
+/// The address after hashing with pedersen, needs to be taken with a specific modulo value (2^251 - 256)
+/// For details see:
+/// <https://docs.starknet.io/documentation/architecture_and_concepts/Smart_Contracts/contract-storage>
+#[must_use]
+pub fn map_storage_address(address: FieldElement) -> FieldElement {
+    let modulus = Felt252::from(2).pow(251) - Felt252::from(256);
+    address % modulus.into_()
+}

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/cheatcodes/mod.rs

@@ -8,6 +8,7 @@ pub mod deploy;
 pub mod elect;
 pub mod get_class_hash;
 pub mod l1_handler_execute;
+pub mod load;
 pub mod mock_call;
 pub mod prank;
 pub mod precalculate_address;

crates/cheatnet/tests/cheatcodes/load.rs

@@ -0,0 +1,248 @@
+use crate::common::state::{create_cached_state, create_cheatnet_state};
+use crate::common::{felt_selector_from_name, get_contracts};
+use blockifier::abi::abi_utils::starknet_keccak;
+use cairo_felt::Felt252;
+use cheatnet::runtime_extensions::call_to_blockifier_runtime_extension::rpc::call_contract;
+use cheatnet::runtime_extensions::forge_runtime_extension::cheatcodes::deploy::deploy;
+use cheatnet::runtime_extensions::forge_runtime_extension::cheatcodes::load::{
+    load, map_storage_address,
+};
+use conversions::felt252::FromShortString;
+use conversions::IntoConv;
+use starknet::core::crypto::pedersen_hash;
+
+#[test]
+fn load_simple_state() {
+    let mut cached_state = create_cached_state();
+    let (mut blockifier_state, mut cheatnet_state) = create_cheatnet_state(&mut cached_state);
+
+    let contract = Felt252::from_short_string("HelloStarknet").unwrap();
+    let contracts = get_contracts();
+
+    let class_hash = blockifier_state.declare(&contract, &contracts).unwrap();
+
+    let contract_address = deploy(&mut blockifier_state, &mut cheatnet_state, &class_hash, &[])
+        .unwrap()
+        .contract_address;
+
+    let selector = felt_selector_from_name("increase_balance");
+
+    call_contract(
+        &mut blockifier_state,
+        &mut cheatnet_state,
+        &contract_address,
+        &selector,
+        &[Felt252::from(420)],
+    )
+    .unwrap();
+
+    let balance_value = load(
+        &mut blockifier_state,
+        contract_address,
+        &starknet_keccak("balance".as_ref()),
+        &Felt252::from(1),
+    )
+    .unwrap();
+
+    assert_eq!(balance_value.len(), 1, "Wrong data amount was returned");
+    let returned_balance_value = balance_value[0].clone();
+    assert_eq!(
+        returned_balance_value,
+        Felt252::from(420),
+        "Wrong data value was returned: {returned_balance_value}"
+    );
+}
+
+#[test]
+fn load_state_map_simple_value() {
+    let mut cached_state = create_cached_state();
+    let (mut blockifier_state, mut cheatnet_state) = create_cheatnet_state(&mut cached_state);
+
+    let contract = Felt252::from_short_string("MapSimpleValueSimpleKey").unwrap();
+    let contracts = get_contracts();
+
+    let class_hash = blockifier_state.declare(&contract, &contracts).unwrap();
+
+    let contract_address = deploy(&mut blockifier_state, &mut cheatnet_state, &class_hash, &[])
+        .unwrap()
+        .contract_address;
+
+    let selector = felt_selector_from_name("insert");
+
+    let map_key = Felt252::from(420);
+    let inserted_value = Felt252::from(69);
+    call_contract(
+        &mut blockifier_state,
+        &mut cheatnet_state,
+        &contract_address,
+        &selector,
+        &[map_key.clone(), inserted_value.clone()],
+    )
+    .unwrap();
+
+    // This can be abstracted in cairo code itself, without compromising load's interface simplicity
+    let variable_name_hashed = starknet_keccak("values".as_ref());
+    let value_address = pedersen_hash(&variable_name_hashed.into_(), &map_key.into_());
+    let value_address = map_storage_address(value_address);
+
+    let map_value = load(
+        &mut blockifier_state,
+        contract_address,
+        &value_address.into_(),
+        &Felt252::from(1),
+    )
+    .unwrap();
+
+    assert_eq!(map_value.len(), 1, "Wrong data amount was returned");
+    let returned_map_value = map_value[0].clone();
+    assert_eq!(
+        returned_map_value, inserted_value,
+        "Wrong data value was returned: {returned_map_value}"
+    );
+}
+
+#[test]
+fn load_state_map_complex_value() {
+    let mut cached_state = create_cached_state();
+    let (mut blockifier_state, mut cheatnet_state) = create_cheatnet_state(&mut cached_state);
+
+    let contract = Felt252::from_short_string("MapComplexValueSimpleKey").unwrap();
+    let contracts = get_contracts();
+
+    let class_hash = blockifier_state.declare(&contract, &contracts).unwrap();
+
+    let contract_address = deploy(&mut blockifier_state, &mut cheatnet_state, &class_hash, &[])
+        .unwrap()
+        .contract_address;
+
+    let selector = felt_selector_from_name("insert");
+
+    let map_key = Felt252::from(420);
+    let inserted_values = vec![Felt252::from(68), Felt252::from(69)];
+    let mut calldata = vec![map_key.clone()];
+    calldata.append(&mut inserted_values.clone());
+    call_contract(
+        &mut blockifier_state,
+        &mut cheatnet_state,
+        &contract_address,
+        &selector,
+        &calldata,
+    )
+    .unwrap();
+
+    // This can be abstracted in cairo code itself, without compromising load's interface simplicity
+    let variable_name_hashed = starknet_keccak("values".as_ref());
+    let value_address = pedersen_hash(&variable_name_hashed.into_(), &map_key.into_());
+    let value_address = map_storage_address(value_address);
+
+    let map_value = load(
+        &mut blockifier_state,
+        contract_address,
+        &value_address.into_(),
+        &Felt252::from(2),
+    )
+    .unwrap();
+
+    assert_eq!(map_value.len(), 2, "Wrong data amount was returned");
+
+    assert_eq!(
+        map_value, inserted_values,
+        "Wrong data value was returned: {map_value:?}"
+    );
+}
+
+#[test]
+fn load_state_map_complex_key() {
+    let mut cached_state = create_cached_state();
+    let (mut blockifier_state, mut cheatnet_state) = create_cheatnet_state(&mut cached_state);
+
+    let contract = Felt252::from_short_string("MapSimpleValueComplexKey").unwrap();
+    let contracts = get_contracts();
+
+    let class_hash = blockifier_state.declare(&contract, &contracts).unwrap();
+
+    let contract_address = deploy(&mut blockifier_state, &mut cheatnet_state, &class_hash, &[])
+        .unwrap()
+        .contract_address;
+
+    let selector = felt_selector_from_name("insert");
+
+    let map_key = vec![Felt252::from(68), Felt252::from(69)];
+    let mut calldata = map_key.clone();
+    let inserted_value = Felt252::from(420);
+    calldata.push(inserted_value.clone());
+    call_contract(
+        &mut blockifier_state,
+        &mut cheatnet_state,
+        &contract_address,
+        &selector,
+        &calldata,
+    )
+    .unwrap();
+
+    // This can be abstracted in cairo code itself, without compromising load's interface simplicity
+    let variable_name_hashed = starknet_keccak("values".as_ref());
+    let value_address = pedersen_hash(
+        &pedersen_hash(&variable_name_hashed.into_(), &map_key[0].clone().into_()),
+        &map_key[1].clone().into_(),
+    );
+    let value_address = map_storage_address(value_address);
+
+    let map_value = load(
+        &mut blockifier_state,
+        contract_address,
+        &value_address.into_(),
+        &Felt252::from(1),
+    )
+    .unwrap();
+
+    assert_eq!(map_value.len(), 1, "Wrong data amount was returned");
+    let returned_map_value = map_value[0].clone();
+    assert_eq!(
+        inserted_value, returned_map_value,
+        "Wrong data value was returned: {returned_map_value}"
+    );
+}
+
+#[test]
+fn load_state_struct() {
+    let mut cached_state = create_cached_state();
+    let (mut blockifier_state, mut cheatnet_state) = create_cheatnet_state(&mut cached_state);
+
+    let contract = Felt252::from_short_string("FlatStateStruct").unwrap();
+    let contracts = get_contracts();
+
+    let class_hash = blockifier_state.declare(&contract, &contracts).unwrap();
+
+    let contract_address = deploy(&mut blockifier_state, &mut cheatnet_state, &class_hash, &[])
+        .unwrap()
+        .contract_address;
+
+    let selector = felt_selector_from_name("insert");
+
+    let calldata = vec![Felt252::from(68), Felt252::from(69)];
+    call_contract(
+        &mut blockifier_state,
+        &mut cheatnet_state,
+        &contract_address,
+        &selector,
+        &calldata,
+    )
+    .unwrap();
+
+    let variable_name_hashed = starknet_keccak("value".as_ref());
+    let struct_value = load(
+        &mut blockifier_state,
+        contract_address,
+        &variable_name_hashed,
+        &Felt252::from(2),
+    )
+    .unwrap();
+
+    assert_eq!(struct_value.len(), 2, "Wrong data amount was returned");
+
+    assert_eq!(
+        calldata, struct_value,
+        "Wrong data value was returned: {struct_value:?}"
+    );
+}

crates/cheatnet/tests/cheatcodes/mod.rs

@@ -2,6 +2,7 @@ mod declare;
 mod deploy;
 mod elect;
 mod get_class_hash;
+mod load;
 mod mock_call;
 mod prank;
 mod precalculate_address;

crates/cheatnet/tests/contracts/src/lib.cairo

@@ -9,3 +9,4 @@ mod elect;
 mod starknet;
 mod warp;
 mod segment_arena_user;
+mod store_load;

crates/cheatnet/tests/contracts/src/store_load.cairo

@@ -0,0 +1,4 @@
+mod map_simple_value_simple_key;
+mod map_complex_value_simple_key;
+mod map_simple_value_complex_key;
+mod flat_state_struct;

crates/cheatnet/tests/contracts/src/store_load/flat_state_struct.cairo

@@ -0,0 +1,28 @@
+#[starknet::contract]
+mod FlatStateStruct {
+    #[derive(Serde, Drop, starknet::Store)]
+    struct NestedStructure {
+        c: felt252
+    }
+    #[derive(Serde, Drop, starknet::Store)]
+    struct StoredStructure {
+        a: felt252,
+        b: NestedStructure,
+    }
+
+
+    #[storage]
+    struct Storage {
+        value: StoredStructure,
+    }
+
+    #[external(v0)]
+    fn insert(ref self: ContractState, value: StoredStructure) {
+        self.value.write(value);
+    }
+
+    #[external(v0)]
+    fn read(self: @ContractState) -> StoredStructure {
+        self.value.read()
+    }
+}

crates/cheatnet/tests/contracts/src/store_load/map_complex_value_simple_key.cairo

@@ -0,0 +1,28 @@
+#[starknet::contract]
+mod MapComplexValueSimpleKey {
+    #[derive(Serde, Drop, starknet::Store)]
+    struct NestedStructure {
+        c: felt252
+    }
+    #[derive(Serde, Drop, starknet::Store)]
+    struct StoredStructure {
+        a: felt252,
+        b: NestedStructure,
+    }
+
+
+    #[storage]
+    struct Storage {
+        values: LegacyMap<felt252, StoredStructure>,
+    }
+
+    #[external(v0)]
+    fn insert(ref self: ContractState, key: felt252, value: StoredStructure) {
+        self.values.write(key, value);
+    }
+
+    #[external(v0)]
+    fn read(self: @ContractState, key: felt252) -> StoredStructure {
+        self.values.read(key)
+    }
+}