Add polkadotskip cheatcode (#448)

Author: smiasojed

crates/cheatcodes/assets/cheatcodes.json

@@ -2906,6 +2906,12 @@ interface Vm {
     #[cheatcode(group = Utilities)]
     function pvm(bool enabled) external;
 
+    /// When running in PVM context, skips the next CREATE or CALL, executing it on the EVM instead.
+    /// All `CREATE`s executed within this skip, will automatically have `CALL`s to their target addresses
+    /// executed in the EVM, and need not be marked with this cheatcode at every usage location.
+    #[cheatcode(group = Testing, safety = Safe)]
+    function polkadotSkip() external pure;
+
     /// Generates the hash of the canonical EIP-712 type representation.
     ///
     /// Supports 2 different inputs:

crates/cheatcodes/spec/src/vm.rs

@@ -950,6 +950,14 @@ impl Cheatcode for pvmCall {
     }
 }
 
+impl Cheatcode for polkadotSkipCall {
+    fn apply_stateful(&self, _ccx: &mut CheatsCtxt) -> Result {
+        // Does nothing by default.
+        // PVM-related logic is implemented in the corresponding strategy object.
+        Ok(Default::default())
+    }
+}
+
 pub(super) fn get_nonce(ccx: &mut CheatsCtxt, address: &Address) -> Result {
     let account = ccx.ecx.journaled_state.load_account(*address)?;
     Ok(account.info.nonce.abi_encode())

crates/cheatcodes/src/evm.rs

@@ -937,6 +937,7 @@ impl Cheatcodes {
         }
 
         self.strategy.runner.revive_remove_duplicate_account_access(self);
+        self.strategy.runner.revive_record_create_address(self, outcome);
     }
 
     // Tells whether PVM is enabled or not.

crates/cheatcodes/src/inspector.rs

@@ -267,6 +267,14 @@ pub trait CheatcodeInspectorStrategyExt {
 
     // Remove duplicate accesses in storage_recorder
     fn revive_remove_duplicate_account_access(&self, _state: &mut crate::Cheatcodes) {}
+
+    // Record create address for skip_pvm_addresses tracking
+    fn revive_record_create_address(
+        &self,
+        _state: &mut crate::Cheatcodes,
+        _outcome: &revm::interpreter::CreateOutcome,
+    ) {
+    }
 }
 
 // Legacy type aliases for backward compatibility

crates/cheatcodes/src/strategy.rs

@@ -0,0 +1,15 @@
+use crate::{config::*, test_helpers::TEST_DATA_REVIVE};
+use foundry_test_utils::Filter;
+use revive_strategy::ReviveRuntimeMode;
+use revm::primitives::hardfork::SpecId;
+use rstest::rstest;
+
+#[rstest]
+#[case::evm(ReviveRuntimeMode::Evm)]
+#[tokio::test(flavor = "multi_thread")]
+async fn test_polkadot_skip(#[case] runtime_mode: ReviveRuntimeMode) {
+    let runner: forge::MultiContractRunner = TEST_DATA_REVIVE.runner_revive(runtime_mode);
+    let filter = Filter::new(".*", "PolkadotSkipTest", ".*/revive/PolkadotSkip.t.sol");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}

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

@@ -5,6 +5,7 @@ pub mod cheat_gas_metering;
 pub mod cheat_mock_call;
 pub mod cheat_mock_calls;
 pub mod cheat_mock_functions;
+pub mod cheat_polkadot_skip;
 pub mod cheat_prank;
 mod cheat_snapshot;
 pub mod cheat_store;

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

@@ -8,9 +8,9 @@ use foundry_cheatcodes::{
     CheatcodeInspectorStrategyContext, CheatcodeInspectorStrategyRunner, CheatsConfig, CheatsCtxt,
     CommonCreateInput, Ecx, EvmCheatcodeInspectorStrategyRunner, Result,
     Vm::{
-        chainIdCall, coinbaseCall, dealCall, etchCall, getNonce_0Call, loadCall, pvmCall,
-        resetNonceCall, revertToStateAndDeleteCall, revertToStateCall, rollCall, setNonceCall,
-        setNonceUnsafeCall, snapshotStateCall, storeCall, warpCall,
+        chainIdCall, coinbaseCall, dealCall, etchCall, getNonce_0Call, loadCall, polkadotSkipCall,
+        pvmCall, resetNonceCall, revertToStateAndDeleteCall, revertToStateCall, rollCall,
+        setNonceCall, setNonceUnsafeCall, snapshotStateCall, storeCall, warpCall,
     },
     journaled_account, precompile_error,
 };
@@ -42,6 +42,7 @@ use crate::{
     tracing::{Tracer, storage_tracer::AccountAccess},
 };
 use foundry_cheatcodes::Vm::{AccountAccess as FAccountAccess, ChainInfo};
+use polkadot_sdk::pallet_revive::tracing::Tracing;
 
 use revm::{
     bytecode::opcode as op,
@@ -113,6 +114,14 @@ impl PvmStartupMigration {
 pub struct PvmCheatcodeInspectorStrategyContext {
     /// Whether we're currently using pallet-revive (migrated from REVM)
     pub using_pvm: bool,
+    /// When in PVM context, execute the next CALL or CREATE in the EVM instead.
+    pub skip_pvm: bool,
+    /// Any contracts that were deployed in `skip_pvm` step.
+    /// This makes it easier to dispatch calls to any of these addresses in PVM context, directly
+    /// to EVM. Alternatively, we'd need to add `vm.polkadotSkip()` to these calls manually.
+    pub skip_pvm_addresses: std::collections::HashSet<Address>,
+    /// Records the next create address for `skip_pvm_addresses`.
+    pub record_next_create_address: bool,
     /// Controls automatic migration to pallet-revive
     pub pvm_startup_migration: PvmStartupMigration,
     pub dual_compiled_contracts: DualCompiledContracts,
@@ -131,6 +140,9 @@ impl PvmCheatcodeInspectorStrategyContext {
         Self {
             // Start in REVM mode by default
             using_pvm: false,
+            skip_pvm: false,
+            skip_pvm_addresses: Default::default(),
+            record_next_create_address: Default::default(),
             // Will be set to Allow when test contract deploys
             pvm_startup_migration: PvmStartupMigration::Defer,
             dual_compiled_contracts,
@@ -275,6 +287,12 @@ impl CheatcodeInspectorStrategyRunner for PvmCheatcodeInspectorStrategyRunner {
                 }
                 Ok(Default::default())
             }
+            t if is::<polkadotSkipCall>(t) => {
+                let polkadotSkipCall { .. } = cheatcode.as_any().downcast_ref().unwrap();
+                let ctx = get_context_ref_mut(ccx.state.strategy.context.as_mut());
+                ctx.skip_pvm = true;
+                Ok(Default::default())
+            }
             t if using_pvm && is::<dealCall>(t) => {
                 tracing::info!(cheatcode = ?cheatcode.as_debug() , using_pvm = ?using_pvm);
                 let dealCall { account, newBalance } = cheatcode.as_any().downcast_ref().unwrap();
@@ -783,6 +801,13 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
             return None;
         }
 
+        if ctx.skip_pvm {
+            ctx.skip_pvm = false; // handled the skip, reset flag
+            ctx.record_next_create_address = true;
+            tracing::info!("running create in EVM, instead of pallet-revive (skipped)");
+            return None;
+        }
+
         if let Some(CreateScheme::Create) = input.scheme() {
             let caller = input.caller();
             let nonce = ecx
@@ -832,9 +857,12 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
         let gas_price_pvm =
             sp_core::U256::from_little_endian(&U256::from(ecx.tx.gas_price).as_le_bytes());
         let mut tracer = Tracer::new(true);
+        let caller_h160 = H160::from_slice(input.caller().as_slice());
+
         let res = ctx.externalities.execute_with(|| {
+            tracer.watch_address(&caller_h160);
+
             tracer.trace(|| {
-                let caller_h160 = H160::from_slice(input.caller().as_slice());
                 let origin_account_id = AccountId::to_fallback_account_id(&caller_h160);
                 let origin = OriginFor::<Runtime>::signed(origin_account_id.clone());
                 let evm_value = sp_core::U256::from_little_endian(&input.value().as_le_bytes());
@@ -956,6 +984,12 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
             return None;
         }
 
+        if ctx.skip_pvm || ctx.skip_pvm_addresses.contains(&call.target_address) {
+            ctx.skip_pvm = false; // handled the skip, reset flag
+            tracing::info!("running call in EVM, instead of pallet-revive (skipped)");
+            return None;
+        }
+
         if ecx
             .journaled_state
             .database
@@ -991,6 +1025,9 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
 
         let mut tracer = Tracer::new(true);
         let res = ctx.externalities.execute_with(|| {
+            // Watch the caller's address so its nonce changes get tracked in prestate trace
+            tracer.watch_address(&caller_h160);
+
             tracer.trace(|| {
                 let origin =
                     OriginFor::<Runtime>::signed(AccountId::to_fallback_account_id(&caller_h160));
@@ -1121,6 +1158,25 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
 
         apply_revm_storage_diff(ctx, ecx, call.target_address);
     }
+
+    fn revive_record_create_address(
+        &self,
+        state: &mut foundry_cheatcodes::Cheatcodes,
+        outcome: &CreateOutcome,
+    ) {
+        let ctx = get_context_ref_mut(state.strategy.context.as_mut());
+
+        if ctx.record_next_create_address {
+            ctx.record_next_create_address = false;
+            if let Some(address) = outcome.address {
+                ctx.skip_pvm_addresses.insert(address);
+                tracing::info!(
+                    "recorded address {:?} for skip execution in the pallet-revive",
+                    address
+                );
+            }
+        }
+    }
 }
 
 fn post_exec(

crates/revive-strategy/src/cheatcodes/mod.rs

@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity ^0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+
+contract Calculator {
+    event Added(uint8 indexed sum);
+
+    function add(uint8 a, uint8 b) public returns (uint8) {
+        uint8 sum = a + b;
+        emit Added(sum);
+        return sum;
+    }
+}
+
+contract EvmTargetContract is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    event Added(uint8 indexed sum);
+
+    function exec() public {
+        emit Added(3);
+
+        Calculator calc = new Calculator();
+        uint8 sum = calc.add(1, 2);
+        assertEq(3, sum);
+        vm.setNonce(address(this), 10);
+    }
+}
+
+contract PolkadotSkipTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    EvmTargetContract helper;
+
+    function setUp() external {
+        assertEq(vm.getNonce(address(this)), 1);
+        helper = new EvmTargetContract();
+        assertEq(vm.getNonce(address(this)), 2);
+
+        // ensure we can call cheatcodes from the helper
+        vm.allowCheatcodes(address(helper));
+    }
+
+    function testUseCheatcodesInEvmWithSkip() external {
+        vm.polkadotSkip();
+        helper.exec();
+        assertEq(vm.getNonce(address(helper)), 10);
+    }
+
+    function testAutoSkipAfterDeployInEvmWithSkip() external {
+        assertEq(vm.getNonce(address(this)), 2);
+        vm.polkadotSkip();
+        EvmTargetContract helper2 = new EvmTargetContract();
+        // this should auto execute in EVM
+        helper2.exec();
+        assertEq(vm.getNonce(address(helper2)), 10);
+    }
+
+    function testreviveWhenUseCheatcodeWithoutSkip() external {
+        uint256 nonceBefore = vm.getNonce(address(helper));
+        helper.exec();
+        assertEq(vm.getNonce(address(helper)), nonceBefore + 1);
+    }
+}