feat: mock_call with dynamic return data

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Rust is no longer required to use `snforge` if using Scarb >= 2.10.0 on supported platforms - precompiled `snforge_scarb_plugin` plugin binaries are now published to [package registry](https://scarbs.xyz) for new versions. 
 - Added a suggestion for using the `--max-n-steps` flag when the Cairo VM returns the error: `Could not reach the end of the program. RunResources has no remaining steps`.
+- `mock_call_when`, `start_mock_call_when`, `stop_mock_call_when` cheatcodes.
 
 #### Fixed
 

crates/cheatnet/Cargo.toml

@@ -13,6 +13,7 @@ bimap.workspace = true
 camino.workspace = true
 starknet_api.workspace = true
 starknet-types-core.workspace = true
+starknet-crypto.workspace = true
 cairo-lang-casm.workspace = true
 cairo-lang-runner.workspace = true
 cairo-lang-utils.workspace = true

crates/cheatnet/src/runtime_extensions/call_to_blockifier_runtime_extension/execution/entry_point.rs

@@ -2,7 +2,7 @@ use std::cell::RefCell;
 use super::cairo1_execution::execute_entry_point_call_cairo1;
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::execution::deprecated::cairo0_execution::execute_entry_point_call_cairo0;
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::CheatnetState;
-use crate::state::{CallTrace, CallTraceNode, CheatStatus, EncounteredError};
+use crate::state::{CallTrace, CallTraceNode, CheatSpan, CheatStatus, EncounteredError};
 use blockifier::execution::call_info::{CallExecution, Retdata};
 use blockifier::{
     execution::{
@@ -18,11 +18,13 @@ use blockifier::{
     state::state_api::State,
 };
 use cairo_vm::vm::runners::cairo_runner::{CairoRunner, ExecutionResources};
+use num_traits::Zero;
 use starknet_api::{
     core::ClassHash,
     deprecated_contract_class::EntryPointType,
     transaction::{Calldata, TransactionVersion},
 };
+use starknet_crypto::poseidon_hash_many;
 use std::collections::HashSet;
 use std::rc::Rc;
 use blockifier::execution::deprecated_syscalls::hint_processor::SyscallCounter;
@@ -268,11 +270,25 @@ fn get_mocked_function_cheat_status<'a>(
     if call.call_type == CallType::Delegate {
         return None;
     }
-
-    cheatnet_state
+    match cheatnet_state
         .mocked_functions
         .get_mut(&call.storage_address)
-        .and_then(|contract_functions| contract_functions.get_mut(&call.entry_point_selector))
+    {
+        None => None,
+        Some(contract_functions) => {
+            let calldata_hash = poseidon_hash_many(call.calldata.0.iter());
+            let key = (call.entry_point_selector, calldata_hash);
+            let key_zero = (call.entry_point_selector, Felt::zero());
+
+            match contract_functions.get(&key) {
+                Some(CheatStatus::Cheated(_, CheatSpan::TargetCalls(0))) => {
+                    contract_functions.get_mut(&key_zero)
+                }
+                Some(CheatStatus::Cheated(_, _)) => contract_functions.get_mut(&key),
+                _ => contract_functions.get_mut(&key_zero),
+            }
+        }
+    }
 }
 
 fn mocked_call_info(call: CallEntryPoint, ret_data: Vec<Felt>) -> CallInfo {

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

@@ -1,6 +1,8 @@
-use crate::state::{CheatSpan, CheatStatus};
+use crate::state::{CheatSpan, CheatStatus, MockCalldata};
 use crate::CheatnetState;
+use num_traits::Zero;
 use starknet_api::core::{ContractAddress, EntryPointSelector};
+use starknet_crypto::poseidon_hash_many;
 use starknet_types_core::felt::Felt;
 use std::collections::hash_map::Entry;
 
@@ -9,26 +11,30 @@ impl CheatnetState {
         &mut self,
         contract_address: ContractAddress,
         function_selector: EntryPointSelector,
+        calldata: MockCalldata,
         ret_data: &[Felt],
         span: CheatSpan,
     ) {
         let contract_mocked_functions = self.mocked_functions.entry(contract_address).or_default();
-
-        contract_mocked_functions.insert(
-            function_selector,
-            CheatStatus::Cheated(ret_data.to_vec(), span),
-        );
+        let calldata_hash = match calldata {
+            MockCalldata::Values(data) => poseidon_hash_many(data.iter()),
+            MockCalldata::Any => Felt::zero(),
+        };
+        let key = (function_selector, calldata_hash);
+        contract_mocked_functions.insert(key, CheatStatus::Cheated(ret_data.to_vec(), span));
     }
 
     pub fn start_mock_call(
         &mut self,
         contract_address: ContractAddress,
         function_selector: EntryPointSelector,
+        calldata: MockCalldata,
         ret_data: &[Felt],
     ) {
         self.mock_call(
             contract_address,
             function_selector,
+            calldata,
             ret_data,
             CheatSpan::Indefinite,
         );
@@ -38,10 +44,15 @@ impl CheatnetState {
         &mut self,
         contract_address: ContractAddress,
         function_selector: EntryPointSelector,
+        calldata: MockCalldata,
     ) {
         if let Entry::Occupied(mut e) = self.mocked_functions.entry(contract_address) {
             let contract_mocked_functions = e.get_mut();
-            contract_mocked_functions.remove(&function_selector);
+            let calldata_hash = match calldata {
+                MockCalldata::Values(data) => poseidon_hash_many(data.iter()),
+                MockCalldata::Any => Felt::zero(),
+            };
+            contract_mocked_functions.remove(&(function_selector, calldata_hash));
         }
     }
 }

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

@@ -87,26 +87,31 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
             "mock_call" => {
                 let contract_address = input_reader.read()?;
                 let function_selector = input_reader.read()?;
+                let calldata = input_reader.read()?;
                 let span = input_reader.read()?;
-
                 let ret_data: Vec<_> = input_reader.read()?;
-
                 extended_runtime
                     .extended_runtime
                     .extension
                     .cheatnet_state
-                    .mock_call(contract_address, function_selector, &ret_data, span);
+                    .mock_call(
+                        contract_address,
+                        function_selector,
+                        calldata,
+                        &ret_data,
+                        span,
+                    );
                 Ok(CheatcodeHandlingResult::from_serializable(()))
             }
             "stop_mock_call" => {
                 let contract_address = input_reader.read()?;
                 let function_selector = input_reader.read()?;
-
+                let calldata = input_reader.read()?;
                 extended_runtime
                     .extended_runtime
                     .extension
                     .cheatnet_state
-                    .stop_mock_call(contract_address, function_selector);
+                    .stop_mock_call(contract_address, function_selector, calldata);
                 Ok(CheatcodeHandlingResult::from_serializable(()))
             }
             "replace_bytecode" => {

crates/cheatnet/src/state.rs

@@ -41,6 +41,12 @@ pub enum CheatSpan {
     TargetCalls(usize),
 }
 
+#[derive(CairoDeserialize, Clone, Debug, PartialEq, Eq)]
+pub enum MockCalldata {
+    Any,
+    Values(Vec<Felt>),
+}
+
 #[derive(Debug)]
 pub struct ExtendedStateReader {
     pub dict_state_reader: DictStateReader,
@@ -326,12 +332,13 @@ pub struct EncounteredError {
     pub class_hash: ClassHash,
 }
 
+type MockedFunctionKey = (EntryPointSelector, Felt);
 pub struct CheatnetState {
     pub cheated_execution_info_contracts: HashMap<ContractAddress, ExecutionInfoMock>,
     pub global_cheated_execution_info: ExecutionInfoMock,
 
     pub mocked_functions:
-        HashMap<ContractAddress, HashMap<EntryPointSelector, CheatStatus<Vec<Felt>>>>,
+        HashMap<ContractAddress, HashMap<MockedFunctionKey, CheatStatus<Vec<Felt>>>>,
     pub replaced_bytecode_contracts: HashMap<ContractAddress, ClassHash>,
     pub detected_events: Vec<Event>,
     pub detected_messages_to_l1: Vec<MessageToL1>,

crates/cheatnet/tests/cheatcodes/mock_call.rs

@@ -8,7 +8,7 @@ use crate::{
     common::{deploy_contract, get_contracts},
 };
 use cheatnet::runtime_extensions::forge_runtime_extension::cheatcodes::declare::declare;
-use cheatnet::state::{CheatSpan, CheatnetState};
+use cheatnet::state::{CheatSpan, CheatnetState, MockCalldata};
 use conversions::IntoConv;
 use starknet::core::utils::get_selector_from_name;
 use starknet_api::core::ContractAddress;
@@ -38,15 +38,19 @@ impl MockCallTrait for TestEnvironment {
         self.cheatnet_state.mock_call(
             *contract_address,
             function_selector.into_(),
+            MockCalldata::Any,
             &ret_data,
             span,
         );
     }
 
     fn stop_mock_call(&mut self, contract_address: &ContractAddress, function_name: &str) {
         let function_selector = get_selector_from_name(function_name).unwrap();
-        self.cheatnet_state
-            .stop_mock_call(*contract_address, function_selector.into_());
+        self.cheatnet_state.stop_mock_call(
+            *contract_address,
+            function_selector.into_(),
+            MockCalldata::Any,
+        );
     }
 }
 
@@ -68,6 +72,7 @@ fn mock_call_simple() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        MockCalldata::Any,
         &ret_data,
     );
 
@@ -100,6 +105,7 @@ fn mock_call_stop() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        MockCalldata::Any,
         &ret_data,
     );
 
@@ -113,7 +119,11 @@ fn mock_call_stop() {
 
     assert_success(output, &ret_data);
 
-    cheatnet_state.stop_mock_call(contract_address, felt_selector_from_name("get_thing"));
+    cheatnet_state.stop_mock_call(
+        contract_address,
+        felt_selector_from_name("get_thing"),
+        MockCalldata::Any,
+    );
 
     let output = call_contract(
         &mut cached_state,
@@ -140,7 +150,11 @@ fn mock_call_stop_no_start() {
 
     let selector = felt_selector_from_name("get_thing");
 
-    cheatnet_state.stop_mock_call(contract_address, felt_selector_from_name("get_thing"));
+    cheatnet_state.stop_mock_call(
+        contract_address,
+        felt_selector_from_name("get_thing"),
+        MockCalldata::Any,
+    );
 
     let output = call_contract(
         &mut cached_state,
@@ -168,10 +182,10 @@ fn mock_call_double() {
     let selector = felt_selector_from_name("get_thing");
 
     let ret_data = [Felt::from(123)];
-    cheatnet_state.start_mock_call(contract_address, selector, &ret_data);
+    cheatnet_state.start_mock_call(contract_address, selector, MockCalldata::Any, &ret_data);
 
     let ret_data = [Felt::from(999)];
-    cheatnet_state.start_mock_call(contract_address, selector, &ret_data);
+    cheatnet_state.start_mock_call(contract_address, selector, MockCalldata::Any, &ret_data);
 
     let output = call_contract(
         &mut cached_state,
@@ -183,7 +197,7 @@ fn mock_call_double() {
 
     assert_success(output, &ret_data);
 
-    cheatnet_state.stop_mock_call(contract_address, selector);
+    cheatnet_state.stop_mock_call(contract_address, selector, MockCalldata::Any);
 
     let output = call_contract(
         &mut cached_state,
@@ -214,6 +228,7 @@ fn mock_call_double_call() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        MockCalldata::Any,
         &ret_data,
     );
 
@@ -255,6 +270,7 @@ fn mock_call_proxy() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        MockCalldata::Any,
         &ret_data,
     );
 
@@ -303,6 +319,7 @@ fn mock_call_proxy_with_other_syscall() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        MockCalldata::Any,
         &ret_data,
     );
 
@@ -352,6 +369,7 @@ fn mock_call_inner_call_no_effect() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        MockCalldata::Any,
         &ret_data,
     );
 
@@ -407,6 +425,7 @@ fn mock_call_library_call_no_effect() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_constant_thing"),
+        MockCalldata::Any,
         &ret_data,
     );
 
@@ -440,6 +459,7 @@ fn mock_call_before_deployment() {
     cheatnet_state.start_mock_call(
         precalculated_address,
         felt_selector_from_name("get_thing"),
+        MockCalldata::Any,
         &ret_data,
     );
 
@@ -482,6 +502,7 @@ fn mock_call_not_implemented() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing_not_implemented"),
+        MockCalldata::Any,
         &ret_data,
     );
 
@@ -512,6 +533,7 @@ fn mock_call_in_constructor() {
     cheatnet_state.start_mock_call(
         balance_contract_address,
         felt_selector_from_name("get_balance"),
+        MockCalldata::Any,
         &ret_data,
     );
 
@@ -559,12 +581,14 @@ fn mock_call_two_methods() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        MockCalldata::Any,
         &ret_data,
     );
 
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_constant_thing"),
+        MockCalldata::Any,
         &ret_data,
     );
 
@@ -602,6 +626,7 @@ fn mock_call_nonexisting_contract() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        MockCalldata::Any,
         &ret_data,
     );