feat: activeFork cheatcode implementation (#220)

matter-labs · Jan 8, 2024 · cbcf7de · cbcf7de
1 parent 03906c5
commit cbcf7de
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 1 deletion.
diff --git a/crates/era-cheatcodes/src/cheatcodes.rs b/crates/era-cheatcodes/src/cheatcodes.rs
@@ -838,6 +838,23 @@ impl CheatcodeTracer {
         use Vm::{VmCalls::*, *};
 
         match call {
+            activeFork(activeForkCall {}) => {
+                tracing::info!("👷 Getting active fork");
+                let handle: &ForkStorage<RevmDatabaseForEra<S>> =
+                    &storage.borrow_mut().storage_handle;
+                let fork_storage = handle.inner.read().unwrap();
+                let fork_id = fork_storage
+                    .fork
+                    .as_ref()
+                    .unwrap()
+                    .fork_source
+                    .db
+                    .lock()
+                    .unwrap()
+                    .active_fork_id();
+                assert!(fork_id.is_some(), "No active fork found. Please create a fork first.");
+                self.return_data = Some(vec![fork_id.unwrap().to_u256()]);
+            }
             addr(addrCall { privateKey: private_key }) => {
                 tracing::info!("👷 Getting address for private key");
                 let Ok(address) = zksync_types::PackedEthSignature::address_from_private_key(

diff --git a/crates/era-cheatcodes/tests/src/cheatcodes/Fork.t.sol b/crates/era-cheatcodes/tests/src/cheatcodes/Fork.t.sol
@@ -42,7 +42,7 @@ contract ForkTest is Test {
         console.log("decimals_after", decimals_after);
         require(
             decimals_after == TOKEN_DECIMALS,
-            "Contract dosent exists in fork"
+            "Contract doesn't exists in fork"
         );
         require(
             block.number == FORK_BLOCK + 1,
@@ -83,4 +83,28 @@ contract ForkTest is Test {
             "ENV for blocks is not set correctly"
         );
     }
+
+    function testActiveFork() public {
+        (bool success, bytes memory data) = Constants.CHEATCODE_ADDRESS.call(
+            abi.encodeWithSignature(
+                "createFork(string,uint256)",
+                "mainnet",
+                FORK_BLOCK + 100
+            )
+        );
+        require(success, "fork failed");
+
+        uint256 forkId = uint256(bytes32(data));
+        (bool success1, ) = Constants.CHEATCODE_ADDRESS.call(
+            abi.encodeWithSignature("selectFork(uint256)", forkId)
+        );
+        require(success1, "select fork failed");
+
+        (bool success3, bytes memory activeFork) = Constants
+            .CHEATCODE_ADDRESS
+            .call(abi.encodeWithSignature("activeFork()"));
+
+        console.log("activeFork", uint256(bytes32(activeFork)));
+        require(success3, "active fork failed");
+    }
 }