add prank* and mock* cheatcodes family support

Cargo.toml

@@ -266,7 +266,7 @@ alloy-sol-types = "1.2.1"
 
 alloy-chains = "0.2"
 alloy-rlp = "0.3"
-alloy-trie = "0.9"
+alloy-trie = "0.9.1"
 
 ## op-alloy
 op-alloy-consensus = "0.17.2"

crates/anvil-polkadot/Cargo.toml

@@ -23,8 +23,8 @@ codec = { version = "3.7.5", default-features = true, package = "parity-scale-co
 substrate-runtime = { path = "substrate-runtime" }
 secp256k1 = { version = "0.28.0", default-features = false }
 libsecp256k1 = { version = "0.7.0", default-features = false }
-sp-runtime-interface = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "master", default-features = false }
-polkadot-sdk = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "master", default-features = false, features = [
+sp-runtime-interface = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/inject_state", default-features = false }
+polkadot-sdk = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/inject_state", default-features = false, features = [
     "sc-allocator",
 	"sc-basic-authorship",
 	"sc-block-builder",
@@ -64,7 +64,6 @@ polkadot-sdk = { git = "https://github.com/paritytech/polkadot-sdk.git", branch
 	"substrate-frame-rpc-system",
 	"substrate-rpc-client",
 	"substrate-wasm-builder",
-	"pallet-revive-eth-rpc"
 ] }
 anvil.workspace = true
 anvil-core.workspace = true

crates/anvil-polkadot/substrate-runtime/Cargo.toml

@@ -12,7 +12,7 @@ license.workspace = true
 [dependencies]
 array-bytes = { version = "6.2.2", default-features = false }
 codec = { version = "3.7.5", default-features = false, package = "parity-scale-codec" }
-polkadot-sdk = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "master", default-features = false, features = [
+polkadot-sdk = {git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/inject_state", default-features = false, features = [
 	"pallet-balances",
 	"pallet-revive",
 	"pallet-sudo",
@@ -27,7 +27,7 @@ scale-info = { version = "2.11.6", default-features = false }
 serde_json = { version = "1.0", default-features = false, features = ["alloc"] }
 
 [build-dependencies]
-polkadot-sdk = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "master", default-features = false, optional = true, features = ["substrate-wasm-builder"] }
+polkadot-sdk = {git = "https://github.com/paritytech/polkadot-sdk.git", branch = "alexggh/inject_state", default-features = false, optional = true, features = ["substrate-wasm-builder"] }
 
 [features]
 default = ["std"]

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

@@ -926,6 +926,11 @@ impl Cheatcodes {
         }
     }
 
+    // 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,
@@ -1009,40 +1014,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

@@ -913,8 +913,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(|| {

crates/forge/tests/it/revive/cheat_mock_call.rs

@@ -0,0 +1,171 @@
+use crate::{config::*, test_helpers::TEST_DATA_REVIVE};
+use foundry_test_utils::Filter;
+use revm::primitives::hardfork::SpecId;
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_getters() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockGetters", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_clear_mock_reverted_calls() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testClearMockRevertedCalls", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_call_empty_account_revert() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallEmptyAccountRevert", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_call_revert() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallResetsMockCallRevert", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_call_reverts_partial_match() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallRevertPartialMatch", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_call_revert_resets_mock_call() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallRevertResetsMockCall", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_call_revert_with_call() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallRevertWithCall", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_call_revert_with_value() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallRevertWithValue", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_call_data_revert() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCalldataRevert", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_getters_revert() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockGettersRevert", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_nested_revert() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockNestedRevert", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_revert_with_custom_error() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockRevertWithCustomError", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_clear_mocked_calls() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testClearMockedCalls", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_call_empty_account() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallEmptyAccount", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_call_multiple_partial_match() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallMultiplePartialMatch", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_call_with_value() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallWithValue", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_call_with_value_calldata_precedence() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallWithValueCalldataPrecedence", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_calldata() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCalldata", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_nested() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockNested", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_nested_delegate() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockNestedDelegate", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_selector() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockSelector", "MockCall", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}

crates/forge/tests/it/revive/cheat_mock_calls.rs

@@ -0,0 +1,27 @@
+use crate::{config::*, test_helpers::TEST_DATA_REVIVE};
+use foundry_test_utils::Filter;
+use revm::primitives::hardfork::SpecId;
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_calls() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCalls", "MockCalls", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_calls_last_should_persist() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallsLastShouldPersist", "MockCalls", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_mock_calls_with_value() {
+    let runner = TEST_DATA_REVIVE.runner_revive();
+    let filter = Filter::new("testMockCallsWithValue", "MockCalls", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}