[issues/306] Added support for cheat code 'vm.wrap' in PVM mode

crates/forge/tests/cli/revive_vm.rs

@@ -267,6 +267,66 @@ Ran 1 test suite [ELAPSED]: 1 tests passed, 0 failed, 0 skipped (1 total tests)
 "#]]);
 });
 
+forgetest!(warp_revive, |prj, cmd| {
+    prj.insert_ds_test();
+    prj.insert_vm();
+    prj.insert_console();
+    prj.add_source(
+        "Warp.t.sol",
+        r#"
+import "./test.sol";
+import "./Vm.sol";
+import {console} from "./console.sol";
+
+contract Warp is DSTest {
+  Vm constant vm = Vm(HEVM_ADDRESS);
+
+  function test_Warp() public {
+      vm.pvm(true);
+      uint256 original = block.timestamp;
+      vm.warp(100);
+      uint256 newValue = block.timestamp;
+      assert(original != newValue);
+      assertEq(newValue, 100);
+  }
+}
+"#,
+    )
+    .unwrap();
+
+    cmd.env("RUST_LOG", "revive_strategy");
+    let res = cmd.args(["test", "--resolc", "-vvv", "--resolc-startup"]).assert_success();
+    res.stderr_eq(str![""]).stdout_eq(str![[r#"
+[COMPILING_FILES] with [SOLC_VERSION]
+[SOLC_VERSION] [ELAPSED]
+Compiler run successful!
+[COMPILING_FILES] with [RESOLC_VERSION]
+[RESOLC_VERSION] [ELAPSED]
+Compiler run successful with warnings:
+Warning: Warning: Your code or one of its dependencies uses the 'extcodesize' instruction, which is
+usually needed in the following cases:
+  1. To detect whether an address belongs to a smart contract.
+  2. To detect whether the deploy code execution has finished.
+Polkadot comes with native account abstraction support (so smart contracts are just accounts
+coverned by code), and you should avoid differentiating between contracts and non-contract
+addresses.
+[FILE]
+[..] INFO revive_strategy::cheatcodes: startup PVM migration initiated
+[..] INFO revive_strategy::cheatcodes: switching to PVM
+[..] INFO revive_strategy::cheatcodes: startup PVM migration completed
+[..] INFO revive_strategy::cheatcodes: cheatcode=pvmCall { enabled: true } using_pvm=true
+[..] INFO revive_strategy::cheatcodes: already in PVM
+[..] INFO revive_strategy::cheatcodes: cheatcode=warpCall { newTimestamp: 100 } using_pvm=true
+
+Ran 1 test for src/Warp.t.sol:Warp
+[PASS] test_Warp() ([GAS])
+Suite result: ok. 1 passed; 0 failed; 0 skipped; [ELAPSED]
+
+Ran 1 test suite [ELAPSED]: 1 tests passed, 0 failed, 0 skipped (1 total tests)
+
+"#]]);
+});
+
 forgetest!(deal, |prj, cmd| {
     prj.insert_ds_test();
     prj.insert_vm();

crates/revive-env/src/lib.rs

@@ -16,7 +16,7 @@ use polkadot_sdk::{
     sp_tracing,
 };
 
-pub use crate::runtime::{AccountId, Balance, Runtime, System};
+pub use crate::runtime::{AccountId, Balance, Runtime, System, Timestamp};
 
 mod runtime;
 

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

@@ -10,11 +10,11 @@ use foundry_cheatcodes::{
     Broadcast, BroadcastableTransactions, CheatcodeInspectorStrategy,
     CheatcodeInspectorStrategyContext, CheatcodeInspectorStrategyRunner, CheatsConfig, CheatsCtxt,
     CommonCreateInput, DealRecord, Ecx, EvmCheatcodeInspectorStrategyRunner, InnerEcx, Result,
-    Vm::{dealCall, getNonce_0Call, pvmCall, rollCall, setNonceCall, setNonceUnsafeCall},
+    Vm::{dealCall, getNonce_0Call, pvmCall, rollCall, setNonceCall, setNonceUnsafeCall, warpCall},
 };
 use foundry_common::sh_err;
 use foundry_compilers::resolc::dual_compiled_contracts::DualCompiledContracts;
-use revive_env::{AccountId, Runtime, System};
+use revive_env::{AccountId, Runtime, System, Timestamp};
 
 use polkadot_sdk::{
     frame_support::traits::{fungible::Mutate, Currency},
@@ -152,6 +152,18 @@ fn set_block_number(new_height: U256, ecx: InnerEcx<'_, '_, '_>) {
     });
 }
 
+fn set_timestamp(new_timestamp: U256, ecx: InnerEcx<'_, '_, '_>) {
+    // Set timestamp in EVM context.
+    ecx.env.block.timestamp = new_timestamp;
+
+    // Set timestamp in pallet-revive runtime.
+    execute_with_externalities(|externalities| {
+        externalities.execute_with(|| {
+            Timestamp::set_timestamp(new_timestamp.try_into().expect("Timestamp exceeds u64"));
+        })
+    });
+}
+
 /// Implements [CheatcodeInspectorStrategyRunner] for PVM.
 #[derive(Debug, Default, Clone)]
 pub struct PvmCheatcodeInspectorStrategyRunner;
@@ -228,6 +240,14 @@ impl CheatcodeInspectorStrategyRunner for PvmCheatcodeInspectorStrategyRunner {
 
                 Ok(Default::default())
             }
+            t if using_pvm && is::<warpCall>(t) => {
+                let &warpCall { newTimestamp } = cheatcode.as_any().downcast_ref().unwrap();
+
+                tracing::info!(cheatcode = ?cheatcode.as_debug() , using_pvm = ?using_pvm);
+                set_timestamp(newTimestamp, ccx.ecx);
+
+                Ok(Default::default())
+            }
             // Not custom, just invoke the default behavior
             _ => cheatcode.dyn_apply(ccx, executor),
         }