Add interact_with_state closure cheatcode

Towards #3331

commit-id:0b8b8c6b

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,18 @@
+#[starknet::contract]
+pub mod StorageNodeContract {
+    use starknet::ContractAddress;
+    use starknet::storage::Map;
+
+    #[starknet::storage_node]
+    struct RandomData {
+        title: felt252,
+        description: felt252,
+        counter: u32,
+        data: Map<(ContractAddress, u16), ByteArray>,
+    }
+
+    #[storage]
+    struct Storage {
+        pub random_data: RandomData,
+    }
+}

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

@@ -0,0 +1,164 @@
+use contract_state::balance::HelloStarknetExtended::InternalFunctionsTrait;
+use contract_state::balance::{
+    HelloStarknetExtended, IHelloStarknetExtendedDispatcher, IHelloStarknetExtendedDispatcherTrait,
+};
+use snforge_std::{ContractClassTrait, DeclareResultTrait, declare, interact_with_state};
+use starknet::ContractAddress;
+use starknet::storage::{
+    MutableVecTrait, StorageMapWriteAccess, StoragePointerReadAccess, StoragePointerWriteAccess,
+};
+
+fn deploy_contract(name: ByteArray) -> ContractAddress {
+    let contract = declare(name).unwrap().contract_class();
+    let mut calldata = array![];
+    calldata.append('Name');
+
+    let (contract_address, _) = contract.deploy(@calldata).unwrap();
+    contract_address
+}
+
+#[test]
+fn test_interact_with_state() {
+    let contract_address = deploy_contract("HelloStarknetExtended");
+    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");
+    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");
+    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");
+    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");
+    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");
+
+    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');
+}

crates/forge/tests/data/contracts/hello_starknet_extended.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);
+        }
+    }
+}