add prank* and mock* cheatcodes family support

crates/cheatcodes/src/evm.rs

@@ -39,7 +39,7 @@ use serde::Serialize;
 
 mod fork;
 pub(crate) mod mapping;
-pub(crate) mod mock;
+pub mod mock;
 pub(crate) mod prank;
 
 /// Records storage slots reads and writes.

crates/cheatcodes/src/inspector.rs

@@ -939,6 +939,11 @@ impl Cheatcodes {
         self.strategy.runner.revive_remove_duplicate_account_access(self);
     }
 
+    // Tells whether PVM is enabled or not.
+    pub fn is_pvm_enabled(&mut self) -> bool {
+        self.strategy.runner.is_pvm_enabled(self)
+    }
+
     pub fn call_with_executor(
         &mut self,
         ecx: Ecx,
@@ -1022,40 +1027,51 @@ impl Cheatcodes {
             }
         }
 
-        // Handle mocked calls
-        if let Some(mocks) = self.mocked_calls.get_mut(&call.bytecode_address) {
-            let ctx = MockCallDataContext {
-                calldata: call.input.bytes(ecx),
-                value: call.transfer_value(),
-            };
+        // Do not handle mocked calls if PVM is enabled and let the revive call handle it.
+        // There is literally no problem with handling mocked calls with PVM enabled here as well,
+        // but the downside is that if call a mocked call from the test it will not exercise the
+        // paths in revive that handle mocked calls and only nested mocks will be handle by the
+        // revive specific calls.
+        // This is undesirable because conformity tests could accidentally pass and the revive code
+        // paths be broken.
+        if !self.is_pvm_enabled() {
+            // Handle mocked calls
+            if let Some(mocks) = self.mocked_calls.get_mut(&call.bytecode_address) {
+                let ctx = MockCallDataContext {
+                    calldata: call.input.bytes(ecx),
+                    value: call.transfer_value(),
+                };
 
-            if let Some(return_data_queue) = match mocks.get_mut(&ctx) {
-                Some(queue) => Some(queue),
-                None => mocks
-                    .iter_mut()
-                    .find(|(mock, _)| {
-                        call.input.bytes(ecx).get(..mock.calldata.len()) == Some(&mock.calldata[..])
-                            && mock.value.is_none_or(|value| Some(value) == call.transfer_value())
-                    })
-                    .map(|(_, v)| v),
-            } && let Some(return_data) = if return_data_queue.len() == 1 {
-                // If the mocked calls stack has a single element in it, don't empty it
-                return_data_queue.front().map(|x| x.to_owned())
-            } else {
-                // Else, we pop the front element
-                return_data_queue.pop_front()
-            } {
-                return Some(CallOutcome {
-                    result: InterpreterResult {
-                        result: return_data.ret_type,
-                        output: return_data.data,
-                        gas,
-                    },
-                    memory_offset: call.return_memory_offset.clone(),
-                });
+                if let Some(return_data_queue) = match mocks.get_mut(&ctx) {
+                    Some(queue) => Some(queue),
+                    None => mocks
+                        .iter_mut()
+                        .find(|(mock, _)| {
+                            call.input.bytes(ecx).get(..mock.calldata.len())
+                                == Some(&mock.calldata[..])
+                                && mock
+                                    .value
+                                    .is_none_or(|value| Some(value) == call.transfer_value())
+                        })
+                        .map(|(_, v)| v),
+                } && let Some(return_data) = if return_data_queue.len() == 1 {
+                    // If the mocked calls stack has a single element in it, don't empty it
+                    return_data_queue.front().map(|x| x.to_owned())
+                } else {
+                    // Else, we pop the front element
+                    return_data_queue.pop_front()
+                } {
+                    return Some(CallOutcome {
+                        result: InterpreterResult {
+                            result: return_data.ret_type,
+                            output: return_data.data,
+                            gas,
+                        },
+                        memory_offset: call.return_memory_offset.clone(),
+                    });
+                }
             }
         }
-
         // Apply our prank
         if let Some(prank) = &self.get_prank(curr_depth) {
             // Apply delegate call, `call.caller`` will not equal `prank.prank_caller`

crates/cheatcodes/src/lib.rs

@@ -57,6 +57,7 @@ mod script;
 pub use script::{Broadcast, Wallets, WalletsInner};
 
 mod strategy;
+pub use evm::mock::{MockCallDataContext, MockCallReturnData};
 pub use strategy::{
     CheatcodeInspectorStrategy, CheatcodeInspectorStrategyContext, CheatcodeInspectorStrategyExt,
     CheatcodeInspectorStrategyRunner, CheatcodesStrategy, EvmCheatcodeInspectorStrategyRunner,

crates/cheatcodes/src/strategy.rs

@@ -238,6 +238,10 @@ impl Clone for CheatcodeInspectorStrategy {
 
 /// Defined in revive-strategy
 pub trait CheatcodeInspectorStrategyExt {
+    fn is_pvm_enabled(&self, _state: &mut crate::Cheatcodes) -> bool {
+        false
+    }
+
     fn revive_try_create(
         &self,
         _state: &mut crate::Cheatcodes,

crates/evm/evm/src/inspectors/stack.rs

@@ -914,8 +914,19 @@ impl Inspector<EthEvmContext<&mut dyn DatabaseExt>> for InspectorStackRefMut<'_>
         );
 
         if let Some(cheatcodes) = self.cheatcodes.as_deref_mut() {
+            let is_pvm_enabled = cheatcodes.is_pvm_enabled();
             // Handle mocked functions, replace bytecode address with mock if matched.
-            if let Some(mocks) = cheatcodes.mocked_functions.get(&call.target_address) {
+
+            // Do not handle mocked functions if PVM is enabled and let the revive call handle it.
+            // There is literally no problem with handling mocked functions with PVM enabled here as
+            // well, but the downside is that if we call a mocked functions from the test it
+            // will not exercise the paths in revive that handle mocked calls and only
+            // nested mocks will be handle by the revive specific calls.
+            // This is undesirable because conformity tests could accidentally pass and the revive
+            // code paths be broken.
+            if let Some(mocks) = cheatcodes.mocked_functions.get(&call.target_address)
+                && !is_pvm_enabled
+            {
                 // Check if any mock function set for call data or if catch-all mock function set
                 // for selector.
                 if let Some(target) = mocks.get(&call.input.bytes(ecx)).or_else(|| {