Merge branch 'main' into jr-fear-rpcUrls-cheatcodes

Cargo.toml

@@ -159,7 +159,7 @@ alloy-rlp = "0.3.3"
 solang-parser = "=0.3.3"
 
 ## zksync 
-era_test_node = { git="https://github.com/Moonsong-Labs/era-test-node.git" , branch = "nish-deniallugo-move-storage-view"}
+era_test_node = { git="https://github.com/matter-labs/era-test-node.git" , rev = "6ee7d29e876b75506f58355218e1ea755a315d17" }
 zksync_basic_types = { git = "https://github.com/matter-labs/zksync-era.git", rev = "3c669e7b0caf6515ad865f4cba5ea6fb36c33811" }
 zksync_types = { git = "https://github.com/matter-labs/zksync-era.git", rev = "3c669e7b0caf6515ad865f4cba5ea6fb36c33811" }
 zksync_state = { git = "https://github.com/matter-labs/zksync-era.git", rev = "3c669e7b0caf6515ad865f4cba5ea6fb36c33811" }

crates/era-cheatcodes/src/cheatcodes.rs

@@ -29,7 +29,7 @@ use multivm::{
     },
 };
 use revm::{
-    primitives::{BlockEnv, CfgEnv, Env, SpecId, U256 as rU256},
+    primitives::{ruint::Uint, BlockEnv, CfgEnv, Env, SpecId, U256 as rU256},
     JournaledState,
 };
 use serde::Serialize;
@@ -119,6 +119,12 @@ pub struct CheatcodeTracer {
     recording_timestamp: u32,
     expected_calls: ExpectedCallsTracker,
     test_status: FoundryTestState,
+    saved_snapshots: HashMap<U256, SavedSnapshot>,
+}
+
+#[derive(Debug, Clone)]
+pub struct SavedSnapshot {
+    modified_storage: HashMap<StorageKey, H256>,
 }
 
 #[derive(Debug, Clone, Serialize, Eq, Hash, PartialEq)]
@@ -141,6 +147,8 @@ enum FinishCycleOneTimeActions {
     CreateSelectFork { url_or_alias: String, block_number: Option<u64> },
     CreateFork { url_or_alias: String, block_number: Option<u64> },
     SelectFork { fork_id: U256 },
+    RevertToSnapshot { snapshot_id: U256 },
+    Snapshot,
 }
 
 #[derive(Debug, Default, Clone)]
@@ -447,6 +455,84 @@ impl<S: DatabaseExt + Send, H: HistoryMode> VmTracer<EraDb<S>, H> for CheatcodeT
 
                     self.return_data = Some(vec![fork_id]);
                 }
+                FinishCycleOneTimeActions::RevertToSnapshot { snapshot_id } => {
+                    let mut storage = storage.borrow_mut();
+
+                    let modified_storage: HashMap<StorageKey, H256> = storage
+                        .modified_storage_keys()
+                        .clone()
+                        .into_iter()
+                        .filter(|(key, _)| key.address() != &zksync_types::SYSTEM_CONTEXT_ADDRESS)
+                        .collect();
+
+                    storage.clean_cache();
+
+                    {
+                        let handle: &ForkStorage<RevmDatabaseForEra<S>> = &storage.storage_handle;
+                        let mut fork_storage = handle.inner.write().unwrap();
+                        fork_storage.value_read_cache.clear();
+                        let era_db = fork_storage.fork.as_ref().unwrap().fork_source.clone();
+                        let bytecodes = bootloader_state
+                            .get_last_tx_compressed_bytecodes()
+                            .iter()
+                            .map(|b| bytecode_to_factory_dep(b.original.clone()))
+                            .collect();
+
+                        let mut journaled_state = JournaledState::new(SpecId::LATEST, vec![]);
+                        let state = storage_to_state(&era_db, &modified_storage, bytecodes);
+                        *journaled_state.state() = state;
+
+                        let mut db = era_db.db.lock().unwrap();
+                        let era_env = self.env.get().unwrap();
+                        let mut env = into_revm_env(era_env);
+                        db.revert(Uint::from_limbs(snapshot_id.0), &journaled_state, &mut env);
+                    }
+
+                    storage.modified_storage_keys =
+                        self.saved_snapshots.remove(&snapshot_id).unwrap().modified_storage;
+                }
+                FinishCycleOneTimeActions::Snapshot => {
+                    let mut storage = storage.borrow_mut();
+
+                    let modified_storage: HashMap<StorageKey, H256> = storage
+                        .modified_storage_keys()
+                        .clone()
+                        .into_iter()
+                        .filter(|(key, _)| key.address() != &zksync_types::SYSTEM_CONTEXT_ADDRESS)
+                        .collect();
+
+                    storage.clean_cache();
+
+                    let snapshot_id = {
+                        let handle: &ForkStorage<RevmDatabaseForEra<S>> = &storage.storage_handle;
+                        let mut fork_storage = handle.inner.write().unwrap();
+                        fork_storage.value_read_cache.clear();
+                        let era_db = fork_storage.fork.as_ref().unwrap().fork_source.clone();
+                        let bytecodes = bootloader_state
+                            .get_last_tx_compressed_bytecodes()
+                            .iter()
+                            .map(|b| bytecode_to_factory_dep(b.original.clone()))
+                            .collect();
+
+                        let mut journaled_state = JournaledState::new(SpecId::LATEST, vec![]);
+                        let state = storage_to_state(&era_db, &modified_storage, bytecodes);
+                        *journaled_state.state() = state;
+
+                        let mut db = era_db.db.lock().unwrap();
+                        let era_env = self.env.get().unwrap();
+                        let env = into_revm_env(era_env);
+                        let snapshot_id = db.snapshot(&journaled_state, &env);
+
+                        self.saved_snapshots.insert(
+                            snapshot_id.to_u256(),
+                            SavedSnapshot { modified_storage: modified_storage.clone() },
+                        );
+                        snapshot_id
+                    };
+
+                    storage.modified_storage_keys = modified_storage;
+                    self.return_data = Some(vec![snapshot_id.to_u256()]);
+                }
             }
         }
 
