feat: add function to calculate deployment addresses

nomoixyz · Aug 19, 2023 · 745b90b · 745b90b
1 parent f011bb7
commit 745b90b
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 1 deletion.
diff --git a/src/_modules/Accounts.sol b/src/_modules/Accounts.sol
@@ -102,6 +102,43 @@ library accountsSafe {
 
         return label(addr, lbl);
     }
+
+    /// @dev Calculates the deployment address of `who` with nonce `nonce`.
+    /// @param who The deployer address.
+    /// @param nonce The deployer nonce.
+    function getDeploymentAddress(address who, uint64 nonce) internal pure returns (address) {
+        bytes memory data;
+
+        if (nonce == 0x00) {
+            data = abi.encodePacked(bytes1(0xd6), bytes1(0x94), who, bytes1(0x80));
+        } else if (nonce <= 0x7f) {
+            data = abi.encodePacked(bytes1(0xd6), bytes1(0x94), who, uint8(nonce));
+        } else if (nonce <= 0xff) {
+            data = abi.encodePacked(bytes1(0xd7), bytes1(0x94), who, bytes1(0x81), uint8(nonce));
+        } else if (nonce <= 0xffff) {
+            data = abi.encodePacked(bytes1(0xd8), bytes1(0x94), who, bytes1(0x82), uint16(nonce));
+        } else if (nonce <= 0xffffff) {
+            data = abi.encodePacked(bytes1(0xd9), bytes1(0x94), who, bytes1(0x83), uint24(nonce));
+        } else if (nonce <= 0xffffffff) {
+            data = abi.encodePacked(bytes1(0xda), bytes1(0x94), who, bytes1(0x84), uint32(nonce));
+        } else if (nonce <= 0xffffffffff) {
+            data = abi.encodePacked(bytes1(0xdb), bytes1(0x94), who, bytes1(0x85), uint40(nonce));
+        } else if (nonce <= 0xffffffffffff) {
+            data = abi.encodePacked(bytes1(0xdc), bytes1(0x94), who, bytes1(0x86), uint48(nonce));
+        } else if (nonce <= 0xffffffffffffff) {
+            data = abi.encodePacked(bytes1(0xdd), bytes1(0x94), who, bytes1(0x87), uint56(nonce));
+        } else if (nonce <= 0xffffffffffffffff) {
+            data = abi.encodePacked(bytes1(0xde), bytes1(0x94), who, bytes1(0x88), uint64(nonce));
+        }
+
+        return address(uint160(uint256(keccak256(data))));
+    }
+
+    /// @dev Calculates the deployment address of `who` with the current nonce.
+    /// @param who The deployer address.
+    function getDeploymentAddress(address who) internal view returns (address) {
+        return getDeploymentAddress(who, getNonce(who));
+    }
 }
 
 library accounts {
@@ -167,6 +204,14 @@ library accounts {
         return accountsSafe.create(name, lbl);
     }
 
+    function getDeploymentAddress(address who, uint64 nonce) internal pure returns (address) {
+        return accountsSafe.getDeploymentAddress(who, nonce);
+    }
+
+    function getDeploymentAddress(address who) internal view returns (address) {
+        return accountsSafe.getDeploymentAddress(who);
+    }
+
     /// @dev Sets the specified `slot` in the storage of the given `self` address to the provided `value`.
     /// @param self The address to modify the storage of.
     /// @param slot The storage slot to set.

diff --git a/test/_modules/Accounts.t.sol b/test/_modules/Accounts.t.sol
@@ -1,6 +1,6 @@
 pragma solidity >=0.8.13 <0.9.0;
 
-import {Test, expect, commands, accounts, console} from "../../src/test.sol";
+import {Test, expect, commands, accounts, console, ctx} from "../../src/test.sol";
 import {Sender} from "../mocks/Sender.sol";
 
 contract AccountsTest is Test {
@@ -202,6 +202,32 @@ contract AccountsTest is Test {
         expect(token.balanceOf(user)).toEqual(balance - firstBurn - secondBurn);
         expect(token.totalSupply()).toEqual(totalSupply - firstBurn - secondBurn);
     }
+
+    function testGetDeploymentAddressWithCurrentNonce(address user, uint64 nonce) external {
+        ctx.assume(nonce < type(uint64).max);
+
+        address deploymentAddress = user.setNonce(uint64(nonce)).getDeploymentAddress();
+
+        user.impersonateOnce();
+
+        address deployedAddress = address(new TestToken());
+
+        expect(deployedAddress).toEqual(deploymentAddress);
+    }
+
+    function testGetDeploymentAddress(address user, uint64 nonce) external {
+        ctx.assume(nonce < type(uint64).max);
+
+        address deploymentAddress = user.getDeploymentAddress(nonce);
+
+        user.setNonce(nonce);
+
+        user.impersonateOnce();
+
+        address deployedAddress = address(new TestToken());
+
+        expect(deployedAddress).toEqual(deploymentAddress);
+    }
 }
 
 contract TestToken {