From 62a878e3feeb7629ca1f223854c2eb4467dcb3f4 Mon Sep 17 00:00:00 2001
From: aon <21188659+aon@users.noreply.github.com>
Date: Thu, 30 Nov 2023 13:46:11 -0300
Subject: [PATCH] feat: add vm.addr cheatcode

---
 e2e-tests/contracts/TestCheatcodes.sol |  7 +++++++
 e2e-tests/test/cheatcodes.test.ts      | 17 +++++++++++++++++
 src/cheatcodes.rs                      | 13 ++++++++++++-
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/e2e-tests/contracts/TestCheatcodes.sol b/e2e-tests/contracts/TestCheatcodes.sol
index 1dc6d7ed..aae84c4c 100644
--- a/e2e-tests/contracts/TestCheatcodes.sol
+++ b/e2e-tests/contracts/TestCheatcodes.sol
@@ -5,6 +5,13 @@ pragma solidity ^0.8.0;
 contract TestCheatcodes {
   address constant CHEATCODE_ADDRESS = 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D;
 
+  function testAddr(uint256 privateKey, address addr) external {
+    (bool success, bytes memory data) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("addr(uint256)", privateKey));
+    require(success, "addr failed");
+    address recovered = abi.decode(data, (address));
+    require(recovered == addr, "address mismatch");
+  }
+
   function testDeal(address account, uint256 amount) external {
     uint balanceBefore = address(account).balance;
     (bool success, ) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("deal(address,uint256)", account, amount));
diff --git a/e2e-tests/test/cheatcodes.test.ts b/e2e-tests/test/cheatcodes.test.ts
index 65c05369..776eb415 100644
--- a/e2e-tests/test/cheatcodes.test.ts
+++ b/e2e-tests/test/cheatcodes.test.ts
@@ -8,6 +8,23 @@ import { deployContract, getTestProvider } from "../helpers/utils";
 const provider = getTestProvider();
 
 describe("Cheatcodes", function () {
+  it("Should test vm.addr", async function () {
+    // Arrange
+    const wallet = new Wallet(RichAccounts[0].PrivateKey);
+    const deployer = new Deployer(hre, wallet);
+    const randomWallet = Wallet.createRandom().connect(provider);
+
+    // Act
+    const greeter = await deployContract(deployer, "TestCheatcodes", []);
+    const tx = await greeter.testAddr(randomWallet.privateKey, randomWallet.address, {
+      gasLimit: 1000000,
+    });
+    const receipt = await tx.wait();
+
+    // Assert
+    expect(receipt.status).to.eq(1);
+  });
+
   it("Should test vm.deal", async function () {
     // Arrange
     const wallet = new Wallet(RichAccounts[0].PrivateKey);
diff --git a/src/cheatcodes.rs b/src/cheatcodes.rs
index 76ef466e..ec41e85e 100644
--- a/src/cheatcodes.rs
+++ b/src/cheatcodes.rs
@@ -52,6 +52,7 @@ pub trait NodeCtx {
 abigen!(
     CheatcodeContract,
     r#"[
+        function addr(uint256 privateKey)
         function deal(address who, uint256 newBalance)
         function etch(address who, bytes calldata code)
         function setNonce(address account, uint64 nonce)
@@ -170,6 +171,16 @@ impl<F: NodeCtx> CheatcodeTracer<F> {
     ) {
         use CheatcodeContractCalls::*;
         match call {
+            Addr(AddrCall { private_key }) => {
+                tracing::info!("Getting address for private key");
+                let Ok(address) = zksync_types::PackedEthSignature::address_from_private_key(
+                    &u256_to_h256(private_key),
+                ) else {
+                    tracing::error!("Failed generating address for private key");
+                    return;
+                };
+                self.returndata = Some(vec![h256_to_u256(address.into())]);
+            }
             Deal(DealCall { who, new_balance }) => {
                 tracing::info!("Setting balance for {who:?} to {new_balance}");
                 storage
@@ -205,7 +216,7 @@ impl<F: NodeCtx> CheatcodeTracer<F> {
                     account_nonce.as_u64()
                 );
                 tracing::info!("👷 Setting returndata",);
-                self.returndata = Some(vec![account_nonce.into()]);
+                self.returndata = Some(vec![account_nonce]);
             }
             SetNonce(SetNonceCall { account, nonce }) => {
                 tracing::info!("Setting nonce for {account:?} to {nonce}");