From 43f4dcc5eea8745a18740e83c74d7bec81028f39 Mon Sep 17 00:00:00 2001
From: Bogdan Opanchuk <bogdan@opanchuk.net>
Date: Sun, 10 Nov 2024 11:27:32 -0800
Subject: [PATCH] Add `get_create2_address()`

---
 docs/api.rst        |  6 ++++++
 docs/changelog.rst  |  2 ++
 pons/__init__.py    |  2 ++
 pons/_utils.py      | 14 ++++++++++++++
 tests/TestUtils.sol | 45 +++++++++++++++++++++++++++++++++++++++++++++
 tests/test_utils.py | 31 +++++++++++++++++++++++++++++++
 6 files changed, 100 insertions(+)
 create mode 100644 pons/_utils.py
 create mode 100644 tests/TestUtils.sol
 create mode 100644 tests/test_utils.py

diff --git a/docs/api.rst b/docs/api.rst
index f0b73a7..0fa3617 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -206,6 +206,12 @@ Utility classes
    :members:
 
 
+Utility methods
+---------------
+
+.. autofunction:: pons.get_create2_address
+
+
 Compiled and deployed contracts
 -------------------------------
 
diff --git a/docs/changelog.rst b/docs/changelog.rst
index e0c67f5..59042de 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -39,6 +39,7 @@ Added
 - ``ClientSession.eth_get_logs()`` and ``eth_get_filter_logs()``. (PR_68_)
 - Support for a custom block number in gas estimation methods. (PR_70_)
 - ``LocalProvider`` accepts an ``evm_version`` parameter. (PR_78_)
+- ``get_create2_address()``. (PR_80_)
 
 
 Fixed
@@ -73,6 +74,7 @@ Fixed
 .. _PR_76: https://github.com/fjarri-eth/pons/pull/76
 .. _PR_77: https://github.com/fjarri-eth/pons/pull/77
 .. _PR_78: https://github.com/fjarri-eth/pons/pull/78
+.. _PR_80: https://github.com/fjarri-eth/pons/pull/80
 
 
 0.7.0 (09-07-2023)
diff --git a/pons/__init__.py b/pons/__init__.py
index 7ee6a7a..72fed7e 100644
--- a/pons/__init__.py
+++ b/pons/__init__.py
@@ -49,6 +49,7 @@
 from ._local_provider import LocalProvider, SnapshotID
 from ._provider import HTTPProvider, ProtocolError, Provider, Unreachable
 from ._signer import AccountSigner, Signer
+from ._utils import get_create2_address
 
 __all__ = [
     "ABIDecodingError",
@@ -98,4 +99,5 @@
     "Unreachable",
     "abi",
     "compile_contract_file",
+    "get_create2_address",
 ]
diff --git a/pons/_utils.py b/pons/_utils.py
new file mode 100644
index 0000000..37581e3
--- /dev/null
+++ b/pons/_utils.py
@@ -0,0 +1,14 @@
+from ethereum_rpc import Address, keccak
+
+
+def get_create2_address(deployer: Address, init_code: bytes, salt: int) -> Address:
+    """
+    Returns the deterministic deployed contract address as produced by ``CREATE2`` opcode.
+    Here `deployer` is the contract address invoking ``CREATE2``
+    (**not** the transaction initiator),
+    ``init_code`` is the deployment code (see :py:attr:`~pons.BoundConstructorCall.data_bytes`),
+    and ``salt`` is an integer up to 256 bits in length.
+    """
+    salt_bytes = salt.to_bytes(32, byteorder="big")
+    contract_address = keccak(b"\xff" + bytes(deployer) + salt_bytes + keccak(init_code))[-20:]
+    return Address(contract_address)
diff --git a/tests/TestUtils.sol b/tests/TestUtils.sol
new file mode 100644
index 0000000..98c80ec
--- /dev/null
+++ b/tests/TestUtils.sol
@@ -0,0 +1,45 @@
+pragma solidity >=0.8.0 <0.9.0;
+
+
+contract ToDeploy {
+    uint256 public state;
+
+    constructor(uint256 _state) {
+        state = _state;
+    }
+
+    function getState() public view returns (uint256) {
+        return state;
+    }
+}
+
+
+contract Create2Deployer {
+    event Deployed(
+        address deployedAddress
+    );
+
+    function deploy(bytes memory bytecode, uint256 _salt) public payable {
+        address addr;
+        bool success = true;
+
+        assembly {
+            addr := create2(
+                callvalue(),
+                add(bytecode, 0x20), // Skip the first 32 bytes, which is the size of `bytecode`
+                mload(bytecode), // Load the size of code contained in the first 32 bytes
+                _salt
+            )
+
+            if iszero(extcodesize(addr)) {
+                success := false
+            }
+        }
+
+        if (!success) {
+            revert("Failed to deploy the contract");
+        }
+
+        emit Deployed(addr);
+    }
+}
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..53db8e8
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,31 @@
+import secrets
+from pathlib import Path
+
+import pytest
+
+from pons import EVMVersion, compile_contract_file, get_create2_address
+
+
+@pytest.fixture
+def compiled_contracts():
+    path = Path(__file__).resolve().parent / "TestUtils.sol"
+    return compile_contract_file(path, evm_version=EVMVersion.CANCUN)
+
+
+async def test_create2(session, root_signer, compiled_contracts):
+    compiled_deployer = compiled_contracts["Create2Deployer"]
+    compiled_to_deploy = compiled_contracts["ToDeploy"]
+
+    deployer = await session.deploy(root_signer, compiled_deployer.constructor())
+
+    salt = secrets.randbelow(1 << 256)
+    to_deploy = compiled_to_deploy.constructor(123)
+    events = await session.transact(
+        root_signer,
+        deployer.method.deploy(to_deploy.data_bytes, salt),
+        return_events=[deployer.event.Deployed],
+    )
+    assert len(events[deployer.event.Deployed]) == 1
+    assert events[deployer.event.Deployed][0] == dict(
+        deployedAddress=get_create2_address(deployer.address, to_deploy.data_bytes, salt)
+    )