Add interact_with_state closure cheatcode (#3526)

Towards #3331

commit-id:0b8b8c6b

---

**Stack**:
- #3539
- #3527
- #3526 ⬅


⚠️ *Part of a stack created by [spr](https://github.com/ejoffe/spr). Do
not merge manually using the UI - doing so may have unexpected results.*

Author: ddoktorski

.github/workflows/ci.yml

@@ -159,6 +159,22 @@ jobs:
 
       - run: cargo test --package forge --features scarb_since_2_10 sierra_gas
 
+  # TODO(#3212): Closures in Cairo are fully supported since version 2.11
+  test-interact-with-state:
+    name: Test interact with state
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6
+      - uses: software-mansion/setup-scarb@v1
+        with:
+          scarb-version: "2.11.0"
+      - uses: software-mansion/setup-universal-sierra-compiler@v1
+
+      - run: cargo test --release --package forge --features interact-with-state --test main integration::interact_with_state
+      - run: cargo test --release --package forge --features interact-with-state --test main e2e::running::test_interact_with_state
+
   test-forge-runner:
     name: Test Forge Runner
     runs-on: ubuntu-latest

crates/forge/Cargo.toml

@@ -13,6 +13,7 @@ no_scarb_installed = []
 debugging = []
 assert_non_exact_gas = ["test_utils/assert_non_exact_gas"]
 supports-panic-backtrace = []
+interact-with-state = []
 
 [dependencies]
 anyhow.workspace = true

crates/forge/tests/data/contract_state/Scarb.toml

@@ -0,0 +1,13 @@
+[package]
+name = "contract_state"
+version = "0.1.0"
+edition = "2024_07"
+
+[dependencies]
+starknet = "2.11.0"
+
+[dev-dependencies]
+snforge_std = { path = "../../../../../snforge_std" }
+assert_macros = "2.11.0"
+
+[[target.starknet-contract]]

crates/forge/tests/data/contract_state/src/balance.cairo

@@ -0,0 +1,75 @@
+#[starknet::interface]
+pub trait IHelloStarknetExtended<TContractState> {
+    fn increase_balance(ref self: TContractState, amount: u256);
+    fn get_balance(self: @TContractState) -> u256;
+    fn get_caller_info(self: @TContractState, address: starknet::ContractAddress) -> u256;
+    fn get_balance_at(self: @TContractState, index: u64) -> u256;
+}
+
+
+#[starknet::contract]
+pub mod HelloStarknetExtended {
+    use starknet::storage::{
+        Map, MutableVecTrait, StoragePathEntry, StoragePointerReadAccess, StoragePointerWriteAccess,
+        Vec, VecTrait,
+    };
+    use starknet::{ContractAddress, get_caller_address};
+
+    #[derive(starknet::Store)]
+    struct Owner {
+        pub address: ContractAddress,
+        pub name: felt252,
+    }
+
+    #[storage]
+    struct Storage {
+        pub owner: Owner,
+        pub balance: u256,
+        pub balance_records: Vec<u256>,
+        pub callers: Map<ContractAddress, u256>,
+    }
+
+    #[constructor]
+    fn constructor(ref self: ContractState, owner_name: felt252) {
+        self
+            ._set_owner(
+                starknet::get_execution_info().tx_info.account_contract_address, owner_name,
+            );
+        self.balance_records.push(0);
+    }
+
+    #[abi(embed_v0)]
+    impl HelloStarknetExtendedImpl of super::IHelloStarknetExtended<ContractState> {
+        fn increase_balance(ref self: ContractState, amount: u256) {
+            let caller = get_caller_address();
+            let value_before = self.callers.entry(caller).read();
+
+            assert(amount != 0, 'Amount cannot be 0');
+
+            self.balance.write(self.balance.read() + amount);
+            self.callers.entry(caller).write(value_before + amount);
+            self.balance_records.push(self.balance.read());
+        }
+
+        fn get_balance(self: @ContractState) -> u256 {
+            self.balance.read()
+        }
+
+        fn get_caller_info(self: @ContractState, address: ContractAddress) -> u256 {
+            self.callers.entry(address).read()
+        }
+
+        fn get_balance_at(self: @ContractState, index: u64) -> u256 {
+            assert(index < self.balance_records.len(), 'Index out of range');
+            self.balance_records.at(index).read()
+        }
+    }
+
+    #[generate_trait]
+    pub impl InternalFunctions of InternalFunctionsTrait {
+        fn _set_owner(ref self: ContractState, address: ContractAddress, name: felt252) {
+            self.owner.address.write(address);
+            self.owner.name.write(name);
+        }
+    }
+}

crates/forge/tests/data/contract_state/src/lib.cairo

@@ -0,0 +1,2 @@
+pub mod balance;
+pub mod storage_node;

crates/forge/tests/data/contract_state/src/storage_node.cairo

@@ -0,0 +1,37 @@
+#[starknet::interface]
+pub trait IStorageNodeContract<TContractState> {
+    fn get_description_at(self: @TContractState, index: u64) -> felt252;
+    fn get_data_at(
+        self: @TContractState, index: u64, address: starknet::ContractAddress, key: u16,
+    ) -> ByteArray;
+}
+
+#[starknet::contract]
+pub mod StorageNodeContract {
+    use starknet::ContractAddress;
+    use starknet::storage::{StoragePointerReadAccess, StoragePathEntry, Map};
+
+    #[starknet::storage_node]
+    pub struct RandomData {
+        pub description: felt252,
+        pub data: Map<(ContractAddress, u16), ByteArray>,
+    }
+
+    #[storage]
+    pub struct Storage {
+        pub random_data: Map<u64, RandomData>,
+    }
+
+    #[abi(embed_v0)]
+    impl IStorageNodeContractImpl of super::IStorageNodeContract<ContractState> {
+        fn get_description_at(self: @ContractState, index: u64) -> felt252 {
+            self.random_data.entry(index).description.read()
+        }
+
+        fn get_data_at(
+            self: @ContractState, index: u64, address: ContractAddress, key: u16,
+        ) -> ByteArray {
+            self.random_data.entry(index).data.entry((address, key)).read()
+        }
+    }
+}

crates/forge/tests/data/contract_state/tests/test_fork.cairo

@@ -0,0 +1,50 @@
+#[starknet::interface]
+trait IMap<TMapState> {
+    fn put(ref self: TMapState, key: felt252, value: felt252);
+    fn get(self: @TMapState, key: felt252) -> felt252;
+}
+
+#[starknet::contract]
+mod Map {
+    use starknet::storage::{Map, StorageMapReadAccess, StoragePathEntry, StoragePointerWriteAccess};
+    #[storage]
+    pub struct Storage {
+        pub storage: Map<felt252, felt252>,
+    }
+
+    #[abi(embed_v0)]
+    impl MapImpl of super::IMap<ContractState> {
+        fn put(ref self: ContractState, key: felt252, value: felt252) {
+            self.storage.entry(key).write(value);
+        }
+
+        fn get(self: @ContractState, key: felt252) -> felt252 {
+            self.storage.read(key)
+        }
+    }
+}
+use snforge_std::interact_with_state;
+use starknet::ContractAddress;
+use starknet::storage::StorageMapWriteAccess;
+
+#[test]
+#[fork(url: "{{ NODE_RPC_URL }}", block_number: 900_000)]
+fn test_fork_contract() {
+    let contract_address: ContractAddress =
+        0x00cd8f9ab31324bb93251837e4efb4223ee195454f6304fcfcb277e277653008
+        .try_into()
+        .unwrap();
+    let dispatcher = IMapDispatcher { contract_address };
+
+    assert(dispatcher.get(1) == 2, 'Wrong value');
+
+    interact_with_state(
+        contract_address,
+        || {
+            let mut state = Map::contract_state_for_testing();
+            state.storage.write(1, 13579)
+        },
+    );
+
+    assert(dispatcher.get(1) == 13579, 'Wrong value');
+}

crates/forge/tests/data/contract_state/tests/test_state.cairo

@@ -0,0 +1,156 @@
+use contract_state::balance::HelloStarknetExtended::InternalFunctionsTrait;
+use contract_state::balance::{
+    HelloStarknetExtended, IHelloStarknetExtendedDispatcher, IHelloStarknetExtendedDispatcherTrait,
+};
+use snforge_std::interact_with_state;
+use starknet::ContractAddress;
+use starknet::storage::{
+    MutableVecTrait, StorageMapWriteAccess, StoragePointerReadAccess, StoragePointerWriteAccess,
+};
+use crate::utils::deploy_contract;
+
+#[test]
+fn test_interact_with_state() {
+    let contract_address = deploy_contract("HelloStarknetExtended", array!['Name']);
+    let dispatcher = IHelloStarknetExtendedDispatcher { contract_address };
+
+    assert(dispatcher.get_balance() == 0, 'Wrong balance');
+
+    interact_with_state(
+        contract_address,
+        || {
+            let mut state = HelloStarknetExtended::contract_state_for_testing();
+            state.balance.write(987);
+        },
+    );
+
+    assert(dispatcher.get_balance() == 987, 'Wrong balance');
+    dispatcher.increase_balance(13);
+    assert(dispatcher.get_balance() == 1000, 'Wrong balance');
+}
+
+#[test]
+fn test_interact_with_state_return() {
+    let contract_address = deploy_contract("HelloStarknetExtended", array!['Name']);
+    let dispatcher = IHelloStarknetExtendedDispatcher { contract_address };
+
+    assert(dispatcher.get_balance() == 0, 'Wrong balance');
+
+    let res = interact_with_state(
+        contract_address,
+        || -> u256 {
+            let mut state = HelloStarknetExtended::contract_state_for_testing();
+            state.balance.write(111);
+            state.balance.read()
+        },
+    );
+
+    assert(res == 111, 'Wrong balance');
+}
+
+#[test]
+fn test_interact_with_initialized_state() {
+    let contract_address = deploy_contract("HelloStarknetExtended", array!['Name']);
+    let dispatcher = IHelloStarknetExtendedDispatcher { contract_address };
+
+    dispatcher.increase_balance(199);
+
+    interact_with_state(
+        contract_address,
+        || {
+            let mut state = HelloStarknetExtended::contract_state_for_testing();
+            assert(state.balance.read() == 199, 'Wrong balance');
+            state.balance.write(1);
+        },
+    );
+
+    assert(dispatcher.get_balance() == 1, 'Wrong balance');
+}
+
+#[test]
+fn test_interact_with_state_vec() {
+    let contract_address = deploy_contract("HelloStarknetExtended", array!['Name']);
+    let dispatcher = IHelloStarknetExtendedDispatcher { contract_address };
+
+    dispatcher.increase_balance(1);
+    dispatcher.increase_balance(1);
+    dispatcher.increase_balance(1);
+
+    interact_with_state(
+        contract_address,
+        || {
+            let mut state = HelloStarknetExtended::contract_state_for_testing();
+            assert(state.balance_records.len() == 4, 'Wrong length');
+            state.balance_records.push(10);
+        },
+    );
+
+    assert(dispatcher.get_balance_at(0) == 0, 'Wrong balance');
+    assert(dispatcher.get_balance_at(2) == 2, 'Wrong balance');
+    assert(dispatcher.get_balance_at(4) == 10, 'Wrong balance');
+}
+
+#[test]
+fn test_interact_with_state_map() {
+    let contract_address = deploy_contract("HelloStarknetExtended", array!['Name']);
+    let dispatcher = IHelloStarknetExtendedDispatcher { contract_address };
+
+    dispatcher.increase_balance(1);
+
+    interact_with_state(
+        contract_address,
+        || {
+            let mut state = HelloStarknetExtended::contract_state_for_testing();
+            state.callers.write(0x123.try_into().unwrap(), 1000);
+            state.callers.write(0x321.try_into().unwrap(), 2000);
+        },
+    );
+
+    assert(
+        dispatcher.get_caller_info(0x123.try_into().unwrap()) == 1000,
+        'Wrong data for address 0x123',
+    );
+    assert(
+        dispatcher.get_caller_info(0x321.try_into().unwrap()) == 2000,
+        'Wrong data for address 0x321',
+    );
+    assert(
+        dispatcher.get_caller_info(0x12345.try_into().unwrap()) == 0,
+        'Wrong data for address 0x12345',
+    );
+}
+
+#[test]
+fn test_interact_with_state_internal_function() {
+    let contract_address = deploy_contract("HelloStarknetExtended", array!['Name']);
+
+    let get_owner = 
+        || -> (
+            ContractAddress, felt252,
+        ) {
+            interact_with_state(
+                contract_address,
+                || -> (
+                    ContractAddress, felt252,
+                ) {
+                    let mut state = HelloStarknetExtended::contract_state_for_testing();
+                    (state.owner.address.read(), state.owner.name.read())
+                },
+            )
+        };
+    let (owner_address, owner_name) = get_owner();
+    assert(owner_address == 0.try_into().unwrap(), 'Incorrect owner address');
+    assert(owner_name == 'Name', 'Incorrect owner name');
+
+    interact_with_state(
+        contract_address,
+        || {
+            let mut state = HelloStarknetExtended::contract_state_for_testing();
+            state._set_owner(0x777.try_into().unwrap(), 'New name');
+        },
+    );
+    let (owner_address, owner_name) = get_owner();
+
+    assert(owner_address == 0x777.try_into().unwrap(), 'Incorrect owner address');
+    assert(owner_name == 'New name', 'Incorrect owner name');
+}