feat(cheatnet): add "dynamic" return value fir mock_call

Author: ptisserand

Cargo.lock

@@ -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/Cargo.toml

@@ -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,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 call_data_hash = poseidon_hash_many(call.calldata.0.iter());
+            let key = (call.entry_point_selector, call_data_hash);
+            if contract_functions.contains_key(&key) {
+                contract_functions.get_mut(&key)
+            } else {
+                let key_zero = (call.entry_point_selector, Felt::zero());
+                contract_functions.get_mut(&key_zero)
+            }
+        }
+    }
 }
 
 fn mocked_call_info(call: CallEntryPoint, ret_data: Vec<Felt>) -> CallInfo {

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

@@ -1,6 +1,8 @@
 use crate::state::{CheatSpan, CheatStatus};
 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,
+        call_data: Option<Vec<Felt>>,
         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 call_data_hash = match call_data {
+            Some(data) => poseidon_hash_many(data.iter()),
+            None => Felt::zero(),
+        };
+        let key = (function_selector, call_data_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,
+        call_data: Option<Vec<Felt>>,
         ret_data: &[Felt],
     ) {
         self.mock_call(
             contract_address,
             function_selector,
+            call_data,
             ret_data,
             CheatSpan::Indefinite,
         );
@@ -38,10 +44,15 @@ impl CheatnetState {
         &mut self,
         contract_address: ContractAddress,
         function_selector: EntryPointSelector,
+        call_data: Option<Vec<Felt>>,
     ) {
         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 call_data_hash = match call_data {
+                Some(data) => poseidon_hash_many(data.iter()),
+                None => Felt::zero(),
+            };
+            contract_mocked_functions.remove(&(function_selector, call_data_hash));
         }
     }
 }

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/cheatcodes/mock_call.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 call_data = 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,
+                        call_data,
+                        &ret_data,
+                        span,
+                    );
                 Ok(CheatcodeHandlingResult::from_serializable(()))
             }
             "stop_mock_call" => {
                 let contract_address = input_reader.read()?;
                 let function_selector = input_reader.read()?;
-
+                let call_data = input_reader.read()?;
                 extended_runtime
                     .extended_runtime
                     .extension
                     .cheatnet_state
-                    .stop_mock_call(contract_address, function_selector);
+                    .stop_mock_call(contract_address, function_selector, call_data);
                 Ok(CheatcodeHandlingResult::from_serializable(()))
             }
             "replace_bytecode" => {

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

@@ -331,7 +331,7 @@ pub struct CheatnetState {
     pub global_cheated_execution_info: ExecutionInfoMock,
 
     pub mocked_functions:
-        HashMap<ContractAddress, HashMap<EntryPointSelector, CheatStatus<Vec<Felt>>>>,
+        HashMap<ContractAddress, HashMap<(EntryPointSelector, Felt), CheatStatus<Vec<Felt>>>>,
     pub replaced_bytecode_contracts: HashMap<ContractAddress, ClassHash>,
     pub detected_events: Vec<Event>,
     pub detected_messages_to_l1: Vec<MessageToL1>,

crates/cheatnet/src/state.rs

@@ -68,6 +68,7 @@ fn mock_call_simple() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        None,
         &ret_data,
     );
 
@@ -100,6 +101,7 @@ fn mock_call_stop() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        None,
         &ret_data,
     );
 
@@ -168,10 +170,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, None, &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, None, &ret_data);
 
     let output = call_contract(
         &mut cached_state,
@@ -214,6 +216,7 @@ fn mock_call_double_call() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        None,
         &ret_data,
     );
 
@@ -255,6 +258,7 @@ fn mock_call_proxy() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        None,
         &ret_data,
     );
 
@@ -303,6 +307,7 @@ fn mock_call_proxy_with_other_syscall() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        None,
         &ret_data,
     );
 
@@ -352,6 +357,7 @@ fn mock_call_inner_call_no_effect() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        None,
         &ret_data,
     );
 
@@ -407,6 +413,7 @@ fn mock_call_library_call_no_effect() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_constant_thing"),
+        None,
         &ret_data,
     );
 
@@ -440,6 +447,7 @@ fn mock_call_before_deployment() {
     cheatnet_state.start_mock_call(
         precalculated_address,
         felt_selector_from_name("get_thing"),
+        None,
         &ret_data,
     );
 
@@ -482,6 +490,7 @@ fn mock_call_not_implemented() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing_not_implemented"),
+        None,
         &ret_data,
     );
 
@@ -512,6 +521,7 @@ fn mock_call_in_constructor() {
     cheatnet_state.start_mock_call(
         balance_contract_address,
         felt_selector_from_name("get_balance"),
+        None,
         &ret_data,
     );
 
@@ -559,12 +569,14 @@ fn mock_call_two_methods() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        None,
         &ret_data,
     );
 
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_constant_thing"),
+        None,
         &ret_data,
     );
 
@@ -602,6 +614,7 @@ fn mock_call_nonexisting_contract() {
     cheatnet_state.start_mock_call(
         contract_address,
         felt_selector_from_name("get_thing"),
+        None,
         &ret_data,
     );