Merge pull request #20 from Moonsong-Labs/aon/feat-add-vm-roll-cheatcode

feat: add vm.roll cheatcode
matter-labs · Nov 23, 2023 · e8294ec · e8294ec
2 parents a2e6888 + 128df7d
commit e8294ec
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 6 deletions.
diff --git a/e2e-tests/contracts/TestCheatcodes.sol b/e2e-tests/contracts/TestCheatcodes.sol
@@ -21,12 +21,23 @@ contract TestCheatcodes {
     require(success, "setGreeting failed");
   }
 
-  function setNonce(address account, uint256 nonce) external {
+  function testSetNonce(address account, uint256 nonce) external {
     (bool success, ) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("setNonce(address,uint64)", account, nonce));
     require(success, "setNonce failed");
   }
 
-  function warp(uint256 timestamp) external {
+  function testRoll(uint256 blockNumber) external {
+    uint256 initialBlockNumber = block.number;
+    require(blockNumber != initialBlockNumber, "block number must be different than current block number");
+
+    (bool success, ) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("roll(uint256)", blockNumber));
+    require(success, "roll failed");
+
+    uint256 finalBlockNumber = block.number;
+    require(finalBlockNumber == blockNumber, "block number was not changed");
+  }
+
+  function testWarp(uint256 timestamp) external {
     uint256 initialTimestamp = block.timestamp;
     require(timestamp != initialTimestamp, "timestamp must be different than current block timestamp");
 

diff --git a/e2e-tests/test/cheatcodes.test.ts b/e2e-tests/test/cheatcodes.test.ts
@@ -56,7 +56,7 @@ describe("Cheatcodes", function () {
     // Act
     const cheatcodes = await deployContract(deployer, "TestCheatcodes", []);
     const initialNonce = await provider.getTransactionCount(randomWallet.address);
-    const tx = await cheatcodes.setNonce(randomWallet.address, 1234);
+    const tx = await cheatcodes.testSetNonce(randomWallet.address, 1234);
     const receipt = await tx.wait();
 
     // Assert
@@ -71,10 +71,12 @@ describe("Cheatcodes", function () {
     const wallet = new Wallet(RichAccounts[0].PrivateKey);
     const deployer = new Deployer(hre, wallet);
     const contract = await deployContract(deployer, "TestCheatcodes", []);
-    const randomWallet = Wallet.createRandom().connect(provider);
+
+    const blockNumber = await provider.getBlockNumber();
+    const newBlockNumber = blockNumber + 345;
 
     // Act
-    const tx = await contract.testRoll({ gasLimit: 1000000 });
+    const tx = await contract.testRoll(newBlockNumber, { gasLimit: 1000000 });
     const receipt = await tx.wait();
 
     // Assert
@@ -92,7 +94,7 @@ describe("Cheatcodes", function () {
 
     // Act
     const cheatcodes = await deployContract(deployer, "TestCheatcodes", []);
-    const tx = await cheatcodes.warp(expectedTimestamp, {
+    const tx = await cheatcodes.testWarp(expectedTimestamp, {
       gasLimit: 1000000,
     });
     expectedTimestamp += 2; // New transaction will add two blocks

diff --git a/src/cheatcodes.rs b/src/cheatcodes.rs
@@ -47,6 +47,7 @@ abigen!(
         function deal(address who, uint256 newBalance)
         function etch(address who, bytes calldata code)
         function setNonce(address account, uint64 nonce)
+        function roll(uint256 blockNumber)
         function warp(uint256 timestamp)
     ]"#
 );
@@ -167,6 +168,21 @@ impl<F: NodeCtx> CheatcodeTracer<F> {
                 );
                 storage.set_value(nonce_key, u256_to_h256(enforced_full_nonce));
             }
+            Roll(RollCall { block_number }) => {
+                tracing::info!("Setting block number to {}", block_number);
+
+                let key = StorageKey::new(
+                    AccountTreeId::new(zksync_types::SYSTEM_CONTEXT_ADDRESS),
+                    zksync_types::CURRENT_VIRTUAL_BLOCK_INFO_POSITION,
+                );
+                let mut storage = storage.borrow_mut();
+                let (_, block_timestamp) =
+                    unpack_block_info(h256_to_u256(storage.read_value(&key)));
+                storage.set_value(
+                    key,
+                    u256_to_h256(pack_block_info(block_number.as_u64(), block_timestamp)),
+                );
+            }
             Warp(WarpCall { timestamp }) => {
                 tracing::info!("Setting block timestamp {}", timestamp);
                 self.node_ctx.set_time(timestamp.as_u64());