Merge pull request #27 from Moonsong-Labs/jr/feat-add-vm-load-cheatcode

feat: add vm load cheatcode
Moonsong-Labs · Dec 4, 2023 · f02534c · f02534c
2 parents 6ef8f07 + 00b7f1c
commit f02534c
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 28 deletions.
diff --git a/e2e-tests/contracts/TestCheatcodes.sol b/e2e-tests/contracts/TestCheatcodes.sol
@@ -148,6 +148,18 @@ contract TestCheatcodes {
 
     testStoreInstance.testStoredValue(value);
   }
+
+  function testLoad(bytes32 slot) external {
+    TestLoadTarget testLoadTarget = new TestLoadTarget();
+    (bool success, bytes memory data) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("load(address,bytes32)", address(testLoadTarget), slot));
+    require(success, "load failed");
+    bytes32 loadedValue = abi.decode(data, (bytes32));
+    require(loadedValue == bytes32(uint256(1337)), "address mismatch");
+  }
+}
+
+contract TestLoadTarget {
+  bytes32 public testValue = bytes32(uint256(1337)); //slot 0
 }
 
 contract testStoreTarget {

diff --git a/e2e-tests/test/cheatcodes.test.ts b/e2e-tests/test/cheatcodes.test.ts
@@ -64,6 +64,23 @@ describe("Cheatcodes", function () {
     expect(finalRandomWalletCode).to.not.eq(initialRandomWalletCode);
   });
 
+  it("Should test vm.load", async function () {
+    // Arrange
+    const wallet = new Wallet(RichAccounts[0].PrivateKey);
+    const deployer = new Deployer(hre, wallet);
+
+    // Act
+    const cheatcodes = await deployContract(deployer, "TestCheatcodes", []);
+    const slot = hre.ethers.constants.HashZero;
+    const tx = await cheatcodes.testLoad(slot, {
+      gasLimit: 10000000,
+    });
+    const receipt = await tx.wait();
+
+    // Assert
+    expect(receipt.status).to.eq(1);
+  });
+
   it("Should test vm.roll", async function () {
     // Arrange
     const wallet = new Wallet(RichAccounts[0].PrivateKey);
@@ -140,44 +157,44 @@ describe("Cheatcodes", function () {
     expect(receipt2.status).to.eq(1);
   });
 
-  it("Should test vm.warp", async function () {
+  it("Should test vm.store", async function () {
     // Arrange
     const wallet = new Wallet(RichAccounts[0].PrivateKey);
     const deployer = new Deployer(hre, wallet);
 
-    const timeIncreaseInMS = 123;
-    let expectedTimestamp: number = await provider.send("config_getCurrentTimestamp", []);
-    expectedTimestamp += timeIncreaseInMS;
-
     // Act
     const cheatcodes = await deployContract(deployer, "TestCheatcodes", []);
-    const tx = await cheatcodes.testWarp(expectedTimestamp, {
-      gasLimit: 1000000,
+    const slot = hre.ethers.constants.HashZero;
+    const value = hre.ethers.constants.MaxUint256;
+    const tx = await cheatcodes.testStore(slot, value, {
+      gasLimit: 10000000,
     });
-    expectedTimestamp += 2; // New transaction will add two blocks
     const receipt = await tx.wait();
 
     // Assert
     expect(receipt.status).to.eq(1);
-    const newBlockTimestamp = (await provider.getBlock("latest")).timestamp;
-    expect(newBlockTimestamp).to.equal(expectedTimestamp);
   });
 
-  it("Should test vm.store", async function () {
+  it("Should test vm.warp", async function () {
     // Arrange
     const wallet = new Wallet(RichAccounts[0].PrivateKey);
     const deployer = new Deployer(hre, wallet);
 
+    const timeIncreaseInMS = 123;
+    let expectedTimestamp: number = await provider.send("config_getCurrentTimestamp", []);
+    expectedTimestamp += timeIncreaseInMS;
+
     // Act
     const cheatcodes = await deployContract(deployer, "TestCheatcodes", []);
-    const slot = hre.ethers.constants.HashZero;
-    const value = hre.ethers.BigNumber.from(hre.ethers.utils.randomBytes(32));
-    const tx = await cheatcodes.testStore(slot, value, {
-      gasLimit: 10000000,
+    const tx = await cheatcodes.testWarp(expectedTimestamp, {
+      gasLimit: 1000000,
     });
+    expectedTimestamp += 2; // New transaction will add two blocks
     const receipt = await tx.wait();
 
     // Assert
     expect(receipt.status).to.eq(1);
+    const newBlockTimestamp = (await provider.getBlock("latest")).timestamp;
+    expect(newBlockTimestamp).to.equal(expectedTimestamp);
   });
 });
diff --git a/src/cheatcodes.rs b/src/cheatcodes.rs
@@ -94,6 +94,7 @@ abigen!(
         function deal(address who, uint256 newBalance)
         function etch(address who, bytes calldata code)
         function getNonce(address account)
+        function load(address account, bytes32 slot)
         function roll(uint256 blockNumber)
         function setNonce(address account, uint64 nonce)
         function startPrank(address sender)
@@ -303,6 +304,13 @@ impl<F: NodeCtx> CheatcodeTracer<F> {
                 tracing::info!("👷 Setting returndata",);
                 self.returndata = Some(vec![account_nonce]);
             }
+            Load(LoadCall { account, slot }) => {
+                tracing::info!("Getting storage slot {:?} for account {:?}", slot, account);
+                let key = StorageKey::new(AccountTreeId::new(account), H256(slot));
+                let mut storage = storage.borrow_mut();
+                let value = storage.read_value(&key);
+                self.returndata = Some(vec![h256_to_u256(value)]);
+            }
             SetNonce(SetNonceCall { account, nonce }) => {
                 tracing::info!("Setting nonce for {account:?} to {nonce}");
                 let nonce_key = get_nonce_key(&account);
@@ -370,19 +378,6 @@ impl<F: NodeCtx> CheatcodeTracer<F> {
 
                 self.start_prank_opts = None;
             }
-            Warp(WarpCall { timestamp }) => {
-                tracing::info!("Setting block timestamp {}", timestamp);
-                self.node_ctx.set_time(timestamp.as_u64());
-
-                let key = StorageKey::new(
-                    AccountTreeId::new(zksync_types::SYSTEM_CONTEXT_ADDRESS),
-                    zksync_types::CURRENT_VIRTUAL_BLOCK_INFO_POSITION,
-                );
-                let (block_number, _) =
-                    unpack_block_info(h256_to_u256(storage.borrow_mut().read_value(&key)));
-                let value = u256_to_h256(pack_block_info(block_number, timestamp.as_u64()));
-                self.write_storage(key, value, storage);
-            }
             Store(StoreCall {
                 account,
                 slot,
@@ -397,6 +392,21 @@ impl<F: NodeCtx> CheatcodeTracer<F> {
                 let key = StorageKey::new(AccountTreeId::new(account), H256(slot));
                 self.write_storage(key, H256(value), storage);
             }
+            Warp(WarpCall { timestamp }) => {
+                tracing::info!("Setting block timestamp {}", timestamp);
+                self.node_ctx.set_time(timestamp.as_u64());
+
+                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_number, _) = unpack_block_info(h256_to_u256(storage.read_value(&key)));
+                storage.set_value(
+                    key,
+                    u256_to_h256(pack_block_info(block_number, timestamp.as_u64())),
+                );
+            }
         };
     }