Add polkadotskip cheatcode

crates/cheatcodes/spec/src/vm.rs

@@ -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/src/evm.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/inspector.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/strategy.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/forge/tests/it/revive/cheat_polkadot_skip.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_mode(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/mod.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;
 pub mod cheat_store;
 pub mod migration;

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

@@ -8,8 +8,8 @@ use foundry_cheatcodes::{
     CheatcodeInspectorStrategyContext, CheatcodeInspectorStrategyRunner, CheatsConfig, CheatsCtxt,
     CommonCreateInput, Ecx, EvmCheatcodeInspectorStrategyRunner, Result,
     Vm::{
-        chainIdCall, dealCall, etchCall, getNonce_0Call, loadCall, pvmCall, resetNonceCall,
-        rollCall, setNonceCall, setNonceUnsafeCall, storeCall, warpCall,
+        chainIdCall, dealCall, etchCall, getNonce_0Call, loadCall, polkadotSkipCall, pvmCall,
+        resetNonceCall, rollCall, setNonceCall, setNonceUnsafeCall, storeCall, warpCall,
     },
     journaled_account, precompile_error,
 };
@@ -112,6 +112,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,
@@ -130,6 +138,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,
@@ -273,6 +284,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();
@@ -756,6 +773,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
@@ -926,6 +950,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
@@ -1085,6 +1115,22 @@ 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(

testdata/default/revive/PolkadotSkip.t.sol

@@ -0,0 +1,58 @@
+// 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 {
+        // We emit the event we expect to see.
+        vm.expectEmit();
+        emit Added(3);
+
+        Calculator calc = new Calculator();
+        uint8 sum = calc.add(1, 2);
+        assertEq(3, sum);
+    }
+}
+
+contract PolkadotSkipTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    EvmTargetContract helper;
+
+    function setUp() external {
+        vm.pvm(true);
+        helper = new EvmTargetContract();
+        // ensure we can call cheatcodes from the helper
+        vm.allowCheatcodes(address(helper));
+        // and that the contract is kept between vm switches
+        vm.makePersistent(address(helper));
+    }
+
+    function testUseCheatcodesInEvmWithSkip() external {
+        vm.polkadotSkip();
+        helper.exec();
+    }
+
+    function testAutoSkipAfterDeployInEvmWithSkip() external {
+        vm.polkadotSkip();
+        EvmTargetContract helper2 = new EvmTargetContract();
+
+        // this should auto execute in EVM
+        helper2.exec();
+    }
+}

testdata/default/revive/PolkadotSkipInvariant.t.sol

@@ -0,0 +1,145 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity ^0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+import {ReviveDetector, RequiresRevive} from "./ReviveDetector.sol";
+import "../../default/logs/console.sol";
+
+contract Counter is RequiresRevive {
+    uint256 public number;
+
+    function inc() public onlyRevive {
+        number += 1;
+    }
+
+    function reset() public onlyRevive {
+        number = 0;
+    }
+}
+
+contract CounterHandler is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    uint256 public incCounter;
+    uint256 public resetCounter;
+    bool public isResetLast;
+    Counter public counter;
+
+    constructor(Counter _counter) {
+        counter = _counter;
+    }
+
+    function inc() public {
+        console.log("inc");
+        incCounter += 1;
+        isResetLast = false;
+
+        vm.deal(tx.origin, 1 ether); // ensure caller has funds
+        counter.inc();
+    }
+
+    function reset() public {
+        console.log("reset");
+        resetCounter += 1;
+        isResetLast = true;
+
+        vm.deal(tx.origin, 1 ether); // ensure caller has funds
+        counter.reset();
+    }
+}
+
+// partial from forge-std/StdInvariant.sol
+abstract contract StdInvariant {
+    struct FuzzSelector {
+        address addr;
+        bytes4[] selectors;
+    }
+
+    address[] internal _targetedContracts;
+
+    function targetContracts() public view returns (address[] memory) {
+        return _targetedContracts;
+    }
+
+    FuzzSelector[] internal _targetedSelectors;
+
+    function targetSelectors() public view returns (FuzzSelector[] memory) {
+        return _targetedSelectors;
+    }
+
+    address[] internal _targetedSenders;
+
+    function targetSenders() public view returns (address[] memory) {
+        return _targetedSenders;
+    }
+}
+
+contract PolkadotSkipInvariantWithHandler is DSTest, StdInvariant {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    Counter cnt;
+    CounterHandler handler;
+
+    function setUp() public {
+        vm.pvm(true);
+        cnt = new Counter();
+        vm.makePersistent(address(cnt));
+
+        vm.polkadotSkip();
+        handler = new CounterHandler(cnt);
+
+        // add the handler selectors to the fuzzing targets
+        bytes4[] memory selectors = new bytes4[](2);
+        selectors[0] = CounterHandler.inc.selector;
+        selectors[1] = CounterHandler.reset.selector;
+
+        _targetedContracts.push(address(handler));
+        _targetedSelectors.push(FuzzSelector({addr: address(handler), selectors: selectors}));
+    }
+
+    /// forge-config: default.invariant.fail-on-revert = true
+    function invariant_ghostVariables() external {
+        uint256 num = cnt.number();
+
+        if (handler.resetCounter() == 0) {
+            assert(handler.incCounter() == num);
+        } else if (handler.isResetLast()) {
+            assert(num == 0);
+        } else {
+            assert(num != 0);
+        }
+    }
+}
+
+contract PolkadotSkipInvariantWithoutHandler is DSTest, StdInvariant {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    Counter cnt;
+
+    uint256 constant dealAmount = 1 ether;
+
+    function setUp() public {
+        vm.pvm(true);
+        cnt = new Counter();
+
+        // so we can fund them ahead of time for fees
+        _targetedSenders.push(address(65536 + 1));
+        _targetedSenders.push(address(65536 + 12));
+        _targetedSenders.push(address(65536 + 123));
+        _targetedSenders.push(address(65536 + 1234));
+
+        for (uint256 i = 0; i < _targetedSenders.length; i++) {
+            vm.deal(_targetedSenders[i], dealAmount);
+        }
+
+        // add the counter selectors to the fuzzing targets
+        bytes4[] memory selectors = new bytes4[](2);
+        selectors[0] = Counter.inc.selector;
+        selectors[1] = Counter.reset.selector;
+
+        _targetedContracts.push(address(cnt));
+        _targetedSelectors.push(FuzzSelector({addr: address(cnt), selectors: selectors}));
+    }
+
+    /// forge-config: default.invariant.fail-on-revert = true
+    function invariant_itWorks() external {}
+}