Merge branch 'grw/feat-cheatcode-tracer' of github.com:Moonsong-Labs/…

…era-test-node into aon/feat-add-vm-to-string-cheatcode

e2e-tests/contracts/TestCheatcodes.sol

@@ -28,6 +28,16 @@ contract TestCheatcodes {
     require(success, "setGreeting failed");
   }
 
+  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");
+  }
+
   function testRoll(uint256 blockNumber) external {
     uint256 initialBlockNumber = block.number;
     require(blockNumber != initialBlockNumber, "block number must be different than current block number");
@@ -126,6 +136,18 @@ contract TestCheatcodes {
     );
   }
 
+  function testStore(bytes32 slot, bytes32 value) external {
+    testStoreTarget testStoreInstance = new testStoreTarget();
+    testStoreInstance.testStoredValue(0);
+
+    (bool success, ) = CHEATCODE_ADDRESS.call(
+      abi.encodeWithSignature("store(address,bytes32,bytes32)", address(testStoreInstance), slot, value)
+    );
+    require(success, "store failed");
+
+    testStoreInstance.testStoredValue(value);
+  }
+
   function testToStringFromAddress() external {
     address testAddress = 0x413D15117be7a498e68A64FcfdB22C6e2AaE1808;
     (bool success, bytes memory rawData) = CHEATCODE_ADDRESS.call(
@@ -204,18 +226,6 @@ contract TestCheatcodes {
     require(finalTimestamp == timestamp, "timestamp was not changed");
   }
 
-  function testStore(bytes32 slot, bytes32 value) external {
-    testStoreTarget testStoreInstance = new testStoreTarget();
-    testStoreInstance.testStoredValue(0);
-
-    (bool success, ) = CHEATCODE_ADDRESS.call(
-      abi.encodeWithSignature("store(address,bytes32,bytes32)", address(testStoreInstance), slot, value)
-    );
-    require(success, "store failed");
-
-    testStoreInstance.testStoredValue(value);
-  }
-
   function trimReturnBytes(bytes memory rawData) internal pure returns (bytes memory) {
     uint256 lengthStartingPos = rawData.length - 32;
     bytes memory lengthSlice = new bytes(32);
@@ -231,6 +241,10 @@ contract TestCheatcodes {
   }
 }
 
+contract TestLoadTarget {
+  bytes32 public testValue = bytes32(uint256(1337)); //slot 0
+}
+
 contract testStoreTarget {
   bytes32 public testValue = bytes32(uint256(0)); //slot 0
 

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,6 +157,24 @@ describe("Cheatcodes", function () {
     expect(receipt2.status).to.eq(1);
   });
 
+  it("Should test vm.store", async function () {
+    // Arrange
+    const wallet = new Wallet(RichAccounts[0].PrivateKey);
+    const deployer = new Deployer(hre, wallet);
+    const slot = hre.ethers.constants.HashZero;
+    const value = hre.ethers.constants.MaxUint256;
+
+    // Act
+    const cheatcodes = await deployContract(deployer, "TestCheatcodes", []);
+    const tx = await cheatcodes.testStore(slot, value, {
+      gasLimit: 10000000,
+    });
+    const receipt = await tx.wait();
+
+    // Assert
+    expect(receipt.status).to.eq(1);
+  });
+
   it("Should test vm.toString from address", async function () {
     // Arrange
     const wallet = new Wallet(RichAccounts[0].PrivateKey);
@@ -222,7 +257,6 @@ describe("Cheatcodes", function () {
 
     // Act
     const cheatcodes = await deployContract(deployer, "TestCheatcodes", []);
-
     const tx = await cheatcodes.testToStringFromBytes({ gasLimit: 100000000 });
     const receipt = await tx.wait();
 
@@ -252,22 +286,4 @@ describe("Cheatcodes", function () {
     const newBlockTimestamp = (await provider.getBlock("latest")).timestamp;
     expect(newBlockTimestamp).to.equal(expectedTimestamp);
   });
-
-  it("Should test vm.store", 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 value = hre.ethers.BigNumber.from(hre.ethers.utils.randomBytes(32));
-    const tx = await cheatcodes.testStore(slot, value, {
-      gasLimit: 10000000,
-    });
-    const receipt = await tx.wait();
-
-    // Assert
-    expect(receipt.status).to.eq(1);
-  });
 });

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)
@@ -309,6 +310,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);
@@ -376,6 +384,20 @@ impl<F: NodeCtx> CheatcodeTracer<F> {
 
                 self.start_prank_opts = None;
             }
+            Store(StoreCall {
+                account,
+                slot,
+                value,
+            }) => {
+                tracing::info!(
+                    "👷 Setting storage slot {:?} for account {:?} to {:?}",
+                    slot,
+                    account,
+                    value
+                );
+                let key = StorageKey::new(AccountTreeId::new(account), H256(slot));
+                self.write_storage(key, H256(value), storage);
+            }
             ToString0(ToString0Call { value }) => {
                 tracing::info!("Converting address into string");
                 let address = Address::from(value);
@@ -420,20 +442,6 @@ impl<F: NodeCtx> CheatcodeTracer<F> {
                 let value = u256_to_h256(pack_block_info(block_number, timestamp.as_u64()));
                 self.write_storage(key, value, storage);
             }
-            Store(StoreCall {
-                account,
-                slot,
-                value,
-            }) => {
-                tracing::info!(
-                    "👷 Setting storage slot {:?} for account {:?} to {:?}",
-                    slot,
-                    account,
-                    value
-                );
-                let key = StorageKey::new(AccountTreeId::new(account), H256(slot));
-                self.write_storage(key, H256(value), storage);
-            }
         };
     }