feat: mock_call with dynamic return data

CHANGELOG.md

@@ -324,6 +324,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-utils.workspace = true
 cairo-lang-starknet-classes.workspace = true

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

@@ -31,12 +31,14 @@ use cairo_vm::vm::runners::cairo_runner::{CairoRunner, ExecutionResources};
 use cairo_vm::vm::trace::trace_entry::RelocatedTraceEntry;
 use conversions::FromConv;
 use conversions::string::TryFromHexStr;
+use num_traits::Zero;
 use shared::vm::VirtualMachineExt;
 use starknet_api::{
     contract_class::EntryPointType,
     core::ClassHash,
     transaction::{TransactionVersion, fields::Calldata},
 };
+use starknet_crypto::poseidon_hash_many;
 use starknet_types_core::felt::Felt;
 use std::collections::{HashMap, HashSet};
 use thiserror::Error;
@@ -415,11 +417,22 @@ 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(_, _)) => contract_functions.get_mut(&key),
+                _ => contract_functions.get_mut(&key_zero),
+            }
+        }
+    }
 }
 
 fn mocked_call_info(

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

@@ -1,6 +1,8 @@
 use crate::CheatnetState;
-use crate::state::{CheatSpan, CheatStatus};
+use crate::state::{CheatSpan, CheatStatus, MockCalldata};
+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

@@ -95,26 +95,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

@@ -47,6 +47,12 @@ pub enum CheatSpan {
     TargetCalls(NonZeroUsize),
 }
 
+#[derive(CairoDeserialize, Clone, Debug, PartialEq, Eq)]
+pub enum MockCalldata {
+    Any,
+    Values(Vec<Felt>),
+}
+
 #[derive(Debug)]
 pub struct ExtendedStateReader {
     pub dict_state_reader: DictStateReader,
@@ -369,12 +375,13 @@ pub struct TraceData {
     pub is_vm_trace_needed: bool,
 }
 
+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>,