@@ -598,8 +684,11 @@ impl CheatcodeTracer {
                     tracing::error!("Failed to run ffi: no args");
                     return
                 };
-                // TODO: set directory to root
-                let Ok(output) = Command::new(first_arg).args(&command_input[1..]).output() else {
+                let Ok(output) = Command::new(first_arg)
+                    .args(&command_input[1..])
+                    .current_dir(&self.config.root)
+                    .output()
+                else {
                     tracing::error!("Failed to run ffi");
                     return
                 };
@@ -722,6 +811,12 @@ impl CheatcodeTracer {
                 };
                 self.add_trimmed_return_data(&data);
             }
+            revertTo(revertToCall { snapshotId }) => {
+                tracing::info!("👷 Reverting to snapshot {}", snapshotId);
+                self.one_time_actions.push(FinishCycleOneTimeActions::RevertToSnapshot {
+                    snapshot_id: snapshotId.to_u256(),
+                });
+            }
             roll(rollCall { newHeight: new_height }) => {
                 tracing::info!("👷 Setting block number to {}", new_height);
                 let key = StorageKey::new(
@@ -858,6 +953,10 @@ impl CheatcodeTracer {
                 );
                 self.write_storage(nonce_key, u256_to_h256(enforced_full_nonce), &mut storage);
             }
+            snapshot(snapshotCall {}) => {
+                tracing::info!("👷 Creating snapshot");
+                self.one_time_actions.push(FinishCycleOneTimeActions::Snapshot);
+            }
             startPrank_0(startPrank_0Call { msgSender: msg_sender }) => {
                 tracing::info!("👷 Starting prank to {msg_sender:?}");
                 self.permanent_actions.start_prank =
@@ -940,8 +1039,11 @@ impl CheatcodeTracer {
                     tracing::error!("Failed to run ffi: no args");
                     return
                 };
-                // TODO: set directory to root
-                let Ok(output) = Command::new(first_arg).args(&command_input[1..]).output() else {
+                let Ok(output) = Command::new(first_arg)
+                    .args(&command_input[1..])
+                    .current_dir(&self.config.root)
+                    .output()
+                else {
                     tracing::error!("Failed to run ffi");
                     return
                 };

crates/era-cheatcodes/tests/src/cheatcodes/Snapshot.t.sol

@@ -0,0 +1,97 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import {Test, Vm, console2 as console} from "../../lib/forge-std/src/Test.sol";
+import {Constants} from "./Constants.sol";
+import {Utils} from "./Utils.sol";
+
+struct Storage {
+    uint256 slot0;
+    uint256 slot1;
+}
+
+contract SnapshotTest is Test {
+    Storage store;
+
+    function setUp() public {
+        store.slot0 = 10;
+        store.slot1 = 20;
+    }
+
+    function testSnapshot() public {
+        console.log("calling snapshot");
+
+        store.slot0 = 10;
+        store.slot1 = 20;
+
+        (bool success, bytes memory data) = Constants.CHEATCODE_ADDRESS.call(
+            abi.encodeWithSignature("snapshot()")
+        );
+        require(success, "snapshot failed");
+
+        uint256 snapshot = abi.decode(data, (uint256));
+
+        console.log("snapshot ", snapshot);
+
+        console.log("store values: ", store.slot0, store.slot1);
+
+        store.slot0 = 300;
+        store.slot1 = 400;
+
+        assertEq(store.slot0, 300);
+        assertEq(store.slot1, 400);
+
+        console.log("store values: ", store.slot0, store.slot1);
+        console.log("calling revertTo");
+
+        (success, ) = Constants.CHEATCODE_ADDRESS.call(
+            abi.encodeWithSignature("revertTo(uint256)", snapshot)
+        );
+        require(success, "revertTo failed");
+
+        console.log("store values: ", store.slot0, store.slot1);
+
+        assertEq(store.slot0, 10, "snapshot revert for slot 0 unsuccessful");
+        assertEq(store.slot1, 20, "snapshot revert for slot 1 unsuccessful");
+    }
+
+    function testBlockValues() public {
+        uint256 num = block.number;
+        uint256 time = block.timestamp;
+
+        (bool success, bytes memory data) = Constants.CHEATCODE_ADDRESS.call(
+            abi.encodeWithSignature("snapshot()")
+        );
+        require(success, "snapshot failed");
+
+        uint256 snapshot = abi.decode(data, (uint256));
+
+        (success, ) = Constants.CHEATCODE_ADDRESS.call(
+            abi.encodeWithSignature("warp(uint256)", 1337)
+        );
+        require(success, "warp failed");
+        assertEq(block.timestamp, 1337);
+
+        (success, ) = Constants.CHEATCODE_ADDRESS.call(
+            abi.encodeWithSignature("roll(uint256)", 99)
+        );
+        require(success, "roll failed");
+        assertEq(block.number, 99);
+
+        (success, ) = Constants.CHEATCODE_ADDRESS.call(
+            abi.encodeWithSignature("revertTo(uint256)", snapshot)
+        );
+        require(success, "revertTo failed");
+
+        assertEq(
+            block.number,
+            num,
+            "snapshot revert for block.number unsuccessful"
+        );
+        assertEq(
+            block.timestamp,
+            time,
+            "snapshot revert for block.timestamp unsuccessful"
+        );
+    }
+}